from __future__ import annotations

from amsdal_server.configs.main import settings
from fastapi import FastAPI
from fastapi import Request
from starlette.datastructures import MutableHeaders
from starlette.status import HTTP_401_UNAUTHORIZED
from starlette.types import ASGIApp
from starlette.types import Message
from starlette.types import Receive
from starlette.types import Scope
from starlette.types import Send

from amsdal_ml.mcp_server.mcp_oauth.services.types import APP_STATE_OAUTH_WWW_AUTHENTICATE_INSTALLED

_MIDDLEWARE_STATE_FLAG = APP_STATE_OAUTH_WWW_AUTHENTICATE_INSTALLED


class MCPAuthMiddleware:
    def __init__(self, app: ASGIApp):
        self.app = app

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope.get('type') != 'http':
            await self.app(scope, receive, send)
            return

        path = scope.get('path', '')
        auth_header = _scope_authorization_header(scope)

        async def send_with_www_authenticate(message: Message) -> None:
            if message.get('type') == 'http.response.start':
                status_code = int(message.get('status', 0))
                if status_code == HTTP_401_UNAUTHORIZED and _is_mcp_resource_path(path):
                    mutable_headers = MutableHeaders(scope=message)
                    if mutable_headers.get('www-authenticate') is None:
                        resource_metadata = _resource_metadata_url(Request(scope))
                        error_description = (
                            'The access token is invalid or expired.'
                            if auth_header
                            else 'Authentication required for this resource.'
                        )
                        mutable_headers['WWW-Authenticate'] = (
                            f'Bearer resource_metadata="{resource_metadata}", '
                            f'error="invalid_token", '
                            f'error_description="{error_description}"'
                        )

            await send(message)

        await self.app(scope, receive, send_with_www_authenticate)


def install_www_authenticate_middleware(app: FastAPI) -> None:
    if getattr(app.state, _MIDDLEWARE_STATE_FLAG, False):
        return

    app.add_middleware(MCPAuthMiddleware)
    setattr(app.state, _MIDDLEWARE_STATE_FLAG, True)


def build_www_authenticate_challenge(request: Request) -> str:
    resource_metadata = _resource_metadata_url(request)
    auth_header = request.headers.get(settings.AUTHORIZATION_HEADER)
    error_description = (
        'The access token is invalid or expired.' if auth_header else 'Authentication required for this resource.'
    )
    return (
        f'Bearer resource_metadata="{resource_metadata}", '
        f'error="invalid_token", '
        f'error_description="{error_description}"'
    )


def _is_mcp_resource_path(path: str) -> bool:
    normalized_path = path.rstrip('/') or '/'
    if normalized_path == '/mcp':
        return True
    return normalized_path.startswith('/mcp/')


def _resource_metadata_url(request: Request) -> str:
    base_url = str(request.base_url).rstrip('/')
    return f'{base_url}/.well-known/oauth-protected-resource'


def _scope_authorization_header(scope: Scope) -> str | None:
    header_name = settings.AUTHORIZATION_HEADER.lower().encode('latin1')
    for raw_name, raw_value in scope.get('headers', []):
        if raw_name.lower() == header_name:
            return raw_value.decode('latin1')
    return None
