"""AttestMCP transparent proxy.

Accepts MCP Streamable HTTP requests from agent clients,
forwards them unchanged to the upstream memory server,
and writes a tamper-evident audit record for every operation.

Architecture:
  - Downstream: FastAPI POST /mcp endpoint
  - Upstream: httpx async client forwarding raw HTTP
  - Audit: synchronous SQLite append before returning response

The proxy is transparent: request and response payloads are
never modified. The agent and memory server behave identically
with or without the proxy in the path.

v1 handles JSON-RPC request/response (single POST → JSON).
SSE streaming responses are deferred to v2.
"""

from __future__ import annotations

import json
import logging
import time
from contextlib import asynccontextmanager
from typing import Any

import httpx
from fastapi import FastAPI, Request, Response
from prometheus_client import CONTENT_TYPE_LATEST, generate_latest

from attestmcp import metrics
from attestmcp.canonical import compute_params_hash, compute_result_hash
from attestmcp.config import Settings
from attestmcp.identity import IdentityExtractor
from attestmcp.models import AgentIdentity, AuditRecord, PolicyDecision
from attestmcp.store import SQLiteAuditStore

logger = logging.getLogger(__name__)


def create_app(
    settings: Settings | None = None,
    *,
    store: SQLiteAuditStore | None = None,
    identity_extractor: IdentityExtractor | None = None,
    http_client: httpx.AsyncClient | None = None,
) -> FastAPI:
    """Create the AttestMCP proxy FastAPI application.

    Args:
        settings: Configuration. If None, loads from environment variables.
        store: Audit store. If None, creates from settings.
        identity_extractor: Identity extractor. If None, creates from settings.
        http_client: HTTP client for upstream. If None, creates in lifespan.
            Injecting this is useful for testing with mock transports.

    Returns:
        Configured FastAPI application.
    """
    if settings is None:
        settings = Settings()

    if store is None:
        store = SQLiteAuditStore(db_path=settings.db_path)

    if identity_extractor is None:
        identity_extractor = IdentityExtractor(
            jwks_url=settings.jwks_url,
            audience=settings.jwt_audience,
            issuer=settings.jwt_issuer,
        )

    # Track whether we own the client (and must close it)
    _owns_client = http_client is None

    @asynccontextmanager
    async def lifespan(app: FastAPI):  # type: ignore[no-untyped-def]
        nonlocal http_client
        if http_client is None:
            http_client = httpx.AsyncClient(timeout=30.0)
        logger.info(
            "AttestMCP proxy started — upstream: %s, db: %s",
            settings.upstream_url,
            settings.db_path,
        )
        yield
        if _owns_client and http_client is not None:
            await http_client.aclose()

    app = FastAPI(
        title="AttestMCP",
        description="Transparent MCP proxy with tamper-evident audit trails",
        version="0.1.0",
        lifespan=lifespan,
    )

    # Store refs for external access (e.g., test assertions)
    app.state.store = store
    app.state.settings = settings

    @app.get("/health")
    async def health() -> dict[str, Any]:
        """Health check endpoint."""
        return {
            "status": "ok",
            "upstream": settings.upstream_url,
            "audit_records": store.count(),
            "chain_head": store.get_chain_head(),
        }

    @app.get("/metrics")
    async def prometheus_metrics() -> Response:
        """Prometheus metrics endpoint."""
        return Response(
            content=generate_latest(),
            media_type=CONTENT_TYPE_LATEST,
        )

    @app.post("/mcp")
    async def mcp_proxy(request: Request) -> Response:
        """Transparent MCP proxy endpoint."""
        nonlocal http_client
        assert http_client is not None, "HTTP client not initialized"

        start_time = time.monotonic()

        # Read the raw request body
        body = await request.body()

        # Parse JSON-RPC for audit purposes (non-destructive)
        jsonrpc_method: str = "unknown"
        tool_name: str | None = None
        params: dict[str, Any] | None = None
        request_id: Any = None

        try:
            parsed = json.loads(body)
            jsonrpc_method = parsed.get("method", "unknown")
            params = parsed.get("params")
            request_id = parsed.get("id")

            if jsonrpc_method == "tools/call" and isinstance(params, dict):
                tool_name = params.get("name")
        except (json.JSONDecodeError, AttributeError):
            logger.warning("Failed to parse JSON-RPC body for audit")

        # Extract identity from headers
        authorization = request.headers.get("authorization")
        session_id = request.headers.get("mcp-session-id")
        identity: AgentIdentity = identity_extractor.extract(
            authorization, session_id=session_id
        )

        # Build upstream headers
        upstream_headers: dict[str, str] = {
            "content-type": request.headers.get(
                "content-type", "application/json"
            ),
            "accept": request.headers.get("accept", "application/json"),
        }
        if session_id:
            upstream_headers["mcp-session-id"] = session_id

        upstream_error: str | None = None
        upstream_response: httpx.Response | None = None
        response_body: bytes = b""

        try:
            with metrics.upstream_latency.time():
                upstream_response = await http_client.post(
                    settings.upstream_url,
                    content=body,
                    headers=upstream_headers,
                )
            response_body = upstream_response.content

        except httpx.HTTPError as e:
            upstream_error = f"Upstream error: {type(e).__name__}: {e}"
            logger.error("Upstream request failed: %s", upstream_error)

        # Compute duration
        duration_ms = int((time.monotonic() - start_time) * 1000)

        # Build and write audit record — SYNCHRONOUS before returning
        record = AuditRecord(
            session_id=session_id,
            agent_identity=identity,
            method=jsonrpc_method,
            tool_name=tool_name,
            params_hash=compute_params_hash(params),
            policy_decision=PolicyDecision.ALLOW,
            result_hash=compute_result_hash(
                response_body if response_body else None
            ),
            duration_ms=duration_ms,
            error=upstream_error,
        )

        try:
            store.append(record)
            metrics.audit_records_written.labels(
                policy_decision=record.policy_decision.value
            ).inc()
        except Exception:
            logger.exception("CRITICAL: Failed to write audit record")
            metrics.chain_verification_failures.inc()

        metrics.requests_proxied.labels(method=jsonrpc_method).inc()

        # Return upstream response unchanged, or error
        if upstream_response is not None:
            response_headers: dict[str, str] = {}
            if "mcp-session-id" in upstream_response.headers:
                response_headers["mcp-session-id"] = upstream_response.headers[
                    "mcp-session-id"
                ]
            if "content-type" in upstream_response.headers:
                response_headers["content-type"] = upstream_response.headers[
                    "content-type"
                ]

            return Response(
                content=response_body,
                status_code=upstream_response.status_code,
                headers=response_headers,
            )

        # Upstream failed — return JSON-RPC error
        error_response = {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32603,
                "message": "Upstream memory server unavailable",
            },
        }
        return Response(
            content=json.dumps(error_response),
            status_code=502,
            media_type="application/json",
        )

    return app
