createUser backend

import json
import logging
import boto3
from botocore.exceptions import ClientError
from decimal import Decimal
from utils.dynamodb_utils import initializeDynamodb, getTableName
from utils.constants import EMOTIONS
from utils.validation_utils import validate_user_id
from utils.logging_utils import setup_logger

# Initialize DynamoDB resource
dynamodb = initializeDynamodb()

# Get the appropriate table names
USER_TABLE_NAME = getTableName('User')
ORG_TABLE_NAME = getTableName('Org')
BATCH_ANALYTICS_TABLE_NAME = getTableName('BatchAnalytics')
USER_ANALYTICS_TABLE_NAME = getTableName('UserAnalytics')
INVITE_TABLE_NAME = getTableName('PendingInvite')

# Configure logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)


def get_org_id_from_join_code(join_code):
    """
    Retrieve org_id using the join_code from the PendingInvite table.
    Mark the invite as 'used' if valid.
    """
    invite_table = dynamodb.Table(INVITE_TABLE_NAME)
    try:
        logger.info({"message": "Querying join_code", "join_code": join_code})
        response = invite_table.query(
            IndexName='CodeIndex',
            KeyConditionExpression='join_code = :join_code',
            FilterExpression='#type = :user_type',
            ExpressionAttributeNames={'#type': 'type'},
            ExpressionAttributeValues={':join_code': join_code, ':user_type': 'user'}
        )

        if response['Items'] and not response['Items'][0].get('used'):
            item = response['Items'][0]
            org_id = item.get('org_id')
            email = item['email']

            # Mark invite as used
            invite_table.update_item(
                Key={'email': email},
                UpdateExpression="SET used = :used",
                ExpressionAttributeValues={':used': True}
            )
            logger.info({"message": "Join code validated and marked as used", "org_id": org_id, "email": email})
            return org_id
        else:
            logger.warning({"message": "Invalid or already used join code", "join_code": join_code})
            return None
    except ClientError as e:
        logger.error({"message": "Error querying join_code", "error": str(e), "join_code": join_code})
        raise

def lambda_handler(event, context):
    """Main Lambda function handler."""
    setup_logger(
        logger,
        aws_request_id=context.aws_request_id,
        function_name=context.function_name,
    )
    logger.info({"message": "Lambda function invoked", "event": event})

    try:
        body = json.loads(event['body'])
        user_id = body['user_id']
        first_name = body.get('first_name')
        last_name = body.get('last_name')
        phone_number = body.get('phone_number')
        join_code = body.get('join_code')

        # Access the Cognito sub from the event context
        cognito_sub = event['requestContext']['authorizer']['claims']['sub']
        validate_user_id(user_id, cognito_sub)
        logger.info({"message": "User ID validated", "user_id": user_id, "cognito_sub": cognito_sub})

        # Reference to the User table
        user_table = dynamodb.Table(USER_TABLE_NAME)

        # Define additional fields for the User item
        update_fields = {
            'first_name': first_name,
            'last_name': last_name,
            'phone_number': phone_number,
            'is_paid_user': False
        }

        # Handle join_code if provided
        if join_code:
            org_id = get_org_id_from_join_code(join_code)
            if not org_id:
                logger.warning({"message": "Invalid join code", "join_code": join_code})
                return {
                    'statusCode': 400,
                    'body': json.dumps({'error': 'Invalid join code'})
                }
        else:
            org_id = "CarthaUsers"
        update_fields['org_id'] = org_id
        update_fields['join_code'] = join_code

        # Update the org table to increment num_users
        org_table = dynamodb.Table(ORG_TABLE_NAME)
        org_table.update_item(
            Key={'org_id': org_id},
            UpdateExpression='ADD num_users :val',
            ExpressionAttributeValues={':val': 1}
        )
        logger.info({"message": "Org table updated", "org_id": org_id})
        
        # Update the user item in the User table
        update_expression = "SET " + ", ".join(f"{k} = :{k}" for k in update_fields)
        expression_attribute_values = {f":{k}": v for k, v in update_fields.items()}

        user_table.update_item(
            Key={'user_id': user_id},
            UpdateExpression=update_expression,
            ExpressionAttributeValues=expression_attribute_values
        )
        logger.info({"message": "User table updated", "user_id": user_id, "update_fields": update_fields})

        # Prepare BatchAnalytics item
        batch_analytics_item = {
            'user_id': user_id,
            'user_message_count': 0,
            'prompt_responses': 0,
            'delayed_responses': 0,
            'satisfied_responses': 0,
            'unsatisfied_responses': 0,
            'total_response_time': Decimal(0),
            'total_sentiment_score': Decimal(0),
            'total_session_time': Decimal(0),
            'total_active_session_time': Decimal(0),
            'number_of_sessions': 1,
            'weekly_active_session_time_count': [0] * 7,
            'weekly_session_time_count': [0] * 7,
            'status': 'Normal'
        }
        batch_analytics_item.update({emotion: 0 for emotion in EMOTIONS})
        if 'org_id' in update_fields:
            batch_analytics_item['org_id'] = update_fields['org_id']
        
        batch_analytics_table = dynamodb.Table(BATCH_ANALYTICS_TABLE_NAME)
        batch_analytics_table.put_item(Item=batch_analytics_item)
        logger.info({"message": "Batch analytics item saved", "user_id": user_id})

        # Prepare UserAnalytics item
        user_analytics_item = {
            'user_id': user_id,
            'active_session_time': [],
            'total_session_time': [],
            'avg_response_times': [],
            'weekly_message_count': [0] * 7
        }
        
        user_analytics_table = dynamodb.Table(USER_ANALYTICS_TABLE_NAME)
        user_analytics_table.put_item(Item=user_analytics_item)
        logger.info({"message": "User analytics item saved", "user_id": user_id})

        return {
            'statusCode': 200,
            'body': json.dumps({'message': 'User info added successfully'})
        }
    except ClientError as e:
        logger.error({"message": "DynamoDB ClientError occurred", "error": str(e)})
        return {
            'statusCode': 500,
            'body': json.dumps({'error': str(e)})
        }
    except Exception as e:
        logger.exception({"message": "Unexpected error occurred", "error": str(e)})
        return {
            'statusCode': 500,
            'body': json.dumps({'error': f'Unexpected error: {str(e)}'})
        }