"""
LockStock FastAPI Middleware
-----------------------------
Production middleware for securing FastAPI inference servers with LockStock chain auth.

Architecture:
- Sequential chain advancement using asyncio.Lock
- Minimal lock hold time (only Guard daemon calls)
- Concurrent inference processing
- Timeout protection against hung daemon
- X-LockStock-* response headers for audit trail

Phase 3 Features:
- Fork recovery: Auto-resync on /verify rejection and retry once
- Detailed rejection handling by error type
- X-LockStock-Status header on failures
"""

import asyncio
import logging
import os
from typing import Optional, Dict, Any

from fastapi import Request, Response
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from starlette.concurrency import run_in_threadpool
import httpx

# Module-level lock for serializing chain operations
# Chain advancement is sequential by design - lock ensures orderly progression
_chain_lock = asyncio.Lock()

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


class LockStockMiddleware(BaseHTTPMiddleware):
    """
    FastAPI middleware for LockStock chain-based authentication.

    Serializes chain advancement operations using asyncio.Lock while allowing
    concurrent inference processing. Each request advances the chain sequentially,
    then releases the lock before calling inference.

    Phase 3: Includes fork recovery - on /verify rejection, automatically resync
    from server and retry once before failing.

    Args:
        app: FastAPI application instance
        agent_id: LockStock agent identifier (from provisioning)
        guard_socket: Path to Guard daemon Unix socket
        task: Task name for inference operations (default: "network")
        block_on_failure: Reject requests on auth failure (default: False)
        timeout: Guard daemon call timeout in seconds (default: 5.0)
        paths: List of paths to protect (default: all paths)
        server_url: LockStock API endpoint (default: from env or official server)
        server_sync: Enable server sync via /verify endpoint (default: True)

    Example:
        >>> from fastapi import FastAPI
        >>> from lockstock_fastapi import LockStockMiddleware
        >>>
        >>> app = FastAPI()
        >>> app.add_middleware(
        ...     LockStockMiddleware,
        ...     agent_id="agent_abc123",
        ...     guard_socket="/var/run/lockstock-guard/guard.sock"
        ... )

    Environment Variables:
        LOCKSTOCK_AGENT_ID: Agent identifier
        LOCKSTOCK_SOCKET: Guard daemon socket path
        LOCKSTOCK_TASK: Task name for operations
        LOCKSTOCK_BLOCK_ON_FAILURE: Reject on failure (true/false)
        LOCKSTOCK_TIMEOUT: Guard daemon timeout (seconds)
        LOCKSTOCK_SERVER_URL: LockStock API endpoint
        LOCKSTOCK_SERVER_SYNC: Enable server sync (true/false, default: true)
    """

    def __init__(
        self,
        app,
        agent_id: Optional[str] = None,
        guard_socket: Optional[str] = None,
        task: str = "network",
        block_on_failure: bool = False,
        timeout: float = 5.0,
        paths: Optional[list] = None,
        path_tasks: Optional[dict] = None,
        server_url: Optional[str] = None,
        server_sync: bool = True,
    ):
        super().__init__(app)

        # Configuration from args or environment
        self.agent_id = agent_id or os.getenv("LOCKSTOCK_AGENT_ID")
        self.guard_socket = guard_socket or os.getenv(
            "LOCKSTOCK_SOCKET", "/var/run/lockstock-guard/guard.sock"
        )
        self.task = task or os.getenv("LOCKSTOCK_TASK", "network")
        self.block_on_failure = block_on_failure or (
            os.getenv("LOCKSTOCK_BLOCK_ON_FAILURE", "false").lower() == "true"
        )
        self.timeout = timeout or float(os.getenv("LOCKSTOCK_TIMEOUT", "5.0"))
        self.paths = paths  # None = protect all paths
        self.path_tasks = path_tasks or {}  # Path prefix -> task mapping

        # Response task mapping: inbound task -> outbound response task
        # Only tasks with a response counterpart get auto-stamped on successful response
        self.response_tasks = {
            "CHATCOMPLETION": "CHATCOMPLETIONRESPONSE",
        }

        # Server sync configuration
        self.server_url = server_url or os.getenv(
            "LOCKSTOCK_SERVER_URL", "https://lockstock-api-i9kp.onrender.com"
        )
        self.server_sync = server_sync and (
            os.getenv("LOCKSTOCK_SERVER_SYNC", "true").lower() == "true"
        )

        if not self.agent_id:
            raise ValueError(
                "agent_id is required. Provide via parameter or LOCKSTOCK_AGENT_ID env var."
            )

        # Import Guard client
        try:
            from lockstock_guard.client import sign_and_advance, get_chain_state, confirm_advance
            self._sign_and_advance = sign_and_advance
            self._get_chain_state = get_chain_state
            self._confirm_advance = confirm_advance
        except ImportError as e:
            raise RuntimeError(
                f"lockstock_guard.client not available: {e}\n"
                f"Install with: pip install lockstock-integrations[guard]\n"
                f"Or ensure lockstock_guard is in PYTHONPATH."
            ) from e

        # HTTP client for server sync
        self._http = httpx.AsyncClient(timeout=30.0) if self.server_sync else None

        logger.info(
            f"LockStock middleware initialized: agent={self.agent_id}, "
            f"socket={self.guard_socket}, timeout={self.timeout}s, "
            f"server_sync={self.server_sync}, path_tasks={list(self.path_tasks.keys())}"
        )

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

        Returns:
            Task string for chain entry, or None to skip chain write entirely.

        Note: Uses longest-prefix-match to prevent shorter prefixes from
        shadowing longer ones (e.g., /v1/chat/completions shadowing
        /v1/chat/completions/response).
        """
        matches = [(prefix, task) for prefix, task in self.path_tasks.items()
                   if path.startswith(prefix)]
        if matches:
            # Return longest matching prefix - prevents shorter prefixes from shadowing longer ones
            return max(matches, key=lambda x: len(x[0]))[1]
        return self.task  # Fall back to default task

    async def _resync_from_server(self) -> bool:
        """
        Tell Guard daemon to resync chain state from LockStock server.

        This is called when /verify rejects with a recoverable error.
        Uses the daemon's sync_chain() which:
        1. Fetches current state from server
        2. Updates daemon's in-memory state
        3. Persists to chain_state.json

        IMPORTANT: Writing chain_state.json directly does NOT work because
        the daemon caches state in memory and only loads from file at startup.

        Returns:
            True if resync successful, False otherwise.
        """
        try:
            from lockstock_guard.client import sync_chain

            # sync_chain is synchronous, run in threadpool to avoid blocking
            result = await run_in_threadpool(sync_chain, self.agent_id)

            if result.get("success"):
                logger.info(
                    f"Daemon resync successful: seq={result.get('sequence')} "
                    f"hash={result.get('current_hash', '')[:16]}..."
                )
                return True
            else:
                logger.error(f"Daemon resync failed: {result}")
                return False

        except Exception as e:
            logger.error(f"Failed to resync daemon from server: {e}")
            return False

    async def dispatch(self, request: Request, call_next):
        """
        Process request with LockStock chain authentication (TWO-PHASE COMMIT).

        Flow:
        1. Check if path should be protected
        2. Acquire lock (serialize chain advancement)
        3. Get current chain state from Guard daemon
        4. Call sign_and_advance (computes tentative state, NOT persisted)
        5. Call /verify with signature
        6. On /verify success: call confirm_advance (persists to disk)
        7. On /verify rejection: attempt recovery (resync + retry once)
        8. On /verify failure: don't confirm, pending state expires, next attempt retries
        9. Release lock
        10. Call downstream handler (inference - concurrent)
        11. Add X-LockStock-* headers to response

        Two-phase commit prevents "burned agent" problem where daemon advances locally
        during transient server failure (502, timeout) and can never resync.

        Returns:
            Response with X-LockStock-* headers added
        """
        # Check if this path should be protected
        if self.paths is not None:
            path_match = any(request.url.path.startswith(p) for p in self.paths)
            if not path_match:
                return await call_next(request)

        # Resolve path to task (None means exclude from chain)
        resolved_task = self._resolve_task(request.url.path)
        if resolved_task is None:
            # Excluded path - pass through without chain write
            return await call_next(request)

        sign_result = None
        retry_attempted = False

        try:
            # CRITICAL SECTION: Serialize chain advancement with timeout
            # Two-phase commit: sign_and_advance computes tentative state,
            # confirm_advance persists ONLY after /verify succeeds
            async with _chain_lock:
                sign_result, verify_result = await self._attempt_chain_advance(
                    resolved_task, retry_attempted=False
                )

                if verify_result and not verify_result.get("authorized"):
                    # /verify rejected - attempt recovery
                    rejection_reason = verify_result.get("reason", "unknown")
                    logger.warning(f"Server /verify rejected: {rejection_reason}")

                    # Check if this is a recoverable error
                    if self._is_recoverable_rejection(rejection_reason):
                        logger.info("Attempting recovery: resync from server...")

                        resync_ok = await self._resync_from_server()
                        if resync_ok:
                            logger.info("Resync successful, retrying chain advance...")
                            retry_attempted = True

                            # Retry the chain advance
                            sign_result, verify_result = await self._attempt_chain_advance(
                                resolved_task, retry_attempted=True
                            )

                            if verify_result and not verify_result.get("authorized"):
                                # Retry also failed
                                retry_reason = verify_result.get("reason", "unknown")
                                logger.critical(
                                    f"Retry after resync also rejected: {retry_reason}"
                                )

                                # Check if agent is bricked
                                if self._is_bricked(retry_reason):
                                    return self._bricked_response(retry_reason)

                                return self._rejection_response(
                                    retry_reason,
                                    retry_attempted=True
                                )
                        else:
                            # Resync failed
                            logger.error("Resync failed, cannot recover")
                            return self._rejection_response(
                                rejection_reason,
                                resync_failed=True
                            )
                    else:
                        # Not recoverable (e.g., agent revoked)
                        logger.critical(f"Non-recoverable rejection: {rejection_reason}")
                        return self._rejection_response(rejection_reason)

            # Lock released - inference proceeds concurrently

        except asyncio.TimeoutError:
            error_msg = f"Guard daemon timeout after {self.timeout}s"
            logger.error(f"{error_msg} for agent {self.agent_id}")

            if self.block_on_failure:
                return JSONResponse(
                    status_code=503,  # Service Unavailable
                    content={
                        "error": "guard_timeout",
                        "message": "Chain authentication service timeout",
                        "agent_id": self.agent_id,
                    },
                    headers={"X-LockStock-Status": "timeout"},
                )
            else:
                logger.warning(
                    f"Proceeding without chain auth (block_on_failure=False)"
                )
                return await call_next(request)

        except Exception as e:
            error_type = type(e).__name__
            error_msg = str(e)[:100]  # Limit error details
            logger.error(
                f"LockStock chain auth failed: {error_type} - {error_msg}",
                exc_info=True,
            )

            if self.block_on_failure:
                return JSONResponse(
                    status_code=403,  # Forbidden
                    content={
                        "error": "chain_auth_failed",
                        "message": "Chain authentication failed",
                        "agent_id": self.agent_id,
                    },
                    headers={"X-LockStock-Status": "error"},
                )
            else:
                logger.warning(
                    f"Proceeding without chain auth (block_on_failure=False)"
                )
                return await call_next(request)

        # Call downstream handler (inference) - CONCURRENT, lock released
        response = await call_next(request)

        # Add LockStock headers to response (provenance for request stamp)
        if sign_result:
            response.headers["X-LockStock-Agent-Id"] = self.agent_id
            response.headers["X-LockStock-Seq"] = str(sign_result["new_sequence"])
            response.headers["X-LockStock-Hash"] = sign_result["new_hash"]
            response.headers["X-LockStock-Sig"] = sign_result["signature"]
            response.headers["X-LockStock-Task"] = resolved_task
            response.headers["X-LockStock-Status"] = "accepted"

        # AUTO-STAMP OUTBOUND RESPONSE if task has a response counterpart
        # This creates a second chain entry for the response, completing the audit trail
        response_task = self.response_tasks.get(resolved_task)
        if response_task and response.status_code >= 200 and response.status_code < 300:
            try:
                async with _chain_lock:
                    response_sign_result, response_verify_result = await self._attempt_chain_advance(
                        response_task, retry_attempted=False
                    )

                    if response_verify_result and response_verify_result.get("authorized"):
                        # Add response stamp headers
                        response.headers["X-LockStock-Response-Seq"] = str(response_sign_result["new_sequence"])
                        response.headers["X-LockStock-Response-Hash"] = response_sign_result["new_hash"]
                        response.headers["X-LockStock-Response-Task"] = response_task
                        logger.info(
                            f"Response stamped: task={response_task} "
                            f"seq={response_sign_result['new_sequence']} "
                            f"hash={response_sign_result['new_hash'][:16]}..."
                        )
                    else:
                        # Response stamp failed - log but don't fail the request
                        # The request was already successful
                        logger.warning(
                            f"Response stamp failed: task={response_task} "
                            f"reason={response_verify_result.get('reason') if response_verify_result else 'unknown'}"
                        )
                        response.headers["X-LockStock-Response-Status"] = "stamp_failed"

            except Exception as e:
                # Response stamp failed - log but don't fail the request
                logger.warning(f"Response stamp exception: {e}")
                response.headers["X-LockStock-Response-Status"] = "error"

        return response

    async def _attempt_chain_advance(
        self,
        resolved_task: str,
        retry_attempted: bool
    ) -> tuple:
        """
        Attempt a single chain advance cycle: sign -> verify -> confirm.

        Returns:
            Tuple of (sign_result, verify_result)
            sign_result: Dict with signature, new_hash, etc. or None on failure
            verify_result: Dict with authorized, reason, state_hash or None
        """
        # Get current chain state (async wrapper prevents event loop blocking)
        chain_state = await asyncio.wait_for(
            run_in_threadpool(self._get_chain_state, self.agent_id),
            timeout=self.timeout,
        )
        parent_hash = chain_state.get("current_hash")

        if not parent_hash:
            raise ValueError(
                f"Chain state not initialized for {self.agent_id}. "
                f"Run: lockstock-guard sync --agent {self.agent_id}"
            )

        # PHASE 1: Sign and compute tentative advance (NOT persisted yet)
        # Guard validates parent_hash matches chain head (prevents forks)
        sign_result = await asyncio.wait_for(
            run_in_threadpool(
                self._sign_and_advance,
                self.agent_id,
                resolved_task,
                parent_hash,
            ),
            timeout=self.timeout,
        )

        logger.debug(
            f"Tentative advance: seq={sign_result['new_sequence']} "
            f"hash={sign_result['new_hash'][:16]}... task={resolved_task} "
            f"(awaiting /verify){' [RETRY]' if retry_attempted else ''}"
        )

        # Server sync with two-phase commit
        # CRITICAL: Only persist (confirm_advance) after /verify succeeds
        verify_result = None

        if self.server_sync:
            try:
                verify_result = await self._verify_with_server(
                    task=resolved_task,
                    signature=sign_result["signature"],
                    parent_hash=parent_hash,
                    state_matrix=sign_result["new_matrix"],
                    timestamp=sign_result["timestamp"],
                    sequence=sign_result["new_sequence"]
                )

                if verify_result.get("authorized"):
                    # PHASE 2: /verify succeeded - persist to disk
                    # Pass server's state_hash so daemon uses it as authoritative chain head
                    server_hash = verify_result.get("state_hash")
                    await asyncio.wait_for(
                        run_in_threadpool(
                            self._confirm_advance,
                            self.agent_id,
                            server_hash,
                        ),
                        timeout=self.timeout,
                    )
                    logger.debug(
                        f"Chain confirmed: seq={sign_result['new_sequence']} "
                        f"hash={server_hash[:16] if server_hash else 'local'}... "
                        f"(persisted after /verify success)"
                    )
                # If not authorized, don't persist - return verify_result for handling

            except httpx.HTTPStatusError as e:
                # HTTP error from server (4xx, 5xx)
                logger.warning(f"Server /verify HTTP error: {e}")
                verify_result = {
                    "authorized": False,
                    "reason": f"HTTP {e.response.status_code}",
                }
            except httpx.ConnectError as e:
                # Server unreachable
                logger.warning(f"Server unreachable during /verify: {e}")
                verify_result = {
                    "authorized": False,
                    "reason": "server_unreachable",
                }
            except httpx.TimeoutException as e:
                # Server timeout
                logger.warning(f"Server timeout during /verify: {e}")
                verify_result = {
                    "authorized": False,
                    "reason": "server_timeout",
                }
            except Exception as e:
                # Other errors
                logger.warning(f"Server sync failed: {e}")
                verify_result = {
                    "authorized": False,
                    "reason": str(e)[:100],
                }
        else:
            # Server sync disabled - persist immediately (legacy behavior)
            # WARNING: This can still burn agents on local-only mode
            await asyncio.wait_for(
                run_in_threadpool(
                    self._confirm_advance,
                    self.agent_id,
                ),
                timeout=self.timeout,
            )
            verify_result = {"authorized": True}

        return sign_result, verify_result

    def _is_recoverable_rejection(self, reason: str) -> bool:
        """
        Check if a rejection reason is recoverable via resync.

        Recoverable errors:
        - Sequence mismatch: local is ahead/behind server
        - Fork detected / parent_hash mismatch: hash chain diverged
        - Signature invalid: HMAC key (chain hash) doesn't match server

        Non-recoverable errors:
        - Agent revoked
        - Agent not found
        - Unauthorized task
        """
        reason_lower = reason.lower()

        recoverable_patterns = [
            "sequence mismatch",
            "sequence",
            "fork",
            "parent_hash",
            "parent hash",
            "signature invalid",
            "signature",
            "hash mismatch",
            "state mismatch",
        ]

        non_recoverable_patterns = [
            "revoked",
            "not found",
            "unauthorized",
            "forbidden",
            "invalid agent",
        ]

        # Check non-recoverable first
        for pattern in non_recoverable_patterns:
            if pattern in reason_lower:
                return False

        # Check recoverable
        for pattern in recoverable_patterns:
            if pattern in reason_lower:
                return True

        # Default: treat unknown errors as potentially recoverable
        # The retry will fail anyway if it's not actually recoverable
        return True

    def _is_bricked(self, reason: str) -> bool:
        """
        Check if the rejection indicates the agent is permanently bricked.

        A bricked agent cannot be recovered by resync - it needs reprovisioning.
        """
        reason_lower = reason.lower()

        bricked_patterns = [
            "signature invalid",  # After resync, still can't sign correctly
            "hmac",
            "key mismatch",
        ]

        for pattern in bricked_patterns:
            if pattern in reason_lower:
                return True

        return False

    def _rejection_response(
        self,
        reason: str,
        retry_attempted: bool = False,
        resync_failed: bool = False
    ) -> JSONResponse:
        """
        Build a 503 response for /verify rejection.
        """
        status_parts = ["rejected"]
        if retry_attempted:
            status_parts.append("retry_failed")
        if resync_failed:
            status_parts.append("resync_failed")

        return JSONResponse(
            status_code=503,
            content={
                "error": "chain_verification_rejected",
                "message": f"Chain verification rejected: {reason}",
                "agent_id": self.agent_id,
                "retry_attempted": retry_attempted,
                "resync_failed": resync_failed,
            },
            headers={"X-LockStock-Status": "|".join(status_parts)},
        )

    def _bricked_response(self, reason: str) -> JSONResponse:
        """
        Build a 503 response for a bricked agent.

        The agent needs to be reprovisioned - this cannot be fixed by resync.
        """
        return JSONResponse(
            status_code=503,
            content={
                "error": "agent_bricked",
                "message": (
                    f"Agent is cryptographically bricked: {reason}. "
                    "The HMAC signing key (chain hash) does not match server state. "
                    "This agent needs to be revoked and reprovisioned."
                ),
                "agent_id": self.agent_id,
                "recovery": "Contact administrator to revoke this agent and provision a new one.",
            },
            headers={"X-LockStock-Status": "bricked"},
        )

    async def _verify_with_server(
        self,
        task: str,
        signature: str,
        parent_hash: str,
        state_matrix: Dict[str, int],
        timestamp: int,
        sequence: int
    ) -> Dict[str, Any]:
        """
        Send Guard-signed request to LockStock server for verification and audit trail.

        The signature was created by Guard daemon using current_hash as HMAC key.
        Server will verify by:
        1. Looking up its current_hash for this agent
        2. Recalculating HMAC using that hash
        3. Comparing to signature we send
        4. Publishing to guard_traffic Redis channel for audit log ingestion

        Args:
            task: Capability being verified
            signature: HMAC signature from Guard (signed with current_hash)
            parent_hash: Parent hash in chain
            state_matrix: New state matrix
            timestamp: Timestamp used in signature
            sequence: Sequence number

        Returns:
            Dict with authorized (bool), reason (str), and state_hash (str)
        """
        payload = {
            "client_id": self.agent_id,
            "task": task.lower(),
            "parent_hash": parent_hash,
            "state_matrix": state_matrix,
            "signature": signature,
            "timestamp": timestamp,
            "sequence": sequence
        }

        response = await self._http.post(
            f"{self.server_url}/verify",
            json=payload,
            # No admin auth header - HMAC signature IS the authentication
        )

        # DEBUG: Log response status and body to diagnose issues
        logger.info(
            f"Server /verify response: status={response.status_code}, "
            f"body={response.text[:500]}"  # Log first 500 chars
        )

        # If non-2xx, log full details before attempting to parse
        if response.status_code != 200:
            logger.error(
                f"Server /verify failed: status={response.status_code}, "
                f"headers={dict(response.headers)}, "
                f"body={response.text}"
            )

        data = response.json()

        return {
            "authorized": data.get("status") == "accepted",
            "reason": data.get("reason"),
            "state_hash": data.get("state_hash")
        }
