import asyncio
import os
import unittest
from datetime import datetime, timedelta, timezone
from unittest.mock import AsyncMock, patch

from agenttrace import trace


class TraceTest(unittest.TestCase):
    def test_trace_run_sets_current_context(self) -> None:
        with trace.run(name="demo") as ctx:
            self.assertIsNotNone(trace.current())
            self.assertEqual(trace.current().run_id, ctx.run_id)
        self.assertIsNone(trace.current())

    def test_add_span_buffers_in_memory(self) -> None:
        now = datetime.now(timezone.utc)
        with trace.run(name="demo") as ctx:
            ctx.add_span(
                span_type="chain",
                name="step-1",
                start_time=now,
                end_time=now + timedelta(milliseconds=5),
                input_json={"a": 1},
                output_json={"b": 2},
            )
            self.assertEqual(len(ctx.spans), 1)

    @patch("agenttrace.client.request.urlopen")
    def test_flush_on_exit_posts_to_backend(self, mock_urlopen) -> None:
        now = datetime.now(timezone.utc)
        with patch.dict(
            os.environ,
            {"AGENTTRACE_API_URL": "http://localhost:8000", "AGENTTRACE_API_KEY": "secret"},
            clear=False,
        ):
            with trace.run(name="demo") as ctx:
                ctx.add_span(
                    span_type="agent",
                    name="root",
                    start_time=now,
                    end_time=now,
                    input_json={},
                    output_json={},
                )

        self.assertEqual(mock_urlopen.call_count, 1)
        req = mock_urlopen.call_args[0][0]
        self.assertEqual(req.full_url, "http://localhost:8000/v1/ingest")
        self.assertEqual(req.get_header("Authorization"), "Bearer secret")


class AsyncTraceTest(unittest.TestCase):
    def test_async_trace_run_sets_current_context(self) -> None:
        async def _test():
            async with trace.run(name="async-demo") as ctx:
                self.assertIsNotNone(trace.current())
                self.assertEqual(trace.current().run_id, ctx.run_id)
            self.assertIsNone(trace.current())

        asyncio.run(_test())

    def test_async_add_span_buffers_in_memory(self) -> None:
        now = datetime.now(timezone.utc)

        async def _test():
            async with trace.run(name="async-demo") as ctx:
                ctx.add_span(
                    span_type="chain",
                    name="step-1",
                    start_time=now,
                    end_time=now + timedelta(milliseconds=5),
                    input_json={"a": 1},
                    output_json={"b": 2},
                )
                self.assertEqual(len(ctx.spans), 1)

        asyncio.run(_test())

    def test_async_flush_on_exit_posts_to_backend(self) -> None:
        try:
            import aiohttp
        except ImportError:
            self.skipTest("aiohttp not installed")
            return

        now = datetime.now(timezone.utc)

        mock_response = AsyncMock()
        mock_response.__aenter__ = AsyncMock(return_value=mock_response)
        mock_response.__aexit__ = AsyncMock(return_value=False)

        mock_session = AsyncMock()
        mock_session.post = AsyncMock(return_value=mock_response)
        mock_session.__aenter__ = AsyncMock(return_value=mock_session)
        mock_session.__aexit__ = AsyncMock(return_value=False)

        async def _test():
            with patch.dict(
                os.environ,
                {"AGENTTRACE_API_URL": "http://localhost:8000", "AGENTTRACE_API_KEY": "secret"},
                clear=False,
            ), patch.object(aiohttp, "ClientSession", return_value=mock_session):
                async with trace.run(name="async-demo") as ctx:
                    ctx.add_span(
                        span_type="agent",
                        name="root",
                        start_time=now,
                        end_time=now,
                        input_json={},
                        output_json={},
                    )

            mock_session.post.assert_called_once()
            call_kwargs = mock_session.post.call_args
            self.assertEqual(call_kwargs[0][0], "http://localhost:8000/v1/ingest")
            self.assertEqual(call_kwargs[1]["headers"]["Authorization"], "Bearer secret")

        asyncio.run(_test())

    def test_async_exit_captures_error_span(self) -> None:
        async def _test():
            ctx = None
            with self.assertRaises(ValueError):
                async with trace.run(name="async-error") as ctx:
                    raise ValueError("boom")
            error_spans = [s for s in ctx.spans if s.error is not None]
            self.assertEqual(len(error_spans), 1)
            self.assertEqual(error_spans[0].name, "run_error")
            self.assertIn("boom", error_spans[0].error)

        asyncio.run(_test())


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