from __future__ import annotations

import pytest
from sepurux import TraceBuilder, sepurux_trace


def test_trace_builder_schema_shape() -> None:
    builder = TraceBuilder("demo_task", {"prompt": "hello"})
    builder.add_model_step("plan", {"goal": "create issue"}, output={"ok": True})
    builder.add_tool_call("jira.create_issue(commit)", {"summary": "test"})
    builder.add_tool_result("jira.create_issue(commit)", {"issue_id": "OPS-123"})

    trace = builder.to_trace()

    assert trace["trace_version"] == "0.1"
    assert trace["source"] == "sdk"
    assert trace["task"] == {"name": "demo_task", "input": {"prompt": "hello"}}
    assert len(trace["events"]) == 3
    assert trace["events"][0]["type"] == "model"
    assert trace["events"][1]["type"] == "tool_call"
    assert trace["events"][2]["type"] == "tool_result"


def test_trace_context_supports_custom_source() -> None:
    with pytest.warns(DeprecationWarning, match="SepuruxClient.trace"):
        with sepurux_trace(
            "demo_task",
            {"prompt": "hello"},
            source="langsmith-like",
            trace_version="0.2",
        ) as rec:
            rec.model_step("plan", {"goal": "create issue"})

    trace = rec.to_trace()

    assert trace["source"] == "langsmith-like"
    assert trace["trace_version"] == "0.2"


def test_trace_builder_supports_timeline_friendly_step_types() -> None:
    builder = TraceBuilder("customer_support_agent", {"prompt": "Find the balance"})
    builder.add_prompt({"prompt": "Find the balance"})
    builder.add_llm_call(
        "planner", {"goal": "fetch balance"}, {"tool": "payments.get_balance"}, latency_ms=220
    )
    builder.add_mutation("latency_spike", input={"tool": "payments.get_balance"})
    builder.add_retry("payments.get_balance", retries=2)
    builder.add_agent_output({"message": "Balance is 42"})
    builder.add_error("tool timed out", tool="payments.get_balance")

    trace = builder.to_trace()

    assert [event["type"] for event in trace["events"]] == [
        "prompt",
        "llm_call",
        "mutation",
        "retry",
        "agent_output",
        "failure",
    ]
    assert trace["events"][1]["latency_ms"] == 220
    assert trace["events"][5]["status"] == "error"
