from __future__ import annotations

import logging
import uuid
from datetime import UTC
from datetime import datetime
from datetime import timedelta
from typing import Annotated
from typing import Literal
from urllib.parse import parse_qs
from urllib.parse import urlencode
from urllib.parse import urlparse
from urllib.parse import urlunparse

import jwt
from fastapi import Depends
from fastapi import HTTPException
from fastapi import Query
from fastapi import Request
from fastapi.responses import RedirectResponse

from amsdal_ml.mcp_server.mcp_oauth import oauth_router
from amsdal_ml.mcp_server.mcp_oauth.services.dependencies import get_current_user_sub
from amsdal_ml.mcp_server.mcp_oauth.services.models import AuthorizeResponse
from amsdal_ml.mcp_server.mcp_oauth.services.types import JWT_ALGORITHM
from amsdal_ml.mcp_server.mcp_oauth.services.types import ClientIdTokenClaims
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.types import ResponseTypes
from amsdal_ml.mcp_server.mcp_oauth.services.utils import get_signing_key
from amsdal_ml.mcp_server.mcp_oauth.services.utils import store_authorization_code_jti
from amsdal_ml.ml_config import ml_config

logger = logging.getLogger(__name__)


@oauth_router.get('/oauth/authorize')
async def authorize(
    request: Request,
    user_sub: Annotated[str | None, Depends(get_current_user_sub)] = None,
    response_type: Literal[ResponseTypes.CODE] = Query(...),
    client_id: str = Query(...),
    redirect_uri: str | None = Query(None),
    state: str = Query(...),
    code_challenge: str = Query(...),
    code_challenge_method: Literal[OAuthConstants.CODE_CHALLENGE_METHOD_S256] = Query(...),
    scope: str | None = Query(None),
    resource: str | None = Query(None),
):
    _ = response_type

    signing_key = get_signing_key()

    try:
        client_payload = jwt.decode(
            client_id,
            signing_key,
            algorithms=[JWT_ALGORITHM],
        )
    except jwt.ExpiredSignatureError as e:
        raise HTTPException(status_code=400, detail='client_id expired') from e
    except Exception as exc:
        logger.exception('Failed to decode client_id')
        raise HTTPException(status_code=400, detail='invalid client_id') from exc

    if client_payload.get(ClientIdTokenClaims.ISSUER.value) != ml_config.oauth_issuer:
        raise HTTPException(status_code=400, detail='invalid client issuer')

    allowed = client_payload.get(ClientIdTokenClaims.REDIRECT_URIS.value) or []
    if allowed and redirect_uri not in allowed:
        raise HTTPException(status_code=400, detail='redirect_uri not allowed')

    if not user_sub:
        next_url = str(request.url)
        login_query = urlencode({'next': next_url})
        login_url = f'{ml_config.oauth_login_path}?{login_query}'
        return RedirectResponse(login_url)

    now = datetime.now(tz=UTC)
    iat = int(now.timestamp())
    exp_time = now + timedelta(minutes=ml_config.oauth_code_exp_minutes)
    exp_ts = int(exp_time.timestamp())
    jti = str(uuid.uuid4())

    code_payload = {
        JwtClaims.CLIENT_ID.value: client_id,
        JwtClaims.CODE_CHALLENGE.value: code_challenge,
        JwtClaims.CODE_CHALLENGE_METHOD.value: code_challenge_method,
        JwtClaims.SUBJECT.value: user_sub,
        JwtClaims.REDIRECT_URI.value: redirect_uri,
        JwtClaims.SCOPE.value: scope,
        JwtClaims.ISSUED_AT.value: iat,
        JwtClaims.EXPIRATION.value: exp_ts,
        JwtClaims.JWT_ID.value: jti,
        JwtClaims.RESOURCE.value: resource,
        JwtClaims.STATE.value: state,
        JwtClaims.ISSUER.value: ml_config.oauth_issuer,
    }

    code = jwt.encode(
        code_payload,
        signing_key,
        algorithm=JWT_ALGORITHM,
    )

    store_authorization_code_jti(jti, exp_ts)

    if redirect_uri:
        parsed = urlparse(redirect_uri)
        query = parse_qs(parsed.query)
        query[JwtClaims.CODE.value] = [code]
        query[JwtClaims.STATE.value] = [state]
        new_query = urlencode(query, doseq=True)
        redirect_to = urlunparse(parsed._replace(query=new_query))
        return RedirectResponse(url=redirect_to)

    return AuthorizeResponse(code=code, state=state)
