from __future__ import annotations

from datetime import UTC
from datetime import datetime
from datetime import timedelta
from typing import Any

import jwt
from amsdal_server.apps.common.events.auth import AuthenticateContext
from starlette.requests import HTTPConnection

from amsdal_ml.mcp_server.mcp_oauth.event_listeners.auth_bridge_listener import OAuthAuthenticateBridgeListener
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.ml_config import ml_config


def _build_connection(base_url: str) -> HTTPConnection:
    scope = {
        'type': 'http',
        'method': 'GET',
        'path': '/',
        'raw_path': b'/',
        'query_string': b'',
        'headers': [],
        'scheme': 'https',
        'server': ('mcp.example.com', 443),
        'client': ('127.0.0.1', 12345),
    }
    connection = HTTPConnection(scope)
    if base_url != 'https://mcp.example.com/':
        alt_scope = {
            **scope,
            'scheme': 'https',
            'server': (base_url.replace('https://', '').rstrip('/'), 443),
        }
        connection = HTTPConnection(alt_scope)
    return connection


def _build_oauth_access_token(signing_key: str, *, sub: str, resource: str, issuer: str) -> str:
    now = datetime.now(tz=UTC)
    payload: dict[str, Any] = {
        JwtClaims.ISSUER.value: issuer,
        JwtClaims.SUBJECT.value: sub,
        JwtClaims.AUDIENCE.value: resource,
        JwtClaims.RESOURCE.value: resource,
        JwtClaims.TOKEN_TYPE.value: OAuthConstants.TOKEN_TYPE_ACCESS.value,
        JwtClaims.ISSUED_AT.value: int(now.timestamp()),
        JwtClaims.EXPIRATION.value: int((now + timedelta(minutes=30)).timestamp()),
    }
    return jwt.encode(payload, signing_key, algorithm=JWT_ALGORITHM)


def test_bridge_passes_through_non_bearer():
    listener = OAuthAuthenticateBridgeListener()
    context = AuthenticateContext(
        auth_header='raw-jwt-token',
        connection=_build_connection('https://mcp.example.com/'),
    )

    result = listener.handle(context, lambda ctx: ctx)

    assert result.auth_header == 'raw-jwt-token'


def test_bridge_translates_valid_oauth_access_bearer_to_internal_token(oauth_signing_key, monkeypatch):
    monkeypatch.setattr(ml_config, 'oauth_issuer', 'https://issuer.example')
    listener = OAuthAuthenticateBridgeListener()
    oauth_access = _build_oauth_access_token(
        oauth_signing_key,
        sub='user@example.com',
        resource='https://mcp.example.com',
        issuer='https://issuer.example',
    )
    context = AuthenticateContext(
        auth_header=f'Bearer {oauth_access}',
        connection=_build_connection('https://mcp.example.com/'),
    )

    result = listener.handle(context, lambda ctx: ctx)

    assert result.auth_header != context.auth_header
    internal_payload = jwt.decode(
        result.auth_header,
        oauth_signing_key,
        algorithms=[JWT_ALGORITHM],
    )
    assert internal_payload['email'] == 'user@example.com'


def test_bridge_passes_through_non_oauth_bearer_token(oauth_signing_key):
    now = datetime.now(tz=UTC)
    raw_internal_token = jwt.encode(
        {
            'email': 'user@example.com',
            JwtClaims.ISSUED_AT.value: int(now.timestamp()),
            JwtClaims.EXPIRATION.value: int((now + timedelta(minutes=30)).timestamp()),
        },
        oauth_signing_key,
        algorithm=JWT_ALGORITHM,
    )
    listener = OAuthAuthenticateBridgeListener()
    context = AuthenticateContext(
        auth_header=f'Bearer {raw_internal_token}',
        connection=_build_connection('https://mcp.example.com/'),
    )

    result = listener.handle(context, lambda ctx: ctx)

    assert result.auth_header == context.auth_header


def test_bridge_passes_through_oauth_token_with_resource_mismatch(oauth_signing_key, monkeypatch):
    monkeypatch.setattr(ml_config, 'oauth_issuer', 'https://issuer.example')
    listener = OAuthAuthenticateBridgeListener()
    oauth_access = _build_oauth_access_token(
        oauth_signing_key,
        sub='user@example.com',
        resource='https://other.example.com',
        issuer='https://issuer.example',
    )
    context = AuthenticateContext(
        auth_header=f'Bearer {oauth_access}',
        connection=_build_connection('https://mcp.example.com/'),
    )

    result = listener.handle(context, lambda ctx: ctx)

    assert result.auth_header == context.auth_header


def test_bridge_translates_when_audience_list_contains_resource(oauth_signing_key, monkeypatch):
    monkeypatch.setattr(ml_config, 'oauth_issuer', 'https://issuer.example')
    now = datetime.now(tz=UTC)
    token_payload = {
        JwtClaims.ISSUER.value: 'https://issuer.example',
        JwtClaims.SUBJECT.value: 'user@example.com',
        JwtClaims.AUDIENCE.value: ['https://other.example.com', 'https://mcp.example.com'],
        JwtClaims.TOKEN_TYPE.value: OAuthConstants.TOKEN_TYPE_ACCESS.value,
        JwtClaims.ISSUED_AT.value: int(now.timestamp()),
        JwtClaims.EXPIRATION.value: int((now + timedelta(minutes=30)).timestamp()),
    }
    oauth_access = jwt.encode(token_payload, oauth_signing_key, algorithm=JWT_ALGORITHM)

    listener = OAuthAuthenticateBridgeListener()
    context = AuthenticateContext(
        auth_header=f'Bearer {oauth_access}',
        connection=_build_connection('https://mcp.example.com/'),
    )

    result = listener.handle(context, lambda ctx: ctx)

    internal_payload = jwt.decode(
        result.auth_header,
        oauth_signing_key,
        algorithms=[JWT_ALGORITHM],
    )
    assert internal_payload['email'] == 'user@example.com'


def test_bridge_async_handle_translates_token(oauth_signing_key, monkeypatch):
    monkeypatch.setattr(ml_config, 'oauth_issuer', 'https://issuer.example')
    listener = OAuthAuthenticateBridgeListener()
    oauth_access = _build_oauth_access_token(
        oauth_signing_key,
        sub='user@example.com',
        resource='https://mcp.example.com',
        issuer='https://issuer.example',
    )
    context = AuthenticateContext(
        auth_header=f'Bearer {oauth_access}',
        connection=_build_connection('https://mcp.example.com/'),
    )

    async def _next(ctx: AuthenticateContext) -> AuthenticateContext:
        return ctx

    import asyncio

    result = asyncio.run(listener.ahandle(context, _next))

    internal_payload = jwt.decode(
        result.auth_header,
        oauth_signing_key,
        algorithms=[JWT_ALGORITHM],
    )
    assert internal_payload['email'] == 'user@example.com'
