import unittest
from unittest.mock import MagicMock, patch
import sys

# Mock anthropic module before importing provider
mock_anthropic = MagicMock()
sys.modules["anthropic"] = mock_anthropic

from orionagent.models.anthropic_provider import Anthropic

class TestAnthropicProvider(unittest.TestCase):
    def setUp(self):
        self.api_key = "test_key"
        self.provider = Anthropic(api_key=self.api_key, token_count=True)
        self.mock_client = mock_anthropic.Anthropic.return_value

    def test_init(self):
        self.assertEqual(self.provider.api_key, self.api_key)
        self.assertEqual(self.provider.model_name, "claude-3-5-sonnet-20240620")

    def test_generate_basic(self):
        mock_response = MagicMock()
        mock_response.content = [MagicMock(type="text", text="Hello world")]
        mock_response.stop_reason = "end_turn"
        mock_response.usage = MagicMock(input_tokens=10, output_tokens=20)
        
        self.mock_client.messages.create.return_value = mock_response
        
        res = self.provider.generate("Hi")
        self.assertEqual(res, "Hello world")
        self.assertEqual(self.provider.session_input_tokens, 10)
        self.assertEqual(self.provider.session_output_tokens, 20)

    @patch('orionagent.tools.tool_executor.ToolExecutor.execute')
    def test_generate_with_tools(self, mock_execute):
        # First call returns tool_use
        mock_tool_use = MagicMock(type="tool_use", id="to_1", name="get_weather", input={"location": "SF"})
        response1 = MagicMock()
        response1.content = [mock_tool_use]
        response1.stop_reason = "tool_use"
        response1.usage = MagicMock(input_tokens=5, output_tokens=5)

        # Second call returns internal text
        response2 = MagicMock()
        response2.content = [MagicMock(type="text", text="It is sunny")]
        response2.stop_reason = "end_turn"
        response2.usage = MagicMock(input_tokens=5, output_tokens=5)

        self.mock_client.messages.create.side_effect = [response1, response2]
        mock_execute.return_value = "Sunny, 72F"

        mock_tool = MagicMock()
        mock_tool.name = "get_weather"
        mock_tool.description = "Get weather"
        mock_tool.parameters = {"type": "object", "properties": {}}

        res = self.provider.generate("What is the weather?", tools=[mock_tool])
        
        self.assertEqual(res, "It is sunny")
        self.assertEqual(self.provider.session_total_tokens, 20)
        self.assertTrue(mock_execute.called)

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