lambda_function.py

import json
import db  # Implements the DynamoDB helper methods we built earlier
import ai  # Implements the Amazon Bedrock AI helper methods we built earlier

def lambda_handler(event, context):
    """
    Main entry point for the HabStrack serverless backend API.
    Handles routing based on the incoming HTTP method and request path.
    """
    # 1. Handle CORS Preflight requests smoothly
    http_context = event.get("requestContext", {}).get("http", {})
    method = http_context.get("method", "GET")
    path = http_context.get("path", "/")
    
    if method == "OPTIONS":
        return build_response(200, {"message": "CORS preflight OK"})

    try:
        # Extract query parameters and body safely
        query_params = event.get("queryStringParameters") or {}
        body_raw = event.get("body", "")
        body = json.loads(body_raw) if body_raw else {}

        # 2. Extract user authorization metadata 
        # For a weekend hackathon, passing a simple 'x-user-id' header works perfectly.
        headers = event.get("headers") or {}
        user_id = headers.get("x-user-id") or query_params.get("userId")
        
        if not user_id and path != "/health":
            return build_response(400, {"error": "Missing Required Header or Query Parameter: x-user-id"})

        # 3. Application Route Engine
        # Route: Health Check
        if path == "/health":
            return build_response(200, {"status": "healthy", "service": "HabStrack Core"})

        # Route: Save Profile Onboarding
        elif path == "/user/profile" and method == "POST":
            result = db.save_user_profile(user_id, body)
            return build_response(200, result)

        # Route: Get Profile Onboarding
        elif path == "/user/profile" and method == "GET":
            profile = db.get_user_profile(user_id)
            return build_response(200, {"profile": profile})

        # Route: Fetch Todos (Handles specific calendar dates via ?date=YYYY-MM-DD)
        elif path == "/todos" and method == "GET":
            date_str = query_params.get("date")
            if not date_str:
                return build_response(400, {"error": "Missing parameter 'date' (Format: YYYY-MM-DD)"})
            todos = db.get_todos_by_date(user_id, date_str)
            return build_response(200, {"todos": todos})

        # Route: Create New Todo Task
        elif path == "/todos" and method == "POST":
            date_str = body.get("date")
            title = body.get("title")
            if not date_str or not title:
                return build_response(400, {"error": "Missing 'date' or 'title' in request payload"})
            new_todo = db.create_todo(user_id, date_str, title)
            return build_response(201, new_todo)

        # Route: Toggle Todo Completion (Tracker Checkmarks)
        elif path == "/todos/toggle" and method == "PATCH":
            date_str = body.get("date")
            todo_id = body.get("todoId")
            completed = body.get("completed", False)
            if not date_str or not todo_id:
                return build_response(400, {"error": "Missing 'date' or 'todoId' in request payload"})
            
            db.update_todo_status(user_id, date_str, todo_id, completed)
            return build_response(200, {"status": "updated", "todoId": todo_id, "completed": completed})
        
        # Route: The Magic AI Stacking Button
        elif path == "/todos/stack" and method == "POST":
            date_str = body.get("date")
            if not date_str:
                return build_response(400, {"error": "Missing 'date' parameter."})
                
            # 1. Gather data
            profile = db.get_user_profile(user_id)
            todos = db.get_todos_by_date(user_id, date_str)
            
            unstacked = [t for t in todos if not t.get("is_stacked")]
            if not unstacked:
                return build_response(200, {"message": "All tasks are already stacked!"})
                
            # 2. Call Bedrock
            try:
                stacked_results = ai.generate_habit_stacks(profile, unstacked)
            except Exception as e:
                print(f"Bedrock Error: {e}")
                return build_response(502, {"error": "AI processing failed."})
                
            # 3. Update the database with the new AI groups
            for item in stacked_results:
                # We reuse the dynamodb update method logic here, or add a specific update_stack method to db.py
                db.table.update_item(
                    Key={"PK": f"USER#{user_id}", "SK": f"TODO#{date_str}#{item['todo_id']}"},
                    UpdateExpression="SET is_stacked = :is_stacked, stack_group = :grp, stack_reason = :rsn",
                    ExpressionAttributeValues={
                        ":is_stacked": True,
                        ":grp": item["stack_group"],
                        ":rsn": item["stack_reason"]
                    }
                )
                
            return build_response(200, {"status": "success", "stacked_items": stacked_results})

        # Fallback route for unmatched requests
        else:
            return build_response(404, {"error": f"Route not found: {method} {path}"})

    except Exception as e:
        print(f"Server Error occurred: {str(e)}")
        return build_response(500, {"error": "Internal Server Error", "details": str(e)})


def build_response(status_code, body_content):
    """Generates standardized HTTP responses complete with mandatory CORS headers."""
    return {
        "statusCode": status_code,
        "headers": {
            "Content-Type": "application/json",
            "Access-Control-Allow-Origin": "*",  # Wildcard allows seamless local development testing
            "Access-Control-Allow-Headers": "Content-Type,x-user-id",
            "Access-Control-Allow-Methods": "OPTIONS,GET,POST,PATCH"
        },
        "body": json.dumps(body_content)
    }