import unittest
from unittest.mock import MagicMock, patch
from app import chat_agent_stream
import re
import json

class TestAccumulationBug(unittest.TestCase):
    
    @patch('app.get_llm')
    @patch('app.TextIteratorStreamer')
    @patch('app.retrieve_relevant_chunks')
    @patch('app.detect_language')
    def test_multi_turn_accumulation(self, mock_detect, mock_rag, mock_streamer, mock_llm):
        """
        Simulates:
        Turn 1: "Thinking..." + Tool Call
        Turn 2: "Here is the answer."
        Expectation: Final yield should contain BOTH strings.
        """
        mock_detect.return_value = "English"
        mock_rag.return_value = []
        mock_llm.return_value = (MagicMock(), MagicMock())
        
        # Determine behavior manually to mock the turns
        # We need the streamer to yield different things on subsequent calls
        
        # Turn 1: "Thinking about it... " + <tool_call>...
        turn1_tokens = ["Thinking", " about", " it...", 
                        ' <tool_call>{"name": "oracle_consultation", "arguments": {"topic": "life"}}</tool_call>']
        
        # Turn 2: "The answer is 42."
        turn2_tokens = ["The", " answer", " is", " 42."]
        
        mock_inst = mock_streamer.return_value
        # side_effect allows us to return different iterators for each call (turn)
        mock_inst.__iter__.side_effect = [
            iter(turn1_tokens),
            iter(turn2_tokens)
        ]
        
        # Mock get_oracle_data so tool call succeeds (if app tries to import it)
        with patch('app.get_oracle_data', return_value={"wisdom": "42"}):
            gen = chat_agent_stream("query", [], None, None)
            
            yields = []
            for y in gen:
                yields.append(y)
                
            final_output = yields[-1]
            print(f"\nFINAL OUTPUT SEEN BY UI: {final_output}")
            
            # Check conditions
            self.assertIn("Thinking about it...", final_output, "Turn 1 text was lost!")
            self.assertIn("The answer is 42.", final_output, "Turn 2 text is missing!")
            
            # Also ensure the tool call xml is NOT visible (cleaned)
            self.assertNotIn("<tool_call>", final_output)

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