from __future__ import annotations

from fastapi import FastAPI
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient
from starlette.status import HTTP_401_UNAUTHORIZED

from amsdal_ml.mcp_server.mcp_oauth.middlwares.www_authenticate import install_www_authenticate_middleware

HTTP_OK = 200


def test_www_authenticate_added_for_mcp_401() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/mcp/protected')
    def mcp_protected():
        return JSONResponse(status_code=401, content={'detail': 'unauthorized'})

    client = TestClient(app)
    response = client.get('/mcp/protected')

    assert response.status_code == HTTP_401_UNAUTHORIZED
    header_value = response.headers.get('WWW-Authenticate')
    assert header_value is not None
    assert 'Bearer ' in header_value
    assert 'resource_metadata=' in header_value
    assert '/.well-known/oauth-protected-resource' in header_value
    assert 'error="invalid_token"' in header_value


def test_www_authenticate_added_for_mcp_sse_401() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/mcp/sse')
    def mcp_sse():
        return JSONResponse(status_code=401, content={'detail': 'unauthorized'})

    client = TestClient(app)
    response = client.get('/mcp/sse')

    assert response.status_code == HTTP_401_UNAUTHORIZED
    assert response.headers.get('WWW-Authenticate') is not None


def test_www_authenticate_not_added_for_prefixed_mcp_path_401() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/api/mcp/sse')
    def prefixed_mcp_sse():
        return JSONResponse(status_code=401, content={'detail': 'unauthorized'})

    client = TestClient(app)
    response = client.get('/api/mcp/sse')

    assert response.status_code == HTTP_401_UNAUTHORIZED
    assert response.headers.get('WWW-Authenticate') is None


def test_www_authenticate_not_added_for_non_mcp_401() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/health')
    def health():
        return JSONResponse(status_code=401, content={'detail': 'unauthorized'})

    client = TestClient(app)
    response = client.get('/health')

    assert response.status_code == HTTP_401_UNAUTHORIZED
    assert response.headers.get('WWW-Authenticate') is None


def test_www_authenticate_not_added_for_mcp_non_401() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/mcp/protected')
    def mcp_protected_ok():
        return JSONResponse(status_code=200, content={'ok': True})

    client = TestClient(app)
    response = client.get('/mcp/protected')

    assert response.status_code == HTTP_OK
    assert response.headers.get('WWW-Authenticate') is None


def test_www_authenticate_not_overridden_when_present() -> None:
    app = FastAPI()
    install_www_authenticate_middleware(app)

    @app.get('/mcp/protected')
    def mcp_protected():
        return JSONResponse(
            status_code=401,
            content={'detail': 'unauthorized'},
            headers={'WWW-Authenticate': 'Basic realm="test"'},
        )

    client = TestClient(app)
    response = client.get('/mcp/protected')

    assert response.status_code == HTTP_401_UNAUTHORIZED
    assert response.headers.get('WWW-Authenticate') == 'Basic realm="test"'
