import hashlib
import logging
import os
import time
from collections import OrderedDict
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)
current_expert_id: ContextVar[str | None] = ContextVar("current_expert_id", 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,
        options={"verify_aud": False, "verify_iss": True, "verify_exp": True},
    )

    # Accept any audience under the MCP server URL (e.g. /u/{slug}/)
    aud = claims.get("aud", "")
    audiences = aud if isinstance(aud, list) else [aud]
    base = mcp_server_url.rstrip("/")
    if not any(a.rstrip("/").startswith(base) for a in audiences):
        raise jwt.InvalidTokenError(f"Invalid audience: {aud}")

    return claims


INTROSPECT_CACHE_TTL = 60  # seconds
INTROSPECT_CACHE_MAX = 256


class _TTLCache:
    """Simple LRU+TTL cache bounded by size."""

    def __init__(self, maxsize: int, ttl: float):
        self._maxsize = maxsize
        self._ttl = ttl
        self._data: OrderedDict[str, tuple[dict, float]] = OrderedDict()

    def get(self, key: str) -> dict | None:
        entry = self._data.get(key)
        if entry is None:
            return None
        value, expiry = entry
        if time.monotonic() > expiry:
            del self._data[key]
            return None
        # Move to end (most recently used)
        self._data.move_to_end(key)
        return value

    def put(self, key: str, value: dict) -> None:
        self._data[key] = (value, time.monotonic() + self._ttl)
        # Evict oldest if over capacity
        while len(self._data) > self._maxsize:
            self._data.pop(next(iter(self._data)))


_introspect_cache = _TTLCache(maxsize=INTROSPECT_CACHE_MAX, ttl=INTROSPECT_CACHE_TTL)

EXPERT_SLUG_CACHE_TTL = 300  # 5 minutes
EXPERT_SLUG_CACHE_MAX = 200
_expert_slug_cache = _TTLCache(maxsize=EXPERT_SLUG_CACHE_MAX, ttl=EXPERT_SLUG_CACHE_TTL)


def _token_hash(token: str) -> str:
    return hashlib.sha256(token.encode()).hexdigest()


async def verify_token_opaque(token: str) -> dict:
    """
    Validate an opaque access token via the OAuth introspection endpoint (RFC 7662).
    Results are cached (LRU, 60s TTL) to avoid repeated introspection calls.
    """
    key = _token_hash(token)
    cached = _introspect_cache.get(key)
    if cached is not None:
        return cached

    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")

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

    _introspect_cache.put(key, claims)
    return claims


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(request)

        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(request)

        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(request)

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

        # Resolve expert from slug (set by ExpertPathRewriteMiddleware)
        expert_id = None
        slug = request.scope.get("expert_slug")
        if slug:
            expert_id = await self._resolve_expert(token, slug)
            if expert_id is None:
                return JSONResponse({"error": "Expert not found"}, status_code=404)

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

        return response

    @staticmethod
    async def _resolve_expert(token: str, slug: str) -> str | None:
        """Resolve an expert slug to an expert ID via the Django API.

        Results are cached per user (token hash) + slug for 5 minutes.
        """
        cache_key = f"{_token_hash(token)}:{slug}"
        cached = _expert_slug_cache.get(cache_key)
        if cached is not None:
            return cached.get("id")

        api_url = os.environ.get("OBRIS_API_URL", "https://api.obris.ai")
        try:
            async with httpx.AsyncClient() as client:
                resp = await client.get(
                    f"{api_url}/v1/experts",
                    params={"slug": slug},
                    headers={"Authorization": f"Bearer {token}"},
                    timeout=10.0,
                )
                if resp.status_code == 404:
                    return None
                resp.raise_for_status()
                expert_id = resp.json().get("id")
                _expert_slug_cache.put(cache_key, {"id": expert_id})
                return expert_id
        except Exception as e:
            logger.warning("Failed to resolve expert slug %s: %s", slug, e)
        return None

    @staticmethod
    def _unauthorized_response(request: Request):
        mcp_server_url = os.environ.get("MCP_SERVER_URL", "https://mcp.obris.ai")
        slug = request.scope.get("expert_slug")
        base = mcp_server_url.rstrip("/")
        resource_base = f"{base}/u/{slug}" if slug else base
        return JSONResponse(
            {"error": "unauthorized"},
            status_code=401,
            headers={
                "WWW-Authenticate": (
                    f'Bearer resource_metadata="{resource_base}/.well-known/oauth-protected-resource"'
                ),
            },
        )
