import unittest
from unittest.mock import MagicMock, patch
from orionagent import Agent, OpenAI, MemoryConfig

class TestMemoryAmnesia(unittest.TestCase):

    def test_memory_persistence_between_turns(self):
        # 1. Setup Agent with session memory
        llm = MagicMock()
        # Mock streaming response
        llm.generate_stream.side_effect = [
            iter(["Hello ", "Samir! ", "I am Jervis."]), # Turn 1
            iter(["You are ", "Samir, ", "a developer."]) # Turn 2
        ]
        llm.verbose = False
        llm.debug = False

        agent = Agent(
            name="Jervis",
            model=llm,
            memory="session"
        )

        # Turn 1 - Consume stream
        res = agent.ask("My name is Samir and you are Jervis.")
        if hasattr(res, "__iter__") and not isinstance(res, str):
            "".join(list(res))
        
        # Verify message 1 (User) and message 2 (Assistant) are in cache
        self.assertEqual(len(agent._active_session.messages), 2)
        self.assertEqual(agent._active_session.messages[1]["content"], "Hello Samir! I am Jervis.")

        # Turn 2 - Verify history is built into prompt
        # We'll patch generate_stream to see what prompt it receives
        with patch.object(llm, 'generate_stream', return_value=iter(["OK"])) as mock_gen:
            res2 = agent.ask("Who am I?")
            list(res2) # consume
            
            # The prompt sent to the LLM should contain the previous turns
            sent_prompt = mock_gen.call_args[1].get("prompt")
            self.assertIn("My name is Samir", sent_prompt)
            self.assertIn("Hello Samir! I am Jervis.", sent_prompt)
            self.assertIn("ACTIVE TASK", sent_prompt)

    def test_session_saving_every_turn(self):
        # Verify that the session is saved to disk on every turn now
        llm = MagicMock()
        llm.generate_stream.return_value = iter(["Done"])
        agent = Agent(name="Tester", model=llm, memory="session")
        
        # Clear any existing sessions for this user/agent
        import shutil
        import os
        if os.path.exists(agent.memory_config.storage_path):
            shutil.rmtree(agent.memory_config.storage_path)

        # Turn 1
        res = agent.ask("Turn 1")
        list(res) # consume
        
        # Check disk
        path = os.path.join(agent.memory_config.storage_path, "sessions", agent.user_id, agent.name)
        files = os.listdir(path)
        self.assertEqual(len(files), 1)
        
        with open(os.path.join(path, files[0]), 'r') as f:
            data = f.read()
            self.assertIn("Turn 1", data)

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