from __future__ import annotations

from leanbench.typeindex import ROLE_GOAL, ROLE_INPUT, ROLE_OWN, index_types, parse_signature


def test_theorem_is_indexed_by_the_types_of_its_inputs() -> None:
    code = "theorem add_dvd (a b c : Int) (h : a ∣ b) (k : a ∣ c) : a ∣ b + c := by\n  sorry"
    signature = parse_signature(code)
    assert signature is not None
    assert signature.is_proof
    assert signature.name == "add_dvd"
    assert "Int" in signature.inputs
    pairs = index_types(code)
    assert ("Int", ROLE_INPUT) in pairs


def test_definition_is_indexed_by_its_own_type() -> None:
    code = "def divides (a b : Int) : Prop := ∃ k : Int, b = a * k"
    signature = parse_signature(code)
    assert signature is not None
    assert not signature.is_proof
    pairs = index_types(code)
    assert ("Int", ROLE_INPUT) in pairs
    assert ("Prop", ROLE_OWN) in pairs, "a predicate's own type is the point of it"
    assert all(role != ROLE_GOAL for _, role in pairs)


def test_a_proof_is_not_indexed_by_Prop() -> None:
    pairs = index_types("theorem t (n : Nat) : n = n := rfl")
    assert all(name != "Prop" for name, _ in pairs)


def test_definition_of_a_concrete_type_records_it() -> None:
    pairs = index_types("def double (n : Nat) : Nat := 2 * n")
    assert ("Nat", ROLE_OWN) in pairs


def test_instance_binders_count_as_inputs() -> None:
    code = "theorem foo {G : Type*} [Group G] (x : G) : x * 1 = x := by simp"
    signature = parse_signature(code)
    assert signature is not None
    assert "Group" in signature.inputs


def test_unicode_number_types_are_normalised() -> None:
    code = "theorem sq_nonneg' (n : ℤ) : 0 ≤ n ^ 2 := by positivity"
    pairs = dict(index_types(code))
    assert pairs.get("Int") == ROLE_INPUT


def test_structure_records_its_own_type() -> None:
    code = "structure Monoid' (M : Type) : Type where\n  one : M"
    signature = parse_signature(code)
    assert signature is not None
    assert signature.keyword == "structure"
    assert signature.own or signature.inputs


def test_model_supplied_types_are_merged_without_duplicates() -> None:
    code = "theorem t (n : Nat) : n = n := rfl"
    pairs = index_types(code, ["Nat", "Polynomial", "  "])
    names = [name for name, _ in pairs]
    assert names.count("Nat") == 1
    assert "Polynomial" in names


def test_comments_do_not_leak_into_the_index() -> None:
    code = "-- about Matrix stuff\n/- Group -/\ntheorem t (n : Nat) : n = n := rfl"
    names = [name for name, _ in index_types(code)]
    assert "Matrix" not in names and "Group" not in names


def test_declaration_keyword_detected_after_attributes() -> None:
    code = "@[simp]\nprivate theorem t (n : Nat) : n = n := rfl"
    signature = parse_signature(code)
    assert signature is not None and signature.keyword == "theorem"


def test_unparsable_code_yields_nothing_rather_than_nonsense() -> None:
    assert parse_signature("this is not lean") is None
    assert index_types("this is not lean") == []