import unittest
from unittest.mock import patch, MagicMock
from orionagent.agents.manager import Manager
from orionagent.agents.base_agent import Agent

class TestManagerHITL(unittest.TestCase):
    def setUp(self):
        # Use a real model mock that returns strings, not other mocks
        self.mock_model = MagicMock()
        self.mock_model.generate.return_value = "Mocked Response"
        self.mock_model.verbose = False
        self.mock_model.debug = False
        
        self.agent = Agent(name="TestAgent", role="Tester", model=self.mock_model)
        self.manager = Manager(agents=[self.agent], model=self.mock_model, strategy="planning", hitl=True)


    @patch('builtins.input', return_value='y')
    @patch('orionagent.agents.strategies.direct.DirectStrategy._approve_direct')
    def test_hitl_approve_direct(self, mock_approve, mock_input):
        # A short task should trigger DirectStrategy
        res = self.manager.ask("Hello", stream=False)
        self.assertEqual(res, "Mocked Response")
        self.assertTrue(mock_approve.called)

    @patch('builtins.input', return_value='n')
    def test_hitl_reject_direct(self, mock_input):
        # Should raise InterruptedError if rejected
        with self.assertRaises(InterruptedError):
            self.manager.ask("Hello", stream=False)

    @patch('builtins.input', return_value='y')
    @patch('orionagent.agents.strategies.planning.PlanningStrategy._approve_plan')
    @patch('orionagent.agents.strategies.base.BaseStrategy.is_complex_task', return_value=True)
    def test_hitl_approve_planning(self, mock_is_complex, mock_approve, mock_input):
        # Force PlanningStrategy by mocking is_complex_task
        # Need to mock the planner's response to be a valid JSON plan
        self.mock_model.generate.return_value = '[{"a": "TestAgent", "s": "Do something"}]'
        
        # Mocking the agent's ask to return a string for the second call
        self.agent.ask = MagicMock(return_value="Final Answer")
        
        res = self.manager.ask("Complex task", stream=False)
        self.assertEqual(res, "Final Answer")
        self.assertTrue(mock_approve.called)


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