"""API authentication and rate limiting -- stdlib only.

Provides:
- API key validation (static keys from config or env)
- JWT-like HMAC tokens (HS256 via stdlib hmac/hashlib)
- Token bucket rate limiter (per-key)
- Middleware for http.server integration

All stdlib -- no PyJWT or external auth libraries required.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import json
import secrets
import threading
import time
from dataclasses import dataclass, field
from typing import Any

# ── API Keys ──────────────────────────────────────────────────


@dataclass(frozen=True)
class APIKey:
    """An API key with metadata."""

    key_id: str
    key_hash: str  # SHA-256 hash of the actual key
    name: str = ""
    scopes: frozenset[str] = frozenset()
    created_at: float = 0.0
    expires_at: float = 0.0  # 0 = never


class APIKeyStore:
    """Thread-safe store for API keys.

    Keys are stored as SHA-256 hashes. The actual key value is only
    returned once during generation and never persisted.
    """

    def __init__(self) -> None:
        self._keys: dict[str, APIKey] = {}
        self._hash_index: dict[str, str] = {}  # hash -> key_id
        self._lock = threading.Lock()

    def generate_key(
        self,
        name: str = "",
        scopes: frozenset[str] | None = None,
        ttl_seconds: float = 0,
    ) -> tuple[str, APIKey]:
        """Generate a new API key. Returns (raw_key, APIKey).

        The raw key is returned exactly once. Store it securely.
        """
        raw_key = f"orxaq_{secrets.token_urlsafe(32)}"
        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
        key_id = f"key_{secrets.token_hex(8)}"
        now = time.time()

        record = APIKey(
            key_id=key_id,
            key_hash=key_hash,
            name=name,
            scopes=scopes or frozenset(),
            created_at=now,
            expires_at=now + ttl_seconds if ttl_seconds > 0 else 0.0,
        )

        with self._lock:
            self._keys[key_id] = record
            self._hash_index[key_hash] = key_id

        return raw_key, record

    def validate(self, raw_key: str) -> APIKey | None:
        """Validate a raw API key. Returns the APIKey or None."""
        key_hash = hashlib.sha256(raw_key.encode()).hexdigest()
        with self._lock:
            key_id = self._hash_index.get(key_hash)
            if key_id is None:
                return None
            record = self._keys.get(key_id)

        if record is None:
            return None

        # Check expiry.
        if record.expires_at > 0 and time.time() > record.expires_at:
            return None

        return record

    def revoke(self, key_id: str) -> bool:
        """Revoke an API key by its ID."""
        with self._lock:
            record = self._keys.pop(key_id, None)
            if record is None:
                return False
            self._hash_index.pop(record.key_hash, None)
            return True

    def list_keys(self) -> list[APIKey]:
        """List all active API keys (without raw values)."""
        now = time.time()
        with self._lock:
            return [k for k in self._keys.values() if k.expires_at == 0 or k.expires_at > now]

    def key_count(self) -> int:
        """Number of registered keys."""
        with self._lock:
            return len(self._keys)


# ── JWT-like HMAC Tokens ─────────────────────────────────────


def _b64url_encode(data: bytes) -> str:
    """URL-safe base64 encoding without padding."""
    return base64.urlsafe_b64encode(data).rstrip(b"=").decode()


def _b64url_decode(s: str) -> bytes:
    """URL-safe base64 decoding with padding restoration."""
    padding = 4 - len(s) % 4
    if padding != 4:
        s += "=" * padding
    return base64.urlsafe_b64decode(s)


@dataclass(frozen=True)
class TokenClaims:
    """Claims extracted from a validated token."""

    subject: str
    scopes: frozenset[str]
    issued_at: float
    expires_at: float
    extra: dict[str, Any] = field(default_factory=dict)

    @property
    def is_expired(self) -> bool:
        return time.time() > self.expires_at


class TokenProvider:
    """HMAC-SHA256 token provider (JWT-compatible structure).

    Generates tokens with header.payload.signature format using
    stdlib hmac. Compatible with HS256 JWT validation.
    """

    def __init__(self, secret: str | None = None) -> None:
        self._secret = (secret or secrets.token_hex(32)).encode()

    def issue(
        self,
        subject: str,
        scopes: frozenset[str] | None = None,
        ttl_seconds: float = 3600,
        extra: dict[str, Any] | None = None,
    ) -> str:
        """Issue a new signed token."""
        now = time.time()
        header = {"alg": "HS256", "typ": "JWT"}
        payload: dict[str, Any] = {
            "sub": subject,
            "iat": now,
            "exp": now + ttl_seconds,
            "scopes": sorted(scopes) if scopes else [],
        }
        if extra:
            payload["extra"] = extra

        header_b64 = _b64url_encode(json.dumps(header).encode())
        payload_b64 = _b64url_encode(json.dumps(payload).encode())
        signing_input = f"{header_b64}.{payload_b64}"

        signature = hmac.new(self._secret, signing_input.encode(), hashlib.sha256).digest()
        sig_b64 = _b64url_encode(signature)

        return f"{header_b64}.{payload_b64}.{sig_b64}"

    def validate(self, token: str) -> TokenClaims | None:
        """Validate a token and return claims, or None if invalid."""
        parts = token.split(".")
        if len(parts) != 3:
            return None

        header_b64, payload_b64, sig_b64 = parts

        # Verify signature.
        signing_input = f"{header_b64}.{payload_b64}"
        expected_sig = hmac.new(self._secret, signing_input.encode(), hashlib.sha256).digest()

        try:
            actual_sig = _b64url_decode(sig_b64)
        except Exception:  # noqa: BLE001
            return None

        if not hmac.compare_digest(expected_sig, actual_sig):
            return None

        # Decode payload.
        try:
            payload = json.loads(_b64url_decode(payload_b64))
        except (ValueError, UnicodeDecodeError):
            return None

        claims = TokenClaims(
            subject=payload.get("sub", ""),
            scopes=frozenset(payload.get("scopes", [])),
            issued_at=payload.get("iat", 0),
            expires_at=payload.get("exp", 0),
            extra=payload.get("extra", {}),
        )

        # Check expiry.
        if claims.is_expired:
            return None

        return claims


# ── Rate Limiter ─────────────────────────────────────────────


@dataclass
class _Bucket:
    """Token bucket for rate limiting."""

    tokens: float
    last_refill: float
    capacity: float
    refill_rate: float  # tokens per second


class RateLimiter:
    """Token bucket rate limiter, keyed by client identifier.

    Each client gets a bucket with `capacity` tokens that refills
    at `refill_rate` tokens per second. A request consumes one token.
    """

    def __init__(
        self,
        capacity: float = 60.0,
        refill_rate: float = 1.0,
    ) -> None:
        self._capacity = capacity
        self._refill_rate = refill_rate
        self._buckets: dict[str, _Bucket] = {}
        self._lock = threading.Lock()

    def allow(self, client_id: str) -> bool:
        """Check if a request is allowed. Consumes one token if yes."""
        now = time.time()
        with self._lock:
            bucket = self._buckets.get(client_id)
            if bucket is None:
                bucket = _Bucket(
                    tokens=self._capacity - 1,
                    last_refill=now,
                    capacity=self._capacity,
                    refill_rate=self._refill_rate,
                )
                self._buckets[client_id] = bucket
                return True

            # Refill tokens based on elapsed time.
            elapsed = now - bucket.last_refill
            bucket.tokens = min(
                bucket.capacity,
                bucket.tokens + elapsed * bucket.refill_rate,
            )
            bucket.last_refill = now

            if bucket.tokens >= 1.0:
                bucket.tokens -= 1.0
                return True

            return False

    def remaining(self, client_id: str) -> float:
        """Return remaining tokens for a client."""
        with self._lock:
            bucket = self._buckets.get(client_id)
            if bucket is None:
                return self._capacity
            now = time.time()
            elapsed = now - bucket.last_refill
            return min(
                bucket.capacity,
                bucket.tokens + elapsed * bucket.refill_rate,
            )

    def reset(self, client_id: str) -> None:
        """Reset a client's rate limit bucket."""
        with self._lock:
            self._buckets.pop(client_id, None)


# ── Auth Result ──────────────────────────────────────────────


@dataclass(frozen=True)
class AuthResult:
    """Result of authenticating a request."""

    authenticated: bool
    identity: str = ""
    scopes: frozenset[str] = frozenset()
    error: str = ""
    rate_limited: bool = False


# ── Auth Middleware ───────────────────────────────────────────


class AuthMiddleware:
    """Authentication middleware for the HTTP server.

    Supports:
    - API key via X-API-Key header or ?api_key query parameter
    - Bearer token via Authorization header
    - Rate limiting per client identity
    - Configurable path exclusions (e.g., /health, /ready)
    """

    def __init__(
        self,
        key_store: APIKeyStore | None = None,
        token_provider: TokenProvider | None = None,
        rate_limiter: RateLimiter | None = None,
        excluded_paths: frozenset[str] | None = None,
    ) -> None:
        self.key_store = key_store or APIKeyStore()
        self.token_provider = token_provider or TokenProvider()
        self.rate_limiter = rate_limiter
        self.excluded_paths = excluded_paths or frozenset(
            {
                "/health",
                "/ready",
                "/",
            }
        )

    def authenticate(self, headers: dict[str, str], path: str) -> AuthResult:
        """Authenticate a request from its headers and path.

        Returns AuthResult indicating success/failure and identity.
        """
        # Skip auth for excluded paths.
        if path in self.excluded_paths:
            return AuthResult(authenticated=True, identity="anonymous")

        # Try API key first.
        key_header = headers.get("X-API-Key", "") or headers.get("x-api-key", "")
        if key_header:
            record = self.key_store.validate(key_header)
            if record is None:
                return AuthResult(authenticated=False, error="invalid_api_key")

            identity = record.key_id
            if self.rate_limiter and not self.rate_limiter.allow(identity):
                return AuthResult(
                    authenticated=True,
                    identity=identity,
                    scopes=record.scopes,
                    rate_limited=True,
                )
            return AuthResult(
                authenticated=True,
                identity=identity,
                scopes=record.scopes,
            )

        # Try Bearer token.
        auth_header = headers.get("Authorization", "") or headers.get("authorization", "")
        if auth_header.startswith("Bearer "):
            token_str = auth_header[7:]
            claims = self.token_provider.validate(token_str)
            if claims is None:
                return AuthResult(authenticated=False, error="invalid_token")

            identity = claims.subject
            if self.rate_limiter and not self.rate_limiter.allow(identity):
                return AuthResult(
                    authenticated=True,
                    identity=identity,
                    scopes=claims.scopes,
                    rate_limited=True,
                )
            return AuthResult(
                authenticated=True,
                identity=identity,
                scopes=claims.scopes,
            )

        return AuthResult(authenticated=False, error="no_credentials")
