import unittest
from uuid import uuid4

from agenttrace import AgentTraceCallbackHandler, trace


class LangChainCallbackHandlerTest(unittest.TestCase):
    def test_builds_parent_child_spans(self) -> None:
        callback = AgentTraceCallbackHandler()
        chain_run_id = uuid4()
        llm_run_id = uuid4()

        with trace.run(name="demo") as ctx:
            callback.on_chain_start(
                {"name": "root-chain"},
                {"input": "hello"},
                run_id=chain_run_id,
                parent_run_id=None,
            )
            callback.on_llm_start(
                {"name": "chat-openai"},
                ["hello"],
                run_id=llm_run_id,
                parent_run_id=chain_run_id,
            )
            callback.on_llm_end({"text": "hi"}, run_id=llm_run_id)
            callback.on_chain_end({"output": "hi"}, run_id=chain_run_id)

            spans = ctx.spans
            self.assertEqual(len(spans), 2)
            chain_span = next(span for span in spans if span.name == "root-chain")
            llm_span = next(span for span in spans if span.name == "chat-openai")
            self.assertEqual(str(llm_span.parent_span_id), str(chain_span.span_id))

    def test_captures_error_on_tool_span(self) -> None:
        callback = AgentTraceCallbackHandler()
        tool_run_id = uuid4()

        with trace.run(name="demo") as ctx:
            callback.on_tool_start(
                {"name": "search-tool"},
                "find weather",
                run_id=tool_run_id,
                parent_run_id=None,
            )
            callback.on_tool_error(RuntimeError("boom"), run_id=tool_run_id)

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].span_type, "tool")
            self.assertIn("RuntimeError: boom", spans[0].error or "")

    def test_chat_start_and_non_serializable_payloads(self) -> None:
        callback = AgentTraceCallbackHandler()
        run_id = uuid4()

        class Unserializable:
            pass

        with trace.run(name="demo") as ctx:
            callback.on_chat_model_start(
                {"name": "ollama-chat"},
                [[{"role": "user", "content": "hello"}, Unserializable()]],
                run_id=run_id,
                parent_run_id=None,
            )
            callback.on_llm_new_token("x", run_id=run_id)
            callback.on_chat_model_end({"raw": Unserializable()}, run_id=run_id)

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].name, "ollama-chat")


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