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

# Add project root to path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

# Mock Transformers
with patch('transformers.AutoProcessor.from_pretrained'), \
     patch('transformers.AutoModelForCausalLM.from_pretrained'), \
     patch('transformers.AutoTokenizer.from_pretrained'):
    import agent_module
    import oracle_module

class TestAgentScenarios(unittest.TestCase):

    def setUp(self):
        self.mock_llm_patcher = patch('agent_module.get_llm')
        self.mock_llm = self.mock_llm_patcher.start()
        self.mock_model = MagicMock()
        self.mock_model.device = "cpu"
        self.mock_processor = MagicMock()
        self.mock_llm.return_value = (self.mock_model, self.mock_processor)
        
        self.mock_oracle_patcher = patch('agent_module.get_oracle_data')
        self.mock_oracle = self.mock_oracle_patcher.start()
        self.mock_oracle.return_value = {"wisdom": "Mock Wisdom"}

        self.mock_streamer_patcher = patch('agent_module.TextIteratorStreamer')
        self.mock_streamer = self.mock_streamer_patcher.start()

    def tearDown(self):
        self.mock_llm_patcher.stop()
        self.mock_oracle_patcher.stop()
        self.mock_streamer_patcher.stop()

    def run_chat(self, user_input, llm_output_stream):
        """Helper to run chat with mocked LLM stream"""
        self.mock_streamer.return_value = iter(llm_output_stream)
        gen = agent_module.chat_agent_stream(user_input, [])
        return list(gen)

    def test_name_only(self):
        """Scenario: User says 'I am Julian'"""
        tool_call = {
            "name": "oracle_consultation",
            "arguments": {"topic": "General", "name": "Julian", "date_str": "today"}
        }
        
        # LLM output simulation
        llm_output = [
            "Greetings Julian. ", 
            f"<tool_call>{json.dumps(tool_call)}</tool_call>"
        ]
        
        results = self.run_chat("I am Julian", llm_output)
        
        # Verify Oracle called correctly
        self.mock_oracle.assert_called_with(name="Julian", topic="General", date_str="today")
        self.assertTrue("*(Consulting the Oracle...)*" in results)

    def test_topic_only(self):
        """Scenario: User says 'Thema: Liebe'"""
        tool_call = {
            "name": "oracle_consultation",
            "arguments": {"topic": "Liebe", "name": "Seeker", "date_str": "today"}
        }
        
        llm_output = [
            "Let us check the Oracle for Love. ",
            f"<tool_call>{json.dumps(tool_call)}</tool_call>"
        ]
        
        results = self.run_chat("Thema: Liebe", llm_output)
        
        self.mock_oracle.assert_called_with(name="Seeker", topic="Liebe", date_str="today")
        self.assertTrue("*(Consulting the Oracle...)*" in results)

    def test_date_only(self):
        """Scenario: User says 'Reading for 2025-01-01'"""
        tool_call = {
            "name": "oracle_consultation",
            "arguments": {"topic": "General", "name": "Seeker", "date_str": "2025-01-01"}
        }
        
        llm_output = [
            "Consulting for that date. ",
            f"<tool_call>{json.dumps(tool_call)}</tool_call>"
        ]
        
        results = self.run_chat("Reading for 2025-01-01", llm_output)
        
        self.mock_oracle.assert_called_with(name="Seeker", topic="General", date_str="2025-01-01")

    def test_name_and_date(self):
        """Scenario: User says 'I am Julian, date 2025-01-01'"""
        tool_call = {
            "name": "oracle_consultation",
            "arguments": {"topic": "General", "name": "Julian", "date_str": "2025-01-01"}
        }
        
        llm_output = [
            f"<tool_call>{json.dumps(tool_call)}</tool_call>"
        ]
        
        results = self.run_chat("I am Julian, date 2025-01-01", llm_output)
        
        self.mock_oracle.assert_called_with(name="Julian", topic="General", date_str="2025-01-01")

    def test_general_conversation(self):
        """Scenario: User says 'Hi' (No tool)"""
        llm_output = ["Hello there! How can I help you?"]
        
        results = self.run_chat("Hi", llm_output)
        
        self.mock_oracle.assert_not_called()
        self.assertIn("Hello there! How can I help you?", results)

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