"""
cortex/api/auth.py — Authentication module for Cortex hosted service.

Supports two auth methods (checked in order):
    1. Supabase JWT (Authorization: Bearer <jwt>) — primary for hosted users
    2. API Key (X-Cortex-API-Key or X-API-Key) — fallback for self-host + power users

JWT validation (stdlib only — zero external dependencies):
    - Verify HS256 signature with SUPABASE_JWT_SECRET
    - Check expiration (exp claim)
    - Extract subject (sub claim) → owner_id
    - Resolve owner_id via flyio-sync lookup (with in-memory TTL cache)
    - Fall back to sub UUID when flyio-sync returns null or errors

The existing CORTEX_API_KEY env var remains for self-host mode.
When neither SUPABASE_JWT_SECRET nor CORTEX_API_KEY is set → open access (local dev).

Owner ID resolution (runtime, #737):
    JWT sub (UUID) → POST flyio-sync get_cortex_owner_id → "andrew-main" (or null)
    Result cached per-sub with TTL=3600s to avoid per-request roundtrips.
    Fallback: sub UUID (safe, unique, backward compatible).

    This replaces the earlier app_metadata / JWT hook approach, which required
    Supabase auth settings unavailable in Lovable cloud deployments.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import logging
import os
import time
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

import httpx

logger = logging.getLogger(__name__)


@dataclass
class AuthResult:
    """Result of a successful authentication check."""

    owner_id: str
    """Resolved owner ID: mapped friendly name (e.g. 'andrew-main') or Supabase UUID."""

    method: str
    """Auth method used: 'jwt' | 'api_key' | 'open'."""

    plan: Optional[str] = None
    """Subscription plan from JWT custom claims: 'free' | 'pro' | 'scale' | None."""


class AuthError(Exception):
    """Raised when a provided credential fails validation."""


# ---------------------------------------------------------------------------
# Base64url helpers (no external dependency)
# ---------------------------------------------------------------------------

def _b64url_decode(s: str) -> bytes:
    """Base64url-decode a string, adding padding as needed."""
    padding = (4 - len(s) % 4) % 4
    return base64.urlsafe_b64decode(s + "=" * padding)


def _b64url_encode(data: bytes) -> str:
    """Base64url-encode bytes, stripping trailing padding."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


# ---------------------------------------------------------------------------
# flyio-sync owner ID lookup + session-level cache
# ---------------------------------------------------------------------------

def _get_flyio_sync_url() -> Optional[str]:
    """Return the flyio-sync endpoint URL.

    Resolution order:
      1. FLYIO_SYNC_URL env var (explicit override — staging / self-hosted)
      2. SUPABASE_URL env var + "/functions/v1/flyio-sync"
      3. Hardcoded production fallback (last resort — avoids silent misconfiguration)
    """
    explicit = os.getenv("FLYIO_SYNC_URL")
    if explicit:
        return explicit
    supabase_base = os.getenv("SUPABASE_URL", "").rstrip("/")
    if supabase_base:
        return f"{supabase_base}/functions/v1/flyio-sync"
    logger.warning(
        "Neither FLYIO_SYNC_URL nor SUPABASE_URL is set — "
        "flyio-sync owner ID lookup will be skipped. "
        "Set SUPABASE_URL (or FLYIO_SYNC_URL) for hosted deployments."
    )
    return None


_OWNER_ID_CACHE_TTL = 3600  # seconds — successful (non-UUID) resolutions
_OWNER_ID_CACHE_TTL_FAILURE = 60  # seconds — failed / UUID-fallback resolutions

# Cache: sub UUID → (owner_id, cached_at timestamp)
_owner_id_cache: Dict[str, Tuple[str, float]] = {}


def _is_uuid(value: str) -> bool:
    """Return True if *value* looks like a UUID (8-4-4-4-12 hex groups)."""
    import re
    _UUID_RE = re.compile(
        r"^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$",
        re.IGNORECASE,
    )
    return bool(_UUID_RE.match(value))


def _lookup_cortex_owner_id(sub: str, flyio_sync_secret: str) -> Optional[str]:
    """
    Call flyio-sync to resolve a Supabase user UUID to a cortex_owner_id.

    POST /flyio-sync {"action": "get_cortex_owner_id", "user_id": sub}
    → {"cortex_owner_id": "andrew-main"} or {"cortex_owner_id": null}

    Returns the cortex_owner_id string, or None on null/error.
    Uses stdlib urllib — zero external dependencies.
    """
    sync_url = _get_flyio_sync_url()
    if sync_url is None:
        return None
    body = json.dumps({"action": "get_cortex_owner_id", "user_id": sub}).encode()
    req = urllib.request.Request(
        sync_url,
        data=body,
        headers={
            "Content-Type": "application/json",
            "x-flyio-sync-secret": flyio_sync_secret,
        },
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=5) as resp:
            data = json.loads(resp.read())
            owner_id = data.get("cortex_owner_id")
            return str(owner_id) if owner_id else None
    except Exception:
        # Network errors, timeouts, non-200 responses → fall back gracefully
        return None


def _resolve_owner_id(sub: str, flyio_sync_secret: Optional[str]) -> str:
    """
    Resolve a Supabase JWT sub (UUID) to a cortex_owner_id.

    Resolution order:
      1. In-memory cache (TTL = 3600s)
      2. flyio-sync lookup (if FLYIO_SYNC_SECRET configured)
      3. Fallback → sub UUID (always safe and unique)

    Args:
        sub:               Supabase user UUID from JWT 'sub' claim.
        flyio_sync_secret: Value of FLYIO_SYNC_SECRET env var, or None.

    Returns:
        Resolved owner_id string.
    """
    now = time.time()

    # 1. Cache hit?
    cached = _owner_id_cache.get(sub)
    if cached is not None:
        owner_id, cached_at = cached
        # Use short TTL for UUID fallbacks (transient failures), full TTL for resolved names
        ttl = _OWNER_ID_CACHE_TTL_FAILURE if _is_uuid(owner_id) else _OWNER_ID_CACHE_TTL
        if now - cached_at < ttl:
            return owner_id
        # Expired — evict
        del _owner_id_cache[sub]

    # 2. flyio-sync lookup (if secret available)
    resolved: Optional[str] = None
    if flyio_sync_secret:
        resolved = _lookup_cortex_owner_id(sub, flyio_sync_secret)

    # 3. Fallback to sub UUID
    owner_id = resolved if resolved else str(sub)

    # Cache the result:
    # - Resolved non-UUID names → full TTL (3600s)
    # - UUID fallbacks / null resolutions → short TTL (60s) so transient
    #   flyio-sync failures don't stay sticky for a full hour.
    _owner_id_cache[sub] = (owner_id, now)
    return owner_id


# ---------------------------------------------------------------------------
# AuthProvider
# ---------------------------------------------------------------------------

class AuthProvider:
    """
    Validates Supabase JWTs and API keys.

    Auth priority in :meth:`authenticate`:
        1. ``Authorization: Bearer <jwt>`` → JWT validation (if jwt_secret set)
        2. ``X-Cortex-API-Key: <key>`` → API key validation (if api_key set)
        3. ``X-API-Key: <key>`` → Legacy API key (backward compat)
        4. Open access (owner_id = "default")

    JWT owner resolution uses the flyio-sync edge function to map a Supabase
    user UUID to a friendly cortex_owner_id (e.g. "andrew-main").  Results
    are cached per-session with a 1-hour TTL.  On null/error, falls back to
    the raw Supabase UUID.
    """

    def __init__(
        self,
        jwt_secret: Optional[str] = None,
        api_key: Optional[str] = None,
        flyio_sync_secret: Optional[str] = None,
        api_key_owner: Optional[str] = None,
    ) -> None:
        # Store None rather than empty string so bool checks work
        self._jwt_secret: Optional[str] = jwt_secret or None
        self._api_key: Optional[str] = api_key or None
        self._flyio_sync_secret: Optional[str] = flyio_sync_secret or None
        self._api_key_owner: Optional[str] = api_key_owner or None

    # ------------------------------------------------------------------
    # JWT validation
    # ------------------------------------------------------------------

    def validate_jwt(self, token: str) -> AuthResult:
        """
        Validate a Supabase HS256 JWT and extract owner_id.

        Supabase JWTs are HS256-signed with the project's JWT secret.
        The ``sub`` claim contains the Supabase user UUID, which is then
        resolved to a cortex_owner_id via flyio-sync (with 1h cache).

        Args:
            token: Raw JWT string (without 'Bearer ' prefix).

        Returns:
            :class:`AuthResult` with ``method='jwt'``.

        Raises:
            :class:`AuthError`: Invalid format, bad signature, or expired token.
        """
        if not self._jwt_secret:
            raise AuthError("JWT secret not configured")

        # Step 1 — split into header.payload.signature
        parts = token.split(".")
        if len(parts) != 3:
            raise AuthError("Invalid JWT format: expected 3 dot-separated parts")

        header_b64, payload_b64, signature_b64 = parts

        # Step 2 — verify HMAC-SHA256 signature
        signing_input = f"{header_b64}.{payload_b64}".encode()
        expected_sig = hmac.new(
            self._jwt_secret.encode(),
            signing_input,
            hashlib.sha256,
        ).digest()

        try:
            actual_sig = _b64url_decode(signature_b64)
        except Exception:
            raise AuthError("Invalid JWT signature encoding")

        # Constant-time comparison prevents timing-based attacks
        if not hmac.compare_digest(expected_sig, actual_sig):
            raise AuthError("JWT signature verification failed")

        # Step 3 — decode payload
        try:
            payload = json.loads(_b64url_decode(payload_b64))
        except Exception:
            raise AuthError("JWT payload could not be decoded")

        # Step 4 — check expiration
        exp = payload.get("exp")
        if exp is None:
            raise AuthError("JWT missing exp claim")
        if time.time() > float(exp):
            raise AuthError("JWT has expired")

        # Step 5 — extract sub claim
        sub = payload.get("sub")
        if not sub:
            raise AuthError("JWT missing sub claim")

        # Step 6 — resolve owner_id via flyio-sync lookup (cached, with UUID fallback)
        # Mapping: Supabase UUID → cortex_owner_id (e.g. "andrew-main")
        # Data lives in Lovable/Supabase edge function; no SQL migration needed.
        owner_id = _resolve_owner_id(str(sub), self._flyio_sync_secret)

        # Optional plan from custom claims (set via Supabase hooks / RLS)
        plan: Optional[str] = payload.get("plan") or payload.get("user_plan") or None

        return AuthResult(owner_id=owner_id, method="jwt", plan=plan)

    # ------------------------------------------------------------------
    # API key validation
    # ------------------------------------------------------------------

    async def validate_api_key(self, key: str) -> AuthResult:
        """
        Validate an API key.

        Admin/self-host key (CORTEX_API_KEY) → ``owner_id='default'`` (backward
        compat).  All other keys → SHA-256 hashed and verified against the
        flyio-sync Supabase edge function which returns the real owner_id.

        Uses constant-time comparison for the admin key to prevent timing attacks.

        Args:
            key: API key string from request header.

        Returns:
            :class:`AuthResult` with ``method='api_key'``.

        Raises:
            :class:`AuthError`: Key does not match, lookup fails, or no key
                                configured.
        """
        if not self._api_key:
            raise AuthError("API key not configured")

        # Admin/self-host key — bypass flyio-sync, return configured master owner
        if hmac.compare_digest(key, self._api_key):
            return AuthResult(owner_id=self._api_key_owner or "default", method="api_key")

        # Non-admin key — resolve owner via flyio-sync edge function
        return await self._verify_key_via_flyio_sync(key)

    async def _verify_key_via_flyio_sync(self, key: str) -> AuthResult:
        """Look up the owner of a non-admin API key via the flyio-sync edge function.

        Args:
            key: Raw API key string.

        Returns:
            :class:`AuthResult` with the resolved ``owner_id``.

        Raises:
            :class:`AuthError`: Lookup failed (network, key not found, bad response).
        """
        sync_secret = os.environ.get("FLYIO_SYNC_SECRET")
        if not sync_secret:
            raise AuthError("flyio-sync secret not configured")

        key_hash = hashlib.sha256(key.encode()).hexdigest()
        url = _get_flyio_sync_url()
        if url is None:
            raise AuthError("flyio-sync URL not configured (SUPABASE_URL or FLYIO_SYNC_URL must be set)")

        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.post(
                    url,
                    json={"action": "verify_api_key", "key_hash": key_hash},
                    headers={
                        "Content-Type": "application/json",
                        "x-flyio-sync-secret": sync_secret,
                    },
                )
        except httpx.HTTPError as exc:
            logger.error("flyio-sync request failed: %s", exc)
            raise AuthError("API key verification service unavailable") from exc

        if resp.status_code != 200:
            logger.warning(
                "flyio-sync returned %d for key verification", resp.status_code
            )
            raise AuthError("Invalid API key")

        try:
            data = resp.json()
        except Exception:
            raise AuthError("Invalid response from key verification service")

        # Explicit check: flyio-sync returns {valid: true/false} as the
        # canonical signal.  We must not proceed to owner extraction when
        # the key is invalid/revoked — even if owner_id happens to be set
        # in the response (e.g. for audit logging).  (#1589 audit F-01)
        if not data.get("valid"):
            raise AuthError("Invalid API key")

        # Prefer the resolved cortex_owner_id (friendly namespace) returned
        # by flyio-sync's verify_api_key action.  Fall back to user_id (raw
        # Supabase UUID) only when no mapping exists.  This keeps API-key
        # auth consistent with JWT auth, which resolves the same way via
        # _resolve_owner_id().
        raw_user_id = data.get("user_id")
        owner_id = (
            data.get("owner_id")
            or data.get("cortex_owner_id")
            or raw_user_id
        )
        if not owner_id:
            raise AuthError("Invalid API key")

        # If flyio-sync returned only a raw UUID, attempt to resolve the
        # friendly cortex_owner_id via the same lookup used by JWT auth.
        # This provides a graceful path for older flyio-sync deployments
        # that do not yet return owner_id in the verify_api_key response.
        sync_secret_for_resolve = os.environ.get("FLYIO_SYNC_SECRET")
        if _is_uuid(str(owner_id)) and sync_secret_for_resolve:
            resolved = _lookup_cortex_owner_id(str(owner_id), sync_secret_for_resolve)
            if resolved:
                owner_id = resolved

        plan = data.get("plan")
        return AuthResult(owner_id=str(owner_id), method="api_key", plan=plan)

    # ------------------------------------------------------------------
    # Main entry point
    # ------------------------------------------------------------------

    async def authenticate(self, headers: object) -> AuthResult:
        """
        Authenticate a request by inspecting its headers.

        Tries credentials in priority order:
            1. ``Authorization: Bearer <token>`` → JWT (when jwt_secret is set)
            2. ``X-Cortex-API-Key`` → API key (when api_key is set)
            3. ``X-API-Key`` → Legacy API key, backward compat (when api_key is set)
            4. Open access → returns ``AuthResult(owner_id='default', method='open')``

        A credential that is *present but invalid* (bad signature, expired,
        wrong key) always raises :class:`AuthError`.  Missing credentials
        fall through to the next method.

        Args:
            headers: Any object supporting ``.get(name)`` header lookup.
                     Starlette ``Headers`` (case-insensitive) works directly.

        Returns:
            :class:`AuthResult` on success.

        Raises:
            :class:`AuthError`: A credential was provided but failed validation.
        """

        def _get(name: str) -> Optional[str]:
            """Try header lookup, falling back to lowercase key."""
            val = getattr(headers, "get", None)
            if val is not None:
                return headers.get(name) or headers.get(name.lower())  # type: ignore[union-attr]
            # Plain dict fallback
            if isinstance(headers, dict):
                return headers.get(name) or headers.get(name.lower())
            return None

        # 1. Bearer JWT
        auth_header = _get("Authorization")
        if auth_header and auth_header.startswith("Bearer "):
            token = auth_header[7:].strip()
            if token and self._jwt_secret:
                # JWT secret is configured — validate; any failure is a hard error
                return self.validate_jwt(token)
            # No jwt_secret configured — can't validate; fall through

        # 2. X-Cortex-API-Key (preferred new header)
        cortex_key = _get("X-Cortex-API-Key")
        if cortex_key and self._api_key:
            return await self.validate_api_key(cortex_key)

        # 3. X-API-Key (legacy header — backward compat for existing self-host)
        legacy_key = _get("X-API-Key")
        if legacy_key and self._api_key:
            return await self.validate_api_key(legacy_key)

        # 4. Open access — no usable credential provided
        return AuthResult(owner_id="default", method="open")
