from __future__ import annotations

from datetime import UTC
from datetime import datetime

import jwt
from amsdal.contrib.auth.errors import AuthenticationError
from amsdal_server.apps.common.events.auth import AuthenticateContext
from amsdal_utils.events import AsyncNextFn
from amsdal_utils.events import EventListener
from amsdal_utils.events import NextFn

from amsdal_ml.mcp_server.mcp_oauth.services.models import InternalAuthTokenPayload
from amsdal_ml.mcp_server.mcp_oauth.services.models import JwtDecodedPayload
from amsdal_ml.mcp_server.mcp_oauth.services.types import JWT_ALGORITHM
from amsdal_ml.mcp_server.mcp_oauth.services.types import OAuthConstants
from amsdal_ml.mcp_server.mcp_oauth.services.utils import decode_jwt_payload
from amsdal_ml.mcp_server.mcp_oauth.services.utils import get_signing_key
from amsdal_ml.mcp_server.mcp_oauth.services.utils import normalize_resource
from amsdal_ml.mcp_server.mcp_oauth.services.utils import resources_match
from amsdal_ml.ml_config import ml_config


class OAuthAuthenticateBridgeListener(EventListener[AuthenticateContext]):
    def handle(
        self,
        context: AuthenticateContext,
        next_fn: NextFn[AuthenticateContext],
    ) -> AuthenticateContext:
        next_context = self._translate_context(context)
        return next_fn(next_context)

    async def ahandle(
        self,
        context: AuthenticateContext,
        next_fn: AsyncNextFn[AuthenticateContext],
    ) -> AuthenticateContext:
        next_context = self._translate_context(context)
        return await next_fn(next_context)

    def _translate_context(self, context: AuthenticateContext) -> AuthenticateContext:
        if not ml_config.oauth_enabled:
            return context

        auth_header = context.auth_header
        if not auth_header.lower().startswith(OAuthConstants.BEARER_PREFIX):
            return context

        token_parts = auth_header.split(None, 1)
        expected_parts_count = 2
        if len(token_parts) != expected_parts_count:
            return context
        token_value = token_parts[1]

        try:
            signing_key = get_signing_key()
        except (ImportError, AuthenticationError):
            return context

        try:
            token_payload = decode_jwt_payload(token_value, signing_key)
        except Exception:
            return context

        token_type = token_payload.token_type
        token_issuer = token_payload.iss
        is_oauth_access_token = (
            token_type == OAuthConstants.TOKEN_TYPE_ACCESS and token_issuer == ml_config.oauth_issuer
        )
        if not is_oauth_access_token:
            return context

        if not self._matches_resource(token_payload, context):
            return context

        subject = token_payload.sub
        if not isinstance(subject, str) or not subject:
            return context

        internal_auth_token = self._build_internal_auth_token(subject=subject, signing_key=signing_key)
        if internal_auth_token is None:
            return context

        return context.create_next(listener_id=self.listener_id, auth_header=internal_auth_token)

    def _matches_resource(self, token_payload: JwtDecodedPayload, context: AuthenticateContext) -> bool:
        request_resource = normalize_resource(str(context.connection.base_url))
        token_resource = token_payload.resource
        token_audience = token_payload.aud

        if resources_match(token_resource, request_resource):
            return True

        if isinstance(token_audience, str) and resources_match(token_audience, request_resource):
            return True

        if isinstance(token_audience, list):
            return any(isinstance(item, str) and resources_match(item, request_resource) for item in token_audience)

        return False

    def _build_internal_auth_token(self, *, subject: str, signing_key: str) -> str | None:
        try:
            from amsdal.contrib.auth.settings import auth_settings
        except ImportError:
            return None

        now_ts = int(datetime.now(tz=UTC).timestamp())
        exp_ts = now_ts + auth_settings.AUTH_TOKEN_EXPIRATION
        internal_auth_payload = InternalAuthTokenPayload(email=subject, iat=now_ts, exp=exp_ts)
        return jwt.encode(
            internal_auth_payload.model_dump(),
            signing_key,
            algorithm=JWT_ALGORITHM,
        )
