from __future__ import annotations

import logging
from datetime import UTC
from datetime import datetime
from typing import Annotated
from typing import NoReturn

import jwt
from fastapi import Form
from jwt.types import Options

from amsdal_ml.mcp_server.mcp_oauth import oauth_router
from amsdal_ml.mcp_server.mcp_oauth.services.errors import OAuthBadRequestError
from amsdal_ml.mcp_server.mcp_oauth.services.models import AuthorizationCodeFormData
from amsdal_ml.mcp_server.mcp_oauth.services.models import RefreshTokenFormData
from amsdal_ml.mcp_server.mcp_oauth.services.models import TokenFromData
from amsdal_ml.mcp_server.mcp_oauth.services.models import TokenResponse
from amsdal_ml.mcp_server.mcp_oauth.services.types import JWT_ALGORITHM
from amsdal_ml.mcp_server.mcp_oauth.services.types import JwtClaims
from amsdal_ml.mcp_server.mcp_oauth.services.types import OAuthConstants
from amsdal_ml.mcp_server.mcp_oauth.services.utils import access_token_expires_in_seconds
from amsdal_ml.mcp_server.mcp_oauth.services.utils import build_access_payload
from amsdal_ml.mcp_server.mcp_oauth.services.utils import build_refresh_payload
from amsdal_ml.mcp_server.mcp_oauth.services.utils import consume_authorization_code_jti
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 pkce_s256
from amsdal_ml.mcp_server.mcp_oauth.services.utils import resources_match
from amsdal_ml.ml_config import ml_config

logger = logging.getLogger(__name__)


@oauth_router.post('/oauth/token', response_model=TokenResponse)
async def token(token_data: Annotated[TokenFromData, Form(...)]):
    signing_key = get_signing_key()
    now = datetime.now(tz=UTC)

    if isinstance(token_data, AuthorizationCodeFormData):
        return _handle_authorization_code(token_data, signing_key, now)

    return _handle_refresh_token(token_data, signing_key, now)


def _handle_authorization_code(
    data: AuthorizationCodeFormData,
    signing_key: str,
    now: datetime,
) -> TokenResponse:

    try:
        code_payload = jwt.decode(
            data.code,
            signing_key,
            algorithms=[JWT_ALGORITHM],
        )
    except jwt.ExpiredSignatureError:
        _raise_bad_request('authorization code expired', grant_type=data.grant_type, reason='expired_code')
    except Exception:
        _raise_bad_request('invalid authorization code', grant_type=data.grant_type, reason='invalid_code')

    if code_payload.get(JwtClaims.ISSUER.value) != ml_config.oauth_issuer:
        _raise_bad_request('invalid code issuer', grant_type=data.grant_type, reason='issuer_mismatch')

    if code_payload.get(JwtClaims.CLIENT_ID.value) != data.client_id:
        _raise_bad_request('client_id mismatch', grant_type=data.grant_type, reason='client_id_mismatch')

    if code_payload.get(JwtClaims.REDIRECT_URI.value) != data.redirect_uri:
        _raise_bad_request('redirect_uri mismatch', grant_type=data.grant_type, reason='redirect_uri_mismatch')

    code_resource_raw = code_payload.get(JwtClaims.RESOURCE.value)
    code_resource = code_resource_raw if isinstance(code_resource_raw, str) else None
    if data.resource is not None and not resources_match(data.resource, code_resource):
        _raise_bad_request(
            'resource mismatch',
            grant_type=data.grant_type,
            reason='resource_mismatch',
            requested_resource=normalize_resource(data.resource),
            code_resource=normalize_resource(code_resource),
        )

    if code_payload.get(JwtClaims.CODE_CHALLENGE_METHOD.value) != OAuthConstants.CODE_CHALLENGE_METHOD_S256:
        _raise_bad_request(
            'unsupported code_challenge_method',
            grant_type=data.grant_type,
            reason='unsupported_code_challenge_method',
        )

    expected_challenge = code_payload.get(JwtClaims.CODE_CHALLENGE.value)
    if pkce_s256(data.code_verifier) != expected_challenge:
        _raise_bad_request('invalid code_verifier', grant_type=data.grant_type, reason='invalid_code_verifier')

    code_jti = code_payload.get(JwtClaims.JWT_ID.value)
    if not isinstance(code_jti, str) or not consume_authorization_code_jti(code_jti):
        _raise_bad_request(
            'authorization code already used or unknown',
            grant_type=data.grant_type,
            reason='code_jti_not_consumable',
        )

    sub = code_payload.get(JwtClaims.SUBJECT.value)
    if not isinstance(sub, str) or not sub:
        _raise_bad_request('invalid code subject', grant_type=data.grant_type, reason='invalid_subject')

    scope = code_payload.get(JwtClaims.SCOPE.value)
    access_payload = build_access_payload(
        sub=sub,
        resource=code_resource if data.resource is None else data.resource,
        scope=scope,
        now=now,
    )
    refresh_payload = build_refresh_payload(
        sub=sub,
        client_id=data.client_id,
        resource=code_resource if data.resource is None else data.resource,
        scope=scope,
        now=now,
    )

    access_token = jwt.encode(access_payload, signing_key, algorithm=JWT_ALGORITHM)
    new_refresh_token = jwt.encode(refresh_payload, signing_key, algorithm=JWT_ALGORITHM)
    expires_in = access_token_expires_in_seconds()

    return TokenResponse(
        access_token=access_token,
        refresh_token=new_refresh_token,
        token_type=OAuthConstants.TOKEN_TYPE_BEARER,
        expires_in=expires_in,
        scope=scope if isinstance(scope, str) else None,
    )


def _handle_refresh_token(
    data: RefreshTokenFormData,
    signing_key: str,
    now: datetime,
) -> TokenResponse:

    try:
        refresh_payload = jwt.decode(
            data.refresh_token,
            signing_key,
            algorithms=[JWT_ALGORITHM],
            options=Options(verify_aud=False),
        )
    except jwt.ExpiredSignatureError:
        _raise_bad_request('refresh token expired', grant_type=data.grant_type, reason='expired_refresh_token')
    except Exception:
        _raise_bad_request('invalid refresh token', grant_type=data.grant_type, reason='invalid_refresh_token')

    if refresh_payload.get(JwtClaims.ISSUER.value) != ml_config.oauth_issuer:
        _raise_bad_request('invalid refresh issuer', grant_type=data.grant_type, reason='refresh_issuer_mismatch')

    if refresh_payload.get(JwtClaims.TOKEN_TYPE.value) != OAuthConstants.TOKEN_TYPE_REFRESH:
        _raise_bad_request('invalid token type', grant_type=data.grant_type, reason='invalid_refresh_token_type')

    token_client_id = refresh_payload.get(JwtClaims.CLIENT_ID.value)
    if token_client_id != data.client_id:
        _raise_bad_request('client_id mismatch', grant_type=data.grant_type, reason='client_id_mismatch')

    token_resource = refresh_payload.get(JwtClaims.RESOURCE.value)
    normalized_token_resource = token_resource if isinstance(token_resource, str) else None
    if data.resource is not None and not resources_match(normalized_token_resource, data.resource):
        _raise_bad_request(
            'resource mismatch',
            grant_type=data.grant_type,
            reason='resource_mismatch',
            requested_resource=normalize_resource(data.resource),
            token_resource=normalize_resource(normalized_token_resource),
        )

    sub = refresh_payload.get(JwtClaims.SUBJECT.value)
    if not isinstance(sub, str) or not sub:
        _raise_bad_request('invalid refresh subject', grant_type=data.grant_type, reason='invalid_subject')

    scope = refresh_payload.get(JwtClaims.SCOPE.value)
    effective_resource = normalized_token_resource if isinstance(normalized_token_resource, str) else data.resource
    effective_scope = scope if isinstance(scope, str) else None

    access_payload = build_access_payload(
        sub=sub,
        resource=effective_resource,
        scope=effective_scope,
        now=now,
    )

    rotated_refresh_payload = build_refresh_payload(
        sub=sub,
        client_id=data.client_id,
        resource=effective_resource,
        scope=effective_scope,
        now=now,
    )

    access_token = jwt.encode(access_payload, signing_key, algorithm=JWT_ALGORITHM)
    rotated_refresh_token = jwt.encode(rotated_refresh_payload, signing_key, algorithm=JWT_ALGORITHM)
    expires_in = access_token_expires_in_seconds()

    return TokenResponse(
        access_token=access_token,
        refresh_token=rotated_refresh_token,
        token_type=OAuthConstants.TOKEN_TYPE_BEARER,
        expires_in=expires_in,
        scope=effective_scope,
    )


def _raise_bad_request(detail: str, *, grant_type: str, reason: str, **context: object) -> NoReturn:
    logger.warning(
        'OAuth token request rejected',
        extra={
            'grant_type': grant_type,
            'reason': reason,
            **context,
        },
    )
    raise OAuthBadRequestError(detail)
