"""Tests for bandito traces and trace commands."""

from __future__ import annotations

import json
import os

from bandito.cli.observe import list_traces as _list_traces_from_store, show_trace as _show_trace_from_store
from bandito.storage import LocalStore


def _write_trace(base, trace_id, project, cost=None, error=None, spans=None):
    if spans is None:
        span = {
            "span_id": f"sp-{trace_id}",
            "name": "llm-call",
            "kind": "llm",
            "model": "gpt-4o",
            "provider": "openai",
            "started_at": "2026-03-19T12:00:00Z",
            "ended_at": "2026-03-19T12:00:01Z",
            "status": "error" if error else "ok",
            "input_tokens": 100,
            "output_tokens": 50,
        }
        if cost is not None:
            span["cost"] = cost
        if error:
            span["error"] = error
        spans = [span]
    trace = {
        "trace_id": trace_id,
        "project": project,
        "input": "hello",
        "output": "world",
        "started_at": "2026-03-19T12:00:00Z",
        "ended_at": "2026-03-19T12:00:01Z",
        "status": "error" if error else "ok",
        "error": error,
        "spans": spans,
    }
    project_dir = os.path.join(str(base), "traces", project)
    os.makedirs(project_dir, exist_ok=True)
    with open(os.path.join(project_dir, f"{trace_id}.json"), "w") as f:
        json.dump(trace, f)


class TestListTraces:
    def test_no_traces(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        _list_traces_from_store(LocalStore(str(tmp_path)),None)
        captured = capsys.readouterr()
        assert "No traces found" in captured.out

    def test_lists_traces(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        for i in range(3):
            _write_trace(tmp_path, f"t-{i}", "bot", cost=0.01)
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "3 traces" in captured.out
        assert "bot" in captured.out
        assert "t-0" in captured.out

    def test_shows_error_status(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-err", "bot", error="timeout")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "ERROR" in captured.out

    def test_shows_no_grade_indicator(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        _write_trace(tmp_path, "t-1", "bot")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "- grade" in captured.out

    def test_multi_project_prompts(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-1", "proj-a")
        _write_trace(tmp_path, "t-2", "proj-b")
        _list_traces_from_store(LocalStore(str(tmp_path)),None)
        captured = capsys.readouterr()
        assert "projects found" in captured.out
        assert "bandito observe traces --project" in captured.out

    def test_hints_inspect_with_real_id(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-1", "bot")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "bandito observe trace t-1" in captured.out

    def test_shows_input_output_with_prefixes(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-1", "bot")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "input: hello" in captured.out
        assert "output: world" in captured.out

    def test_truncates_long_input_output(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        long_text = "x" * 200
        trace = {
            "trace_id": "t-long", "project": "bot",
            "input": long_text, "output": long_text,
            "started_at": "2026-03-19T12:00:00Z",
            "ended_at": "2026-03-19T12:00:01Z",
            "status": "ok", "spans": [],
        }
        import json
        project_dir = os.path.join(str(tmp_path), "traces", "bot")
        os.makedirs(project_dir, exist_ok=True)
        with open(os.path.join(project_dir, "t-long.json"), "w") as f:
            json.dump(trace, f)
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        # Should be truncated with ellipsis, not full 200 chars
        assert "\u2026" in captured.out
        assert long_text not in captured.out

    def test_shows_separator_between_traces(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-1", "bot")
        _write_trace(tmp_path, "t-2", "bot")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert "\u2500" in captured.out  # horizontal line separator

    def test_shows_full_trace_id(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        long_id = "912d0fac-1234-5678-9abc-6ed900fa"
        _write_trace(tmp_path, long_id, "bot")
        _list_traces_from_store(LocalStore(str(tmp_path)),"bot")
        captured = capsys.readouterr()
        assert long_id in captured.out  # full ID, not truncated


class TestShowTrace:
    def test_not_found(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        _show_trace_from_store(LocalStore(str(tmp_path)),"nope")
        captured = capsys.readouterr()
        assert "not found" in captured.out

    def test_basic_display(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-1", "bot", cost=0.005)
        _show_trace_from_store(LocalStore(str(tmp_path)),"t-1")
        captured = capsys.readouterr()
        assert "trace t-1" in captured.out
        assert "bot" in captured.out
        assert "Input" in captured.out
        assert "Output" in captured.out
        assert "Spans" in captured.out
        assert "gpt-4o" in captured.out

    def test_shows_error(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        _write_trace(tmp_path, "t-err", "bot", error="connection refused")
        _show_trace_from_store(LocalStore(str(tmp_path)),"t-err")
        captured = capsys.readouterr()
        assert "Error" in captured.out
        assert "connection refused" in captured.out

    def test_shows_span_tree(self, tmp_path, capsys):
        (tmp_path / "traces").mkdir()
        (tmp_path / "evaluations").mkdir()
        spans = [
            {
                "span_id": "sp-root", "name": "agent", "kind": "agent",
                "parent_id": None,
                "started_at": "2026-03-19T12:00:00Z",
                "ended_at": "2026-03-19T12:00:05Z",
                "status": "ok",
            },
            {
                "span_id": "sp-llm", "name": "generate", "kind": "llm",
                "parent_id": "sp-root", "model": "gpt-4o",
                "started_at": "2026-03-19T12:00:00.1Z",
                "ended_at": "2026-03-19T12:00:02Z",
                "status": "ok",
                "input_tokens": 100, "output_tokens": 50, "cost": 0.003,
            },
            {
                "span_id": "sp-tool", "name": "search", "kind": "tool",
                "parent_id": "sp-root",
                "started_at": "2026-03-19T12:00:02Z",
                "ended_at": "2026-03-19T12:00:04Z",
                "status": "ok",
            },
        ]
        _write_trace(tmp_path, "t-1", "bot", spans=spans)
        _show_trace_from_store(LocalStore(str(tmp_path)),"t-1")
        captured = capsys.readouterr()
        # Root span at top level
        assert "agent:agent" in captured.out
        # Children indented under root
        assert "llm:generate" in captured.out
        assert "tool:search" in captured.out
