"""
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
    - Optionally extract plan from custom claims

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

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import time
from dataclasses import dataclass
from typing import Optional


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

    owner_id: str
    """Supabase user UUID (from JWT sub) or 'default' for API key / open access."""

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


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

    def __init__(
        self,
        jwt_secret: Optional[str] = None,
        api_key: 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

    # ------------------------------------------------------------------
    # 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.
        Custom claims (``plan``, ``user_plan``) are extracted when present.

        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 owner_id from sub claim
        sub = payload.get("sub")
        if not sub:
            raise AuthError("JWT missing sub claim")

        # 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=str(sub), method="jwt", plan=plan)

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

    def validate_api_key(self, key: str) -> AuthResult:
        """
        Validate an API key against the configured secret.

        Uses constant-time comparison 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 or no key is configured.
        """
        if not self._api_key:
            raise AuthError("API key not configured")

        # Constant-time comparison to prevent timing attacks
        if not hmac.compare_digest(key, self._api_key):
            raise AuthError("Invalid API key")

        return AuthResult(owner_id="default", method="api_key")

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

    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 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 self.validate_api_key(legacy_key)

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