from datetime import datetime from typing import Dict, List, Optional from fastapi import HTTPException, status from models.user import AccountType, UserBase from utils.dynamo_db import DynamoDB class UserService: def __init__(self): self.db = DynamoDB() def get_or_create_user(self, user_info: Dict) -> UserBase: """ Get existing user or create a new one if it doesn't exist """ user_id = user_info.get("sub") email = user_info.get("email") existing_user = self.db.get_item(self.db.user, user_id) # Check if user exists by user_id or email if not existing_user: # Try finding by email existing_user = self.db.get_item(self.db.user, email) if existing_user: return {"userObj": UserBase(**existing_user), "isNewUser": False} # Create new user if doesn't exist first_name = user_info.get("given_name") last_name = user_info.get("family_name") # Handle case where name is a single string if "name" in user_info and not first_name: name_parts = user_info["name"].split(" ", 1) first_name = name_parts[0] last_name = name_parts[1] if len(name_parts) > 1 else None new_user = UserBase( user_id=email, email=user_info.get("email"), username=user_info.get("cognito:username"), first_name=first_name, last_name=last_name, account_type=AccountType.BANKER, ) # Save to DynamoDB self.db.user.put_item(Item=new_user.model_dump()) return {"userObj": new_user, "isNewUser": True} def get_user_by_id(self, user_id: str) -> Optional[UserBase]: """ Get user by ID """ user_data = self.db.get_item(self.db.user, user_id) if user_data: return UserBase(**user_data) return None def update_user(self, user_id: str, user_data: UserBase) -> UserBase: """ Update user details """ existing_user = self.get_user_by_id(user_id) if not existing_user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) user_data.update_modify() self.db.upload_to_dynamodb(self.db.user, user_data.model_dump()) return user_data def update_profile_picture(self, user_id: str, profile_pic_url: str) -> UserBase: """ Update user's profile picture URL """ user = self.get_user_by_id(user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) user.update_profile(profile_pic=profile_pic_url) self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user def update_account_type(self, user_id: str, account_type: AccountType) -> UserBase: """ Update user's account type """ user = self.get_user_by_id(user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) user.update_account_type(account_type) self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user def update_user_status(self, user_id: str, is_active: bool) -> UserBase: """ Activate or deactivate user account """ user = self.get_user_by_id(user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) if is_active: user.activate() else: user.deactivate() self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user def update_last_login(self, user_id: str) -> UserBase: """ Update user's last login timestamp """ user = self.get_user_by_id(user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) user.update_last_login() self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user def get_users_by_team(self, team_id: str) -> List[UserBase]: """ Get all users in a team """ # Note: This is a simplified implementation. In production, you'd want to use proper indexing # and filtering in DynamoDB for better performance all_users = self.db.get_all_table_items(self.db.user) team_users = [ UserBase(**user) for user in all_users if user.get("team_id") == team_id ] return team_users def update_user_team(self, user_id: str, team_id: str) -> UserBase: """ Update user's team ID """ user = self.get_user_by_id(user_id) if not user: raise HTTPException( status_code=status.HTTP_404_NOT_FOUND, detail=f"User with ID {user_id} not found", ) user.team_id = team_id self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user def get_user_team(self, user_id: str) -> Optional[str]: """ Get user's team ID """ user = self.get_user_by_id(user_id) return user.team_id if user else None def save_user(self, user: UserBase) -> UserBase: """ Save a new user to the database """ try: # Update modification timestamp user.update_modify() # Save to DynamoDB self.db.upload_to_dynamodb(self.db.user, user.model_dump()) return user except Exception as e: raise HTTPException( status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Failed to save user: {str(e)}", ) def get_users_by_ids(self, user_ids: List[str]) -> List[UserBase]: """ Get multiple users by their IDs using batch get item """ if not user_ids: return [] # DynamoDB batch_get_item has a limit of 100 items per request batch_size = 100 all_users = [] # Process user IDs in batches for i in range(0, len(user_ids), batch_size): batch_ids = user_ids[i : i + batch_size] keys = [{"user_id": user_id} for user_id in batch_ids] response = self.db.dynamodb.batch_get_item( RequestItems={"user": {"Keys": keys}} ) if "Responses" in response and "user" in response["Responses"]: users = response["Responses"]["user"] all_users.extend([UserBase(**user) for user in users]) return all_users def get_active_users_by_ids(self, user_ids: List[str]) -> List[UserBase]: """ Get multiple active users by their IDs using batch get item Returns only users where is_active=True """ # Get all users first users = self.get_users_by_ids(user_ids) # Filter to return only active users return [user for user in users if user.is_active]