from aws_lambda_powertools.event_handler.api_gateway import (
    APIGatewayHttpResolver,
    Response,
)
from aws_lambda_powertools.event_handler.middlewares import (
    BaseMiddlewareHandler,
    NextMiddleware,
)
from layercake.dateutils import now, ttl
from layercake.dynamodb import ComposeKey, DynamoDBCollection, KeyPair
from pydantic import UUID4, BaseModel, Field


class AuthorizerMiddleware(BaseMiddlewareHandler):
    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.
        authorizer = app.current_event.request_context.authorizer.get_lambda

        if 'user' in authorizer:
            user = authorizer['user']
            app.append_context(authenticated_user=AuthenticatedUser(**user))

        return next_middleware(app)


class AuthenticatedUser(BaseModel):
    id: str = Field(alias='custom:user_id')
    name: str
    email: str
    email_verified: bool
    sub: UUID4


class AuditLogMiddleware(BaseMiddlewareHandler):
    def __init__(self, action: str, /, collect: DynamoDBCollection) -> None:
        self.action = action
        self.collect = collect

    def handler(
        self,
        app: APIGatewayHttpResolver,
        next_middleware: NextMiddleware,
    ) -> Response:
        response = next_middleware(app)
        auth_user = app.context.get('authenticated_user')
        request_ctx = app.current_event.request_context
        ip_addr = request_ctx.http.source_ip

        # Successful request
        if 200 <= response.status_code < 300 and auth_user:
            now_ = now()

            self.collect.put_item(
                key=KeyPair(
                    pk=ComposeKey(auth_user.id, prefix='logs'),
                    sk=now_.isoformat(),
                ),
                action=self.action,
                ip=ip_addr,
                ttl=ttl(start_dt=now_, days=365 * 2),
            )

        return response