from __future__ import annotations

import os
import shutil
import subprocess
import unittest
from types import SimpleNamespace
from unittest.mock import patch

from backend.lean_checker import (
    _sanitize_lean_output,
    _build_lean_command,
    check_conjecture,
    extract_lean_code,
)


class LeanCheckerTests(unittest.TestCase):
    def test_extracts_lean_from_markdown_fence(self) -> None:
        source = """```lean
theorem trivial_truth : True := by
  trivial
```"""
        self.assertEqual(
            extract_lean_code(source),
            "theorem trivial_truth : True := by\n  trivial",
        )

    def test_sanitizes_common_tiny_model_agda_artifacts(self) -> None:
        source = """open import Mathlib

theorem implies_self_rc (P : Prop) : P -> P := by
  <proof>
    intro h
    exact h
  end

--- TESTING ---
test_proof_attempt "implies_self_rc"
"""
        self.assertEqual(
            _sanitize_lean_output(source),
            "import Mathlib\ntheorem implies_self_rc (P : Prop) : P -> P := by\n    intro h\n    exact h",
        )

    def test_preserves_existing_import_mathlib_without_duplication(self) -> None:
        source = "import Mathlib\n\ntheorem nat_refl_rc (n : Nat) : n = n := by\n  rfl"
        self.assertEqual(_sanitize_lean_output(source), source)

    def test_strips_prompt_scaffolding_from_sanitized_output(self) -> None:
        source = """<proof>

--- TASK TYPE ---
proof_attempt

--- GOAL ---
Prove something
"""
        self.assertEqual(_sanitize_lean_output(source), "")

    def test_extract_lean_code_ignores_english_theorem_sentence(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_empty_input_returns_false(self) -> None:
        result = check_conjecture("   ")
        self.assertFalse(result.is_true)
        self.assertEqual(result.verification_status, "no_code")
        self.assertEqual(result.stderr, "No Lean code was provided.")

    def test_placeholder_input_is_rejected_before_running_lean(self) -> None:
        result = check_conjecture("theorem placeholder : True := by\n  sorry")
        self.assertFalse(result.is_true)
        self.assertEqual(result.verification_status, "unsafe_placeholder")
        self.assertIn("placeholder", result.verification_summary.lower())

    @unittest.skipUnless(shutil.which("elan"), "Elan must be installed for compile checks")
    def test_returns_true_for_valid_theorem(self) -> None:
        result = check_conjecture(
            """theorem trivial_truth : True := by
  trivial"""
        )
        self.assertTrue(result.is_true)
        self.assertEqual(result.exit_code, 0)

    @unittest.skipUnless(shutil.which("elan"), "Elan must be installed for compile checks")
    def test_returns_false_for_invalid_theorem(self) -> None:
        result = check_conjecture(
            """theorem broken_truth : True := by
  exact False.elim ?h"""
        )
        self.assertFalse(result.is_true)
        self.assertNotEqual(result.exit_code, 0)

    @patch("backend.lean_checker._build_lean_command", return_value=("lean", ["/fake/lean"], None))
    @patch("backend.lean_checker.subprocess.run", side_effect=subprocess.TimeoutExpired(cmd=["lean"], timeout=1))
    def test_timeout_returns_false(self, _mock_run, _mock_command) -> None:
        result = check_conjecture("theorem foo : True := by\n  trivial", timeout_seconds=1)
        self.assertFalse(result.is_true)
        self.assertEqual(result.exit_code, 124)
        self.assertIn("timed out", result.stderr)

    @patch.dict(os.environ, {"LEAN_PROJECT_DIR": "/tmp/project", "LEAN_TOOLCHAIN": "leanprover/lean4:v4.21.0"}, clear=False)
    @patch("backend.lean_checker.shutil.which")
    def test_build_command_prefers_elan_lake_for_project(self, mock_which) -> None:
        def fake_which(binary: str) -> str | None:
            mapping = {
                "elan": "/usr/local/bin/elan",
                "lake": "/usr/local/bin/lake",
                "lean": "/usr/local/bin/lean",
            }
            return mapping.get(binary)

        mock_which.side_effect = fake_which
        compiler, command, cwd = _build_lean_command()
        self.assertEqual(cwd, "/tmp/project")
        self.assertEqual(compiler, "elan run leanprover/lean4:v4.21.0 lake env lean")
        self.assertEqual(command[:4], ["/usr/local/bin/elan", "run", "leanprover/lean4:v4.21.0", "lake"])

    @patch.dict(os.environ, {}, clear=True)
    @patch("backend.lean_checker._resolve_lean_project_dir", return_value=None)
    @patch("backend.lean_checker.shutil.which", return_value=None)
    def test_build_command_reports_unavailable_when_no_binary_exists(self, _mock_which, _mock_project_dir) -> None:
        compiler, command, cwd = _build_lean_command()
        self.assertEqual(compiler, "unavailable")
        self.assertEqual(command, [])
        self.assertIsNone(cwd)

    @patch("backend.lean_checker._build_lean_command", return_value=("lean", ["/fake/lean"], None))
    @patch(
        "backend.lean_checker.subprocess.run",
        return_value=SimpleNamespace(
            returncode=1,
            stdout="error: unexpected token 'import'\nerror: invalid 'import' command",
            stderr="",
        ),
    )
    def test_stdout_parse_errors_are_classified_as_malformed_output(self, _mock_run, _mock_command) -> None:
        result = check_conjecture("theorem foo : True := by\n  trivial")
        self.assertFalse(result.is_true)
        self.assertEqual(result.verification_status, "malformed_output")

    @patch("backend.lean_checker._build_lean_command", return_value=("lean", ["/fake/lean"], None))
    @patch(
        "backend.lean_checker.subprocess.run",
        side_effect=[
            SimpleNamespace(
                returncode=1,
                stdout="",
                stderr="error: unknown module prefix 'Mathlib'",
            ),
            SimpleNamespace(
                returncode=0,
                stdout="",
                stderr="",
            ),
        ],
    )
    def test_missing_imports_can_fallback_without_mathlib(self, _mock_run, _mock_command) -> None:
        result = check_conjecture("import Mathlib\n\ntheorem foo : True := by\n  trivial")
        self.assertTrue(result.is_true)
        self.assertEqual(result.normalized_code, "theorem foo : True := by\n  trivial")


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