cli.py

"""lean-debug CLI — entry point."""
from __future__ import annotations

import json
from pathlib import Path
from typing import Optional

import typer
from rich.console import Console
from rich.table import Table

from lean_proof_debugger.benchmark import BenchmarkAnalyzer
from lean_proof_debugger.corpus import Corpus
from lean_proof_debugger.state_diff import ProofStateDiff
from lean_proof_debugger.failure_classifier import FailureClassifier
from lean_proof_debugger.parsers import auto_detect_and_parse, get_parser
from lean_proof_debugger.premise_analysis import PremiseAnalyzer
from lean_proof_debugger.proof_tree import ProofTree
from lean_proof_debugger.report import ReportBuilder
from lean_proof_debugger.run_comparison import RunComparison
from lean_proof_debugger.utils import FAILURE_COLORS

app = typer.Typer(name="lean-debug", help="Lean Proof Debugger — analyse and debug theorem proving runs.")
console = Console()


def _load(path: Path, fmt: str = "auto"):
    if fmt == "auto":
        return auto_detect_and_parse(path)
    return get_parser(fmt).parse_file(path)


@app.command()
def parse(
    path: Path = typer.Argument(..., help="Path to trace file or directory"),
    format: str = typer.Option("auto", "--format", "-f", help="leandojo|reprover|proofnet|raw|auto"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write parsed traces to JSON file"),
):
    """Parse a trace file and print a summary table."""
    traces = _load(path, format)
    table = Table(title=f"Parsed Traces: {path}", show_lines=True)
    table.add_column("Theorem", style="cyan", max_width=50)
    table.add_column("Format")
    table.add_column("Steps", justify="right")
    table.add_column("Result")

    for t in traces:
        table.add_row(
            t.theorem_name,
            t.source_format,
            str(len(t.steps)),
            "[green]✓[/green]" if t.succeeded else "[red]✗[/red]",
        )
    console.print(table)
    console.print(f"\n[dim]{len(traces)} trace(s) loaded.[/dim]")

    if output:
        data = [{"theorem": t.theorem_name, "steps": len(t.steps), "succeeded": t.succeeded} for t in traces]
        output.write_text(json.dumps(data, indent=2))
        console.print(f"[dim]Summary written to {output}[/dim]")


@app.command()
def classify(
    path: Path = typer.Argument(..., help="Path to trace file"),
    format: str = typer.Option("auto", "--format", "-f"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write Markdown report"),
):
    """Classify failure modes for all traces in a file."""
    traces = _load(path, format)
    classifier = FailureClassifier()
    diagnoses = classifier.classify_batch(traces)

    table = Table(title="Failure Classification", show_lines=True)
    table.add_column("Theorem", style="cyan", max_width=40)
    table.add_column("Label")
    table.add_column("Confidence", justify="right")
    table.add_column("First Failure Step", justify="right")
    table.add_column("Evidence")

    for trace, diag in zip(traces, diagnoses):
        color = FAILURE_COLORS.get(diag.label.value, "white")
        table.add_row(
            trace.theorem_name[:40],
            f"[{color}]{diag.label.value}[/{color}]",
            f"{diag.confidence:.0%}",
            str(diag.first_failure_step) if diag.first_failure_step is not None else "—",
            diag.evidence[0][:60] if diag.evidence else "",
        )
    console.print(table)

    if output:
        lines = []
        for trace, diag in zip(traces, diagnoses):
            lines.append(ReportBuilder(trace, diag).to_markdown())
            lines.append("\n---\n")
        output.write_text("\n".join(lines))
        console.print(f"[dim]Report written to {output}[/dim]")


@app.command()
def tree(
    path: Path = typer.Argument(..., help="Path to trace file"),
    theorem: Optional[str] = typer.Option(None, "--theorem", "-t", help="Theorem name filter"),
    format: str = typer.Option("auto", "--format", "-f"),
):
    """Print proof tree(s) as ASCII art."""
    traces = _load(path, format)
    if theorem:
        traces = [t for t in traces if theorem in t.theorem_name]
    if not traces:
        console.print("[red]No matching traces found.[/red]")
        raise typer.Exit(1)
    for trace in traces:
        console.rule(f"[bold]{trace.theorem_name}")
        pt = ProofTree(trace)
        console.print(pt.to_ascii())
        dead_ends = pt.get_dead_ends()
        console.print(f"[dim]{len(dead_ends)} dead-end node(s)[/dim]\n")


@app.command()
def premises(
    path: Path = typer.Argument(..., help="Path to trace file"),
    theorem: Optional[str] = typer.Option(None, "--theorem", "-t"),
    k: int = typer.Option(10, "--k", help="Top-k for coverage"),
    format: str = typer.Option("auto", "--format", "-f"),
):
    """Show per-step premise retrieval recall."""
    traces = _load(path, format)
    if theorem:
        traces = [t for t in traces if theorem in t.theorem_name]
    analyzer = PremiseAnalyzer()
    for trace in traces:
        console.rule(f"[bold]{trace.theorem_name}")
        report = analyzer.analyze(trace)
        console.print(f"Overall recall: [bold]{report.overall_recall:.2f}[/bold]  Miss rate: {report.miss_rate:.0%}")
        console.print(analyzer.highlight_misses(report))


@app.command()
def compare(
    paths: list[Path] = typer.Argument(..., help="Two or more trace files (same theorem)"),
    theorem: Optional[str] = typer.Option(None, "--theorem", "-t"),
):
    """Compare multiple proof runs on the same theorem."""
    if len(paths) < 2:
        console.print("[red]Provide at least 2 trace files.[/red]")
        raise typer.Exit(1)
    traces_map = {}
    for i, p in enumerate(paths):
        loaded = _load(p)
        t = loaded[0] if loaded else None
        if t:
            run_id = t.metadata.get("run_id", p.stem)
            traces_map[run_id] = t

    comparison = RunComparison(traces_map)
    summary = comparison.summary()
    console.print_json(json.dumps(summary, indent=2))


@app.command()
def report(
    path: Path = typer.Argument(..., help="Path to trace file"),
    format: str = typer.Option("markdown", "--format", "-f", help="markdown|json|html"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Output directory"),
    fmt_in: str = typer.Option("auto", "--input-format"),
):
    """Full pipeline: parse → classify → write reports."""
    traces = _load(path, fmt_in)
    classifier = FailureClassifier()
    diagnoses = classifier.classify_batch(traces)
    out_dir = output or Path("reports")
    ReportBuilder.batch_report(traces, diagnoses, out_dir, fmt=format)
    console.print(f"[green]{len(traces)} report(s) written to {out_dir}/[/green]")


@app.command()
def benchmark(
    path: Path = typer.Argument(..., help="Directory of trace files"),
    glob: str = typer.Option("**/*.json", "--glob", "-g", help="File glob pattern"),
    output: Optional[Path] = typer.Option(None, "--output", "-o", help="Write report (markdown/json)"),
    fmt: str = typer.Option("markdown", "--format", "-f", help="markdown|json"),
    max_files: Optional[int] = typer.Option(None, "--max-files", help="Cap number of files loaded"),
    compare: Optional[str] = typer.Option(None, "--compare", help="Comma-separated run_name=dir pairs for model comparison"),
):
    """Analyze a full benchmark run directory: pass rate, failure breakdown, hardest theorems."""
    analyzer = BenchmarkAnalyzer()

    if compare:
        # e.g. --compare "reprover=results/reprover,gptf=results/gptf"
        run_dirs = {}
        for pair in compare.split(","):
            name, dir_path = pair.strip().split("=", 1)
            run_dirs[name.strip()] = Path(dir_path.strip())
        df = analyzer.compare_runs(run_dirs, glob=glob)
        console.print(df.to_string())
        if output:
            output.write_text(df.to_csv(index=False))
            console.print(f"[dim]Comparison written to {output}[/dim]")
        return

    result = analyzer.analyze_dir(path, glob=glob, max_files=max_files)
    n = len(result.traces)
    n_ok = sum(1 for t in result.traces if t.succeeded)

    console.print(f"\n[bold]Benchmark:[/bold] {path}")
    console.print(f"Theorems: {n}  |  Pass rate: [bold]{n_ok}/{n} ({result.pass_rate():.1%})[/bold]  |  Mean premise recall: {result.mean_premise_recall():.3f}\n")

    table = Table(title="Failure Breakdown", show_lines=True)
    table.add_column("Label")
    table.add_column("Count", justify="right")
    table.add_column("%", justify="right")
    breakdown = result.failure_breakdown()
    for label, count in breakdown.items():
        from lean_proof_debugger.utils import FAILURE_COLORS
        color = FAILURE_COLORS.get(label, "white")
        table.add_row(f"[{color}]{label}[/{color}]", str(count), f"{count/n*100:.1f}%")
    console.print(table)

    hard = result.hardest_theorems(10)
    if not hard.empty:
        console.print("\n[bold]Hardest theorems (failed, lowest premise recall):[/bold]")
        console.print(hard.to_string(index=False))

    if output:
        if fmt == "json":
            import json as _json
            output.write_text(_json.dumps(result.to_json(), indent=2))
        else:
            output.write_text(result.to_markdown())
        console.print(f"\n[dim]Report written to {output}[/dim]")


@app.command("state-diff")
def state_diff(
    path: Path = typer.Argument(..., help="Trace file"),
    theorem: Optional[str] = typer.Option(None, "--theorem", "-t"),
    format: str = typer.Option("auto", "--format", "-f"),
    output: Optional[Path] = typer.Option(None, "--output", "-o"),
):
    """Show step-by-step proof state changes as a diff timeline."""
    traces = _load(path, format)
    if theorem:
        traces = [t for t in traces if theorem in t.theorem_name]
    if not traces:
        console.print("[red]No matching traces.[/red]")
        raise typer.Exit(1)
    for trace in traces:
        console.rule(f"[bold]{trace.theorem_name}")
        diff = ProofStateDiff(trace)
        import pandas as pd
        df = pd.DataFrame(diff.summary())
        console.print(df.to_string(index=False))
        if output:
            output.write_text(diff.to_markdown())
            console.print(f"[dim]State diff written to {output}[/dim]")


@app.command("corpus-info")
def corpus_info(
    corpus_path: Path = typer.Argument(..., help="Path to corpus.jsonl"),
):
    """Print summary statistics for a LeanDojo corpus.jsonl file."""
    corpus = Corpus.from_jsonl(corpus_path)
    console.print(f"Premises loaded: [bold]{len(corpus)}[/bold]")


@app.command()
def demo():
    """Run the built-in demo on sample traces."""
    from demo.run_demo import run_demo
    run_demo()


if __name__ == "__main__":
    app()