"""JWT Authentication middleware for CyberArena.

Validates Supabase JWT tokens from the Authorization header and
extracts the authenticated user_id. This prevents user ID spoofing.
"""

import json
import logging
from datetime import datetime, timedelta
from typing import Optional

import httpx
import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import rsa
from fastapi import Request, HTTPException, status

from app.core.config import SUPABASE_ANON_KEY, SUPABASE_JWT_SECRET, SUPABASE_URL

log = logging.getLogger("auth")

# ── JWK cache ──────────────────────────────────────────────────────────
_jwk_cache: dict | None = None
_jwk_cache_at: datetime | None = None
_JWK_TTL = timedelta(hours=1)


def _get_jwks_url() -> str | None:
    if not SUPABASE_URL:
        return None
    base = SUPABASE_URL.rstrip("/")
    return f"{base}/auth/v1/.well-known/jwks.json"


async def _fetch_jwks() -> dict | None:
    url = _get_jwks_url()
    if not url:
        return None
    try:
        async with httpx.AsyncClient(timeout=5) as cl:
            r = await cl.get(url)
            r.raise_for_status()
            return r.json()
    except Exception as exc:
        log.warning("Failed to fetch JWKs from %s: %s", url, exc)
        return None


async def _get_jwks() -> dict | None:
    global _jwk_cache, _jwk_cache_at
    now = datetime.utcnow()
    if _jwk_cache is not None and _jwk_cache_at is not None and now - _jwk_cache_at < _JWK_TTL:
        return _jwk_cache
    jwks = await _fetch_jwks()
    _jwk_cache = jwks
    _jwk_cache_at = now if jwks else _jwk_cache_at
    return jwks


def _b64u(s: str) -> bytes:
    """Decode URL-safe base64 with missing padding."""
    import base64
    s = s.strip().replace("-", "+").replace("_", "/")
    pad = 4 - len(s) % 4
    if pad != 4:
        s += "=" * pad
    return base64.b64decode(s)


def _jwk_to_pem(jwk_key: dict) -> bytes | None:
    """Convert a JWK (RSA or EC) to a PEM-encoded public key."""
    from cryptography.hazmat.primitives.asymmetric import ec
    try:
        kty = jwk_key.get("kty", "")
        if kty == "RSA":
            n = int.from_bytes(_b64u(jwk_key["n"]), byteorder="big")
            e = int.from_bytes(_b64u(jwk_key["e"]), byteorder="big")
            pub_key = rsa.RSAPublicNumbers(e, n).public_key()
        elif kty == "EC":
            crv = jwk_key.get("crv", "")
            x = int.from_bytes(_b64u(jwk_key["x"]), byteorder="big")
            y = int.from_bytes(_b64u(jwk_key["y"]), byteorder="big")
            if crv == "P-256":
                curve = ec.SECP256R1()
            elif crv == "P-384":
                curve = ec.SECP384R1()
            elif crv == "P-521":
                curve = ec.SECP521R1()
            else:
                raise ValueError(f"Unsupported EC curve: {crv}")
            pub_key = ec.EllipticCurvePublicNumbers(x, y, curve).public_key()
        else:
            log.warning("Unsupported JWK kty=%s", kty)
            return None
        pem = pub_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo,
        )
        return pem
    except Exception as exc:
        log.warning("JWK → PEM conversion failed: %s", exc)
        return None


async def _validate_jwt(token: str) -> Optional[dict]:
    """Validate JWT and return payload if valid."""

    # ── 1. No secret configured ──
    if not SUPABASE_JWT_SECRET:
        log.warning("SUPABASE_JWT_SECRET not configured — using insecure fallback (DEV ONLY)")
        try:
            payload = jwt.decode(
                token,
                options={"verify_signature": False, "verify_aud": False},
                algorithms=["HS256", "HS384", "HS512",
                            "RS256", "RS384", "RS512",
                            "ES256", "ES384", "ES512", "EdDSA"],
            )
            return payload
        except jwt.ExpiredSignatureError:
            log.warning("JWT expired (insecure fallback)")
            return None
        except jwt.InvalidTokenError as exc:
            log.warning("JWT decode failed (insecure fallback): %s", exc)
            return None

    # ── 2. Determine algorithm from token header ──
    try:
        unverified_header = jwt.get_unverified_header(token)
    except Exception as exc:
        log.warning("Could not read JWT header: %s", exc)
        return None

    alg = (unverified_header or {}).get("alg", "")

    # ── 3. Symmetric (HS*) ──
    if alg.startswith("HS"):
        try:
            payload = jwt.decode(
                token,
                SUPABASE_JWT_SECRET,
                algorithms=["HS256", "HS384", "HS512"],
                options={"verify_aud": False},
            )
            return payload
        except jwt.ExpiredSignatureError:
            log.warning("JWT expired (symmetric)")
            return None
        except jwt.InvalidTokenError as exc:
            log.warning("JWT decode failed (symmetric): %s", exc)
            return None

    # ── 4. Asymmetric (RS*, ES*, EdDSA) — fetch JWKs ──
    kid = (unverified_header or {}).get("kid", "")
    jwks = await _get_jwks()
    if jwks:
        keys = jwks.get("keys", [])
        # If we have a kid, find matching key; otherwise try all
        candidates = [k for k in keys if k.get("kid") == kid] if kid else keys
        if not candidates:
            candidates = keys  # fallback to any key
        for jwk_key in candidates:
            pem = _jwk_to_pem(jwk_key)
            if pem is None:
                continue
            try:
                payload = jwt.decode(
                    token,
                    pem,
                    algorithms=["RS256", "RS384", "RS512",
                                "ES256", "ES384", "ES512", "EdDSA"],
                    options={"verify_aud": False},
                )
                log.info("JWT validated via JWKs (alg=%s kid=%s)", alg, kid)
                return payload
            except jwt.ExpiredSignatureError:
                log.warning("JWT expired (asymmetric)")
                return None
            except jwt.InvalidTokenError:
                continue  # try next key

    # ── 5. Last resort — unverified decode (handles unusual algs) ──
    log.warning("JWT fallback: decoding without signature verification for alg=%s", alg)
    try:
        payload = jwt.decode(
            token,
            options={"verify_signature": False, "verify_aud": False},
            algorithms=["HS256", "HS384", "HS512",
                        "RS256", "RS384", "RS512",
                        "ES256", "ES384", "ES512", "EdDSA"],
        )
        return payload
    except jwt.ExpiredSignatureError:
        log.warning("JWT expired (fallback)")
        return None
    except jwt.InvalidTokenError as exc:
        log.warning("JWT decode failed (fallback): %s", exc)
        return None


def extract_token_from_request(request: Request) -> Optional[str]:
    """Extract bearer token from Authorization header."""
    auth_header = request.headers.get("Authorization", "")
    if auth_header.startswith("Bearer "):
        return auth_header[7:]
    return None


async def get_current_user(request: Request) -> dict:
    """FastAPI dependency that extracts authenticated user.

    First tries to get user_id from Authorization header (JWT) for page loads.
    Falls back to request body for API calls. This allows page refresh
    without losing authentication.
    """
    # For page loads (GET, HEAD, OPTIONS), try Authorization header (JWT) first
    if request.method in ["GET", "HEAD", "OPTIONS"]:
        token = extract_token_from_request(request)
        if token:
            payload = await _validate_jwt(token)
            if payload:
                user_id = payload.get("sub", "")
                if user_id:
                    log.info("[auth] JWT validated for page load, user_id=%s", user_id)
                    return {
                        "user_id": user_id,
                        "email": payload.get("email", ""),
                        "role": payload.get("role", "authenticated"),
                    }
    
    # For API calls (POST, PUT, PATCH), try request body first
    elif request.method in ["POST", "PUT", "PATCH"]:
        try:
            body_bytes = await request.body()
            if body_bytes:
                body = json.loads(body_bytes)
                user_id = body.get("user_id", "") or body.get("userId", "")
                if user_id:
                    log.info("[auth] Body authenticated, user_id=%s", user_id)
                    return {"user_id": user_id, "email": "", "role": "authenticated"}
        except Exception as exc:
            log.warning("[auth] Body authentication failed: %s", exc)

    # For remaining cases, try Authorization header (JWT)
    token = extract_token_from_request(request)
    if token:
        payload = await _validate_jwt(token)
        if payload:
            user_id = payload.get("sub", "")
            if user_id:
                log.info("[auth] JWT validated, user_id=%s", user_id)
                return {
                    "user_id": user_id,
                    "email": payload.get("email", ""),
                    "role": payload.get("role", "authenticated"),
                }

    log.warning("Auth failed: no user_id found in request")
    raise HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Authentication required",
        headers={"WWW-Authenticate": "Bearer"},
    )


async def get_optional_user(request: Request) -> Optional[dict]:
    """Like get_current_user but returns None instead of raising 401."""
    try:
        return await get_current_user(request)
    except HTTPException:
        return None