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

class TestComprehensiveHITL(unittest.TestCase):
    def setUp(self):
        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)

    @patch('builtins.input', return_value='y')
    def test_level_low_always_asks(self, mock_input):
        config = HitlConfig(permission_level="low")
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        manager.ask("Safe task", stream=False)
        self.assertTrue(mock_input.called)

    @patch('builtins.input', return_value='y')
    def test_level_medium_risky(self, mock_input):
        config = HitlConfig(permission_level="medium")
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        manager.ask("delete database", stream=False)
        self.assertTrue(mock_input.called)

    @patch('builtins.input', return_value='y')
    def test_level_medium_safe(self, mock_input):
        config = HitlConfig(permission_level="medium")
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        manager.ask("What is 2+2?", stream=False)
        self.assertFalse(mock_input.called)

    @patch('builtins.input', return_value='y')
    def test_level_high_never_asks(self, mock_input):
        config = HitlConfig(permission_level="high")
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        manager.ask("delete everything", stream=False)
        self.assertFalse(mock_input.called)

    @patch('builtins.input', return_value='y')
    def test_ask_once_persistence(self, mock_input):
        config = HitlConfig(permission_level="low", ask_once=True)
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        
        manager.ask("Task 1", stream=False)
        manager.ask("Task 2", stream=False)
        self.assertEqual(mock_input.call_count, 1)

    @patch('builtins.input', return_value='y')
    @patch('orionagent.agents.strategies.planning.PlanningStrategy._create_plan')
    def test_planning_with_hitl_config(self, mock_create_plan, mock_input):
        # Mock a multi-step plan
        mock_create_plan.return_value = [[{"a": "TestAgent", "s": "Step 1"}]]
        
        config = HitlConfig(permission_level="low", plan_review=True)
        manager = Manager(agents=[self.agent], model=self.mock_model, strategy="planning", hitl=config)
        
        manager.ask("Complex mission", stream=False)
        self.assertTrue(mock_input.called)

    @patch('builtins.input', return_value='n')
    def test_rejection_raises_error(self, mock_input):
        config = HitlConfig(permission_level="low")
        manager = Manager(agents=[self.agent], model=self.mock_model, hitl=config)
        with self.assertRaises(InterruptedError):
            manager.ask("Any task", stream=False)

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