"""
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 base64
import hashlib
import json
import logging
from dataclasses import dataclass
from typing import Optional, Dict, Callable, Awaitable, Any, Tuple

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)
        account_key: Operator Account Key (32 bytes). Enables envelope encryption
                     (Layer 4). When provided with encrypt_payloads=True, request
                     and response bodies are encrypted and stored as blobs.
                     LockStock cannot decrypt these blobs (no Account Key, no DEK).
        encrypt_payloads: If True (requires account_key), encrypt and store payload
                          blobs after each successful stamp. Implies include_payload_hash.
    """

    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"
        account_key: Optional[bytes] = None,
        encrypt_payloads: bool = False,
    ):
        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 {}

        # Envelope encryption (Layer 4)
        self._account_key = account_key
        self.encrypt_payloads = encrypt_payloads
        if encrypt_payloads and account_key is None:
            raise ValueError(
                "encrypt_payloads=True requires account_key (32-byte operator secret)"
            )
        if encrypt_payloads:
            # Envelope encryption implies payload hashing
            include_payload_hash = True

        self.include_payload_hash = include_payload_hash
        self.max_payload_hash_bytes = max_payload_hash_bytes
        self.on_canonicalization_error = on_canonicalization_error

        # DEK state (lazy-initialized on first request)
        self._dek: Optional[bytes] = None
        self._dek_version: int = 1
        self._dek_lock = asyncio.Lock()

        # 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()

        # Token budget tracking (updated from stamp responses)
        self._budget_remaining: Optional[int] = None

    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")

    async def _ensure_dek(self) -> None:
        """
        Lazy-initialize the DEK for envelope encryption.

        On first call:
        1. Derive KEK from account_key + agent_id
        2. Try to fetch existing wrapped DEK from server
        3. If none exists, generate new DEK, wrap it, store on server
        4. Unwrap DEK and hold in memory

        Thread-safe via _dek_lock.
        """
        if self._dek is not None:
            return

        async with self._dek_lock:
            if self._dek is not None:
                return  # Double-check after lock

            from lockstock_core.keyring import (
                derive_kek, generate_dek, wrap_dek, unwrap_dek,
            )

            kek = derive_kek(self._account_key, self.agent_id)
            client = await self._get_client()

            # Try to fetch existing wrapped DEK
            try:
                response = await client.get(
                    f"{self.server_url}/v1/agents/{self.agent_id}/keys/latest",
                    headers={"Authorization": f"Bearer {self.api_key}"},
                )
                if response.status_code == 200:
                    data = response.json()
                    wrapped_dek = base64.b64decode(data["wrapped_dek_b64"])
                    self._dek = unwrap_dek(kek, wrapped_dek)
                    self._dek_version = data["key_version"]
                    logger.info(
                        "DEK loaded from server: agent=%s version=%d",
                        self.agent_id, self._dek_version,
                    )
                    return
            except Exception as e:
                logger.debug("No existing DEK found (will generate): %s", e)

            # Generate new DEK, wrap, store with retry.
            # CRITICAL: If the PUT fails, the DEK is held in memory but lost
            # on reboot. Any blobs encrypted with an unpersisted DEK become
            # permanently undecryptable. Retry aggressively — losing a DEK
            # is catastrophically worse than losing a single blob.
            dek = generate_dek()
            wrapped = wrap_dek(kek, dek)
            wrapped_b64 = base64.b64encode(wrapped).decode("ascii")

            dek_stored = False
            for attempt in range(3):
                try:
                    response = await client.put(
                        f"{self.server_url}/v1/agents/{self.agent_id}/keys/1",
                        json={"wrapped_dek_b64": wrapped_b64},
                        headers={"Authorization": f"Bearer {self.api_key}"},
                    )
                    if response.status_code == 200:
                        dek_stored = True
                        logger.info(
                            "New DEK generated and stored: agent=%s version=1",
                            self.agent_id,
                        )
                        break
                    else:
                        logger.warning(
                            "DEK store attempt %d failed (HTTP %d)",
                            attempt + 1, response.status_code,
                        )
                except Exception as e:
                    logger.warning(
                        "DEK store attempt %d error: %s", attempt + 1, e,
                    )
                if attempt < 2:
                    await asyncio.sleep(0.5 * (2 ** attempt))

            if not dek_stored:
                logger.error(
                    "CRITICAL: Wrapped DEK not persisted to server after 3 attempts "
                    "(agent=%s). DEK is held in memory only. If this process restarts "
                    "before the DEK is stored, all blobs encrypted in this session "
                    "will be permanently undecryptable. Scheduling background retry.",
                    self.agent_id,
                )
                # Schedule background retry that keeps trying until success
                asyncio.create_task(
                    self._retry_dek_store(wrapped_b64)
                )

            self._dek = dek
            self._dek_version = 1

    async def _retry_dek_store(self, wrapped_b64: str) -> None:
        """
        Background retry for wrapped DEK storage.

        Keeps trying with exponential backoff (5s, 10s, 20s, 40s, 60s cap)
        until the DEK is persisted or max attempts (10) reached.
        """
        for attempt in range(10):
            delay = min(5.0 * (2 ** attempt), 60.0)
            await asyncio.sleep(delay)
            try:
                client = await self._get_client()
                response = await client.put(
                    f"{self.server_url}/v1/agents/{self.agent_id}/keys/{self._dek_version}",
                    json={"wrapped_dek_b64": wrapped_b64},
                    headers={"Authorization": f"Bearer {self.api_key}"},
                )
                if response.status_code == 200:
                    logger.info(
                        "Background DEK store succeeded on attempt %d (agent=%s)",
                        attempt + 1, self.agent_id,
                    )
                    return
            except Exception as e:
                logger.debug("Background DEK store attempt %d: %s", attempt + 1, e)

        logger.error(
            "CRITICAL: Wrapped DEK storage permanently failed after 10 background "
            "retries (agent=%s). This DEK exists only in memory. Operator must "
            "manually store the wrapped DEK or reprovision before process restart.",
            self.agent_id,
        )

    def _canonicalize_body(self, body: bytes) -> Optional[Tuple[bytes, str]]:
        """
        Canonicalize body and return (canonical_bytes, payload_hash).

        Returns None if body is empty, too large, or canonicalization fails.
        """
        if not body or len(body) > self.max_payload_hash_bytes:
            return None

        try:
            from lockstock_core.canonicalize import bencode
            parsed = json.loads(body)
            canonical_bytes = bencode(parsed)
            payload_hash = hashlib.sha256(canonical_bytes).hexdigest()
            return (canonical_bytes, payload_hash)
        except Exception as e:
            if self.on_canonicalization_error == "raise":
                raise
            logger.warning(
                "Canonicalization failed (skipping encryption): %s", e
            )
            return None

    async def _fire_and_forget_blob_store(
        self,
        canonical_bytes: bytes,
        stamp_hash: str,
        payload_hash: str,
    ) -> None:
        """
        Encrypt canonical bytes and store as blob (fire-and-forget).

        Uses envelope encryption: DEK + stamp_hash + payload_hash → blob key.
        Errors are logged but never block the response.
        """
        try:
            from lockstock_core.encrypt import encrypt_payload_envelope

            ciphertext = encrypt_payload_envelope(
                canonical_bytes, stamp_hash, payload_hash, self._dek,
            )
            ciphertext_b64 = base64.b64encode(ciphertext).decode("ascii")

            client = await self._get_client()
            response = await client.post(
                f"{self.server_url}/v1/payload",
                json={
                    "payload_hash": payload_hash,
                    "stamp_hash": stamp_hash,
                    "agent_id": self.agent_id,
                    "ciphertext_b64": ciphertext_b64,
                    "plaintext_length": len(canonical_bytes),
                    "key_version": self._dek_version,
                },
                headers={"Authorization": f"Bearer {self.api_key}"},
            )
            if response.status_code == 200:
                logger.debug(
                    "Blob stored: hash=%s... size=%d",
                    payload_hash[:16], len(ciphertext),
                )
            else:
                logger.warning(
                    "Blob store failed (HTTP %d): hash=%s...",
                    response.status_code, payload_hash[:16],
                )
        except Exception as e:
            logger.warning("Blob store error: %s", e)

    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,
        metadata: Optional[Dict[str, Any]] = 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
            if metadata is not None:
                stamp_payload["metadata"] = metadata

            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"):
                # Track budget remaining from server (enables pre-request gate + clamp)
                tbr = data.get("token_budget_remaining")
                if tbr is not None and isinstance(tbr, (int, float)):
                    self._budget_remaining = int(tbr)

                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,
        metadata: Optional[Dict[str, Any]] = 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, metadata=metadata)
            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,
        canonical_bytes: Optional[bytes] = None,
        metadata: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        Stamp a response task in the background (non-blocking).

        If canonical_bytes is provided and encrypt_payloads is enabled,
        also encrypts and stores the response as a blob after stamping.

        Errors are logged but don't affect the response.
        """
        try:
            result = await self._stamp_with_retry(task, payload_hash=payload_hash, metadata=metadata)
            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 "?"
                )
                # Layer 4: Encrypt and store response blob
                if (
                    self.encrypt_payloads
                    and result.state_hash
                    and canonical_bytes is not None
                    and payload_hash is not None
                    and self._dek is not None
                ):
                    await self._fire_and_forget_blob_store(
                        canonical_bytes, result.state_hash, payload_hash,
                    )
            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

    def _extract_usage(self, body: bytes) -> Optional[Dict[str, Any]]:
        """
        Extract token usage from a provider response body.

        Parses the JSON response looking for the standard OpenAI-compatible
        ``usage`` object. Works with OpenAI, Azure OpenAI, Anthropic (which
        uses the same shape), vLLM, and any provider that returns::

            {"usage": {"prompt_tokens": N, "completion_tokens": N, "total_tokens": N}}

        Returns a metadata dict suitable for the stamp request, or None if
        usage data is not present or the body is not parseable JSON.
        """
        try:
            data = json.loads(body)
        except (json.JSONDecodeError, UnicodeDecodeError):
            return None

        usage = data.get("usage") if isinstance(data, dict) else None
        if not isinstance(usage, dict):
            return None

        meta: Dict[str, Any] = {}
        for key in ("total_tokens", "prompt_tokens", "completion_tokens"):
            val = usage.get(key)
            if isinstance(val, int):
                meta[key] = val

        # Anthropic uses input_tokens / output_tokens
        if "input_tokens" in usage and "prompt_tokens" not in meta:
            val = usage["input_tokens"]
            if isinstance(val, int):
                meta["prompt_tokens"] = val
        if "output_tokens" in usage and "completion_tokens" not in meta:
            val = usage["output_tokens"]
            if isinstance(val, int):
                meta["completion_tokens"] = val

        # Compute total if provider gave components but not total
        if "total_tokens" not in meta and "prompt_tokens" in meta and "completion_tokens" in meta:
            meta["total_tokens"] = meta["prompt_tokens"] + meta["completion_tokens"]

        if not meta:
            return None

        # Wrap in a namespace so it's clear where this came from
        return {"usage": meta}

    def _clamp_max_tokens(self, body: bytes, remaining: int) -> bytes:
        """
        Rewrite max_tokens in the request body to min(requested, remaining).

        Prevents a single response from overshooting the token budget.
        If body is not parseable JSON or has no max_tokens, returns unchanged.
        """
        try:
            data = json.loads(body)
            if not isinstance(data, dict):
                return body
            current = data.get("max_tokens")
            if current is None or (isinstance(current, int) and current <= remaining):
                return body
            data["max_tokens"] = remaining
            logger.info(
                "Clamped max_tokens from %s to %d (budget remaining)",
                current, remaining,
            )
            return json.dumps(data).encode("utf-8")
        except (json.JSONDecodeError, UnicodeDecodeError):
            return body

    async def _send_budget_exhausted_response(self, send: Send) -> None:
        """Send a 429 response when token budget is exhausted."""
        body = json.dumps({
            "error": {
                "message": "Token budget exhausted",
                "type": "insufficient_quota",
                "code": "token_budget_exhausted",
            },
            "agent_id": self.agent_id,
        }).encode("utf-8")

        await send({
            "type": "http.response.start",
            "status": 429,
            "headers": [
                (b"content-type", b"application/json"),
                (b"content-length", str(len(body)).encode()),
                (b"x-lockstock-status", b"budget-exhausted"),
            ],
        })
        await send({
            "type": "http.response.body",
            "body": body,
        })

    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

        # Pre-request budget gate: reject immediately if budget is exhausted.
        # This prevents burning compute on a request that will be rejected anyway.
        if self._budget_remaining is not None and self._budget_remaining <= 0:
            logger.warning(
                "Blocking request %s %s - token budget exhausted (agent=%s, remaining=%d)",
                method, path, self.agent_id, self._budget_remaining,
            )
            await self._send_budget_exhausted_response(send)
            return

        # Layer 4: Ensure DEK is initialized (lazy, first request only)
        if self.encrypt_payloads:
            await self._ensure_dek()

        # Always buffer request body: needed for budget clamping (max_tokens
        # rewrite) and optionally for payload hashing / encryption.
        request_body = await self._read_body(receive)

        # Budget clamping: rewrite max_tokens to min(requested, remaining)
        # so a single response can't overshoot the budget.
        if self._budget_remaining is not None and self._budget_remaining > 0:
            request_body = self._clamp_max_tokens(request_body, self._budget_remaining)

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

        if self.include_payload_hash:
            if self.encrypt_payloads:
                # Need both canonical bytes (for encryption) and hash (for stamp)
                canonical_data = self._canonicalize_body(request_body)
                if canonical_data is not None:
                    request_canonical_bytes, request_payload_hash = canonical_data
            else:
                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)

        # Layer 4: Encrypt and store request payload blob (fire-and-forget)
        if (
            self.encrypt_payloads
            and stamp.authorized
            and stamp.state_hash
            and request_canonical_bytes is not None
            and request_payload_hash is not None
            and self._dek is not None
        ):
            asyncio.create_task(
                self._fire_and_forget_blob_store(
                    request_canonical_bytes, stamp.state_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":
                # Always accumulate response body for usage extraction
                # (and optionally for payload hashing / encryption)
                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:
                            response_payload_hash = None
                            response_canonical_bytes = None
                            response_metadata = None

                            if response_body_chunks:
                                response_body = b"".join(response_body_chunks)

                                # Always extract token usage (decoupled from payload hash)
                                response_metadata = self._extract_usage(response_body)

                                # Compute response payload hash if enabled
                                if self.include_payload_hash:
                                    if self.encrypt_payloads:
                                        canonical_data = self._canonicalize_body(response_body)
                                        if canonical_data is not None:
                                            response_canonical_bytes, response_payload_hash = canonical_data
                                    else:
                                        response_payload_hash = self._compute_payload_hash(response_body)

                            # Update local budget tracking from usage
                            if response_metadata and self._budget_remaining is not None:
                                used = response_metadata.get("usage", {}).get("total_tokens", 0)
                                if used > 0:
                                    self._budget_remaining = max(0, self._budget_remaining - used)
                                    logger.debug(
                                        "Budget updated: used=%d remaining=%d (agent=%s)",
                                        used, self._budget_remaining, self.agent_id,
                                    )

                            asyncio.create_task(
                                self._fire_and_forget_response_stamp(
                                    response_task,
                                    payload_hash=response_payload_hash,
                                    canonical_bytes=response_canonical_bytes,
                                    metadata=response_metadata,
                                )
                            )

            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,
        })
