"""
D3cipher Gateway — Multi-agent LockStock reverse proxy.

Wraps the existing LockStockThinMiddleware in a streaming reverse proxy.
Deployed in the customer's infrastructure. All AI agent traffic flows
through it. Only stamps and encrypted blobs reach LockStock's cloud.

Features:
    - Multi-agent: one middleware instance per agent, resolved via X-D3cipher-Agent header
    - Circuit breaker: server-side velocity monitor → 503 rejection → agent stops
    - Transcript auditor: encrypted request/response bodies (Layer 4) when ACCOUNT_KEY set
    - Bounded: ALLOWED_AGENTS allowlist or MAX_AGENTS hard cap prevents DoS
    - Prometheus /metrics, Kubernetes /healthz
    - ASGI lifespan for graceful shutdown
"""

import asyncio
import json
import logging
import os
import time
from collections import defaultdict
from typing import Dict, Optional, Set

import httpx

from lockstock_fastapi.thin import LockStockThinMiddleware

from .proxy import ProxyApp

logger = logging.getLogger("d3cipher.gateway")

# Default path-to-task mapping for common AI providers
DEFAULT_PATH_TASKS = {
    # OpenAI-compatible (OpenAI, Azure, vLLM, Ollama, etc.)
    "/v1/chat/completions": "CHATCOMPLETION",
    "/v1/completions": "COMPLETION",
    "/v1/embeddings": "EMBEDDING",
    "/v1/images/generations": "IMAGEGENERATION",
    "/v1/audio/transcriptions": "TRANSCRIPTION",
    "/v1/audio/speech": "TEXTTOSPEECH",
    # Anthropic
    "/v1/messages": "CHATCOMPLETION",
    # Exclude metadata endpoints from stamping
    "/v1/models": None,
}


class AgentDenied(Exception):
    """Raised when an agent is not allowed by the gateway."""

    def __init__(self, agent_id: str, status: int, reason: str):
        self.agent_id = agent_id
        self.status = status
        self.reason = reason
        super().__init__(reason)


class D3cipherGateway:
    """Multi-agent LockStock gateway with circuit breaker + transcript encryption."""

    def __init__(self):
        # Required config
        self.api_key = os.environ.get("D3CIPHER_API_KEY", "")
        self.upstream_url = os.environ.get("UPSTREAM_URL", "")
        if not self.api_key or not self.upstream_url:
            logger.warning(
                "D3CIPHER_API_KEY and UPSTREAM_URL are required. "
                "Gateway will reject all requests until configured."
            )

        # Optional config
        self.default_agent_id = os.environ.get("AGENT_ID")
        self.lockstock_url = os.environ.get(
            "LOCKSTOCK_URL", "https://lockstock-api-i9kp.onrender.com"
        )
        self.account_key = self._parse_account_key(os.environ.get("ACCOUNT_KEY"))

        # Agent bounding
        allowed_env = os.environ.get("ALLOWED_AGENTS")
        self.allowed_agents: Optional[Set[str]] = None
        if allowed_env:
            self.allowed_agents = {
                a.strip() for a in allowed_env.split(",") if a.strip()
            }
            logger.info(
                "Agent allowlist configured: %d agents", len(self.allowed_agents)
            )
        self.max_agents = int(os.environ.get("MAX_AGENTS", "50"))

        # Shared proxy app (all agents forward to same upstream)
        self.proxy = ProxyApp(self.upstream_url) if self.upstream_url else None

        # Per-agent middleware instances (lazy)
        self.agents: Dict[str, LockStockThinMiddleware] = {}
        self._agents_lock = asyncio.Lock()

        # Metrics (thread-safe: single worker, async-only access)
        self.request_count: Dict[str, int] = defaultdict(int)
        self.rejection_count: Dict[str, int] = defaultdict(int)
        self.latency_sum: Dict[str, float] = defaultdict(float)

        # Graceful shutdown
        self._shutting_down = False
        self._inflight = 0

        # Long-lived health check client
        self._health_client = httpx.AsyncClient(timeout=5.0)

        logger.info(
            "D3cipher Gateway initialized: upstream=%s lockstock=%s "
            "encrypt=%s agents_bound=%s",
            self.upstream_url,
            self.lockstock_url,
            self.account_key is not None,
            len(self.allowed_agents) if self.allowed_agents else f"max {self.max_agents}",
        )

    async def __call__(self, scope, receive, send) -> None:
        # ASGI lifespan protocol (graceful startup/shutdown)
        if scope["type"] == "lifespan":
            while True:
                message = await receive()
                if message["type"] == "lifespan.startup":
                    await send({"type": "lifespan.startup.complete"})
                elif message["type"] == "lifespan.shutdown":
                    await self.shutdown()
                    await send({"type": "lifespan.shutdown.complete"})
                    return
            return

        if scope["type"] != "http":
            return

        path = scope.get("path", "")

        # Internal endpoints (not proxied)
        if path == "/healthz":
            await self._health(send)
            return
        if path == "/metrics":
            await self._metrics(send)
            return

        # Preflight: check gateway is configured
        if not self.api_key or not self.proxy:
            await self._send_json(
                send, 503,
                {"error": "Gateway not configured",
                 "detail": "Set D3CIPHER_API_KEY and UPSTREAM_URL environment variables"},
            )
            return

        # Reject during shutdown
        if self._shutting_down:
            await self._send_json(send, 503, {"error": "Gateway shutting down"})
            return

        # Resolve agent identity
        agent_id = self._get_agent_id(scope)
        if not agent_id:
            await self._send_json(
                send, 400,
                {"error": "Missing agent identity",
                 "detail": "Set X-D3cipher-Agent header or AGENT_ID environment variable"},
            )
            return

        # Get or create middleware for this agent (bounded)
        try:
            middleware = await self._get_middleware(agent_id)
        except AgentDenied as e:
            logger.warning("Agent denied: %s (%d: %s)", e.agent_id, e.status, e.reason)
            await self._send_json(send, e.status, {"error": e.reason, "agent_id": e.agent_id})
            return

        # Track metrics via send wrapper
        self._inflight += 1
        start = time.monotonic()
        try:
            async def tracking_send(message) -> None:
                if message["type"] == "http.response.start":
                    status = message.get("status", 200)
                    if status == 503:
                        self.rejection_count[agent_id] += 1
                    # Check LockStock rejection headers
                    for name, value in message.get("headers", []):
                        if name == b"x-lockstock-status" and value == b"rejected":
                            self.rejection_count[agent_id] += 1
                            break
                await send(message)

            await middleware(scope, receive, tracking_send)
            self.request_count[agent_id] += 1
        finally:
            elapsed = time.monotonic() - start
            self.latency_sum[agent_id] += elapsed
            self._inflight -= 1

    async def _get_middleware(self, agent_id: str) -> LockStockThinMiddleware:
        """Get or create a middleware instance for the given agent."""
        if agent_id in self.agents:
            return self.agents[agent_id]

        # Allowlist check
        if self.allowed_agents is not None and agent_id not in self.allowed_agents:
            raise AgentDenied(agent_id, 403, "Agent not in ALLOWED_AGENTS allowlist")

        async with self._agents_lock:
            # Double-check after acquiring lock
            if agent_id in self.agents:
                return self.agents[agent_id]

            # Hard cap check (when no allowlist)
            if self.allowed_agents is None and len(self.agents) >= self.max_agents:
                raise AgentDenied(
                    agent_id, 429,
                    f"Maximum concurrent agents ({self.max_agents}) reached"
                )

            logger.info("Creating middleware for agent: %s", agent_id)
            middleware = LockStockThinMiddleware(
                self.proxy,
                agent_id=agent_id,
                api_key=self.api_key,
                server_url=self.lockstock_url,
                path_tasks=DEFAULT_PATH_TASKS,
                block_on_rejection=True,
                include_payload_hash=True,
                encrypt_payloads=self.account_key is not None,
                account_key=self.account_key,
            )
            self.agents[agent_id] = middleware
            return middleware

    def _get_agent_id(self, scope) -> Optional[str]:
        """Read X-D3cipher-Agent header, fall back to AGENT_ID env var."""
        for name, value in scope.get("headers", []):
            if name == b"x-d3cipher-agent":
                return value.decode("utf-8").strip()
        return self.default_agent_id

    def _parse_account_key(self, hex_key: Optional[str]) -> Optional[bytes]:
        """Parse hex Account Key to bytes, or None if not set."""
        if not hex_key:
            return None
        hex_key = hex_key.strip()
        if len(hex_key) != 64:
            logger.error(
                "ACCOUNT_KEY must be 64 hex characters (32 bytes), got %d. "
                "Envelope encryption disabled.", len(hex_key)
            )
            return None
        try:
            return bytes.fromhex(hex_key)
        except ValueError:
            logger.error("ACCOUNT_KEY is not valid hex. Envelope encryption disabled.")
            return None

    # --- Internal endpoints ---

    async def _health(self, send) -> None:
        """Liveness probe. Checks LockStock server reachability."""
        try:
            resp = await self._health_client.get(f"{self.lockstock_url}/health")
            healthy = resp.status_code == 200
        except Exception:
            healthy = False

        status = 200 if healthy else 503
        await self._send_json(send, status, {
            "status": "ok" if healthy else "degraded",
            "agents_active": len(self.agents),
            "inflight": self._inflight,
        })

    async def _metrics(self, send) -> None:
        """Prometheus-compatible metrics endpoint."""
        lines = [
            "# HELP d3cipher_requests_total Total requests proxied per agent",
            "# TYPE d3cipher_requests_total counter",
            "# HELP d3cipher_rejections_total Requests rejected (circuit breaker, auth) per agent",
            "# TYPE d3cipher_rejections_total counter",
            "# HELP d3cipher_request_duration_seconds_sum Total request duration per agent",
            "# TYPE d3cipher_request_duration_seconds_sum counter",
            "# HELP d3cipher_request_duration_seconds_count Request count for duration per agent",
            "# TYPE d3cipher_request_duration_seconds_count counter",
            "# HELP d3cipher_inflight Requests currently in flight",
            "# TYPE d3cipher_inflight gauge",
            "# HELP d3cipher_agents_active Number of active agent middleware instances",
            "# TYPE d3cipher_agents_active gauge",
        ]

        all_agents = set(list(self.request_count) + list(self.rejection_count))
        for agent_id in sorted(all_agents):
            reqs = self.request_count.get(agent_id, 0)
            rejs = self.rejection_count.get(agent_id, 0)
            lat = self.latency_sum.get(agent_id, 0.0)
            lines.append(f'd3cipher_requests_total{{agent="{agent_id}"}} {reqs}')
            lines.append(f'd3cipher_rejections_total{{agent="{agent_id}"}} {rejs}')
            lines.append(f'd3cipher_request_duration_seconds_sum{{agent="{agent_id}"}} {lat:.6f}')
            lines.append(f'd3cipher_request_duration_seconds_count{{agent="{agent_id}"}} {reqs}')

        lines.append(f"d3cipher_inflight {self._inflight}")
        lines.append(f"d3cipher_agents_active {len(self.agents)}")

        body = ("\n".join(lines) + "\n").encode("utf-8")
        await send({
            "type": "http.response.start",
            "status": 200,
            "headers": [
                (b"content-type", b"text/plain; charset=utf-8"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body})

    # --- Utilities ---

    async def _send_json(self, send, status: int, data: dict) -> None:
        """Send a JSON response."""
        body = json.dumps(data).encode("utf-8")
        await send({
            "type": "http.response.start",
            "status": status,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
            ],
        })
        await send({"type": "http.response.body", "body": body})

    async def shutdown(self) -> None:
        """Drain in-flight requests and close all connections."""
        logger.info("Shutting down gateway (inflight=%d)...", self._inflight)
        self._shutting_down = True

        # Wait for in-flight requests (max 30 seconds)
        for _ in range(300):
            if self._inflight == 0:
                break
            await asyncio.sleep(0.1)

        if self._inflight > 0:
            logger.warning(
                "Shutdown timeout: %d requests still in flight", self._inflight
            )

        # Close per-agent middleware HTTP clients
        for agent_id, mw in self.agents.items():
            if hasattr(mw, "_client") and mw._client is not None:
                try:
                    await mw._client.aclose()
                except Exception:
                    pass

        # Close proxy and health clients
        if self.proxy:
            await self.proxy.close()
        await self._health_client.aclose()

        logger.info("Gateway shutdown complete.")
