import json
import os
import tempfile
import unittest
from pathlib import Path

from src.index import SelfCorrectingAgent, _safe_eval


class TestSelfCorrectingAgent(unittest.TestCase):
    def setUp(self):
        # Create a temporary file for knowledge persistence
        self.temp_dir = tempfile.TemporaryDirectory()
        self.knowledge_file = Path(self.temp_dir.name) / "knowledge.json"
        self.agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)

    def tearDown(self):
        self.temp_dir.cleanup()

    def test_safe_eval_basic(self):
        self.assertEqual(_safe_eval("2+3*4"), 14)
        self.assertAlmostEqual(_safe_eval("10/4"), 2.5)
        self.assertEqual(_safe_eval("-5 + 2"), -3)

    def test_safe_eval_invalid(self):
        with self.assertRaises(ValueError):
            _safe_eval("import os; os.system('echo hi')")
        with self.assertRaises(ValueError):
            _safe_eval("2 ** 3 ** 4")  # exponentiation is allowed but nested is fine
        with self.assertRaises(ValueError):
            _safe_eval("2 + unknown_var")

    def test_learning_and_persistence(self):
        problem = "1 + 1"
        # Initially unknown, should compute
        self.assertEqual(self.agent.solve(problem), 2)
        # Simulate user correction
        self.agent.knowledge[problem] = 3
        # Now should return learned answer
        self.assertEqual(self.agent.solve(problem), 3)
        # Persist knowledge
        self.agent._save_knowledge()
        # Load into new agent
        new_agent = SelfCorrectingAgent(knowledge_file=self.knowledge_file)
        self.assertEqual(new_agent.solve(problem), 3)

    def test_invalid_expression(self):
        with self.assertRaises(ValueError):
            self.agent.solve("2 + * 3")


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