"""
LockStock Thin Middleware - Enterprise SaaS Edition

Pure ASGI middleware that stamps requests via /v1/stamp without local crypto.
Designed for streaming responses (SSE) - no response buffering.

Usage:
    from lockstock_fastapi.thin import LockStockThinMiddleware

    app = FastAPI()
    app.add_middleware(
        LockStockThinMiddleware,
        agent_id="agent_xxx_name",
        api_key="lsk_admin_...",
        path_tasks={
            "/v1/chat/completions": "chatcompletion",
            "/v1/models": "network",
            "/health": None,  # Excluded from stamping
        },
    )

NO daemon required. NO local crypto. The server handles everything.
"""

import asyncio
import logging
from dataclasses import dataclass
from typing import Optional, Dict, Callable, Awaitable, Any

import httpx
from starlette.datastructures import MutableHeaders
from starlette.types import ASGIApp, Receive, Send, Scope, Message

logger = logging.getLogger("lockstock.thin")


@dataclass
class StampResult:
    """Result from /v1/stamp call."""
    authorized: bool
    sequence: Optional[int] = None
    state_hash: Optional[str] = None
    trace_id: Optional[str] = None
    timestamp: Optional[int] = None
    reason: Optional[str] = None  # Only set if authorized=False
    http_status: Optional[int] = None  # HTTP status from server


class LockStockThinMiddleware:
    """
    Pure ASGI middleware for LockStock chain authentication.

    This middleware:
    1. Intercepts requests matching path_tasks patterns
    2. Calls /v1/stamp on the LockStock server (server does all crypto)
    3. Injects X-LockStock-* headers into the response
    4. Optionally stamps responses (fire-and-forget, non-blocking)

    NO response buffering - streaming (SSE) passes through unchanged.
    NO local crypto - server is authoritative.
    NO daemon required - just an API key.

    Args:
        app: The ASGI application to wrap
        agent_id: The provisioned agent ID
        api_key: API key for /v1/stamp authentication (Bearer token)
        server_url: LockStock server URL (default: production)
        path_tasks: Dict mapping URL path prefixes to task types.
                    Use None as value to exclude a path from stamping.
        response_tasks: Optional dict mapping request tasks to response tasks.
                        Response stamps are fire-and-forget (non-blocking).
        block_on_rejection: If True, return 503 when stamp is rejected.
                            If False, proceed without stamp (logs warning).
        timeout: HTTP timeout for /v1/stamp calls in seconds
        max_retries: Max retries on 409 Conflict (default: 3)
    """

    def __init__(
        self,
        app: ASGIApp,
        *,
        agent_id: str,
        api_key: str,
        server_url: str = "https://lockstock-api-i9kp.onrender.com",
        path_tasks: Optional[Dict[str, Optional[str]]] = None,
        response_tasks: Optional[Dict[str, str]] = None,
        block_on_rejection: bool = False,
        timeout: float = 10.0,
        max_retries: int = 3,
    ):
        self.app = app
        self.agent_id = agent_id
        self.api_key = api_key
        self.server_url = server_url.rstrip("/")
        self.path_tasks = path_tasks or {}
        self.response_tasks = response_tasks or {}
        self.block_on_rejection = block_on_rejection
        self.timeout = timeout
        self.max_retries = max_retries
        self._client: Optional[httpx.AsyncClient] = None
        self._client_lock = asyncio.Lock()

    async def _get_client(self) -> httpx.AsyncClient:
        """Lazy-initialize the HTTP client with thread-safe locking."""
        if self._client is None:
            async with self._client_lock:
                # Double-check after acquiring lock
                if self._client is None:
                    self._client = httpx.AsyncClient(
                        timeout=self.timeout,
                        # Connection pooling for performance
                        limits=httpx.Limits(
                            max_keepalive_connections=5,
                            max_connections=10,
                        ),
                    )
        return self._client

    async def close(self) -> None:
        """Close the HTTP client. Call on application shutdown."""
        if self._client is not None:
            await self._client.aclose()
            self._client = None
            logger.debug("LockStock thin client closed")

    def _resolve_task(self, path: str) -> Optional[str]:
        """
        Resolve URL path to task type using longest-prefix match.

        Returns:
            Task name if path should be stamped
            None if path is explicitly excluded (value is None) or no match
        """
        best_match: Optional[str] = None
        best_length = 0
        found_explicit_none = False

        for prefix, task in self.path_tasks.items():
            if path.startswith(prefix) and len(prefix) > best_length:
                best_match = task
                best_length = len(prefix)
                found_explicit_none = (task is None)

        # If best match is explicitly None, path is excluded
        if found_explicit_none:
            return None

        return best_match

    async def _stamp(self, task: str) -> StampResult:
        """
        Call /v1/stamp once (no retry).

        Returns StampResult with authorization status and chain state.
        """
        client = await self._get_client()

        try:
            response = await client.post(
                f"{self.server_url}/v1/stamp",
                json={"agent_id": self.agent_id, "task": task},
                headers={"Authorization": f"Bearer {self.api_key}"},
            )

            data = response.json()

            if response.status_code == 200 and data.get("authorized"):
                return StampResult(
                    authorized=True,
                    sequence=data.get("sequence"),
                    state_hash=data.get("state_hash"),
                    trace_id=data.get("trace_id"),
                    timestamp=data.get("timestamp"),
                    http_status=response.status_code,
                )
            else:
                return StampResult(
                    authorized=False,
                    reason=data.get("reason", f"HTTP {response.status_code}"),
                    trace_id=data.get("trace_id"),
                    http_status=response.status_code,
                )

        except httpx.RequestError as e:
            logger.error(
                "Connection error stamping %s (agent=%s): %s",
                task, self.agent_id, e
            )
            return StampResult(
                authorized=False,
                reason=f"Connection error: {type(e).__name__}",
            )

    async def _stamp_with_retry(self, task: str) -> StampResult:
        """
        Call /v1/stamp with retry on 409 Conflict.

        409 indicates a concurrent chain advance - the sequence number was
        claimed between our read and write. This is transient and should
        be retried with short backoff.
        """
        last_result: Optional[StampResult] = None

        for attempt in range(self.max_retries):
            result = await self._stamp(task)
            last_result = result

            # Success - return immediately
            if result.authorized:
                if attempt > 0:
                    logger.info(
                        "Stamp succeeded on retry %d for %s (agent=%s)",
                        attempt + 1, task, self.agent_id
                    )
                return result

            # 409 Conflict - retry with backoff
            if result.http_status == 409:
                if attempt < self.max_retries - 1:
                    backoff = 0.1 * (2 ** attempt)  # 0.1s, 0.2s, 0.4s
                    logger.debug(
                        "409 Conflict on stamp attempt %d, retrying in %.1fs (agent=%s)",
                        attempt + 1, backoff, self.agent_id
                    )
                    await asyncio.sleep(backoff)
                    continue

            # Non-retryable error - return immediately
            break

        return last_result or StampResult(authorized=False, reason="No result")

    async def _fire_and_forget_response_stamp(self, task: str) -> None:
        """
        Stamp a response task in the background (non-blocking).

        Errors are logged but don't affect the response.
        """
        try:
            result = await self._stamp_with_retry(task)
            if result.authorized:
                logger.debug(
                    "Response stamp succeeded: task=%s seq=%s hash=%s",
                    task, result.sequence,
                    result.state_hash[:16] if result.state_hash else "?"
                )
            else:
                logger.warning(
                    "Response stamp rejected: task=%s reason=%s (agent=%s)",
                    task, result.reason, self.agent_id
                )
        except Exception as e:
            logger.error(
                "Response stamp failed: task=%s error=%s (agent=%s)",
                task, e, self.agent_id
            )

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """
        ASGI entry point - handles HTTP requests without buffering.
        """
        # Only process HTTP requests
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

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

        # Resolve path to task
        task = self._resolve_task(path)

        if task is None:
            # Path excluded from stamping - pass through unchanged
            await self.app(scope, receive, send)
            return

        # Stamp the request
        stamp = await self._stamp_with_retry(task)

        # Log rejections/failures when not blocking
        if not stamp.authorized and not self.block_on_rejection:
            logger.warning(
                "Stamp rejected for %s %s (agent=%s, reason=%s, trace=%s) - proceeding unstamped",
                method, path, self.agent_id, stamp.reason, stamp.trace_id
            )

        # Block if configured and stamp rejected
        if not stamp.authorized and self.block_on_rejection:
            logger.warning(
                "Blocking request %s %s - stamp rejected (agent=%s, reason=%s)",
                method, path, self.agent_id, stamp.reason
            )
            await self._send_rejection_response(send, stamp)
            return

        # Track response status for response stamping
        response_status: Optional[int] = None
        response_started = False

        async def send_with_headers(message: Message) -> None:
            nonlocal response_status, response_started

            if message["type"] == "http.response.start":
                response_started = True
                response_status = message.get("status", 200)

                # Inject LockStock headers
                headers = MutableHeaders(scope=message)

                if stamp.authorized:
                    headers.append("X-LockStock-Status", "accepted")
                    headers.append("X-LockStock-Seq", str(stamp.sequence))
                    headers.append("X-LockStock-Hash", stamp.state_hash or "")
                    if stamp.trace_id:
                        headers.append("X-LockStock-Trace", stamp.trace_id)
                else:
                    headers.append("X-LockStock-Status", "skipped")
                    headers.append("X-LockStock-Reason", stamp.reason or "unknown")

            elif message["type"] == "http.response.body":
                # Check if this is the final chunk (for response stamping)
                is_final = not message.get("more_body", False)

                if is_final and stamp.authorized and response_status is not None:
                    # Fire-and-forget response stamp for successful responses
                    if response_status < 400:
                        response_task = self.response_tasks.get(task)
                        if response_task:
                            asyncio.create_task(
                                self._fire_and_forget_response_stamp(response_task)
                            )

            await send(message)

        # Call downstream app with header-injecting send wrapper
        await self.app(scope, receive, send_with_headers)

    async def _send_rejection_response(self, send: Send, stamp: StampResult) -> None:
        """Send a 503 rejection response when block_on_rejection is True."""
        import json

        body = json.dumps({
            "error": "LockStock stamp rejected",
            "reason": stamp.reason,
            "trace_id": stamp.trace_id,
        }).encode("utf-8")

        await send({
            "type": "http.response.start",
            "status": 503,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
                (b"x-lockstock-status", b"rejected"),
                (b"x-lockstock-reason", (stamp.reason or "unknown").encode()),
            ],
        })
        await send({
            "type": "http.response.body",
            "body": body,
        })
