import unittest

from agenttrace.llamaindex import AgentTraceLlamaIndexHandler
from agenttrace import trace


class LlamaIndexCallbackTest(unittest.TestCase):
    def test_event_start_end_creates_span(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        with trace.run(name="llama-test") as ctx:
            eid = handler.on_event_start("llm", {"prompt": "hello"}, event_id="e1")
            handler.on_event_end("llm", {"response": "world"}, event_id="e1")

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].span_type, "llm")
            self.assertEqual(spans[0].name, "llamaindex_llm")
            self.assertEqual(spans[0].input_json["prompt"], "hello")
            self.assertEqual(spans[0].output_json["response"], "world")

    def test_retrieve_event_maps_to_retrieval(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        with trace.run(name="llama-test") as ctx:
            handler.on_event_start("retrieve", {}, event_id="r1")
            handler.on_event_end("retrieve", {"nodes": []}, event_id="r1")

            self.assertEqual(ctx.spans[0].span_type, "retrieval")

    def test_embedding_event_maps_to_llm(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        with trace.run(name="llama-test") as ctx:
            handler.on_event_start("embedding", {}, event_id="em1")
            handler.on_event_end("embedding", {}, event_id="em1")

            self.assertEqual(ctx.spans[0].span_type, "llm")

    def test_end_without_start_still_creates_span(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        with trace.run(name="llama-test") as ctx:
            handler.on_event_end("query", {"result": "ok"}, event_id="orphan")

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].span_type, "chain")

    def test_noop_without_context(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        handler.on_event_start("llm", {}, event_id="x1")
        # Should not raise even without a trace context
        handler.on_event_end("llm", {}, event_id="x1")

    def test_start_trace_and_end_trace_are_noops(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        handler.start_trace("t1")
        handler.end_trace("t1", {})

    def test_auto_generates_event_id(self) -> None:
        handler = AgentTraceLlamaIndexHandler()
        with trace.run(name="llama-test") as ctx:
            eid = handler.on_event_start("llm", {})
            self.assertTrue(len(eid) > 0)


if __name__ == "__main__":
    unittest.main()
