"""Litmus SDK shim — one-liner to start recording or replaying."""

from __future__ import annotations

import atexit
import os
import signal
import subprocess
import sys
import time
from pathlib import Path


_proxy_process: subprocess.Popen | None = None
_original_env: dict[str, str | None] = {}

DEFAULT_PORT = 8787
DEFAULT_OUTPUT_DIR = "./traces"


def init(
    mode: str = "record",
    port: int = DEFAULT_PORT,
    output_dir: str = DEFAULT_OUTPUT_DIR,
    trace_file: str | None = None,
    agent_id: str | None = None,
) -> None:
    """Start the Litmus proxy and redirect LLM SDK calls through it.

    Args:
        mode: "record" to capture calls, "replay" to serve recorded responses.
        port: Port for the proxy server.
        output_dir: Directory to save trace files (record mode).
        trace_file: Path to trace file (required for replay mode).
        agent_id: Optional agent identifier for the trace.
    """
    global _proxy_process

    if _proxy_process is not None:
        return  # Already initialized

    # Build CLI command
    cmd = [
        sys.executable, "-m", "litmus.cli.main",
        "proxy",
        "--mode", mode,
        "--port", str(port),
        "--output-dir", output_dir,
    ]
    if trace_file:
        cmd.extend(["--trace-file", trace_file])
    if agent_id:
        cmd.extend(["--agent-id", agent_id])

    # Start proxy as subprocess
    _proxy_process = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )

    # Wait for proxy to be ready
    proxy_base = f"http://localhost:{port}"
    _wait_for_proxy(proxy_base)

    # Patch environment variables so SDKs route through us
    _patch_env(port)

    # Register cleanup
    atexit.register(shutdown)


def _wait_for_proxy(base_url: str, timeout: float = 10.0) -> None:
    """Poll the proxy health endpoint until it's ready."""
    import httpx

    deadline = time.monotonic() + timeout
    while time.monotonic() < deadline:
        try:
            resp = httpx.get(f"{base_url}/health", timeout=1.0)
            if resp.status_code == 200:
                return
        except (httpx.ConnectError, httpx.TimeoutException):
            pass
        time.sleep(0.1)

    raise TimeoutError(f"Litmus proxy did not start within {timeout}s")


def _patch_env(port: int) -> None:
    """Set SDK base URL env vars to point at the proxy.

    Patches env vars for all known providers. Each SDK checks its own env var:
    - Anthropic SDK: ANTHROPIC_BASE_URL
    - OpenAI SDK:    OPENAI_BASE_URL  (also used by Groq, Together, etc.)
    - Mistral SDK:   MISTRAL_API_URL
    - Google SDK:    (requires code-level base_url override)
    - Cohere SDK:    CO_API_URL
    """
    global _original_env

    proxy_base = f"http://localhost:{port}"

    # Map provider names to the env vars their SDKs respect
    env_vars = {
        "ANTHROPIC_BASE_URL": f"{proxy_base}/anthropic",
        "OPENAI_BASE_URL": f"{proxy_base}/openai",
        "MISTRAL_API_URL": f"{proxy_base}/mistral",
        "CO_API_URL": f"{proxy_base}/cohere",
        # OpenAI-compatible providers: users set OPENAI_BASE_URL or
        # pass base_url= in their SDK constructor. For Groq/Together/etc,
        # they can use: client = OpenAI(base_url="http://localhost:{port}/groq")
    }

    for key, value in env_vars.items():
        _original_env[key] = os.environ.get(key)
        os.environ[key] = value


def _restore_env() -> None:
    """Restore original environment variables."""
    for key, original in _original_env.items():
        if original is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = original
    _original_env.clear()


def shutdown() -> None:
    """Stop the proxy and restore environment."""
    global _proxy_process

    _restore_env()

    if _proxy_process is not None:
        _proxy_process.terminate()
        try:
            _proxy_process.wait(timeout=5)
        except subprocess.TimeoutExpired:
            _proxy_process.kill()
        _proxy_process = None
