import unittest

from snowglobe.client.src.models import (
    CompletionFunctionOutputs,
    CompletionRequest,
    SnowglobeData,
    SnowglobeMessage,
)
from snowglobe.client.src.runner import run_completion_fn
from snowglobe.tools.context import ToolCallContext


class TestRunCompletionFn(unittest.IsolatedAsyncioTestCase):
    async def test_disable_tool_mocking_is_propagated_to_context(self):
        observed = {}

        async def completion_fn(
            request: CompletionRequest,
        ) -> CompletionFunctionOutputs:
            observed.update(ToolCallContext.get_context())
            return CompletionFunctionOutputs(response="ok")

        request = CompletionRequest(
            messages=[
                SnowglobeMessage(
                    role="user",
                    content="hello",
                    snowglobe_data=SnowglobeData(
                        conversation_id="conv-1",
                        test_id="test-1",
                    ),
                )
            ]
        )

        result = await run_completion_fn(
            completion_fn=completion_fn,
            app_id="app-1",
            completion_request=request,
            telemetry_context={},
            enable_tool_mocking=False,
        )

        self.assertEqual(result.response, "ok")
        self.assertEqual(observed["conversation_id"], "conv-1")
        self.assertEqual(observed["test_id"], "test-1")
        self.assertEqual(observed["application_id"], "app-1")
        self.assertFalse(observed["mock_enabled"])

    async def test_enable_tool_mocking_is_propagated_to_context(self):
        observed = {}

        async def completion_fn(
            request: CompletionRequest,
        ) -> CompletionFunctionOutputs:
            observed.update(ToolCallContext.get_context())
            return CompletionFunctionOutputs(response="ok")

        request = CompletionRequest(
            messages=[
                SnowglobeMessage(
                    role="user",
                    content="hello",
                    snowglobe_data=SnowglobeData(
                        conversation_id="conv-2",
                        test_id="test-2",
                    ),
                )
            ]
        )

        result = await run_completion_fn(
            completion_fn=completion_fn,
            app_id="app-2",
            completion_request=request,
            telemetry_context={},
            enable_tool_mocking=True,
        )

        self.assertEqual(result.response, "ok")
        self.assertEqual(observed["conversation_id"], "conv-2")
        self.assertEqual(observed["test_id"], "test-2")
        self.assertEqual(observed["application_id"], "app-2")
        self.assertTrue(observed["mock_enabled"])


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