from __future__ import annotations

import unittest

from backend.research_heuristics import _children_from_verified, build_repair_prompt


class ResearchHeuristicsTests(unittest.TestCase):
    def test_seed_episode_generates_concrete_followup_not_recursive_wrapper(self) -> None:
        episode = {
            "id": "episode-1",
            "domain": "logic",
            "strategy_tag": "seed:logic:true_intro",
            "novelty_score": 8.0,
            "interestingness_score": 3.17,
        }

        candidates = _children_from_verified(episode)

        self.assertEqual(len(candidates), 1)
        self.assertEqual(candidates[0].conjecture_text, "Prove in Lean 4 that True implies itself.")
        self.assertFalse(candidates[0].conjecture_text.startswith("Investigate a converse"))

    def test_followup_episode_does_not_spawn_recursive_child(self) -> None:
        episode = {
            "id": "episode-2",
            "domain": "logic",
            "strategy_tag": "followup:logic:true_implies_true",
            "novelty_score": 9.0,
            "interestingness_score": 6.0,
        }

        self.assertEqual(_children_from_verified(episode), [])

    def test_malformed_output_repair_prompt_explicitly_bans_prose(self) -> None:
        prompt = build_repair_prompt(
            {
                "conjecture_text": "Prove in Lean 4 that True is true.",
                "formal_statement": "theorem true_intro_rc : True := by",
            },
            {
                "verification_status": "malformed_output",
                "generated_lean_code": "### Explanation\n\ntheorem true_intro_rc : True := by\n  trivial",
                "lean_stderr": "unexpected token",
            },
        )

        self.assertIn("Do not include explanations", prompt)
        self.assertIn("Start immediately with `import Mathlib`", prompt)


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