import logging
import os
import time
from contextvars import ContextVar

import httpx
import jwt
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse

from .constants import PUBLIC_PATHS

logger = logging.getLogger("mcp_server.auth")

current_token: ContextVar[str | None] = ContextVar("current_token", default=None)

# Cache the JWKS keys in memory with TTL
_jwks_cache: dict | None = None
_jwks_cache_time: float = 0
JWKS_CACHE_TTL = 300  # 5 minutes


async def _fetch_jwks() -> dict:
    """Fetch JWKS from the authorization server, with 5-minute TTL cache."""
    global _jwks_cache, _jwks_cache_time

    if _jwks_cache is not None and (time.monotonic() - _jwks_cache_time) < JWKS_CACHE_TTL:
        return _jwks_cache

    api_url = os.environ.get("OBRIS_API_URL", "https://api.obris.ai")
    jwks_url = f"{api_url}/oauth/.well-known/jwks.json"

    async with httpx.AsyncClient() as client:
        resp = await client.get(jwks_url, timeout=10.0)
        resp.raise_for_status()
        _jwks_cache = resp.json()
        _jwks_cache_time = time.monotonic()
        return _jwks_cache


def clear_jwks_cache():
    """Clear cached JWKS (useful when keys rotate)."""
    global _jwks_cache, _jwks_cache_time
    _jwks_cache = None
    _jwks_cache_time = 0


async def verify_token_jwt(token: str) -> dict:
    """
    Decode and validate a JWT access token using the authorization server's JWKS.
    Returns the decoded claims dict.
    Raises jwt.PyJWTError on any validation failure.
    """
    issuer = os.environ.get("OAUTH_ISSUER", "https://api.obris.ai")
    mcp_server_url = os.environ.get("MCP_SERVER_URL", "https://mcp.obris.ai")

    jwks_data = await _fetch_jwks()
    jwks_client = jwt.PyJWKSet.from_dict(jwks_data)

    # Get the signing key from the token header
    header = jwt.get_unverified_header(token)
    key = None
    for jwk in jwks_client.keys:
        if jwk.key_id == header.get("kid"):
            key = jwk.key
            break

    if key is None:
        # Key not found - clear cache and retry once (key rotation)
        clear_jwks_cache()
        jwks_data = await _fetch_jwks()
        jwks_client = jwt.PyJWKSet.from_dict(jwks_data)
        for jwk in jwks_client.keys:
            if jwk.key_id == header.get("kid"):
                key = jwk.key
                break

    if key is None:
        raise jwt.InvalidTokenError("No matching signing key found")

    claims = jwt.decode(
        token,
        key,
        algorithms=["RS256"],
        issuer=issuer,
        audience=mcp_server_url,
        options={"verify_aud": True, "verify_iss": True, "verify_exp": True},
    )
    return claims


async def verify_token_opaque(token: str) -> dict:
    """
    Validate an opaque access token via the OAuth introspection endpoint (RFC 7662).
    Requires OAUTH_CLIENT_ID and OAUTH_CLIENT_SECRET env vars for the introspection request.
    """
    api_url = os.environ.get("OBRIS_API_URL", "https://api.obris.ai")
    introspect_url = f"{api_url}/oauth/introspect/"
    client_id = os.environ.get("OAUTH_CLIENT_ID", "")
    client_secret = os.environ.get("OAUTH_CLIENT_SECRET", "")

    async with httpx.AsyncClient() as client:
        resp = await client.post(
            introspect_url,
            data={"token": token},
            auth=(client_id, client_secret),
            timeout=10.0,
        )
        resp.raise_for_status()
        data = resp.json()

    if not data.get("active"):
        raise jwt.InvalidTokenError("Token is inactive or expired")

    return {
        "account_id": data.get("account_id"),
        "user_id": data.get("sub"),
        "scope": data.get("scope", ""),
    }


async def verify_token(token: str) -> dict:
    """
    Validate a Bearer token — tries JWT first, falls back to opaque DB lookup.
    """
    # JWT tokens have 3 dot-separated segments; opaque tokens don't
    if token.count(".") == 2:
        return await verify_token_jwt(token)

    logger.info("Token is opaque (not JWT), validating via database lookup")
    return await verify_token_opaque(token)



class OAuthTokenMiddleware(BaseHTTPMiddleware):
    """
    ASGI middleware that extracts and validates OAuth bearer tokens.
    Sets the account_id context var for downstream tool handlers.
    """

    async def dispatch(self, request: Request, call_next):
        # Allow public paths without auth
        if request.url.path in PUBLIC_PATHS:
            return await call_next(request)

        ua = request.headers.get("user-agent", "(none)")[:200]
        logger.info("AUTH %s %s — User-Agent: %s", request.method, request.url.path, ua)

        auth_header = request.headers.get("authorization", "")
        if not auth_header.startswith("Bearer "):
            logger.warning("AUTH REJECTED: no Bearer token — %s %s", request.method, request.url.path)
            return self._unauthorized_response()

        token = auth_header[len("Bearer "):]

        try:
            claims = await verify_token(token)
        except Exception as e:
            logger.warning("AUTH REJECTED: token validation failed — %s: %s", type(e).__name__, e)
            clear_jwks_cache()
            return self._unauthorized_response()

        account_id = claims.get("account_id")
        if not account_id:
            logger.warning("AUTH REJECTED: no account_id in token claims — claims: %s", list(claims.keys()))
            return self._unauthorized_response()

        logger.info("AUTH OK: account_id=%s scopes=%s", account_id, claims.get("scope"))

        ctx_token = current_token.set(token)
        try:
            response = await call_next(request)
        finally:
            current_token.reset(ctx_token)

        return response

    @staticmethod
    def _unauthorized_response():
        mcp_server_url = os.environ.get("MCP_SERVER_URL", "https://mcp.obris.ai")
        return JSONResponse(
            {"error": "unauthorized"},
            status_code=401,
            headers={
                "WWW-Authenticate": (
                    f'Bearer resource_metadata="{mcp_server_url}/.well-known/oauth-protected-resource"'
                ),
            },
        )