"""Backend-owned conjecture generation and repair heuristics."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

SEED_CONJECTURES = (
    {
        "domain": "logic",
        "conjecture_text": "Prove in Lean 4 that True is true.",
        "formal_statement": "theorem true_intro_rc : True := by",
        "strategy_tag": "seed:logic:true_intro",
        "proof_hint": "Use `trivial` or `exact True.intro`.",
    },
    {
        "domain": "logic",
        "conjecture_text": "Prove in Lean 4 that every proposition implies itself.",
        "formal_statement": "theorem implies_self_rc (P : Prop) : P -> P := by",
        "strategy_tag": "seed:logic:implies_self",
        "proof_hint": "Use `intro h` followed by `exact h`.",
    },
    {
        "domain": "naturals",
        "conjecture_text": "Prove in Lean 4 that every natural number equals itself.",
        "formal_statement": "theorem nat_refl_rc (n : Nat) : n = n := by",
        "strategy_tag": "seed:nat:refl",
        "proof_hint": "Use `rfl`.",
    },
    {
        "domain": "lists",
        "conjecture_text": "Prove in Lean 4 that every list equals itself.",
        "formal_statement": "theorem list_refl_rc {α : Type} (xs : List α) : xs = xs := by",
        "strategy_tag": "seed:list:refl",
        "proof_hint": "Use `rfl`.",
    },
    {
        "domain": "naturals",
        "conjecture_text": "Prove in Lean 4 that for every natural number n, n + 0 = n.",
        "formal_statement": "theorem add_zero_right (n : Nat) : n + 0 = n := by",
        "strategy_tag": "seed:nat:add_zero_right",
        "proof_hint": "Use `simpa` or `exact Nat.add_zero n`.",
    },
    {
        "domain": "naturals",
        "conjecture_text": "Prove in Lean 4 that for every natural number n, 0 + n = n.",
        "formal_statement": "theorem zero_add_nat (n : Nat) : 0 + n = n := by",
        "strategy_tag": "seed:nat:zero_add",
        "proof_hint": "Use `simpa` or `exact Nat.zero_add n`.",
    },
    {
        "domain": "logic",
        "conjecture_text": "Prove in Lean 4 that any conjunction implies its left side.",
        "formal_statement": "theorem and_left_rc (P Q : Prop) : P ∧ Q -> P := by",
        "strategy_tag": "seed:logic:and_left",
        "proof_hint": "Use `intro h` followed by `exact h.left`.",
    },
    {
        "domain": "logic",
        "conjecture_text": "Prove in Lean 4 that any conjunction implies its right side.",
        "formal_statement": "theorem and_right_rc (P Q : Prop) : P ∧ Q -> Q := by",
        "strategy_tag": "seed:logic:and_right",
        "proof_hint": "Use `intro h` followed by `exact h.right`.",
    },
    {
        "domain": "naturals",
        "conjecture_text": "Prove in Lean 4 that every natural number is less than its successor.",
        "formal_statement": "theorem lt_succ_self_rc (n : Nat) : n < Nat.succ n := by",
        "strategy_tag": "seed:nat:lt_succ_self",
        "proof_hint": "Use `exact Nat.lt_succ_self n`.",
    },
    {
        "domain": "naturals",
        "conjecture_text": "Prove in Lean 4 that successor is injective on natural numbers.",
        "formal_statement": "theorem succ_injective_rc {a b : Nat} (h : Nat.succ a = Nat.succ b) : a = b := by",
        "strategy_tag": "seed:nat:succ_injective",
        "proof_hint": "Use `exact Nat.succ.inj h`.",
    },
    {
        "domain": "lists",
        "conjecture_text": "Prove in Lean 4 that appending the empty list changes nothing.",
        "formal_statement": "theorem append_nil_rc {α : Type} (xs : List α) : xs ++ [] = xs := by",
        "strategy_tag": "seed:list:append_nil",
        "proof_hint": "Use `simpa using List.append_nil xs`.",
    },
)


FOLLOW_UP_CONJECTURES: dict[str, tuple[dict[str, str], ...]] = {
    "seed:logic:true_intro": (
        {
            "domain": "logic",
            "conjecture_text": "Prove in Lean 4 that True implies itself.",
            "formal_statement": "theorem true_implies_true_rc : True -> True := by",
            "strategy_tag": "followup:logic:true_implies_true",
            "proof_hint": "Use `intro h` followed by `exact h`.",
        },
    ),
    "seed:logic:implies_self": (
        {
            "domain": "logic",
            "conjecture_text": "Prove in Lean 4 that any proposition implies True.",
            "formal_statement": "theorem implies_true_rc (P : Prop) : P -> True := by",
            "strategy_tag": "followup:logic:implies_true",
            "proof_hint": "Use `intro _` followed by `trivial`.",
        },
    ),
    "seed:nat:refl": (
        {
            "domain": "naturals",
            "conjecture_text": "Prove in Lean 4 that successor preserves reflexive equality.",
            "formal_statement": "theorem succ_refl_rc (n : Nat) : Nat.succ n = Nat.succ n := by",
            "strategy_tag": "followup:nat:succ_refl",
            "proof_hint": "Use `rfl`.",
        },
    ),
    "seed:nat:add_zero_right": (
        {
            "domain": "naturals",
            "conjecture_text": "Prove in Lean 4 that 0 + n = n for every natural number n.",
            "formal_statement": "theorem zero_add_followup_rc (n : Nat) : 0 + n = n := by",
            "strategy_tag": "followup:nat:zero_add",
            "proof_hint": "Use `simpa using Nat.zero_add n`.",
        },
    ),
    "seed:nat:zero_add": (
        {
            "domain": "naturals",
            "conjecture_text": "Prove in Lean 4 that successor is injective on natural numbers.",
            "formal_statement": "theorem succ_injective_followup_rc {a b : Nat} (h : Nat.succ a = Nat.succ b) : a = b := by",
            "strategy_tag": "followup:nat:succ_injective",
            "proof_hint": "Use `exact Nat.succ.inj h`.",
        },
    ),
    "seed:list:refl": (
        {
            "domain": "lists",
            "conjecture_text": "Prove in Lean 4 that appending the empty list changes nothing.",
            "formal_statement": "theorem append_nil_followup_rc {α : Type} (xs : List α) : xs ++ [] = xs := by",
            "strategy_tag": "followup:list:append_nil",
            "proof_hint": "Use `simpa using List.append_nil xs`.",
        },
    ),
}


@dataclass
class ConjectureCandidate:
    domain: str
    conjecture_text: str
    formal_statement: str
    strategy_tag: str
    interestingness_score: float
    novelty_score: float
    metadata: dict[str, Any]


def generate_candidates(
    verified_episodes: list[dict[str, Any]],
    desired_domain: str,
    seed_conjecture: str = "",
) -> list[ConjectureCandidate]:
    candidates: list[ConjectureCandidate] = []

    if seed_conjecture.strip():
        candidates.append(_candidate_from_seed(seed_conjecture.strip(), desired_domain, verified_episodes))

    for seed in SEED_CONJECTURES:
        if desired_domain and seed["domain"] != desired_domain:
            continue
        candidates.append(
            ConjectureCandidate(
                domain=seed["domain"],
                conjecture_text=seed["conjecture_text"],
                formal_statement=seed["formal_statement"],
                strategy_tag=seed["strategy_tag"],
                interestingness_score=score_interestingness(seed["conjecture_text"], seed["formal_statement"]),
                novelty_score=score_novelty(seed["conjecture_text"], verified_episodes),
                metadata={"seed": True, "source": "builtin", "proof_hint": seed.get("proof_hint", "")},
            )
        )

    for episode in verified_episodes:
        candidates.extend(_children_from_verified(episode))
    return sorted(candidates, key=lambda item: (item.novelty_score, item.interestingness_score), reverse=True)


def build_initial_attempt_prompt(candidate: ConjectureCandidate) -> str:
    proof_hint = str(candidate.metadata.get("proof_hint", "") or "")
    return (
        "Write the shortest possible valid Lean 4 file.\n"
        "Use Lean 4 syntax only. Never use Agda, Coq, or Isabelle syntax.\n"
        "Specifically, never write `open import`; the first line must be exactly `import Mathlib`.\n"
        "Output one theorem only and keep the file under 12 lines.\n"
        "Start the first line with `import Mathlib` and return no headings, no explanations, and no markdown fences.\n\n"
        f"Conjecture:\n{candidate.conjecture_text}\n\n"
        f"Use this theorem statement verbatim:\n{candidate.formal_statement}\n\n"
        f"Proof hint:\n{proof_hint or 'Prefer the shortest standard Lean proof.'}\n\n"
        "Target template:\n"
        "import Mathlib\n\n"
        f"{candidate.formal_statement}\n"
        "  <proof>\n"
    )


def build_repair_prompt(episode: dict[str, Any], attempt: dict[str, Any]) -> str:
    diagnostics = attempt.get("lean_stderr") or attempt.get("lean_stdout") or ""
    status = str(attempt.get("verification_status") or "")
    syntax_hint = ""
    previous_code = attempt.get("generated_lean_code") or attempt.get("raw_output", "")
    if status == "malformed_output":
        syntax_hint = (
            "The previous output included prose, repeated tactics, or malformed Lean.\n"
            "Start immediately with `import Mathlib`, keep the result under 12 lines, and stop after the proof closes.\n\n"
        )
    if "open import" in previous_code or "invalid 'import' command" in diagnostics or "unexpected token 'import'" in diagnostics:
        syntax_hint += (
            "The previous output used invalid Agda-style import syntax. "
            "In Lean 4 the file must start with exactly `import Mathlib`, and there should be no `open import` lines.\n\n"
        )
    return (
        "Repair this Lean 4 file.\n"
        "Return only a complete Lean 4 file under 12 lines.\n"
        "Use Lean 4 syntax only. Never use `open import`, `module`, `where`, `admit`, `axiom`, or `sorry`.\n"
        "Prefer the simplest proof term that works, such as `trivial`, `exact h`, `rfl`, or `simpa`.\n"
        "Do not include explanations, bullet points, markdown fences, or repeated tactic lines.\n\n"
        f"{syntax_hint}"
        f"Conjecture:\n{episode.get('conjecture_text', '')}\n\n"
        f"Target theorem statement to use verbatim:\n{episode.get('formal_statement') or attempt.get('formal_statement', '')}\n\n"
        f"Previous failed output:\n{previous_code}\n\n"
        f"Lean diagnostics:\n{diagnostics}\n\n"
        "Required template:\n"
        "import Mathlib\n\n"
        f"{episode.get('formal_statement') or attempt.get('formal_statement', '')}\n"
        "  <proof>\n"
    )


def score_interestingness(conjecture_text: str, formal_statement: str) -> float:
    token_count = len((conjecture_text + " " + formal_statement).split())
    return round(min(10.0, 2.0 + token_count / 12.0), 2)


def score_novelty(conjecture_text: str, verified_episodes: list[dict[str, Any]]) -> float:
    text = conjecture_text.strip().lower()
    verified_texts = {str(episode.get("conjecture_text", "")).strip().lower() for episode in verified_episodes}
    if text not in verified_texts:
        return 8.0
    return 1.0


def _children_from_verified(episode: dict[str, Any]) -> list[ConjectureCandidate]:
    strategy_tag = str(episode.get("strategy_tag", "") or "")
    if not strategy_tag or strategy_tag.startswith(("child:", "followup:")):
        return []

    follow_ups = FOLLOW_UP_CONJECTURES.get(strategy_tag, ())
    candidates: list[ConjectureCandidate] = []
    for follow_up in follow_ups:
        candidates.append(
            ConjectureCandidate(
                domain=follow_up["domain"],
                conjecture_text=follow_up["conjecture_text"],
                formal_statement=follow_up["formal_statement"],
                strategy_tag=follow_up["strategy_tag"],
                interestingness_score=score_interestingness(
                    follow_up["conjecture_text"],
                    follow_up["formal_statement"],
                ),
                novelty_score=min(10.0, float(episode.get("novelty_score", 0.0) or 0.0) + 0.5),
                metadata={
                    "parent_episode_id": episode.get("id"),
                    "seed": False,
                    "source": "followup",
                    "proof_hint": follow_up["proof_hint"],
                },
            )
        )
    return candidates


def _candidate_from_seed(
    seed_conjecture: str,
    desired_domain: str,
    verified_episodes: list[dict[str, Any]],
) -> ConjectureCandidate:
    lowered = seed_conjecture.lower()
    formal_statement = ""
    proof_hint = "Prefer the shortest standard Lean proof."
    inferred_domain = desired_domain or "logic"

    if "true" in lowered and "implies" not in lowered:
        formal_statement = "theorem seeded_true_intro_rc : True := by"
        proof_hint = "Use `trivial` or `exact True.intro`."
        inferred_domain = desired_domain or "logic"
    elif "implies itself" in lowered or "p implies p" in lowered:
        formal_statement = "theorem seeded_implies_self_rc (P : Prop) : P -> P := by"
        proof_hint = "Use `intro h` followed by `exact h`."
        inferred_domain = desired_domain or "logic"
    elif "equals itself" in lowered or "n = n" in lowered:
        formal_statement = "theorem seeded_refl_rc (n : Nat) : n = n := by"
        proof_hint = "Use `rfl`."
        inferred_domain = desired_domain or "naturals"

    return ConjectureCandidate(
        domain=inferred_domain,
        conjecture_text=seed_conjecture,
        formal_statement=formal_statement,
        strategy_tag="seed:user",
        interestingness_score=score_interestingness(seed_conjecture, formal_statement),
        novelty_score=score_novelty(seed_conjecture, verified_episodes),
        metadata={"seed": True, "source": "user", "proof_hint": proof_hint},
    )
