from __future__ import annotations

import unittest

from rc.inference.engine import _extract_lean_code, _resolve_max_tokens


class InferenceEngineHelperTests(unittest.TestCase):
    def test_extract_lean_code_skips_leading_prose(self) -> None:
        source = """### Explanation

The proof is trivial.

theorem true_intro_rc : True := by
  trivial
"""
        self.assertEqual(
            _extract_lean_code(source),
            "theorem true_intro_rc : True := by\n  trivial",
        )

    def test_extract_lean_code_does_not_treat_english_theorem_phrase_as_code(self) -> None:
        source = """### Explanation and Proof

The theorem states that for any propositions P and Q, if P /\\ Q holds, then P holds.
"""
        self.assertEqual(_extract_lean_code(source), "")

    def test_resolve_max_tokens_caps_small_research_tasks(self) -> None:
        self.assertEqual(_resolve_max_tokens("proof_attempt", 256), 96)
        self.assertEqual(_resolve_max_tokens("proof_repair", 256), 96)
        self.assertEqual(_resolve_max_tokens("lean_proof_gen", 256), 128)
        self.assertEqual(_resolve_max_tokens("other", 256), 256)


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