"""Direct SDK monkey-patching — no proxy needed.

Usage:
    import litmus
    litmus.record()          # start recording all LLM calls
    # ... your agent code, unchanged ...
    litmus.stop()            # save trace

    litmus.replay("trace.json")  # replay from a trace file
    # ... your agent code, gets recorded responses ...
    litmus.stop()
"""

from __future__ import annotations

import atexit
import functools
import json
import sys
import threading
import time
from pathlib import Path
from typing import Any

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

# ---------- Thread-safe global state ----------

_lock = threading.Lock()
_active_trace: Trace | None = None
_mode: str | None = None  # "record" or "replay"
_replay_trace: Trace | None = None
_replay_index: int = 0
_output_dir: str = "./traces"
_original_transports: dict[str, Any] = {}
_patched: bool = False
_atexit_registered: bool = False


# ---------- Public API ----------

def record(output_dir: str = "./traces", agent_id: str | None = None) -> Trace:
    """Start recording all LLM SDK calls. No proxy, no env vars, no setup.

    Call this before any LLM SDK usage. All calls to Anthropic/OpenAI
    will be intercepted and recorded.

    Returns the Trace object (also saved to disk on stop()).
    """
    global _active_trace, _mode, _output_dir

    with _lock:
        _mode = "record"
        _output_dir = output_dir
        _active_trace = Trace(agent_id=agent_id, metadata=TraceMetadata.auto())

    _patch_sdks()
    _register_atexit()

    return _active_trace


def replay(trace_path: str, faults: list[str] | None = None) -> Trace:
    """Replay a recorded trace. All LLM calls return recorded responses.

    Call this before any LLM SDK usage. The agent's code runs unchanged
    but gets deterministic pre-recorded responses instead of real API calls.

    Args:
        trace_path: Path to the trace file.
        faults: Optional list of fault strings (e.g. ["llm_refuse:step=0"]).

    Raises:
        FileNotFoundError: If trace_path doesn't exist.
        ValueError: If trace file is corrupted or invalid.
    """
    global _active_trace, _mode, _replay_trace, _replay_index

    # Validate trace file exists
    p = Path(trace_path)
    if not p.exists():
        raise FileNotFoundError(f"[litmus] Trace file not found: {trace_path}")

    try:
        loaded = Trace.load(trace_path)
    except Exception as e:
        raise ValueError(f"[litmus] Failed to load trace {trace_path}: {e}") from e

    # Apply faults if any
    if faults:
        from litmus.faults import apply_faults, parse_fault_string
        fault_objs = [parse_fault_string(f) for f in faults]
        loaded = apply_faults(loaded, fault_objs)

    with _lock:
        _mode = "replay"
        _replay_trace = loaded
        _replay_index = 0
        _active_trace = Trace(
            agent_id=_replay_trace.agent_id,
            metadata=TraceMetadata.auto(),
        )

    _patch_sdks()
    _register_atexit()

    return _replay_trace


def stop() -> str | None:
    """Stop recording/replaying and save the trace. Returns the file path."""
    global _active_trace, _mode, _replay_trace, _replay_index

    _unpatch_sdks()

    with _lock:
        if _active_trace is None:
            return None

        path = None
        if _mode == "record":
            _active_trace.complete()
            path = str(Path(_output_dir) / f"{_active_trace.trace_id}.trace.json")
            _active_trace.save(path)

        _active_trace = None
        _mode = None
        _replay_trace = None
        _replay_index = 0
        return path


def get_trace() -> Trace | None:
    """Get the current active trace (during recording or replay)."""
    return _active_trace


# ---------- Internal ----------

def _register_atexit() -> None:
    global _atexit_registered
    if not _atexit_registered:
        atexit.register(stop)
        _atexit_registered = True


def _patch_sdks() -> None:
    """Monkey-patch installed SDK HTTP transports."""
    global _patched
    if _patched:
        return
    _patched = True
    _try_patch_httpx()


def _unpatch_sdks() -> None:
    """Restore original SDK transports."""
    global _patched
    if not _patched:
        return
    _patched = False
    _try_unpatch_httpx()


def _try_patch_httpx() -> None:
    """Patch httpx.Client.send and httpx.AsyncClient.send.

    Both Anthropic and OpenAI Python SDKs use httpx under the hood.
    By patching httpx, we intercept ALL HTTP calls from both SDKs
    without needing to know their internal APIs.
    """
    try:
        import httpx
    except ImportError:
        return

    # Patch sync client
    original_sync_send = httpx.Client.send
    _original_transports["httpx_sync_send"] = original_sync_send

    @functools.wraps(original_sync_send)
    def patched_sync_send(self, request, **kwargs):
        if _should_intercept(str(request.url)):
            return _handle_sync_request(self, request, original_sync_send, **kwargs)
        return original_sync_send(self, request, **kwargs)

    httpx.Client.send = patched_sync_send

    # Patch async client
    original_async_send = httpx.AsyncClient.send
    _original_transports["httpx_async_send"] = original_async_send

    @functools.wraps(original_async_send)
    async def patched_async_send(self, request, **kwargs):
        if _should_intercept(str(request.url)):
            return await _handle_async_request(self, request, original_async_send, **kwargs)
        return await original_async_send(self, request, **kwargs)

    httpx.AsyncClient.send = patched_async_send


def _try_unpatch_httpx() -> None:
    """Restore original httpx send methods."""
    try:
        import httpx
    except ImportError:
        return

    if "httpx_sync_send" in _original_transports:
        httpx.Client.send = _original_transports.pop("httpx_sync_send")
    if "httpx_async_send" in _original_transports:
        httpx.AsyncClient.send = _original_transports.pop("httpx_async_send")


# ---------- Host detection ----------

_LLM_HOSTS = {
    "api.anthropic.com",
    "api.openai.com",
    "api.mistral.ai",
    "api.cohere.com",
    "api.groq.com",
    "api.together.xyz",
    "api.fireworks.ai",
    "api.deepseek.com",
    "api.perplexity.ai",
    "openrouter.ai",
    "generativelanguage.googleapis.com",
}


def _should_intercept(url: str) -> bool:
    """Check if this URL is an LLM API call we should intercept."""
    if _mode is None:
        return False
    from urllib.parse import urlparse
    host = urlparse(url).hostname or ""
    return host in _LLM_HOSTS


def _detect_provider_from_url(url: str) -> str:
    """Map a URL to a provider name."""
    from urllib.parse import urlparse
    host = urlparse(url).hostname or ""
    mapping = {
        "api.anthropic.com": "anthropic",
        "api.openai.com": "openai",
        "api.mistral.ai": "mistral",
        "api.cohere.com": "cohere",
        "api.groq.com": "groq",
        "api.together.xyz": "together",
        "api.fireworks.ai": "fireworks",
        "api.deepseek.com": "deepseek",
        "api.perplexity.ai": "perplexity",
        "openrouter.ai": "openrouter",
        "generativelanguage.googleapis.com": "google",
    }
    return mapping.get(host, "unknown")


# ---------- Helpers ----------

def _extract_model(body: dict | None) -> str | None:
    if not body:
        return None
    return body.get("model")


def _parse_body(content: bytes) -> dict | None:
    try:
        return json.loads(content) if content else None
    except (json.JSONDecodeError, UnicodeDecodeError):
        return None


def _get_request_path(request) -> str:
    """Extract the path string from an httpx request."""
    raw = request.url.raw_path
    if isinstance(raw, bytes):
        return raw.decode("utf-8", errors="replace")
    return str(raw)


def _replay_next_step(request):
    """Thread-safe: get the next replay step and serve the response.

    Returns an httpx.Response or None if exhausted.
    """
    import httpx

    global _replay_index

    with _lock:
        if _replay_trace is None:
            return None
        if _replay_index >= len(_replay_trace.steps):
            total = len(_replay_trace.steps)
            idx = _replay_index
            print(
                f"\n[litmus] Replay exhausted — agent requested step {idx + 1} "
                f"but recording only has {total} steps.\n"
                f"[litmus] The agent's behavior diverged from the original trace.",
                file=sys.stderr,
            )
            return httpx.Response(
                status_code=500,
                content=json.dumps({
                    "error": f"Litmus replay exhausted — recording had {total} steps, "
                             f"agent requested step {idx + 1}"
                }).encode(),
                request=request,
            )

        step = _replay_trace.steps[_replay_index]
        _replay_index += 1

    if not isinstance(step, LLMCall) or step.response is None:
        return httpx.Response(
            status_code=500,
            content=json.dumps({"error": "Recorded step has no response"}).encode(),
            request=request,
        )

    resp_body = json.dumps(step.response.body or {}).encode()

    if step.response.is_stream and step.response.stream_chunks:
        raw_chunks = "".join(step.response.stream_chunks).encode()
        return httpx.Response(
            status_code=step.response.status_code,
            headers={"content-type": "text/event-stream"},
            stream=httpx.ByteStream(raw_chunks),
            request=request,
        )

    return httpx.Response(
        status_code=step.response.status_code,
        headers={"content-type": "application/json"},
        content=resp_body,
        request=request,
    )


def _record_step(llm_call: LLMCall) -> None:
    """Thread-safe: add a step to the active trace."""
    with _lock:
        if _active_trace is not None:
            _active_trace.add_step(llm_call)


# ---------- Sync handler ----------

def _handle_sync_request(client, request, original_send, **kwargs):
    """Intercept a sync HTTP request — record or replay."""
    url = str(request.url)
    provider = _detect_provider_from_url(url)
    body_json = _parse_body(request.content)
    model = _extract_model(body_json)
    is_stream = body_json.get("stream", False) if body_json else False

    # Replay mode — serve recorded response (thread-safe)
    if _mode == "replay":
        resp = _replay_next_step(request)
        if resp is not None:
            return resp

    # Record mode — forward the real request and capture response
    llm_call = LLMCall(
        request=LLMRequest(
            provider=provider,
            model=model,
            method=request.method,
            path=_get_request_path(request),
            body=body_json,
        )
    )

    start = time.monotonic()
    try:
        response = original_send(client, request, **kwargs)
    except Exception as e:
        llm_call.latency_ms = (time.monotonic() - start) * 1000
        llm_call.error = f"[litmus] Recording failed: {e}"
        _record_step(llm_call)
        raise

    llm_call.latency_ms = (time.monotonic() - start) * 1000

    # Streaming — wrap the stream to capture chunks while forwarding
    if is_stream:
        chunks: list[str] = []
        original_stream = response.stream

        def recording_stream():
            try:
                for chunk in original_stream:
                    if isinstance(chunk, bytes):
                        chunks.append(chunk.decode("utf-8", errors="replace"))
                    else:
                        chunks.append(str(chunk))
                    yield chunk
            finally:
                # Record even if stream is interrupted
                llm_call.response = LLMResponse(
                    status_code=response.status_code,
                    is_stream=True,
                    stream_chunks=chunks,
                )
                _record_step(llm_call)

        response.stream = recording_stream()
        return response

    # Non-streaming — read body
    resp_json = None
    try:
        resp_json = response.json()
    except Exception:
        pass

    llm_call.response = LLMResponse(
        status_code=response.status_code,
        body=resp_json,
        is_stream=False,
    )
    _record_step(llm_call)

    return response


# ---------- Async handler ----------

async def _handle_async_request(client, request, original_send, **kwargs):
    """Intercept an async HTTP request — record or replay."""
    url = str(request.url)
    provider = _detect_provider_from_url(url)
    body_json = _parse_body(request.content)
    model = _extract_model(body_json)
    is_stream = body_json.get("stream", False) if body_json else False

    # Replay mode — serve recorded response (thread-safe)
    if _mode == "replay":
        resp = _replay_next_step(request)
        if resp is not None:
            return resp

    # Record mode — forward the real request and capture response
    llm_call = LLMCall(
        request=LLMRequest(
            provider=provider,
            model=model,
            method=request.method,
            path=_get_request_path(request),
            body=body_json,
        )
    )

    start = time.monotonic()
    try:
        response = await original_send(client, request, **kwargs)
    except Exception as e:
        llm_call.latency_ms = (time.monotonic() - start) * 1000
        llm_call.error = f"[litmus] Async recording failed: {e}"
        _record_step(llm_call)
        raise

    llm_call.latency_ms = (time.monotonic() - start) * 1000

    # Async streaming — wrap the stream to capture chunks
    if is_stream:
        chunks: list[str] = []
        original_astream = response.stream

        async def recording_astream():
            try:
                async for chunk in original_astream:
                    if isinstance(chunk, bytes):
                        chunks.append(chunk.decode("utf-8", errors="replace"))
                    else:
                        chunks.append(str(chunk))
                    yield chunk
            finally:
                llm_call.response = LLMResponse(
                    status_code=response.status_code,
                    is_stream=True,
                    stream_chunks=chunks,
                )
                _record_step(llm_call)

        response.stream = recording_astream()
        return response

    # Non-streaming — read body
    resp_json = None
    try:
        resp_json = response.json()
    except Exception:
        pass

    llm_call.response = LLMResponse(
        status_code=response.status_code,
        body=resp_json,
        is_stream=False,
    )
    _record_step(llm_call)

    return response
