import json
import os

import jsonschema
import pytest

from bandito.models import Span, Trace

SHARED_DIR = os.path.join(os.path.dirname(__file__), "..", "..", "shared")
SCHEMAS_DIR = os.path.join(SHARED_DIR, "schemas")
FIXTURES_DIR = os.path.join(SHARED_DIR, "fixtures")


@pytest.fixture
def trace_schema():
    with open(os.path.join(SCHEMAS_DIR, "trace.schema.json")) as f:
        return json.load(f)


@pytest.fixture
def evaluation_schema():
    with open(os.path.join(SCHEMAS_DIR, "evaluation.schema.json")) as f:
        return json.load(f)


class TestFixturesAgainstSchema:
    """Validate that shared fixtures conform to the trace schema."""

    @pytest.mark.parametrize(
        "fixture_name",
        ["simple-llm-trace.json", "rag-trace.json", "agent-loop-trace.json"],
    )
    def test_fixture_validates(self, trace_schema, fixture_name):
        with open(os.path.join(FIXTURES_DIR, fixture_name)) as f:
            data = json.load(f)
        jsonschema.validate(data, trace_schema)

    def test_simple_llm_trace_fields(self):
        with open(os.path.join(FIXTURES_DIR, "simple-llm-trace.json")) as f:
            data = json.load(f)
        assert data["project"] == "my-chatbot"
        assert len(data["spans"]) == 1
        span = data["spans"][0]
        assert span["kind"] == "llm"
        assert span["model"] == "gpt-4o"
        assert "input_tokens" in span

    def test_rag_trace_has_retrieval_and_llm(self):
        with open(os.path.join(FIXTURES_DIR, "rag-trace.json")) as f:
            data = json.load(f)
        kinds = [s["kind"] for s in data["spans"]]
        assert "retrieval" in kinds
        assert "llm" in kinds

    def test_agent_loop_nesting(self):
        with open(os.path.join(FIXTURES_DIR, "agent-loop-trace.json")) as f:
            data = json.load(f)
        spans = {s["span_id"]: s for s in data["spans"]}
        root = next(s for s in data["spans"] if s["parent_id"] is None)
        assert root["kind"] == "agent"
        children = [s for s in data["spans"] if s["parent_id"] == root["span_id"]]
        assert len(children) == 3


class TestSerializedTraceAgainstSchema:
    """Validate that Python-serialized traces conform to the schema."""

    def test_minimal_trace(self, trace_schema):
        t = Trace(project="test")
        t.spans.append(Span(name="s", kind="llm"))
        jsonschema.validate(t.to_dict(), trace_schema)

    def test_full_trace(self, trace_schema):
        t = Trace(
            project="test",
            input="hello",
            output="world",
            session_id="s1",
            user_id="u1",
            metadata={"k": "v"},
            tags=["a"],
        )
        t.spans.append(
            Span(
                name="gen",
                kind="llm",
                model="gpt-4o",
                provider="openai",
                input_tokens=10,
                output_tokens=5,
                cost=0.001,
                input={"messages": []},
                output={"message": "hi"},
            )
        )
        jsonschema.validate(t.to_dict(), trace_schema)

    def test_evaluation_human(self, evaluation_schema):
        from bandito.models import Evaluation

        e = Evaluation(trace_id="t1", human_score=1, human_comment="good")
        jsonschema.validate(e.to_dict(), evaluation_schema)

    def test_evaluation_judge(self, evaluation_schema):
        from bandito.models import Evaluation

        e = Evaluation(
            trace_id="t1",
            judge_score=0,
            judge_critique="bad output",
        )
        jsonschema.validate(e.to_dict(), evaluation_schema)

    def test_evaluation_both(self, evaluation_schema):
        from bandito.models import Evaluation

        e = Evaluation(
            trace_id="t1",
            human_score=1,
            human_span_ids=["sp1"],
            judge_score=0,
            judge_critique="missed details",
        )
        jsonschema.validate(e.to_dict(), evaluation_schema)
