"""Lean conjecture parsing and compilation helpers."""

from __future__ import annotations

import os
import re
import shutil
import subprocess
import tempfile
from dataclasses import dataclass


@dataclass
class LeanCheckResult:
    conjecture: str
    normalized_code: str
    is_true: bool
    verification_status: str
    verification_summary: str
    compiler: str
    exit_code: int
    stdout: str
    stderr: str


def extract_lean_code(text: str) -> str:
    """Strip markdown fences and return the Lean source."""
    stripped = text.strip()

    if "```lean" in stripped:
        parts = stripped.split("```lean", 1)
        if len(parts) > 1:
            return _sanitize_lean_output(parts[1].split("```", 1)[0].strip())

    if "```" in stripped:
        parts = stripped.split("```", 1)
        if len(parts) > 1:
            return _sanitize_lean_output(parts[1].split("```", 1)[0].strip())

    first_code_offset = _find_first_code_offset(stripped)
    if first_code_offset is not None:
        return _sanitize_lean_output(stripped[first_code_offset:].strip())

    return _sanitize_lean_output(stripped)


def check_conjecture(conjecture: str, timeout_seconds: int = 30) -> LeanCheckResult:
    """Compile Lean code and treat successful type-checking as true."""
    normalized_code = extract_lean_code(conjecture)

    if not normalized_code:
        return LeanCheckResult(
            conjecture=conjecture,
            normalized_code="",
            is_true=False,
            verification_status="no_code",
            verification_summary="No Lean code was provided.",
            compiler="unavailable",
            exit_code=1,
            stdout="",
            stderr="No Lean code was provided.",
        )

    if _contains_unsafe_placeholder(normalized_code):
        return LeanCheckResult(
            conjecture=conjecture,
            normalized_code=normalized_code,
            is_true=False,
            verification_status="unsafe_placeholder",
            verification_summary="Proof contains placeholders such as sorry, admit, or axiom.",
            compiler="static-check",
            exit_code=1,
            stdout="",
            stderr="Unsafe placeholder detected in generated Lean code.",
        )

    compiler, command, cwd = _build_lean_command()

    if not command:
        return LeanCheckResult(
            conjecture=conjecture,
            normalized_code=normalized_code,
            is_true=False,
            verification_status="checker_unavailable",
            verification_summary="Lean is not installed or not available on PATH.",
            compiler=compiler,
            exit_code=1,
            stdout="",
            stderr="Lean is not installed or not available on PATH.",
        )

    completed = _run_lean(command=command, cwd=cwd, normalized_code=normalized_code, timeout_seconds=timeout_seconds)
    if completed is None:
        return LeanCheckResult(
            conjecture=conjecture,
            normalized_code=normalized_code,
            is_true=False,
            verification_status="timeout",
            verification_summary=f"Lean compilation timed out after {timeout_seconds} seconds.",
            compiler=compiler,
            exit_code=124,
            stdout="",
            stderr=f"Lean compilation timed out after {timeout_seconds} seconds.",
        )

    verification_status, verification_summary = _classify_result(
        exit_code=completed.returncode,
        stdout=completed.stdout.strip(),
        stderr=completed.stderr.strip(),
    )
    fallback_code = _fallback_without_mathlib(normalized_code, verification_status)
    if fallback_code:
        fallback_completed = _run_lean(
            command=command,
            cwd=cwd,
            normalized_code=fallback_code,
            timeout_seconds=timeout_seconds,
        )
        if fallback_completed is not None and fallback_completed.returncode == 0:
            return LeanCheckResult(
                conjecture=conjecture,
                normalized_code=fallback_code,
                is_true=True,
                verification_status="verified",
                verification_summary="Lean accepted the proof after removing an unresolved Mathlib import.",
                compiler=compiler,
                exit_code=0,
                stdout=fallback_completed.stdout.strip(),
                stderr=fallback_completed.stderr.strip(),
            )

    return LeanCheckResult(
        conjecture=conjecture,
        normalized_code=normalized_code,
        is_true=completed.returncode == 0,
        verification_status=verification_status,
        verification_summary=verification_summary,
        compiler=compiler,
        exit_code=completed.returncode,
        stdout=completed.stdout.strip(),
        stderr=completed.stderr.strip(),
    )


def _contains_unsafe_placeholder(code: str) -> bool:
    return bool(re.search(r"\b(sorry|admit|axiom)\b|\?\w*", code))


def _sanitize_lean_output(code: str) -> str:
    if not code:
        return code

    sanitized_lines: list[str] = []
    saw_import = False
    needs_mathlib_import = False
    started_code = False
    stop_markers = (
        "--- testing ---",
        "--- task type ---",
        "--- goal ---",
        "--- context ---",
        "--- lean 4 code ---",
        "test_proof_attempt",
        "test_",
        "#eval",
        "###",
        "explanation:",
        "proof sketch:",
        "complete lean 4 code:",
    )

    for raw_line in code.splitlines():
        line = raw_line.rstrip()
        stripped = line.strip()
        lowered = stripped.lower()

        if any(lowered.startswith(marker) for marker in stop_markers):
            break

        if stripped.startswith("```"):
            continue

        if lowered.startswith("open import "):
            module_name = stripped[len("open import ") :].strip()
            if module_name == "Mathlib" and not saw_import:
                sanitized_lines.append("import Mathlib")
                saw_import = True
                needs_mathlib_import = True
            continue

        if stripped == "end":
            continue

        if stripped == "<proof>":
            continue

        if stripped == "import Mathlib":
            if saw_import:
                continue
            saw_import = True
            needs_mathlib_import = True

        if not started_code:
            if _line_looks_like_lean_code(stripped):
                started_code = True
            else:
                continue

        sanitized_lines.append(line)

    sanitized = "\n".join(sanitized_lines).strip()
    if not sanitized:
        return sanitized

    if needs_mathlib_import and "import Mathlib" not in {line.strip() for line in sanitized_lines[:2]}:
        sanitized = f"import Mathlib\n\n{sanitized}"

    return sanitized


def _find_first_code_offset(text: str) -> int | None:
    patterns = (
        r"(?m)^import Mathlib\b",
        r"(?m)^theorem\s+[A-Za-z_][\w']*",
        r"(?m)^lemma\s+[A-Za-z_][\w']*",
        r"(?m)^example\b",
        r"(?m)^def\s+[A-Za-z_][\w']*",
    )
    offsets = [
        match.start()
        for pattern in patterns
        if (match := re.search(pattern, text)) is not None
    ]
    if not offsets:
        return None
    return min(offsets)


def _line_looks_like_lean_code(stripped: str) -> bool:
    if not stripped:
        return False
    code_prefixes = (
        "import ",
        "theorem ",
        "lemma ",
        "example ",
        "def ",
        "| ",
        "where",
        "by",
        "open ",
        "namespace ",
        "section ",
        "variable ",
        "variables ",
        "axiom ",
        "inductive ",
        "structure ",
    )
    return stripped.startswith(code_prefixes) or stripped[0].isspace() or stripped.startswith("--")


def _fallback_without_mathlib(code: str, verification_status: str) -> str | None:
    if verification_status != "missing_imports":
        return None
    lines = code.splitlines()
    filtered = [line for line in lines if line.strip() != "import Mathlib"]
    candidate = "\n".join(filtered).strip()
    if not candidate or candidate == code.strip():
        return None
    return candidate


def _classify_result(exit_code: int, stdout: str, stderr: str) -> tuple[str, str]:
    if exit_code == 0:
        return ("verified", "Lean accepted the proof.")

    diagnostics = "\n".join(part for part in (stderr, stdout) if part).lower()
    lowered = diagnostics
    if "unknown module prefix" in lowered or "object file" in lowered or "unknown package" in lowered:
        return ("missing_imports", "Lean could not resolve one or more imports.")
    if "expected" in lowered or "parse error" in lowered or "unexpected token" in lowered or "invalid 'import' command" in lowered:
        return ("malformed_output", "The generated Lean output was malformed.")
    return ("typecheck_error", "Lean rejected the proof during type checking.")


def _run_lean(
    command: list[str],
    cwd: str | None,
    normalized_code: str,
    timeout_seconds: int,
) -> subprocess.CompletedProcess[str] | None:
    temp_parent = cwd if cwd and os.path.isdir(cwd) else None
    with tempfile.TemporaryDirectory(prefix="lean-check-", dir=temp_parent) as temp_dir:
        file_path = os.path.join(temp_dir, "Main.lean")
        with open(file_path, "w", encoding="utf-8") as handle:
            handle.write(normalized_code)
            handle.write("\n")

        try:
            return subprocess.run(
                [*command, file_path],
                cwd=cwd,
                capture_output=True,
                text=True,
                timeout=timeout_seconds,
                check=False,
            )
        except subprocess.TimeoutExpired:
            return None


def _build_lean_command() -> tuple[str, list[str], str | None]:
    project_dir = _resolve_lean_project_dir()
    lake_bin = os.environ.get("LAKE_BIN") or "lake"
    lean_bin = os.environ.get("LEAN_BIN") or "lean"
    elan_bin = os.environ.get("ELAN_BIN") or "elan"
    elan_path = shutil.which(elan_bin)
    installed_toolchain = _get_installed_elan_toolchain(elan_path)

    if project_dir:
        if elan_path and installed_toolchain:
            return (
                f"elan run {installed_toolchain} lake env lean",
                [elan_path, "run", installed_toolchain, "lake", "env", "lean"],
                project_dir,
            )

        lake_path = shutil.which(lake_bin)
        if lake_path:
            return "lake env lean", [lake_path, "env", "lean"], project_dir

    if elan_path and installed_toolchain:
        return (
            f"elan run {installed_toolchain} lean",
            [elan_path, "run", installed_toolchain, "lean"],
            None,
        )

    lean_path = shutil.which(lean_bin)
    if lean_path:
        return "lean", [lean_path], None

    return "unavailable", [], project_dir or None


def _resolve_lean_project_dir() -> str | None:
    configured = os.environ.get("LEAN_PROJECT_DIR", "").strip()
    return configured or None


def _get_installed_elan_toolchain(elan_path: str | None) -> str | None:
    explicit_toolchain = os.environ.get("LEAN_TOOLCHAIN")
    if explicit_toolchain:
        return explicit_toolchain

    if not elan_path:
        return None

    try:
        completed = subprocess.run(
            [elan_path, "toolchain", "list"],
            capture_output=True,
            text=True,
            timeout=5,
            check=False,
        )
    except (OSError, subprocess.SubprocessError):
        return None

    if completed.returncode != 0:
        return None

    for raw_line in completed.stdout.splitlines():
        line = raw_line.strip()
        if not line:
            continue
        return line.split(" ", 1)[0]

    return None
