import unittest
from unittest.mock import MagicMock, patch
import json
from orionagent import Agent, Manager, Gemini, Ollama, Model

class TestThinking(unittest.TestCase):

    def test_ollama_thinking_filter_sync(self):
        # Mock requests.post for Ollama sync call
        with patch('requests.post') as mock_post:
            mock_response = MagicMock()
            mock_response.json.return_value = {
                "response": "<thought>I should check the current directory.</thought> The directory is empty.",
                "prompt_eval_count": 10,
                "eval_count": 20
            }
            mock_response.raise_for_status = MagicMock()
            mock_post.return_value = mock_response

            # 1. Test with show_thinking=True (Default)
            model_show = Ollama(model_name="deepseek-r1", thinking=True, show_thinking=True)
            res_show = model_show.generate("ls")
            self.assertIn("<thought>", res_show)
            self.assertIn("The directory is empty.", res_show)

            # 2. Test with show_thinking=False
            model_hide = Ollama(model_name="deepseek-r1", thinking=True, show_thinking=False)
            res_hide = model_hide.generate("ls")
            self.assertNotIn("<thought>", res_hide)
            self.assertNotIn("I should check", res_hide)
            self.assertEqual(res_hide, "The directory is empty.")

    def test_ollama_thinking_filter_stream(self):
        # Mock requests.post for Ollama stream call
        with patch('requests.post') as mock_post:
            mock_response = MagicMock()
            
            # Simulated chunks from Ollama
            chunks = [
                json.dumps({"response": "<thought>"}).encode(),
                json.dumps({"response": "Thinking..."}).encode(),
                json.dumps({"response": "</thought>"}).encode(),
                json.dumps({"response": " Hello"}).encode(),
                json.dumps({"response": " World"}).encode(),
                json.dumps({"done": True, "prompt_eval_count": 5, "eval_count": 10}).encode()
            ]
            mock_response.iter_lines.return_value = chunks
            mock_response.raise_for_status = MagicMock()
            mock_post.return_value = mock_response

            # Test with show_thinking=False
            model = Ollama(model_name="deepseek-r1", thinking=True, show_thinking=False)
            stream = model.generate_stream("test")
            
            final_text = "".join(list(stream))
            self.assertEqual(final_text.strip(), "Hello World")
            self.assertNotIn("thought", final_text)

    def test_gemini_thinking_model_switch(self):
        # Verify that thinking=True switches the model name for Gemini
        with patch('google.genai.Client') as mock_client:
            model = Gemini(api_key="test", thinking=True)
            self.assertEqual(model.model_name, "gemini-2.0-flash-thinking-exp")
            
            # Verify that providing a custom model name overrides the switch
            model_custom = Gemini(api_key="test", model_name="custom-model", thinking=True)
            self.assertEqual(model_custom.model_name, "custom-model")

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