"""Litmus proxy server — intercepts LLM/tool API calls for recording and replay."""

from __future__ import annotations

import asyncio
import json
import time
from datetime import datetime, timezone

import httpx
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import Response, StreamingResponse
from starlette.routing import Route, Mount

from litmus.models import (
    Trace,
    LLMCall,
    LLMRequest,
    LLMResponse,
    ToolCall,
    ToolRequest,
    ToolResponse,
    TraceMetadata,
)

# Built-in provider base URLs.
# For SDKs that set base_url to include a version prefix (like OpenAI's /v1),
# we include that prefix here so the proxy reconstructs the full URL correctly.
DEFAULT_PROVIDERS: dict[str, str] = {
    # Major providers
    "anthropic": "https://api.anthropic.com",
    "openai": "https://api.openai.com/v1",
    "google": "https://generativelanguage.googleapis.com/v1beta",
    "mistral": "https://api.mistral.ai/v1",
    "cohere": "https://api.cohere.com/v2",
    # OpenAI-compatible inference providers
    "groq": "https://api.groq.com/openai/v1",
    "together": "https://api.together.xyz/v1",
    "fireworks": "https://api.fireworks.ai/inference/v1",
    "deepseek": "https://api.deepseek.com/v1",
    "perplexity": "https://api.perplexity.ai",
    "openrouter": "https://openrouter.ai/api/v1",
    # Self-hosted / local
    "ollama": "http://localhost:11434/v1",
    "vllm": "http://localhost:8000/v1",
    "lmstudio": "http://localhost:1234/v1",
}

# Mutable registry — starts with defaults, users can add custom providers
PROVIDERS: dict[str, str] = dict(DEFAULT_PROVIDERS)


def register_provider(name: str, base_url: str) -> None:
    """Register a custom provider (self-hosted model, fine-tuned endpoint, etc).

    Example:
        register_provider("my-model", "https://my-finetuned-llama.example.com/v1")

    Then route traffic through:
        MY_MODEL_BASE_URL=http://localhost:8787/my-model
    """
    PROVIDERS[name] = base_url.rstrip("/")


# Headers to skip entirely when recording
SKIP_HEADERS = {"host", "content-length", "transfer-encoding", "connection"}

# Headers containing secrets — redact the value but keep the key
SECRET_HEADERS = {"authorization", "x-api-key", "api-key", "x-goog-api-key"}


def _sanitize_headers(headers: dict[str, str]) -> dict[str, str]:
    """Remove hop-by-hop headers and redact secrets for safe trace storage."""
    result = {}
    for k, v in headers.items():
        lower = k.lower()
        if lower in SKIP_HEADERS:
            continue
        if lower in SECRET_HEADERS:
            result[k] = "[REDACTED]"
        else:
            result[k] = v
    return result


def _detect_provider(path: str) -> str | None:
    """Detect which provider a request is for based on the URL prefix."""
    parts = path.strip("/").split("/")
    if parts and parts[0] in PROVIDERS:
        return parts[0]
    return None


def _strip_provider_prefix(path: str, provider: str) -> str:
    """Remove the /anthropic or /openai prefix from the path."""
    prefix = f"/{provider}"
    if path.startswith(prefix):
        return path[len(prefix):]
    return path


class LitmusProxy:
    """The core proxy that intercepts, records, and replays API calls."""

    def __init__(
        self,
        mode: str = "record",  # "record" or "replay"
        trace: Trace | None = None,
        output_dir: str = "./traces",
    ):
        self.mode = mode
        self.output_dir = output_dir

        if mode == "record":
            self.trace = trace or Trace(metadata=TraceMetadata.auto())
            self._replay_index = 0
        elif mode == "replay":
            if trace is None:
                raise ValueError("Must provide a trace for replay mode")
            self.trace = trace
            self._replay_index = 0
        else:
            raise ValueError(f"Unknown mode: {mode}")

        # In replay mode, track what the agent actually requested (for diff)
        self.replay_requests: list[LLMRequest] = []

        self._client: httpx.AsyncClient | None = None

    async def _get_client(self) -> httpx.AsyncClient:
        if self._client is None:
            self._client = httpx.AsyncClient(timeout=120.0)
        return self._client

    async def shutdown(self):
        if self._client:
            await self._client.aclose()
            self._client = None

    async def handle_request(self, request: Request) -> Response:
        """Route incoming request to record or replay handler."""
        provider = _detect_provider(request.url.path)
        if provider is None:
            # Health check / info endpoint
            if request.url.path in ("/", "/health"):
                return Response(
                    json.dumps({
                        "status": "ok",
                        "mode": self.mode,
                        "trace_id": self.trace.trace_id,
                        "steps_recorded": len(self.trace.steps),
                        "providers": list(PROVIDERS.keys()),
                    }),
                    media_type="application/json",
                )

            # Register a custom provider at runtime
            if request.url.path == "/register" and request.method == "POST":
                body = await request.body()
                try:
                    data = json.loads(body)
                    name = data["name"]
                    base_url = data["base_url"]
                    register_provider(name, base_url)
                    return Response(
                        json.dumps({"registered": name, "base_url": base_url}),
                        media_type="application/json",
                    )
                except (json.JSONDecodeError, KeyError) as e:
                    return Response(
                        json.dumps({"error": f"Bad request: {e}. Send {{\"name\": \"...\", \"base_url\": \"...\"}}"}),
                        status_code=400,
                        media_type="application/json",
                    )

            available = ", ".join(sorted(PROVIDERS.keys()))
            return Response(
                json.dumps({"error": f"Unknown provider. Available: {available}. Or POST /register to add custom."}),
                status_code=404,
                media_type="application/json",
            )

        if self.mode == "record":
            return await self._handle_record(request, provider)
        else:
            return await self._handle_replay(request, provider)

    async def _handle_record(self, request: Request, provider: str) -> Response:
        """Forward request to real API, record both request and response."""
        client = await self._get_client()
        target_base = PROVIDERS[provider]
        target_path = _strip_provider_prefix(request.url.path, provider)
        target_url = f"{target_base}{target_path}"

        # Read request body
        body_bytes = await request.body()
        body_json = None
        try:
            body_json = json.loads(body_bytes) if body_bytes else None
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        # Build request headers (forward auth headers from client)
        # Remove accept-encoding so upstream sends uncompressed — avoids double-decompression
        skip_req = SKIP_HEADERS | {"accept-encoding"}
        forward_headers = {}
        for key, value in request.headers.items():
            if key.lower() not in skip_req:
                forward_headers[key] = value

        # Detect if client wants streaming
        is_stream = body_json.get("stream", False) if isinstance(body_json, dict) else False

        # Extract model name
        model = None
        if isinstance(body_json, dict):
            model = body_json.get("model")

        # Record the request
        llm_call = LLMCall(
            request=LLMRequest(
                provider=provider,
                model=model,
                method=request.method,
                path=target_path,
                headers=_sanitize_headers(dict(request.headers)),
                body=body_json,
            )
        )

        start_time = time.monotonic()

        if is_stream:
            return await self._record_streaming(
                client, target_url, forward_headers, body_bytes, llm_call, start_time
            )

        # Non-streaming: simple forward
        try:
            resp = await client.request(
                method=request.method,
                url=target_url,
                headers=forward_headers,
                content=body_bytes,
            )
        except httpx.HTTPError as e:
            llm_call.latency_ms = (time.monotonic() - start_time) * 1000
            llm_call.error = str(e)
            self.trace.add_step(llm_call)
            return Response(f"Upstream error: {e}", status_code=502)

        elapsed = (time.monotonic() - start_time) * 1000

        # Parse response
        resp_json = None
        try:
            resp_json = resp.json()
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        llm_call.response = LLMResponse(
            status_code=resp.status_code,
            headers=_sanitize_headers(dict(resp.headers)),
            body=resp_json,
            is_stream=False,
        )
        llm_call.latency_ms = elapsed
        self.trace.add_step(llm_call)

        # Strip encoding headers since httpx already decompressed
        resp_headers = {
            k: v for k, v in resp.headers.items()
            if k.lower() not in ("content-encoding", "content-length", "transfer-encoding")
        }
        return Response(
            content=resp.content,
            status_code=resp.status_code,
            headers=resp_headers,
        )

    async def _record_streaming(
        self,
        client: httpx.AsyncClient,
        url: str,
        headers: dict,
        body: bytes,
        llm_call: LLMCall,
        start_time: float,
    ) -> StreamingResponse:
        """Record a streaming SSE response while forwarding it to the client."""
        chunks: list[str] = []

        async def stream_and_record():
            try:
                async with client.stream(
                    "POST", url, headers=headers, content=body
                ) as resp:
                    llm_call.response = LLMResponse(
                        status_code=resp.status_code,
                        headers=_sanitize_headers(dict(resp.headers)),
                        is_stream=True,
                        stream_chunks=[],
                    )
                    async for chunk in resp.aiter_text():
                        chunks.append(chunk)
                        yield chunk

                    llm_call.response.stream_chunks = chunks
                    llm_call.latency_ms = (time.monotonic() - start_time) * 1000
                    self.trace.add_step(llm_call)
            except httpx.HTTPError as e:
                llm_call.latency_ms = (time.monotonic() - start_time) * 1000
                llm_call.error = str(e)
                self.trace.add_step(llm_call)

        return StreamingResponse(
            stream_and_record(),
            media_type="text/event-stream",
        )

    async def _handle_replay(self, request: Request, provider: str) -> Response:
        """Serve a recorded response instead of calling the real API."""
        # Capture what the agent is requesting (for divergence detection)
        body_bytes = await request.body()
        body_json = None
        try:
            body_json = json.loads(body_bytes) if body_bytes else None
        except (json.JSONDecodeError, UnicodeDecodeError):
            pass

        target_path = _strip_provider_prefix(request.url.path, provider)
        self.replay_requests.append(LLMRequest(
            provider=provider,
            model=body_json.get("model") if isinstance(body_json, dict) else None,
            method=request.method,
            path=target_path,
            headers=_sanitize_headers(dict(request.headers)),
            body=body_json,
        ))

        if self._replay_index >= len(self.trace.steps):
            return Response(
                json.dumps({"error": "Replay exhausted — no more recorded steps"}),
                status_code=500,
                media_type="application/json",
            )

        step = self.trace.steps[self._replay_index]
        self._replay_index += 1

        # Only replay LLM calls for now
        if not isinstance(step, LLMCall) or step.response is None:
            return Response(
                json.dumps({"error": f"Step {step.step_id} has no recorded response"}),
                status_code=500,
                media_type="application/json",
            )

        resp = step.response

        if resp.is_stream and resp.stream_chunks:
            # Replay streaming response
            async def replay_stream():
                for chunk in resp.stream_chunks:
                    yield chunk
                    # Small delay to simulate real streaming cadence
                    await asyncio.sleep(0.01)

            return StreamingResponse(
                replay_stream(),
                status_code=resp.status_code,
                media_type="text/event-stream",
            )

        # Non-streaming replay
        body = json.dumps(resp.body) if resp.body else ""
        return Response(
            content=body,
            status_code=resp.status_code,
            media_type="application/json",
        )

    def save_trace(self) -> str:
        """Finalize and save the current trace. Returns the file path."""
        from pathlib import Path

        self.trace.complete()
        path = str(Path(self.output_dir) / f"{self.trace.trace_id}.trace.json")
        self.trace.save(path)
        return path


def create_app(proxy: LitmusProxy) -> Starlette:
    """Create the Starlette ASGI app wrapping the proxy."""

    async def catch_all(request: Request) -> Response:
        return await proxy.handle_request(request)

    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def lifespan(app):
        yield
        await proxy.shutdown()

    app = Starlette(
        routes=[
            Route("/{path:path}", catch_all, methods=["GET", "POST", "PUT", "DELETE", "PATCH"]),
            Route("/", catch_all),
        ],
        lifespan=lifespan,
    )
    return app
