from __future__ import annotations

import asyncio
import hashlib
import json
import logging
import re
from abc import ABC, abstractmethod
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Awaitable, Callable, Optional, Sequence

from aiohttp import web

from ..auth.crypto import lookup_digest, verify_token
from ..config_store import ConfigStore
from ..metrics import METRICS
from ..settings import GatewaySettings, get_gateway_settings
from .context import get_request_context

_LOG = logging.getLogger(__name__)


class AuthBackend(ABC):
    @abstractmethod
    def validate(self, token: str) -> Optional[dict[str, str]]:
        """Validate token and return a principal dict or None.

        Principal must include at least: {"org_id": str, "api_key_id": str}
        """


class InMemoryAuthBackend(AuthBackend):
    """Simple env-configured key store for development and tests.

    Keys are sourced from GatewaySettings.auth.issued_keys.
    """

    def __init__(
        self,
        issued_keys: Sequence[str] | None = None,
        *,
        default_tier: str = "T1",
    ) -> None:
        self._map: dict[str, dict[str, str]] = {}
        self.last_reason: Optional[str] = None
        entries = issued_keys or ["dev-key:org_dev"]
        for part in [p.strip() for p in entries if p.strip()]:
            segments = part.split(":")
            key = segments[0]
            org = segments[1] if len(segments) > 1 else "org_dev"
            tier = segments[2] if len(segments) > 2 else default_tier
            role = segments[3] if len(segments) > 3 else "user"
            self._map[key] = {
                "org_id": org,
                "api_key_id": self._hash_key(key),
                "tier": tier,
                "role": role,
            }

    def _hash_key(self, key: str) -> str:
        return hashlib.sha256(key.encode("utf-8")).hexdigest()[:16]

    def validate(self, token: str) -> Optional[dict[str, str]]:
        principal = self._map.get(token)
        self.last_reason = "success" if principal else "not_found"
        return principal

    def register_key(
        self,
        key: str,
        org_id: str,
        tier: str = "T1",
        role: str = "user",
    ) -> None:
        """Register a new API key in the in-memory store.

        Args:
            key: The raw API key token.
            org_id: Organization identifier for the key owner.
            tier: Service tier (default ``"T1"``).
            role: Key role (default ``"user"``).
        """
        self._map[key] = {
            "org_id": org_id,
            "api_key_id": self._hash_key(key),
            "tier": tier,
            "role": role,
        }

    def has_org(self, org_id: str) -> bool:
        """Return True if any key is registered for the given org_id."""
        return any(p["org_id"] == org_id for p in self._map.values())

    def load_persisted_keys(self, path: Path) -> int:
        """Load previously registered keys from a JSONL file.

        Each line must be a JSON object with at minimum ``key`` and
        ``org_id`` fields.

        Args:
            path: Path to the JSONL file.

        Returns:
            Number of keys loaded.
        """
        if not path.exists():
            return 0
        loaded = 0
        with path.open("r", encoding="utf-8") as fh:
            for line_no, line in enumerate(fh, start=1):
                line = line.strip()
                if not line:
                    continue
                try:
                    entry = json.loads(line)
                except json.JSONDecodeError:
                    _LOG.warning(
                        "Skipping malformed line %d in %s", line_no, path
                    )
                    continue
                key = entry.get("key")
                org_id = entry.get("org_id")
                if not key or not org_id:
                    _LOG.warning(
                        "Skipping incomplete entry on line %d in %s",
                        line_no,
                        path,
                    )
                    continue
                self.register_key(
                    key,
                    org_id,
                    tier=entry.get("tier", "T1"),
                    role=entry.get("role", "user"),
                )
                loaded += 1
        if loaded:
            _LOG.info("Loaded %d persisted keys from %s", loaded, path)
        return loaded


_VALID_TABLE_NAME = re.compile(r"^[a-z_][a-z0-9_]*$")


class PostgresAuthBackend(AuthBackend):
    """Postgres-backed API key store leveraging hashed tokens and lookup digests."""

    def __init__(self, conn: Any, *, table: str = "api_keys", pepper: str) -> None:
        if not pepper:
            raise ValueError("gateway.auth.pepper must be configured when using Postgres auth")
        if not _VALID_TABLE_NAME.match(table):
            raise ValueError(
                f"Invalid table name {table!r}: must match [a-z_][a-z0-9_]*"
            )
        self.conn = conn
        self.conn.autocommit = True
        self._table = table
        self._select_by_kid = (
            f"SELECT id, org_id, kid, key_hash, tier, role, status, expires_at "
            f"FROM {table} WHERE kid = %s LIMIT 1"
        )
        self._select_by_hmac = (
            f"SELECT id, org_id, kid, key_hash, tier, role, status, expires_at "
            f"FROM {table} WHERE lookup_hmac = %s LIMIT 1"
        )
        self._update_last_used = f"UPDATE {table} SET last_used_at = now() WHERE id = %s"
        self._pepper = pepper
        self._last_reason: Optional[str] = None

    def _parse_kid(self, token: str) -> Optional[str]:
        parts = token.split("_", 2)
        if len(parts) == 3:
            return parts[1]
        return None

    def validate(self, token: str) -> Optional[dict[str, str]]:
        kid = self._parse_kid(token)
        cur = self.conn.cursor()
        row = None
        try:
            if kid:
                cur.execute(self._select_by_kid, (kid,))
                row = cur.fetchone()
            if not row:
                digest = lookup_digest(token, self._pepper)
                cur.execute(self._select_by_hmac, (digest,))
                row = cur.fetchone()
        except Exception as exc:
            self._last_reason = "db_error"
            raise RuntimeError("Postgres auth backend query failed") from exc

        if not row:
            self._last_reason = "not_found"
            return None

        key_id, org_id, stored_kid, key_hash, tier, role, status, expires_at = row
        if status != "active":
            self._last_reason = "inactive"
            return None
        if expires_at and isinstance(expires_at, datetime):
            if expires_at.tzinfo is None:
                expires_at = expires_at.replace(tzinfo=timezone.utc)
            if expires_at < datetime.now(timezone.utc):
                self._last_reason = "expired"
                return None

        if not verify_token(token, key_hash):
            self._last_reason = "hash_mismatch"
            return None

        try:
            cur.execute(self._update_last_used, (key_id,))
        except Exception as exc:
            _LOG.debug("Failed to update api_keys.last_used_at: %s", exc)

        self._last_reason = "success"

        return {
            "org_id": org_id,
            "api_key_id": str(key_id),
            "kid": stored_kid,
            "tier": tier,
            "role": role,
        }

    @property
    def last_reason(self) -> Optional[str]:
        return self._last_reason


Handler = Callable[[web.Request], Awaitable[web.StreamResponse]]
Middleware = Callable[[web.Request, Handler], Awaitable[web.StreamResponse]]


def build_auth_middleware(
    backend: Optional[AuthBackend] = None,
    config: Optional[ConfigStore] = None,
    *,
    settings: Optional[GatewaySettings] = None,
) -> Middleware:
    settings = settings or get_gateway_settings()

    # Choose backend from settings if not provided
    if backend is None:
        backend_choice = settings.auth.mode.lower()
        if backend_choice == "postgres":
            dsn = settings.stores.postgres_dsn
            pepper = settings.auth.pepper
            if not dsn or not pepper:
                raise RuntimeError(
                    "gateway.auth.mode is 'postgres' but no Postgres DSN or pepper was provided"
                )
            try:
                import psycopg2
            except ImportError as exc:  # pragma: no cover
                raise RuntimeError("psycopg2 is required for Postgres auth mode") from exc
            try:
                conn = psycopg2.connect(dsn)
            except Exception as exc:
                raise RuntimeError("Failed to connect to Postgres for auth backend") from exc
            backend = PostgresAuthBackend(conn, table="api_keys", pepper=pepper)
        else:
            backend = InMemoryAuthBackend(
                settings.auth.issued_keys, default_tier=settings.auth.default_tier
            )

    require_auth = settings.auth.require_auth
    cfg = config or ConfigStore(settings.config.tiers_path)

    # Paths that bypass authentication entirely (health probes, metrics, signup).
    _bypass_paths = frozenset({
        "/healthz", "/v1/health", "/metrics", "/v1/metrics",
        "/v1/auth/register",
    })

    # Sync backends (e.g. PostgresAuthBackend) block the event loop.
    # Run them in a thread pool to keep the server responsive.
    _is_sync_backend = isinstance(backend, PostgresAuthBackend)

    async def _validate(token: str) -> Optional[dict[str, str]]:
        if _is_sync_backend:
            return await asyncio.to_thread(backend.validate, token)
        return backend.validate(token)

    @web.middleware
    async def auth_middleware(request: web.Request, handler: Handler) -> web.StreamResponse:
        if request.path in _bypass_paths:
            return await handler(request)

        # Allow local dev with no auth if not required
        token = None
        auth = request.headers.get("Authorization")
        if auth and auth.lower().startswith("bearer "):
            token = auth[7:].strip()
        # Smart default: also accept X-API-Key for clients using custom header auth
        if not token:
            token = request.headers.get("X-API-Key")

        backend_label = backend.__class__.__name__

        if token:
            try:
                principal = await _validate(token)
            except Exception as exc:
                METRICS.record_auth(backend_label, False, reason="backend_error")
                raise web.HTTPServiceUnavailable(text="Auth backend unavailable") from exc
            reason = getattr(backend, "last_reason", None)
            METRICS.record_auth(
                backend_label,
                bool(principal),
                reason=None if principal else reason,
            )
            if not principal:
                raise web.HTTPUnauthorized(text="Invalid API token")
            org_id_value = principal.get("org_id")
            if org_id_value is None:
                raise web.HTTPServiceUnavailable(text="Auth backend returned invalid principal")
            ctx = get_request_context(request)
            ctx.org_id = org_id_value
            ctx.api_key_id = principal.get("api_key_id")
            ctx.api_key_kid = principal.get("kid")
            ctx.api_key_role = principal.get("role")
            ctx.api_key_tier = principal.get("tier")
            ctx.plan = ctx.api_key_tier or cfg.tier_for_org(
                ctx.org_id,
                default=settings.auth.default_tier,
            )
        else:
            if require_auth:
                METRICS.record_auth(backend_label, False, reason="missing_token")
                raise web.HTTPUnauthorized(text="Missing Authorization header")
            # Dev fallback: allow X-Org-Id or default
            ctx = get_request_context(request)
            ctx.org_id = request.headers.get("X-Org-Id", "org_anon")
            ctx.plan = cfg.tier_for_org(ctx.org_id, default=settings.auth.dev_default_tier)

        return await handler(request)

    return auth_middleware
