import unittest
from unittest.mock import MagicMock, patch
from app import chat_wrapper, chat_agent_stream, get_oracle_data
import json

class TestFinalSuite(unittest.TestCase):

    @patch('app.get_llm')
    @patch('app.TextIteratorStreamer')
    @patch('app.retrieve_relevant_chunks')
    @patch('app.detect_language')
    def test_multi_message_bubbles(self, mock_detect, mock_rag, mock_streamer, mock_llm):
        """
        Verify that multi-turn agent responses result in multiple distinct message bubbles in history.
        """
         # Setup
        mock_detect.return_value = "English"
        mock_rag.return_value = []
        mock_llm.return_value = (MagicMock(), MagicMock())
        
        # Turn 1: "Thinking..." + Tool Call
        # Turn 2: "Final Answer"
        mock_inst = mock_streamer.return_value
        mock_inst.__iter__.side_effect = [
            iter(["Thinking", "...", ' <tool_call>{"name": "oracle_consultation", "arguments": {"topic": "life"}}</tool_call>']),
            iter(["Final", " Answer"])
        ]
        
        # Mock Oracle so it returns something valid
        with patch('app.get_oracle_data', return_value={"wisdom_nodes": [], "els_revelation": "hidden"}):
            # Run the WRAPPER (which manages bubbles)
            # chat_wrapper yields (history, threads, ...)
            # We want to see the FINAL history state.
            
            gen = chat_wrapper("query", [], {}, "tid", None, None)
            
            final_history = []
            for h, _, _, _, _ in gen:
                final_history = h
            
            # Expectation:
            # 1. User: "query"
            # 2. Assistant: "Thinking..."
            # 3. Assistant: "Final Answer"
            # Total 3 messages.
            
            print(f"\nFINAL HISTORY: {final_history}")
            
            self.assertEqual(len(final_history), 3, "Should have 3 messages (User, Bubble1, Bubble2)")
            self.assertEqual(final_history[1]["content"], "Thinking...")
            self.assertEqual(final_history[2]["content"], "Final Answer")
            
    
    @patch('app.get_llm')
    @patch('app.TextIteratorStreamer')
    @patch('app.retrieve_relevant_chunks')
    @patch('app.detect_language')
    @patch('app.get_oracle_data')
    def test_oracle_filtering(self, mock_oracle, mock_detect, mock_rag, mock_streamer, mock_llm):
        """
        Verify that ONLY wisdom_nodes are passed to the tool result string, masking 'els_revelation' etc.
        """
        mock_detect.return_value = "English"
        mock_rag.return_value = []
        mock_llm.return_value = (MagicMock(), MagicMock())
        
        # Mock Oracle returning sensitive data (BOS API / ELS)
        mock_oracle.return_value = {
            "wisdom_nodes": [{"source": "Psalm 1"}],
            "els_revelation": "SECRET_DATA",
            "bos_api": "HIDDEN"
        }
        
        # We need to capture what chat_agent_stream injects into messages
        # We can inspect the logger OR inspect the messages list if we mock it?
        # Actually, let's run the stream and spy on the 'messages' list built inside using a side_effect on generate?
        # A simpler way: The generator yields 'tool_result' into the User Context message. 
        # But 'chat_agent_stream' function local var 'messages' is hard to access.
        # However, we can use `mock_llm` call args! generate() is called with `input_ids`.
        # Wait, build_agent_prompt is called, but that's initial.
        # The tool result is injected in Turn 2 prompt.
        
        # Turn 1 triggers tool. Turn 2 prompt contains the result.
        mock_inst = mock_streamer.return_value
        mock_inst.__iter__.side_effect = [
            iter(['<tool_call>{"name": "oracle_consultation", "arguments": {}}</tool_call>']),
            iter(["Done"]),
            iter([]), iter([]) # Safety padding for extra turns
        ]
        
        gen = chat_agent_stream("query", [], None, None)
        list(gen) # Exhaust
        
        # Now check the args passed to model.generate in Turn 2?
        # Or easier: Check the LOGS if we could.
        # Best: Mock 'json.dumps' inside app?
        
        # Actually, let's verify what mock_oracle was called with, 
        # AND verify logic by importing the code? 
        # No, let's trust the logic if we can verify the messages list.
        # We can patch 'app.messages' list? No it's local.
        
        # Run agent
        with patch('app.json.dumps', side_effect=json.dumps) as mock_json:
             list(chat_agent_stream("query", [], None, None))
             
        # Inspect the messages injected via Apply Chat Template
        mock_proc = mock_llm.return_value[1]
        calls = mock_proc.apply_chat_template.call_args_list
        
        found_filtered = False
        for call in calls:
            msgs = call[0][0]
            # Check the tool result injection message
            for m in msgs:
                if m["role"] == "user" and "<tool_result>" in m["content"][0]["text"]:
                     content = m["content"][0]["text"]
                     if "wisdom_nodes" in content and "els_revelation" not in content and "bos_api" not in content:
                         found_filtered = True
        
        self.assertTrue(found_filtered, "Tool Result did not contain filtered data (or contained forbidden keys).")

    @patch('app.get_llm')
    @patch('app.TextIteratorStreamer')
    @patch('app.retrieve_relevant_chunks')
    @patch('app.detect_language')
    def test_prompt_fluidity_instruction(self, mock_detect, mock_rag, mock_streamer, mock_llm):
        """
        Verify that the injected prompt contains the 'connect smoothly' instruction.
        """
        mock_detect.return_value = "English"
        mock_rag.return_value = []
        mock_llm.return_value = (MagicMock(), MagicMock())
        
        # Turn 1 triggers tool. Turn 2 prompt injection.
        mock_inst = mock_streamer.return_value
        mock_inst.__iter__.side_effect = [
            iter(['<tool_call>{"name": "oracle_consultation", "arguments": {}}</tool_call>']),
            iter(["Done"]),
            iter([]), iter([])
        ]
        
        # We need to spy on 'messages' appended in app.py.
        # Since we can't easily access the local variable, we can mock `model.generate` 
        # inside `chat_agent_stream` (it's called via thread, but `processor.apply_chat_template` is main thread).
        
        # We'll spy on `processor.apply_chat_template`.
        mock_proc = mock_llm.return_value[1]
        
        # Run agent
        with patch('app.get_oracle_data', return_value={"wisdom_nodes": []}):
             list(chat_agent_stream("query", [], None, None))
             
        # Check calls to apply_chat_template. 
        # The LAST call should contain the injected tool result + instruction.
        calls = mock_proc.apply_chat_template.call_args_list
        # Found relevant call?
        found_instruction = False
        for call in calls:
            # call[0][0] is 'messages' list
            msgs = call[0][0]
            last_msg = msgs[-1]
            if last_msg["role"] == "user" and "connect this smoothly" in last_msg["content"][0]["text"].lower():
                found_instruction = True
                break
        
        self.assertTrue(found_instruction, "Did not find the 'connect smoothly' instruction in the prompt injection.")

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