import unittest

from graders import extract_aime_answer, grade_aime, extract_code, run_humaneval_check


class TestAimeExtraction(unittest.TestCase):
    def test_boxed_answer(self):
        self.assertEqual(extract_aime_answer(r"The answer is \boxed{70}."), 70)

    def test_last_boxed_wins(self):
        text = r"First we get \boxed{12}, but correcting: \boxed{588}"
        self.assertEqual(extract_aime_answer(text), 588)

    def test_boxed_with_spaces_and_commas(self):
        self.assertEqual(extract_aime_answer(r"\boxed{ 1,024 }"), 1024)

    def test_no_boxed_falls_back_to_last_integer(self):
        self.assertEqual(extract_aime_answer("So the final answer is 204."), 204)

    def test_no_answer(self):
        self.assertIsNone(extract_aime_answer("I could not solve this."))

    def test_grade(self):
        self.assertTrue(grade_aime(r"\boxed{70}", "70"))
        self.assertFalse(grade_aime(r"\boxed{71}", "70"))
        self.assertFalse(grade_aime("no idea", "70"))


class TestCodeExtraction(unittest.TestCase):
    def test_python_fence(self):
        text = "Here you go:\n```python\ndef f():\n    return 1\n```\nDone."
        self.assertEqual(extract_code(text), "def f():\n    return 1")

    def test_last_fence_wins(self):
        text = "```python\ndef f():\n    return 1\n```\nOops, fix:\n```python\ndef f():\n    return 2\n```"
        self.assertEqual(extract_code(text), "def f():\n    return 2")

    def test_bare_fence(self):
        text = "```\ndef f():\n    return 3\n```"
        self.assertEqual(extract_code(text), "def f():\n    return 3")

    def test_no_fence_returns_raw(self):
        text = "def f():\n    return 4"
        self.assertEqual(extract_code(text), "def f():\n    return 4")


class TestHumanEvalCheck(unittest.TestCase):
    TEST_SRC = (
        "def check(candidate):\n"
        "    assert candidate(2) == 4\n"
        "    assert candidate(3) == 9\n"
    )

    def test_pass(self):
        ok, detail = run_humaneval_check("def sq(x):\n    return x * x", self.TEST_SRC, "sq")
        self.assertTrue(ok, detail)

    def test_wrong_answer_fails(self):
        ok, _ = run_humaneval_check("def sq(x):\n    return x + x", self.TEST_SRC, "sq")
        self.assertFalse(ok)

    def test_exception_fails(self):
        ok, _ = run_humaneval_check("def sq(x):\n    raise ValueError", self.TEST_SRC, "sq")
        self.assertFalse(ok)

    def test_infinite_loop_times_out(self):
        ok, detail = run_humaneval_check(
            "def sq(x):\n    while True:\n        pass", self.TEST_SRC, "sq", timeout=3
        )
        self.assertFalse(ok)
        self.assertIn("timeout", detail)

    def test_typing_names_available(self):
        code = "def ident(x: List[int]) -> List[int]:\n    return x"
        test = "def check(candidate):\n    assert candidate([1]) == [1]\n"
        ok, detail = run_humaneval_check(code, test, "ident")
        self.assertTrue(ok, detail)


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


class TestMcqExtraction(unittest.TestCase):
    def test_answer_prefix(self):
        from graders import extract_mcq_letter
        self.assertEqual(extract_mcq_letter("Some reasoning.\n\nAnswer: C"), "C")

    def test_answer_bold(self):
        from graders import extract_mcq_letter
        self.assertEqual(extract_mcq_letter("The answer is **B**."), "B")

    def test_parenthesized(self):
        from graders import extract_mcq_letter
        self.assertEqual(extract_mcq_letter("Therefore (D) is correct."), "D")

    def test_last_standalone_letter_fallback(self):
        from graders import extract_mcq_letter
        self.assertEqual(extract_mcq_letter("Options A and B fail, so C."), "C")

    def test_no_letter(self):
        from graders import extract_mcq_letter
        self.assertIsNone(extract_mcq_letter("I am not sure."))

    def test_ignores_lowercase_article_a(self):
        from graders import extract_mcq_letter
        self.assertEqual(extract_mcq_letter("This is a tricky one. Answer: B"), "B")