"""
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
        },
        # response_tasks auto-inferred: CHATCOMPLETION->CHATCOMPLETIONRESPONSE, etc.
        # block_on_rejection defaults to True (fail-closed: no stamp = no action)
    )

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

Defaults match daemon behavior:
- response_tasks=None -> auto-inferred from path_tasks (TASK+"RESPONSE")
- block_on_rejection=True -> fail-closed (no stamp = no action)
"""

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

import httpx
from lockstock_core.fingerprint import generate_fingerprint
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: Dict mapping request tasks to response tasks.
                        Controls response auditing behavior:
                        - None (default): Auto-infer from path_tasks using TASK+"RESPONSE"
                          convention. This matches the daemon's implicit behavior where
                          every response was automatically stamped.
                        - Explicit dict: Use the provided mapping as-is.
                        - Empty dict {}: Conscious opt-out. Disables response stamping
                          entirely. A warning will be logged.
        block_on_rejection: If True (default), return 503 when stamp is rejected
                            or server is unreachable. No stamp = no action.
                            If False, proceed without stamp (logs warning).
                            The daemon was fail-closed by network position. This
                            default preserves that guarantee.
        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 = True,
        timeout: float = 10.0,
        max_retries: int = 3,
        include_payload_hash: bool = False,
        max_payload_hash_bytes: int = 2_097_152,  # 2MB ceiling
        on_canonicalization_error: str = "warn",  # "warn" or "raise"
    ):
        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.include_payload_hash = include_payload_hash
        self.max_payload_hash_bytes = max_payload_hash_bytes
        self.on_canonicalization_error = on_canonicalization_error

        # Response task resolution: None = auto-infer, dict = explicit, {} = opt-out
        if response_tasks is None:
            # Auto-infer from path_tasks: TASKNAME -> TASKNAMERESPONSE
            # This matches the daemon's implicit behavior where every response
            # was automatically stamped. The TASK+RESPONSE convention is enforced
            # by the server's agent_registry — it's a system invariant.
            self.response_tasks = {
                task: task + "RESPONSE"
                for task in self.path_tasks.values()
                if task is not None
            }
            if self.response_tasks:
                logger.info(
                    "response_tasks auto-inferred from path_tasks: %d response tasks "
                    "for %d stamped routes (convention: TASK+RESPONSE)",
                    len(self.response_tasks),
                    len(self.response_tasks),
                )
        else:
            self.response_tasks = response_tasks
            # Explicit empty dict = conscious opt-out
            if self.path_tasks and not self.response_tasks:
                logger.warning(
                    "response_tasks is explicitly empty — outbound responses will NOT "
                    "be stamped on chain. Only request stamps will be recorded. "
                    "If this is unintentional, remove the response_tasks parameter "
                    "to enable auto-inference."
                )

        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, payload_hash: Optional[str] = None) -> StampResult:
        """
        Call /v1/stamp once (no retry).

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

        try:
            stamp_payload: Dict[str, Any] = {
                "agent_id": self.agent_id,
                "task": task,
                "request_fingerprint": generate_fingerprint(),
            }
            if payload_hash is not None:
                stamp_payload["payload_hash"] = payload_hash

            response = await client.post(
                f"{self.server_url}/v1/stamp",
                json=stamp_payload,
                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, payload_hash: Optional[str] = None) -> 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, payload_hash=payload_hash)
            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, payload_hash: Optional[str] = None) -> 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, payload_hash=payload_hash)
            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 _read_body(self, receive: Receive) -> bytes:
        """Read the full ASGI request body from receive()."""
        chunks = []
        while True:
            message = await receive()
            body = message.get("body", b"")
            if body:
                chunks.append(body)
            if not message.get("more_body", False):
                break
        return b"".join(chunks)

    def _compute_payload_hash(self, body: bytes) -> Optional[str]:
        """
        Compute canonical payload hash from raw body bytes.

        Returns hex digest string, or None if body is empty, too large,
        or canonicalization fails (depending on on_canonicalization_error).
        """
        if not body:
            return None

        if len(body) > self.max_payload_hash_bytes:
            logger.debug(
                "Skipping payload hash: body size %d exceeds max %d bytes",
                len(body), self.max_payload_hash_bytes,
            )
            return None

        try:
            from lockstock_core.canonicalize import canonical_hash_from_bytes
            return canonical_hash_from_bytes(body)
        except Exception as e:
            if self.on_canonicalization_error == "raise":
                raise
            logger.warning(
                "Canonicalization failed (proceeding without payload_hash): %s", e
            )
            return None

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """
        ASGI entry point.

        When include_payload_hash is True, buffers the request body to compute
        a canonical payload hash, then replays the buffered body to the app.
        """
        # 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

        # Layer 3: Compute request payload hash if enabled
        request_payload_hash: Optional[str] = None
        actual_receive = receive

        if self.include_payload_hash:
            # Buffer the request body so we can hash it AND replay it to the app
            request_body = await self._read_body(receive)
            request_payload_hash = self._compute_payload_hash(request_body)

            # Wrap receive to replay the buffered body
            body_sent = False

            async def buffered_receive() -> Message:
                nonlocal body_sent
                if not body_sent:
                    body_sent = True
                    return {
                        "type": "http.request",
                        "body": request_body,
                        "more_body": False,
                    }
                # Subsequent calls: pass through (for disconnect, etc.)
                return await receive()

            actual_receive = buffered_receive

        # Stamp the request (with payload hash if computed)
        stamp = await self._stamp_with_retry(task, payload_hash=request_payload_hash)

        # 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 and body for response stamping
        response_status: Optional[int] = None
        response_started = False
        response_body_chunks: list = []

        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":
                # Accumulate response body for payload hashing
                if self.include_payload_hash:
                    body = message.get("body", b"")
                    if body:
                        response_body_chunks.append(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:
                            # Compute response payload hash
                            response_payload_hash = None
                            if self.include_payload_hash and response_body_chunks:
                                response_body = b"".join(response_body_chunks)
                                response_payload_hash = self._compute_payload_hash(response_body)

                            asyncio.create_task(
                                self._fire_and_forget_response_stamp(
                                    response_task, payload_hash=response_payload_hash
                                )
                            )

            await send(message)

        # Call downstream app with header-injecting send wrapper
        await self.app(scope, actual_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,
        })
