from aws_lambda_powertools.event_handler.api_gateway import (
    APIGatewayHttpResolver,
    Response,
)
from aws_lambda_powertools.event_handler.middlewares import (
    BaseMiddlewareHandler,
    NextMiddleware,
)
from pydantic import UUID4, BaseModel, EmailStr, Field

from auth import AuthFlowType


class User(BaseModel):
    id: str
    name: str
    email: EmailStr


class CognitoUser(User):
    id: str = Field(alias='custom:user_id')
    email_verified: bool
    sub: UUID4


class AuthenticationMiddleware(BaseMiddlewareHandler):
    """This middleware extracts user authentication details from the Lambda authorizer context
    and makes them available in the application context."""

    def handler(
        self,
        app: APIGatewayHttpResolver,
        next_middleware: NextMiddleware,
    ) -> Response:
        # Gets the Lambda authorizer associated with the current API Gateway event.
        # You can check the file `auth.py` for more details.
        context = app.current_event.request_context.authorizer.get_lambda
        auth_flow_type = context.get('auth_flow_type')

        if not auth_flow_type:
            return next_middleware(app)

        cls = {
            AuthFlowType.USER_AUTH: CognitoUser,
            AuthFlowType.API_AUTH: User,
        }.get(auth_flow_type)

        if cls:
            app.append_context(user=cls(**context['user']))

        return next_middleware(app)