"""Tests for trace data models."""

import json
import tempfile
from pathlib import Path

from litmus.models import (
    Trace,
    LLMCall,
    LLMRequest,
    LLMResponse,
    ToolCall,
    ToolRequest,
    ToolResponse,
    TraceMetadata,
    TraceStatus,
)


def _make_llm_call(provider="anthropic", model="claude-sonnet-4-20250514", status=200, body=None):
    return LLMCall(
        request=LLMRequest(
            provider=provider, model=model, method="POST", path="/v1/messages",
            headers={"content-type": "application/json", "x-api-key": "sk-test-12345"},
            body={
                "model": model,
                "max_tokens": 1024,
                "system": "You are a helpful assistant with detailed instructions that make the payload large.",
                "messages": [{"role": "user", "content": "Hello, this is a test prompt with some content."}],
                "tools": [{"name": "calculate", "description": "Do math", "input_schema": {"type": "object", "properties": {"expr": {"type": "string"}}}}],
            },
        ),
        response=LLMResponse(status_code=status, body=body or {"content": [{"type": "text", "text": "hello"}]}, is_stream=False),
        latency_ms=500,
    )


def _make_trace(n_steps=2):
    trace = Trace(agent_id="test-agent", metadata=TraceMetadata.auto())
    for i in range(n_steps):
        trace.add_step(_make_llm_call())
    trace.complete()
    return trace


class TestTraceModel:
    def test_create_trace(self):
        trace = Trace()
        assert trace.trace_id.startswith("lt-")
        assert trace.status == TraceStatus.RECORDING
        assert trace.steps == []

    def test_add_steps(self):
        trace = Trace()
        call = _make_llm_call()
        trace.add_step(call)
        assert len(trace.steps) == 1
        assert trace.llm_calls() == [call]
        assert trace.tool_calls() == []

    def test_complete(self):
        trace = Trace()
        trace.complete()
        assert trace.status == TraceStatus.COMPLETED
        assert trace.completed_at is not None
        assert trace.duration_ms() is not None
        assert trace.duration_ms() >= 0

    def test_save_load_roundtrip(self):
        trace = _make_trace(3)
        with tempfile.TemporaryDirectory() as tmpdir:
            path = str(Path(tmpdir) / "test.trace.json")
            trace.save(path)

            loaded = Trace.load(path)
            assert loaded.trace_id == trace.trace_id
            assert loaded.agent_id == trace.agent_id
            assert loaded.status == TraceStatus.COMPLETED
            assert len(loaded.steps) == 3
            assert loaded.llm_calls()[0].request.provider == "anthropic"

    def test_save_load_compact(self):
        trace = _make_trace(2)
        with tempfile.TemporaryDirectory() as tmpdir:
            full_path = str(Path(tmpdir) / "full.trace.json")
            compact_path = str(Path(tmpdir) / "compact.trace.json")

            trace.save(full_path, compact=False)
            trace.save(compact_path, compact=True)

            full_size = Path(full_path).stat().st_size
            compact_size = Path(compact_path).stat().st_size
            assert compact_size < full_size

            # Compact trace should still load and have same structure
            loaded = Trace.load(compact_path)
            assert len(loaded.steps) == 2
            assert loaded.llm_calls()[0].response is not None

    def test_metadata_auto(self):
        meta = TraceMetadata.auto()
        assert meta.python_version is not None
        assert meta.platform is not None
        assert meta.litmus_version == "0.1.0"

    def test_json_serialization(self):
        trace = _make_trace(1)
        data = json.loads(trace.model_dump_json())
        assert data["trace_id"].startswith("lt-")
        assert data["status"] == "completed"
        assert len(data["steps"]) == 1
        assert data["steps"][0]["step_type"] == "llm_call"
