"""Lean 4 proof-generation prompt templates."""

from __future__ import annotations

from typing import Any

SYSTEM_PROMPT = """\
You write Lean 4 code only.
Hard rules:
1. Use Lean 4 syntax, not Agda, Coq, Isabelle, or Haskell syntax.
2. Never write `open import`. If imports are needed, use `import Mathlib` at the top of the file.
3. Never use `sorry`, `admit`, `axiom`, placeholders, markdown fences, or explanations outside code.
4. Start immediately with Lean code. Do not add headings, bullet points, or prose before the file.
5. Keep the output short and compile-focused.

Output ONLY one valid Lean 4 file."""

PROOF_GEN_TEMPLATE = """\
{system}

--- TASK TYPE ---
{task_type}

--- GOAL ---
{problem}

--- CONTEXT ---
{context}

--- LEAN 4 CODE ---
{response_prefix}\
"""


def build_lean_prompt(problem: str) -> str:
    """Construct the default prompt for Lean proof generation."""
    return build_task_prompt("lean_proof_gen", problem, {})


def build_task_prompt(
    task_type: str,
    problem: str,
    context: dict[str, Any] | None = None,
) -> str:
    """Construct a prompt tailored to the queued task type."""
    prompt_context = _build_context_block(task_type, context or {})
    response_prefix = _build_response_prefix(task_type, context or {})
    return PROOF_GEN_TEMPLATE.format(
        system=SYSTEM_PROMPT,
        task_type=task_type.strip() or "lean_proof_gen",
        problem=problem.strip(),
        context=prompt_context,
        response_prefix=response_prefix,
    )


def build_chat_messages(
    problem: str,
    task_type: str = "lean_proof_gen",
    context: dict[str, Any] | None = None,
) -> list[dict[str, str]]:
    """Construct chat-format messages for models that support them."""
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": build_task_prompt(task_type=task_type, problem=problem, context=context),
        },
    ]


def _build_context_block(task_type: str, context: dict[str, Any]) -> str:
    base_lines = [
        f"Task type: {task_type}",
        f"Conjecture family: {context.get('conjecture_family', '')}".rstrip(),
        f"Strategy tag: {context.get('strategy_tag', '')}".rstrip(),
    ]

    optional_fields = [
        ("Domain", context.get("domain")),
        ("Conjecture text", context.get("conjecture_text")),
        ("Formal statement", context.get("formal_statement")),
        ("Proof hint", context.get("proof_hint")),
        ("Repair objective", context.get("repair_objective")),
        ("Previous Lean attempt", context.get("previous_code")),
        ("Lean stderr", context.get("lean_stderr")),
        ("Lean stdout", context.get("lean_stdout")),
        ("Task metadata", context.get("task_metadata")),
    ]

    for label, value in optional_fields:
        if value:
            base_lines.append(f"{label}: {value}")

    if task_type == "counterexample_search":
        base_lines.append(
            "If the statement appears false, produce Lean code that demonstrates the failure or isolates the bad assumption."
        )
    elif task_type == "proof_repair":
        base_lines.append("Focus on repairing the previous failed proof attempt using the Lean diagnostics.")
        base_lines.append("Keep the repaired file under 12 lines and start with `import Mathlib`.")
        base_lines.append("Delete any explanations or markdown that appeared in the previous output.")
    elif task_type == "conjecture_formalization":
        base_lines.append("Prioritize precise theorem formalization before proof search.")
    else:
        base_lines.append("Prioritize a complete, compilable proof with no placeholders.")
        base_lines.append("For simple goals, prefer `trivial`, `exact h`, `rfl`, or `simpa` over longer tactic scripts.")
        base_lines.append("Do not explain the proof; return the Lean file only.")

    return "\n".join(line for line in base_lines if line and not line.endswith(": "))


def _build_response_prefix(task_type: str, context: dict[str, Any]) -> str:
    formal_statement = str(context.get("formal_statement", "") or "").strip()
    prefix = "import Mathlib\n\n"
    if formal_statement and task_type in {"lean_proof_gen", "proof_attempt", "proof_repair", "conjecture_formalization"}:
        return f"{prefix}{formal_statement}\n  "
    return prefix
