from __future__ import annotations

from leanbench.detect import KIND_PROBLEM, KIND_STATEMENT, Detection
from leanbench.formalize import formalize, formalize_all, verification_mode
from leanbench.llm import ScriptedClient
from leanbench.verify import MODE_FULL, MODE_STATEMENT

from .conftest import formalize_reply


def make_detection(kind: str = KIND_STATEMENT) -> Detection:
    return Detection(
        kind=kind,
        label="Theorem 1.6.2",
        summary="divisibility is additive",
        start_page=12,
        start_index=100,
        end_page=12,
        end_index=260,
        text="Theorem 1.6.2. If a divides b and a divides c then a divides b + c.",
    )


def test_open_problems_are_only_checked_as_statements() -> None:
    assert verification_mode(make_detection(KIND_PROBLEM)) == MODE_STATEMENT
    assert verification_mode(make_detection(KIND_STATEMENT)) == MODE_FULL


def test_formalize_produces_verified_code_and_types(fake_checker) -> None:
    client = ScriptedClient([formalize_reply()])
    result = formalize(client, make_detection(), checker=fake_checker)
    assert result.lean_code.startswith("theorem test_decl")
    assert result.verified
    assert not result.proved, "a sorry in the code means it is not proved"
    assert ("Int", "input") in result.index_types


def test_a_real_proof_counts_as_proved(fake_checker) -> None:
    client = ScriptedClient(
        [formalize_reply(lean_code="theorem t (n : Nat) : n = n := rfl", has_proof=True)]
    )
    result = formalize(client, make_detection(), checker=fake_checker)
    assert result.has_proof and result.proved


def test_lean_diagnostics_drive_a_repair_attempt(fake_checker) -> None:
    client = ScriptedClient(
        [
            formalize_reply(lean_code="theorem t : BAD := by sorry"),
            formalize_reply(lean_code="theorem t (n : Nat) : n = n := rfl", has_proof=True),
        ]
    )
    result = formalize(client, make_detection(), checker=fake_checker, repair_attempts=1)
    assert result.attempts == 2
    assert result.verified
    assert len(client.calls) == 2
    repair_prompt = client.calls[1][-1]["content"]
    assert "error: bad" in repair_prompt, "the compiler's own words go back to the model"


def test_repair_gives_up_after_the_budget(fake_checker) -> None:
    client = ScriptedClient([formalize_reply(lean_code="theorem t : BAD := by sorry")] * 3)
    result = formalize(client, make_detection(), checker=fake_checker, repair_attempts=1)
    assert result.attempts == 2
    assert not result.verified


def test_model_may_decline_an_unformalizable_passage(fake_checker) -> None:
    client = ScriptedClient(['{"skip": true, "reason": "prose only"}'])
    result = formalize(client, make_detection(), checker=fake_checker)
    assert result.skipped and result.skip_reason == "prose only"
    assert fake_checker.calls == []


def test_imports_inside_the_body_are_stripped(fake_checker) -> None:
    client = ScriptedClient(
        [formalize_reply(lean_code="import Mathlib\ntheorem t (n : Nat) : n = n := rfl")]
    )
    result = formalize(client, make_detection(), checker=fake_checker)
    assert "import" not in result.lean_code


def test_one_bad_item_does_not_stop_the_batch(fake_checker) -> None:
    client = ScriptedClient(["garbage, not json", formalize_reply()])
    results = formalize_all(
        client, [make_detection(), make_detection(KIND_PROBLEM)], checker=fake_checker
    )
    assert len(results) == 2
    assert results[0].skipped and "formalizer failed" in results[0].skip_reason
    assert results[1].verified