import unittest

from gradescope_utils.autograder_utils.decorators import weight, number

try:
    from programming.react import (
        add,
        subtract,
        multiply,
        divide,
        execute_action,
        react,
    )
except ImportError:
    from react import (
        add,
        subtract,
        multiply,
        divide,
        execute_action,
        react,
    )


class TestReact(unittest.TestCase):
    # Tool function tests
    @weight(2)
    @number("3.1")
    def test_add(self):
        """Test that add function correctly adds two numbers."""
        result = add(5, 3)
        self.assertEqual(result, 8)

    @weight(2)
    @number("3.2")
    def test_subtract(self):
        """Test that subtract function correctly subtracts two numbers."""
        result = subtract(10, 4)
        self.assertEqual(result, 6)

    @weight(2)
    @number("3.3")
    def test_multiply(self):
        """Test that multiply function correctly multiplies two numbers."""
        result = multiply(6, 7)
        self.assertEqual(result, 42)

    @weight(2)
    @number("3.4")
    def test_divide(self):
        """Test that divide function correctly divides two numbers."""
        result = divide(20, 4)
        self.assertEqual(result, 5)

    # execute_action tests
    @weight(2)
    @number("3.5")
    def test_execute_action_add(self):
        """Test that execute_action parses and executes add correctly."""
        result = execute_action("add[10, 5]")
        self.assertEqual(result, 15)

    @weight(2)
    @number("3.6")
    def test_execute_action_subtract(self):
        """Test that execute_action parses and executes subtract correctly."""
        result = execute_action("subtract[10, 5]")
        self.assertEqual(result, 5)

    @weight(2)
    @number("3.7")
    def test_execute_action_multiply(self):
        """Test that execute_action parses and executes multiply correctly."""
        result = execute_action("multiply[10, 5]")
        self.assertEqual(result, 50)

    @weight(2)
    @number("3.8")
    def test_execute_action_divide(self):
        """Test that execute_action parses and executes divide correctly."""
        result = execute_action("divide[10, 5]")
        self.assertEqual(result, 2)

    @weight(2)
    @number("3.9")
    def test_execute_action_finish(self):
        """Test that execute_action parses finish action correctly."""
        result = execute_action("finish[42]")
        self.assertEqual(result, "FINISH: 42")

    @weight(2)
    @number("3.10")
    def test_execute_action_unknown(self):
        """Test that execute_action handles unknown actions."""
        result = execute_action("unknown[1, 2]")
        self.assertTrue(result.startswith("Unknown action:"))

    # react function tests
    @weight(2)
    @number("3.11")
    def test_react_simple_addition(self):
        """Test that react solves a simple addition problem."""
        question = "What is 24880 + 220?"
        answer = react(question, max_steps=7, verbose=False)
        self.assertIn("25100", answer)

    @weight(2)
    @number("3.12")
    def test_react_division_and_addition(self):
        """Test that react solves a division and addition problem."""
        question = "What is (24880 / 10) + 220?"
        answer = react(question, max_steps=7, verbose=False)
        self.assertIn("2708", answer)

    @weight(2)
    @number("3.13")
    def test_react_complex_expression(self):
        """Test that react solves a complex multi-step expression."""
        question = "What is ((24880 / 10) + 220) * 1234?"
        answer = react(question, max_steps=7, verbose=False)
        self.assertIn("3341672", answer)

    @weight(2)
    @number("3.14")
    def test_react_expression_with_subtraction(self):
        """Test that react solves an expression with subtraction."""
        question = "What is ((24880 / 10) + 220 - 110) * 1234?"
        answer = react(question, max_steps=10, verbose=False)
        self.assertIn("3205932", answer)

    @weight(2)
    @number("3.15")
    def test_react_word_problem(self):
        """Test that react solves a word problem about splitting a bill."""
        question = "Bob and Charlie go out to dinner and the bill comes out to $50. Both order the same entree. How much should each pay?"
        answer = react(question, max_steps=7, verbose=False)
        self.assertIn("25", answer)


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