import asyncio
import unittest

from src.snowglobe.tools.context import ToolCallContext


class TestToolCallContext(unittest.IsolatedAsyncioTestCase):
    async def test_context_is_isolated_per_task(self):
        results = {}

        async def worker(name: str, conversation_id: str, test_id: str):
            tokens = ToolCallContext.set_context(
                conversation_id=conversation_id,
                test_id=test_id,
                application_id=f"app-{name}",
                mock_enabled=True,
            )
            try:
                await asyncio.sleep(0.01)
                results[name] = ToolCallContext.get_context()
            finally:
                ToolCallContext.clear_context(tokens)

        await asyncio.gather(
            worker("a", "conv-a", "test-a"),
            worker("b", "conv-b", "test-b"),
        )

        self.assertEqual(results["a"]["conversation_id"], "conv-a")
        self.assertEqual(results["a"]["test_id"], "test-a")
        self.assertTrue(results["a"]["mock_enabled"])
        self.assertEqual(results["b"]["conversation_id"], "conv-b")
        self.assertEqual(results["b"]["test_id"], "test-b")
        self.assertTrue(results["b"]["mock_enabled"])

    async def test_clear_context_restores_previous_values(self):
        outer_tokens = ToolCallContext.set_context(
            conversation_id="outer-conv",
            test_id="outer-test",
            application_id="outer-app",
            mock_enabled=False,
        )
        try:
            inner_tokens = ToolCallContext.set_context(
                conversation_id="inner-conv",
                test_id="inner-test",
                application_id="inner-app",
                mock_enabled=True,
            )
            ToolCallContext.clear_context(inner_tokens)
            restored = ToolCallContext.get_context()
            self.assertEqual(restored["conversation_id"], "outer-conv")
            self.assertEqual(restored["test_id"], "outer-test")
            self.assertFalse(restored["mock_enabled"])
        finally:
            ToolCallContext.clear_context(outer_tokens)
