import unittest

from agenttrace.crewai import AgentTraceCrewAIHandler
from agenttrace import trace


class CrewAICallbackTest(unittest.TestCase):
    def test_step_callback_creates_agent_span(self) -> None:
        handler = AgentTraceCrewAIHandler()
        with trace.run(name="crewai-test") as ctx:
            handler.step_callback("Step output text")

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].span_type, "agent")
            self.assertEqual(spans[0].name, "crewai_step")
            self.assertEqual(spans[0].output_json["step_output"], "Step output text")

    def test_task_callback_creates_chain_span(self) -> None:
        handler = AgentTraceCrewAIHandler()
        with trace.run(name="crewai-test") as ctx:
            handler.task_callback("Task result")

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].span_type, "chain")
            self.assertEqual(spans[0].name, "crewai_task")
            self.assertIn("task_output", spans[0].output_json)

    def test_step_callback_noop_without_context(self) -> None:
        handler = AgentTraceCrewAIHandler()
        # Should not raise even without a trace context
        handler.step_callback("orphan step")

    def test_task_callback_extracts_structured_attrs(self) -> None:
        handler = AgentTraceCrewAIHandler()

        class FakeTaskOutput:
            description = "Summarize article"
            summary = "Done"
            raw = "raw text"

            def __str__(self):
                return "FakeTaskOutput"

        with trace.run(name="crewai-test") as ctx:
            handler.task_callback(FakeTaskOutput())

            spans = ctx.spans
            self.assertEqual(len(spans), 1)
            self.assertEqual(spans[0].output_json["description"], "Summarize article")
            self.assertEqual(spans[0].output_json["summary"], "Done")


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