"""LLMHosts proxy authentication -- API key management and auth middleware.

Provides a :class:`ProxyKeyStore` for creating, validating, and revoking
proxy API keys, and an :class:`AuthMiddleware` that enforces Bearer-token
authentication on incoming requests.
"""

from __future__ import annotations

import hashlib
import json
import logging
import secrets
import uuid
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any

from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.responses import JSONResponse

if TYPE_CHECKING:
    from pathlib import Path

    from starlette.requests import Request
    from starlette.responses import Response

logger = logging.getLogger(__name__)

# Paths that never require authentication.
_EXEMPT_PATHS: frozenset[str] = frozenset({"/health", "/docs", "/openapi.json", "/redoc"})

# Path prefixes that never require authentication (management endpoints).
_EXEMPT_PREFIXES: tuple[str, ...] = ("/api/keys",)

PROXY_KEYS_FILENAME = "proxy_keys.json"


# ---------------------------------------------------------------------------
# ProxyKeyStore
# ---------------------------------------------------------------------------


def _hash_key(raw_key: str) -> str:
    """Return the SHA-256 hex digest of *raw_key*."""
    return hashlib.sha256(raw_key.encode()).hexdigest()


class ProxyKeyStore:
    """Manage proxy API keys persisted in ``<data_dir>/proxy_keys.json``.

    Keys are stored as SHA-256 hashes -- the raw key is only returned once
    at creation time and is never persisted.
    """

    def __init__(self, data_dir: Path) -> None:
        self._data_dir = data_dir
        self._keys_file = data_dir / PROXY_KEYS_FILENAME
        self._keys: dict[str, dict[str, Any]] = {}
        self._load()

    # -- persistence --------------------------------------------------------

    def _load(self) -> None:
        """Load key data from the JSON file (if it exists)."""
        if self._keys_file.is_file():
            try:
                self._keys = json.loads(self._keys_file.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError) as exc:
                logger.warning("Failed to load proxy keys from %s: %s", self._keys_file, exc)
                self._keys = {}
        else:
            self._keys = {}

    def _save(self) -> None:
        """Persist current key data to the JSON file."""
        self._data_dir.mkdir(parents=True, exist_ok=True)
        self._keys_file.write_text(json.dumps(self._keys, indent=2), encoding="utf-8")

    # -- public API ---------------------------------------------------------

    def create_key(self, label: str, plan: str = "free") -> tuple[str, str]:
        """Create a new API key with the given *label* and *plan*.

        Returns
        -------
        tuple[str, str]
            ``(key_id, raw_key)`` -- the raw key is shown only once.
        """
        key_id = uuid.uuid4().hex[:12]
        raw_key = "lmh_" + secrets.token_urlsafe(32)
        key_hash = _hash_key(raw_key)
        self._keys[key_id] = {
            "hash": key_hash,
            "label": label,
            "plan": plan,
            "created_at": datetime.now(tz=timezone.utc).isoformat(),
            "last_used": None,
        }
        self._save()
        logger.info("Created proxy key %s (%s) plan=%s", key_id, label, plan)
        return key_id, raw_key

    def validate_key(self, raw_key: str) -> str | None:
        """Validate *raw_key* against stored hashes.

        Returns the ``key_id`` on success, or ``None`` if the key is invalid.
        """
        incoming_hash = _hash_key(raw_key)
        for key_id, meta in self._keys.items():
            if meta["hash"] == incoming_hash:
                meta["last_used"] = datetime.now(tz=timezone.utc).isoformat()
                self._save()
                return key_id
        return None

    def list_keys(self) -> list[dict[str, Any]]:
        """Return metadata for all keys (no raw keys or hashes)."""
        return [
            {
                "id": key_id,
                "label": meta["label"],
                "plan": meta.get("plan", "free"),
                "created_at": meta["created_at"],
                "last_used": meta["last_used"],
            }
            for key_id, meta in self._keys.items()
        ]

    def get_plan_by_id(self, key_id: str) -> str:
        """Return the plan name for the given *key_id*."""
        meta = self._keys.get(key_id)
        return meta.get("plan", "free") if meta else "free"

    def revoke_key(self, key_id: str) -> bool:
        """Revoke the key identified by *key_id*.

        Returns ``True`` if the key was found and removed, ``False`` otherwise.
        """
        if key_id in self._keys:
            del self._keys[key_id]
            self._save()
            logger.info("Revoked proxy key %s", key_id)
            return True
        return False


# ---------------------------------------------------------------------------
# AuthMiddleware
# ---------------------------------------------------------------------------


class AuthMiddleware(BaseHTTPMiddleware):
    """Authenticate requests via ``Authorization: Bearer <key>`` header.

    When *require_auth* is ``False`` every request is allowed through and
    tagged as anonymous.  When ``True``, missing or invalid keys receive a
    ``401`` JSON response (exempt paths like ``/health`` are always allowed).
    """

    def __init__(self, app: Any, key_store: ProxyKeyStore, require_auth: bool = False) -> None:
        super().__init__(app)
        self._key_store = key_store
        self._require_auth = require_auth

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        path = request.url.path

        # Exempt paths always pass through.
        if path in _EXEMPT_PATHS or path.startswith(_EXEMPT_PREFIXES):
            request.state.authenticated = False
            request.state.api_key_id = "anonymous"
            request.state.plan = "free"
            return await call_next(request)

        # Try to extract and validate a Bearer token.
        key_id: str | None = None
        auth_header = request.headers.get("authorization", "")
        if auth_header.startswith("Bearer "):
            raw_key = auth_header[7:]
            key_id = self._key_store.validate_key(raw_key)

        if key_id is not None:
            request.state.authenticated = True
            request.state.api_key_id = key_id
            request.state.plan = self._key_store.get_plan_by_id(key_id)
            return await call_next(request)

        # No valid key found.
        if self._require_auth:
            return JSONResponse(
                status_code=401,
                content={
                    "error": {
                        "message": "Invalid or missing API key",
                        "type": "authentication_error",
                        "code": "invalid_api_key",
                    }
                },
                headers={"WWW-Authenticate": "Bearer"},
            )

        # Auth not required -- allow as anonymous.
        request.state.authenticated = False
        request.state.api_key_id = "anonymous"
        request.state.plan = "free"
        return await call_next(request)
