"""Litmus CLI — record, replay, and view agent traces."""

from __future__ import annotations

import json
import os
import sys

import click
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.syntax import Syntax

console = Console()


@click.group()
@click.version_option(package_name="litmus-trace")
def cli():
    """Litmus — Record and deterministically replay AI agent executions."""
    pass


@cli.command(context_settings={"ignore_unknown_options": True, "allow_extra_args": True})
@click.option("--replay", "replay_trace", default=None, type=click.Path(exists=True), help="Replay from this trace file")
@click.option("--fault", "fault_strs", multiple=True, help="Fault to inject during replay (e.g. llm_refuse:step=0)")
@click.option("--output-dir", default="./traces", help="Directory to save traces")
@click.argument("command", nargs=-1, required=True, type=click.UNPROCESSED)
def run(replay_trace: str | None, fault_strs: tuple[str, ...], output_dir: str, command: tuple[str, ...]):
    """Wrap a command to record or replay all LLM calls. Zero code changes.

    \b
    Record:  litmus run python my_agent.py
    Replay:  litmus run --replay ./traces/lt-abc.trace.json python my_agent.py
    Faults:  litmus run --replay trace.json --fault llm_refuse:step=0 python my_agent.py
    """
    import subprocess
    import tempfile

    if not command:
        console.print("[red]Error:[/red] No command specified.")
        console.print("Usage: litmus run python my_agent.py")
        sys.exit(1)

    cmd_str = list(command)

    if cmd_str[0] not in ("python", "python3") or len(cmd_str) < 2:
        console.print("[yellow]Non-Python command — use the proxy instead:[/yellow]")
        console.print("  litmus proxy --mode record")
        sys.exit(1)

    if fault_strs and not replay_trace:
        console.print("[red]Error:[/red] --fault requires --replay")
        sys.exit(1)

    mode = "replay" if replay_trace else "record"
    script = cmd_str[1]
    script_args = cmd_str[2:]

    if mode == "record":
        console.print(f"[green]Recording[/green] — {' '.join(cmd_str)}")
    else:
        console.print(f"[cyan]Replaying[/cyan] from {replay_trace}")
        for fs in fault_strs:
            console.print(f"  [red]Fault:[/red] {fs}")
    console.print()

    # Build wrapper that uses runpy for proper __file__, __name__, sys.path
    if mode == "replay":
        replay_path = os.path.abspath(replay_trace)
        faults_repr = repr(list(fault_strs)) if fault_strs else "None"
        setup = f"""
import os
os.environ.setdefault('ANTHROPIC_API_KEY', 'litmus-replay')
os.environ.setdefault('OPENAI_API_KEY', 'litmus-replay')
os.environ.setdefault('MISTRAL_API_KEY', 'litmus-replay')
import litmus.patch as _lp
_lp.replay({replay_path!r}, faults={faults_repr})
"""
    else:
        setup = f"""
import litmus.patch as _lp
_lp.record(output_dir={os.path.abspath(output_dir)!r})
"""

    script_abs = os.path.abspath(script)
    wrapper_code = f"""{setup}
import atexit; atexit.register(_lp.stop)
import sys, os
sys.argv = {[script_abs] + script_args!r}
# Add script's directory to sys.path (like normal python execution)
_script_dir = os.path.dirname({script_abs!r})
if _script_dir not in sys.path:
    sys.path.insert(0, _script_dir)
import runpy
runpy.run_path({script_abs!r}, run_name='__main__')
"""

    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False, prefix="litmus_run_") as f:
        f.write(wrapper_code)
        wrapper_file = f.name

    try:
        env = dict(os.environ)
        result = subprocess.run([cmd_str[0], wrapper_file], env=env)
    finally:
        try:
            os.unlink(wrapper_file)
        except OSError:
            pass

    if mode == "record":
        from pathlib import Path
        out = Path(os.path.abspath(output_dir))
        traces = sorted(out.glob("*.trace.json"), key=lambda p: p.stat().st_mtime, reverse=True)
        if traces:
            console.print(f"\n[green]Trace saved:[/green] {traces[0]}")

    sys.exit(result.returncode)


@cli.command()
@click.option("--mode", type=click.Choice(["record", "replay"]), default="record")
@click.option("--port", default=8787, help="Proxy server port")
@click.option("--output-dir", default="./traces", help="Directory to save traces")
@click.option("--trace-file", default=None, help="Trace file for replay mode")
@click.option("--agent-id", default=None, help="Agent identifier")
@click.option("--provider", "custom_providers", multiple=True, help="Register custom provider: name=base_url")
def proxy(mode: str, port: int, output_dir: str, trace_file: str | None, agent_id: str | None, custom_providers: tuple[str, ...]):
    """Start the Litmus proxy server."""
    import uvicorn

    from litmus.models import Trace, TraceMetadata
    from litmus.proxy.server import LitmusProxy, create_app, register_provider

    # Register any custom providers
    for cp in custom_providers:
        if "=" not in cp:
            console.print(f"[red]Error:[/red] --provider must be name=base_url, got: {cp}")
            sys.exit(1)
        name, base_url = cp.split("=", 1)
        register_provider(name, base_url)
        console.print(f"  [cyan]Custom provider:[/cyan] {name} → {base_url}")

    if mode == "replay":
        if not trace_file:
            console.print("[red]Error:[/red] --trace-file is required for replay mode")
            sys.exit(1)
        trace = Trace.load(trace_file)
        console.print(f"[cyan]Replaying trace:[/cyan] {trace.trace_id} ({len(trace.steps)} steps)")
    else:
        trace = Trace(
            agent_id=agent_id,
            metadata=TraceMetadata.auto(),
        )
        console.print(f"[green]Recording new trace:[/green] {trace.trace_id}")

    litmus_proxy = LitmusProxy(mode=mode, trace=trace, output_dir=output_dir)

    from litmus.proxy.server import PROVIDERS

    console.print(f"[bold]Litmus proxy starting on port {port}[/bold]")
    console.print(f"  Providers: [cyan]{len(PROVIDERS)} available[/cyan]")
    console.print(f"  Anthropic: [dim]ANTHROPIC_BASE_URL=http://localhost:{port}/anthropic[/dim]")
    console.print(f"  OpenAI:    [dim]OPENAI_BASE_URL=http://localhost:{port}/openai[/dim]")
    console.print(f"  Custom:    [dim]BASE_URL=http://localhost:{port}/<name>[/dim]")
    console.print()

    def on_shutdown():
        if mode == "record":
            path = litmus_proxy.save_trace()
            console.print(f"\n[green]Trace saved:[/green] {path}")

    import signal
    import asyncio

    def handle_sigint(sig, frame):
        on_shutdown()
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_sigint)
    signal.signal(signal.SIGTERM, handle_sigint)

    app = create_app(litmus_proxy)

    try:
        uvicorn.run(app, host="0.0.0.0", port=port, log_level="warning")
    finally:
        on_shutdown()


@cli.command()
@click.argument("trace_file", type=click.Path(exists=True))
@click.option("--diff/--no-diff", default=False, help="Show diff against original trace")
@click.option("--fault", "fault_strs", multiple=True, help="Fault to inject (e.g. llm_refuse:step=0)")
@click.option("--port", default=8788, help="Proxy port for replay")
def replay(trace_file: str, diff: bool, fault_strs: tuple[str, ...], port: int):
    """Replay a recorded trace deterministically."""
    from litmus.models import Trace
    from litmus.faults import apply_faults, parse_fault_string

    trace = Trace.load(trace_file)
    console.print(f"[cyan]Loaded trace:[/cyan] {trace.trace_id}")
    console.print(f"  Steps: {len(trace.steps)}")
    console.print(f"  LLM calls: {len(trace.llm_calls())}")
    console.print(f"  Tool calls: {len(trace.tool_calls())}")

    # Apply faults if any
    if fault_strs:
        faults = [parse_fault_string(s) for s in fault_strs]
        for f in faults:
            console.print(f"  [red]Fault:[/red] {f.describe()}")
        trace = apply_faults(trace, faults)
        console.print(f"  [yellow]Trace mutated with {len(faults)} fault(s)[/yellow]")

    console.print()
    console.print(
        f"[bold]Starting replay proxy on port {port}...[/bold]\n"
        f"Run your agent with:\n"
        f"  [cyan]ANTHROPIC_BASE_URL=http://localhost:{port}/anthropic[/cyan]\n"
        f"  [cyan]OPENAI_BASE_URL=http://localhost:{port}/openai[/cyan]\n"
    )

    import uvicorn
    from litmus.proxy.server import LitmusProxy, create_app

    litmus_proxy = LitmusProxy(mode="replay", trace=trace)
    app = create_app(litmus_proxy)
    uvicorn.run(app, host="0.0.0.0", port=port, log_level="warning")


@cli.command()
@click.argument("trace_file", type=click.Path(exists=True))
@click.option("--steps/--no-steps", default=True, help="Show individual steps")
def view(trace_file: str, steps: bool):
    """Pretty-print a trace file."""
    from litmus.models import Trace, LLMCall, ToolCall

    trace = Trace.load(trace_file)

    # Header
    console.print(Panel.fit(
        f"[bold cyan]Trace:[/bold cyan] {trace.trace_id}\n"
        f"[bold]Agent:[/bold] {trace.agent_id or 'unknown'}\n"
        f"[bold]Status:[/bold] {trace.status.value}\n"
        f"[bold]Started:[/bold] {trace.started_at}\n"
        f"[bold]Duration:[/bold] {trace.duration_ms():.0f}ms" if trace.duration_ms() else "",
        title="Litmus Trace",
    ))

    if not steps:
        console.print(f"\n  {len(trace.steps)} steps ({len(trace.llm_calls())} LLM, {len(trace.tool_calls())} tool)")
        return

    console.print()

    # Steps table
    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("#", style="dim", width=4)
    table.add_column("Type", width=10)
    table.add_column("Provider/URL", width=20)
    table.add_column("Model", width=25)
    table.add_column("Status", width=8)
    table.add_column("Latency", justify="right", width=10)

    for i, step in enumerate(trace.steps):
        if isinstance(step, LLMCall):
            table.add_row(
                str(i),
                "[magenta]LLM[/magenta]",
                step.request.provider,
                step.request.model or "—",
                str(step.response.status_code) if step.response else "[red]err[/red]",
                f"{step.latency_ms:.0f}ms",
            )
        elif isinstance(step, ToolCall):
            table.add_row(
                str(i),
                "[green]Tool[/green]",
                step.request.url[:20],
                "—",
                str(step.response.status_code) if step.response else "[red]err[/red]",
                f"{step.latency_ms:.0f}ms",
            )

    console.print(table)


@cli.command()
def providers():
    """List all available LLM providers."""
    from litmus.proxy.server import PROVIDERS

    table = Table(show_header=True, header_style="bold cyan")
    table.add_column("Provider", width=15)
    table.add_column("Base URL", width=50)
    table.add_column("Usage", width=40)

    for name, url in sorted(PROVIDERS.items()):
        table.add_row(
            f"[bold]{name}[/bold]",
            f"[dim]{url}[/dim]",
            f"BASE_URL=http://proxy:8787/{name}",
        )

    console.print(table)
    console.print(f"\n[dim]Total: {len(PROVIDERS)} providers. Add custom: litmus proxy --provider name=https://url[/dim]")


@cli.command()
@click.argument("traces_dir", type=click.Path(exists=True))
@click.option("--threshold", default=80.0, help="Minimum average score to pass (0-100)")
@click.option("--json-output", "json_out", default=None, type=click.Path(), help="Write JSON report to file")
@click.option("--verbose", is_flag=True, help="Show per-trace details")
def ci(traces_dir: str, threshold: float, json_out: str | None, verbose: bool):
    """Run CI scoring on a trace directory. Exits non-zero if score < threshold."""
    import json as json_mod
    from litmus.scorer import run_ci

    console.print(f"[bold]Litmus CI[/bold] — scoring traces in {traces_dir}")
    console.print(f"  Threshold: {threshold}")
    console.print()

    report = run_ci(traces_dir, threshold=threshold)

    if report.trace_count == 0:
        console.print("[yellow]No traces found.[/yellow]")
        sys.exit(1)

    # Per-trace results
    if verbose:
        trace_table = Table(show_header=True, header_style="bold cyan")
        trace_table.add_column("Trace ID", width=16)
        trace_table.add_column("Steps", justify="right", width=6)
        trace_table.add_column("Correct", justify="right", width=8)
        trace_table.add_column("Resilient", justify="right", width=9)
        trace_table.add_column("Efficient", justify="right", width=9)
        trace_table.add_column("Overall", justify="right", width=8)
        trace_table.add_column("Errors", width=30)

        for s in report.scores:
            status_color = "green" if s.overall >= threshold else "red"
            trace_table.add_row(
                s.trace_id,
                str(s.steps),
                f"{s.correctness:.0f}",
                f"{s.resilience:.0f}",
                f"{s.efficiency:.0f}",
                f"[{status_color}]{s.overall:.0f}[/{status_color}]",
                "; ".join(s.errors[:2]) if s.errors else "—",
            )

        console.print(trace_table)
        console.print()

    # Summary
    pass_color = "green" if report.passed else "red"
    console.print(Panel.fit(
        f"[bold]Traces scored:[/bold] {report.trace_count}\n"
        f"[bold]Average score:[/bold] [{pass_color}]{report.average_score:.1f}[/{pass_color}]\n"
        f"[bold]Min / Max:[/bold] {report.min_score:.1f} / {report.max_score:.1f}\n"
        f"[bold]Threshold:[/bold] {threshold}\n"
        f"[bold]Result:[/bold] [{'green' if report.passed else 'red bold'}]"
        f"{'PASSED' if report.passed else 'FAILED'}"
        f"[/{'green' if report.passed else 'red bold'}]",
        title="Litmus CI Report",
        border_style=pass_color,
    ))

    # JSON output
    if json_out:
        with open(json_out, "w") as f:
            json_mod.dump(report.to_dict(), f, indent=2)
        console.print(f"\n[dim]JSON report written to {json_out}[/dim]")

    sys.exit(0 if report.passed else 1)


if __name__ == "__main__":
    cli()
