"""Tests for trace normalization."""

from bandito.storage.normalize import normalize_span, normalize_trace


class TestNormalizeTrace:
    def test_canonical_passes_through(self):
        raw = {
            "trace_id": "tr-1",
            "project": "bot",
            "started_at": "2026-03-25T10:00:00Z",
            "status": "ok",
            "spans": [
                {
                    "span_id": "sp-1",
                    "parent_id": None,
                    "name": "llm-call",
                    "kind": "llm",
                    "started_at": "2026-03-25T10:00:00Z",
                    "status": "ok",
                    "model": "gpt-4o",
                    "input_tokens": 100,
                    "output_tokens": 50,
                    "cost": 0.01,
                }
            ],
        }
        result = normalize_trace(raw)
        assert result["trace_id"] == "tr-1"
        assert result["project"] == "bot"
        assert result["status"] == "ok"
        assert len(result["spans"]) == 1
        assert result["spans"][0]["model"] == "gpt-4o"
        assert result["spans"][0]["cost"] == 0.01

    def test_missing_spans_gets_empty_list(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z"}
        result = normalize_trace(raw)
        assert result["spans"] == []

    def test_missing_status_gets_ok(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z", "spans": []}
        result = normalize_trace(raw)
        assert result["status"] == "ok"

    def test_invalid_status_gets_ok(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z", "status": "weird", "spans": []}
        result = normalize_trace(raw)
        assert result["status"] == "ok"

    def test_missing_trace_id_gets_uuid(self):
        raw = {"project": "bot", "started_at": "2026-03-25T10:00:00Z", "spans": []}
        result = normalize_trace(raw)
        assert result["trace_id"]  # non-empty
        assert len(result["trace_id"]) == 36  # UUID format

    def test_missing_project_gets_unknown(self):
        raw = {"trace_id": "tr-1", "started_at": "2026-03-25T10:00:00Z", "spans": []}
        result = normalize_trace(raw)
        assert result["project"] == "unknown"

    def test_missing_started_at_gets_now(self):
        raw = {"trace_id": "tr-1", "project": "bot", "spans": []}
        result = normalize_trace(raw)
        assert result["started_at"]
        assert "Z" in result["started_at"] or "T" in result["started_at"]

    def test_timestamp_plus_offset_normalized_to_z(self):
        raw = {
            "trace_id": "tr-1", "project": "bot",
            "started_at": "2026-03-25T10:00:00+00:00",
            "ended_at": "2026-03-25T10:00:01+00:00",
            "status": "ok", "spans": [],
        }
        result = normalize_trace(raw)
        assert result["started_at"] == "2026-03-25T10:00:00Z"
        assert result["ended_at"] == "2026-03-25T10:00:01Z"

    def test_tags_non_list_gets_empty_list(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z", "spans": [], "tags": "not-a-list"}
        result = normalize_trace(raw)
        assert result["tags"] == []

    def test_tags_list_passes_through(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z", "spans": [], "tags": ["prod", "v2"]}
        result = normalize_trace(raw)
        assert result["tags"] == ["prod", "v2"]

    def test_optional_fields_pass_through(self):
        raw = {
            "trace_id": "tr-1", "project": "bot",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok", "spans": [],
            "input": "hello", "output": "world",
            "user_id": "u-1", "session_id": "s-1",
            "metadata": {"key": "value"},
        }
        result = normalize_trace(raw)
        assert result["input"] == "hello"
        assert result["output"] == "world"
        assert result["user_id"] == "u-1"
        assert result["session_id"] == "s-1"
        assert result["metadata"] == {"key": "value"}

    def test_does_not_mutate_original(self):
        raw = {"trace_id": "tr-1", "project": "bot", "started_at": "2026-03-25T10:00:00Z", "spans": []}
        normalize_trace(raw)
        assert "status" not in raw  # original unchanged


class TestNormalizeSpan:
    def test_canonical_passes_through(self):
        raw = {
            "span_id": "sp-1", "parent_id": None,
            "name": "llm-call", "kind": "llm",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok",
        }
        result = normalize_span(raw)
        assert result == raw

    def test_missing_span_id_gets_uuid(self):
        raw = {"name": "test", "kind": "llm", "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
        result = normalize_span(raw)
        assert result["span_id"]
        assert len(result["span_id"]) == 36

    def test_missing_parent_id_gets_none(self):
        raw = {"span_id": "sp-1", "name": "test", "kind": "llm", "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
        result = normalize_span(raw)
        assert result["parent_id"] is None

    def test_missing_name_gets_kind(self):
        raw = {"span_id": "sp-1", "kind": "tool", "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
        result = normalize_span(raw)
        assert result["name"] == "tool"

    def test_invalid_kind_gets_custom(self):
        raw = {"span_id": "sp-1", "name": "test", "kind": "banana", "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
        result = normalize_span(raw)
        assert result["kind"] == "custom"

    def test_missing_kind_gets_custom(self):
        raw = {"span_id": "sp-1", "name": "test", "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
        result = normalize_span(raw)
        assert result["kind"] == "custom"

    def test_all_valid_kinds(self):
        for kind in ("llm", "tool", "retrieval", "agent", "chain", "embedding", "guardrail", "custom"):
            raw = {"span_id": "sp-1", "name": "t", "kind": kind, "started_at": "2026-03-25T10:00:00Z", "status": "ok"}
            result = normalize_span(raw)
            assert result["kind"] == kind

    def test_invalid_status_gets_ok(self):
        raw = {"span_id": "sp-1", "name": "test", "kind": "llm", "started_at": "2026-03-25T10:00:00Z", "status": "WARNING"}
        result = normalize_span(raw)
        assert result["status"] == "ok"

    def test_token_fields_coerced_to_int(self):
        raw = {
            "span_id": "sp-1", "name": "test", "kind": "llm",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok",
            "input_tokens": "100", "output_tokens": 50.0,
        }
        result = normalize_span(raw)
        assert result["input_tokens"] == 100
        assert isinstance(result["input_tokens"], int)
        assert result["output_tokens"] == 50
        assert isinstance(result["output_tokens"], int)

    def test_cost_coerced_to_float(self):
        raw = {
            "span_id": "sp-1", "name": "test", "kind": "llm",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok",
            "cost": "0.01",
        }
        result = normalize_span(raw)
        assert result["cost"] == 0.01
        assert isinstance(result["cost"], float)

    def test_invalid_token_value_removed(self):
        raw = {
            "span_id": "sp-1", "name": "test", "kind": "llm",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok",
            "input_tokens": "not-a-number",
        }
        result = normalize_span(raw)
        assert "input_tokens" not in result

    def test_llm_fields_pass_through(self):
        raw = {
            "span_id": "sp-1", "parent_id": None, "name": "test", "kind": "llm",
            "started_at": "2026-03-25T10:00:00Z", "status": "ok",
            "model": "gpt-4o", "provider": "openai",
            "response_model": "gpt-4o-2024-05-13",
            "finish_reason": "stop",
            "time_to_first_token": 0.35,
        }
        result = normalize_span(raw)
        assert result["model"] == "gpt-4o"
        assert result["provider"] == "openai"
        assert result["response_model"] == "gpt-4o-2024-05-13"
        assert result["finish_reason"] == "stop"
        assert result["time_to_first_token"] == 0.35

    def test_does_not_mutate_original(self):
        raw = {"span_id": "sp-1", "name": "test", "kind": "llm", "started_at": "2026-03-25T10:00:00Z"}
        normalize_span(raw)
        assert "status" not in raw
