"""``albus trace`` — inspect execution traces stored in .albus_store/."""

from __future__ import annotations

from typing import Annotated

import typer

from albusos.cli import app
from albusos.cli._helpers import err, out, run_async

trace_app = typer.Typer(
    name="trace",
    help="Inspect execution traces.",
    no_args_is_help=True,
)
app.add_typer(trace_app, name="trace")


# ---------------------------------------------------------------------------
# albus trace list
# ---------------------------------------------------------------------------


@trace_app.command("list")
def trace_list_cmd(
    limit: Annotated[int, typer.Option("--limit", "-n", help="Max traces to show.")] = 20,
    store_root: Annotated[str, typer.Option("--store", help="State store root.")] = ".albus_store",
) -> None:
    """List stored execution traces (newest first)."""
    from core.state import FileStateStore

    store = FileStateStore(root=store_root)
    traces = run_async(store.list_traces(limit=limit))

    if not traces:
        err.print("No traces found.")
        raise typer.Exit(0)

    from rich.table import Table

    table = Table(title="Execution Traces", show_lines=False)
    table.add_column("Execution ID", style="cyan", no_wrap=True)
    table.add_column("Pathway", style="green")
    table.add_column("Nodes", justify="right")
    table.add_column("Tokens", justify="right")
    table.add_column("Duration", justify="right")

    for t in traces:
        from core.observability.trace import _format_gantt_duration

        dur = _format_gantt_duration(t.total_duration_ms)
        table.add_row(
            t.execution_id[:16] + "..." if len(t.execution_id) > 19 else t.execution_id,
            t.pathway_id or "-",
            str(len(t.spans)),
            f"{t.total_tokens:,}",
            dur,
        )

    out.print(table)


# ---------------------------------------------------------------------------
# albus trace show
# ---------------------------------------------------------------------------


@trace_app.command("show")
def trace_show_cmd(
    execution_id: Annotated[str, typer.Argument(help="Execution ID (or prefix).")],
    store_root: Annotated[str, typer.Option("--store", help="State store root.")] = ".albus_store",
) -> None:
    """Show details of a single execution trace."""
    trace = _resolve_trace(execution_id, store_root)

    from core.observability.trace import _format_gantt_duration

    err.print(f"\n[bold]Execution:[/bold]  {trace.execution_id}")
    err.print(f"[bold]Pathway:[/bold]    {trace.pathway_id or '-'}")
    err.print(f"[bold]Duration:[/bold]   {_format_gantt_duration(trace.total_duration_ms)}")
    err.print(f"[bold]Tokens:[/bold]     {trace.total_tokens:,}")
    err.print(f"[bold]Nodes:[/bold]      {len(trace.spans)}\n")

    from rich.table import Table

    table = Table(show_lines=False)
    table.add_column("Node", style="cyan")
    table.add_column("Type", style="dim")
    table.add_column("Status")
    table.add_column("Duration", justify="right")
    table.add_column("Tools", justify="right")
    table.add_column("LLM Calls", justify="right")

    for span in trace.spans:
        status_style = {"success": "green", "failure": "red", "skipped": "dim"}.get(span.status, "")
        table.add_row(
            span.node_id,
            span.node_type,
            f"[{status_style}]{span.status}[/{status_style}]",
            _format_gantt_duration(span.duration_ms),
            str(len(span.tool_calls)),
            str(len(span.llm_calls)),
        )

    out.print(table)


# ---------------------------------------------------------------------------
# albus trace export
# ---------------------------------------------------------------------------


@trace_app.command("export")
def trace_export_cmd(
    execution_id: Annotated[str, typer.Argument(help="Execution ID (or prefix).")],
    fmt: Annotated[
        str, typer.Option("--format", "-f", help="Output format: html, mermaid, json.")
    ] = "html",
    output: Annotated[str | None, typer.Option("--output", "-o", help="Output file path.")] = None,
    store_root: Annotated[str, typer.Option("--store", help="State store root.")] = ".albus_store",
) -> None:
    """Export a trace in various formats."""
    trace = _resolve_trace(execution_id, store_root)

    export_handlers: dict[str, object] = {
        "json": lambda: _export_json(trace, output),
        "mermaid": lambda: _export_mermaid(trace, output),
        "html": lambda: _export_html(trace, output),
    }

    handler = export_handlers.get(fmt)
    if handler is None:
        err.print(f"[red]Error:[/red] Unknown format '{fmt}'. Use: html, mermaid, json.")
        raise typer.Exit(1)
    handler()


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _resolve_trace(execution_id: str, store_root: str) -> object:
    """Load a trace by exact ID or prefix match."""
    from pathlib import Path

    from core.observability.trace import ExecutionTrace
    from core.state import FileStateStore

    store = FileStateStore(root=store_root)

    # Try exact match first.
    trace = run_async(store.load_trace(execution_id))
    if trace is not None:
        return trace

    # Prefix search across trace files.
    traces_dir = Path(store_root) / "traces"
    if traces_dir.is_dir():
        for path in traces_dir.glob("*.json"):
            if path.stem.startswith(execution_id):
                try:
                    return ExecutionTrace.model_validate_json(path.read_text())
                except Exception:
                    continue

    err.print(f"[red]Error:[/red] No trace found for '{execution_id}'.")
    raise typer.Exit(1)


def _export_json(trace: object, output: str | None) -> None:
    content = trace.model_dump_json(indent=2)
    if output:
        from pathlib import Path

        Path(output).write_text(content)
        err.print(f"Written to [bold]{output}[/bold]")
    else:
        out.print_json(content)


def _export_mermaid(trace: object, output: str | None) -> None:
    from core.observability.trace import trace_to_mermaid

    content = trace_to_mermaid(trace)
    if output:
        from pathlib import Path

        Path(output).write_text(content)
        err.print(f"Written to [bold]{output}[/bold]")
    else:
        out.print(content)


def _export_html(trace: object, output: str | None) -> None:
    from core.observability.html_viewer import render_trace_html
    from core.pathways.pathway import Pathway

    # Build a minimal stub pathway for DAG rendering when the real one
    # isn't available (trace export shouldn't require a running workspace).
    pathway = Pathway(id=trace.pathway_id or "unknown", nodes={}, connections=[])

    html = render_trace_html(pathway, trace)
    dest = output or f"trace-{trace.execution_id[:12]}.html"

    from pathlib import Path

    Path(dest).parent.mkdir(parents=True, exist_ok=True)
    Path(dest).write_text(html)
    err.print(f"Written to [bold]{dest}[/bold]")
