#!/usr/bin/env python3
"""Generate Lean proofs with the paper's prompt and last-complete-block parser."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from prompt_utils import (
    LEAN4_DEFAULT_HEADER, DEFAULT_GENERATION_TOP_P,
    build_prover_prompt_for_example, default_generation_temperature,
    default_generation_max_tokens, default_lean_stop_sequences,
    parse_generated_lean, split_context_and_command,
)

DEFAULT_MODEL = 'EPFL-LARA/lean-standard-prover-qwen3-8b-reasoning-grpo'
FEEDBACK = False
REASONING = True


def build_prompt(source: str) -> str:
    context, declaration, found = split_context_and_command(source)
    if not found:
        raise ValueError("Supply a Lean theorem, lemma, or example ending in 'by sorry'.")
    example = {"formal_statement": declaration, "header": context or LEAN4_DEFAULT_HEADER}
    _, prompt = build_prover_prompt_for_example(
        example, use_informal_statement=False, is_deepseek=False,
        tgat_prompt=FEEDBACK, reasoning=REASONING,
    )
    return prompt


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("files", nargs="+", type=Path, help="Lean problem files, including imports and a theorem to complete")
    parser.add_argument("--model", default=DEFAULT_MODEL)
    parser.add_argument("--samples", type=int, default=1)
    parser.add_argument("--temperature", type=float)
    parser.add_argument("--top-p", type=float, default=DEFAULT_GENERATION_TOP_P)
    parser.add_argument("--max-new-tokens", type=int, default=default_generation_max_tokens(False))
    parser.add_argument("--tensor-parallel-size", type=int, default=1)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--print-prompt", action="store_true", help="Show the exact raw prompt without loading the model")
    args = parser.parse_args()
    if args.samples < 1 or args.max_new_tokens < 1:
        parser.error("--samples and --max-new-tokens must be positive")
    prompts = [build_prompt(path.read_text(encoding="utf-8")) for path in args.files]
    if args.print_prompt:
        for prompt in prompts:
            print(prompt, end="")
        return

    from vllm import LLM, SamplingParams

    llm = LLM(model=args.model, tokenizer=args.model,
              tensor_parallel_size=args.tensor_parallel_size, dtype="auto",
              max_model_len=32768, trust_remote_code=False, seed=args.seed,
              gpu_memory_utilization=0.90)
    tokenizer = llm.get_tokenizer()
    tokenized = [tokenizer.encode(prompt, add_special_tokens=False) for prompt in prompts]
    if any(len(ids) >= 32768 for ids in tokenized):
        raise ValueError("The problem prompt must leave space in the 32,768-token context for a proof.")
    temperature = args.temperature if args.temperature is not None else default_generation_temperature(args.samples)
    params = [SamplingParams(n=args.samples, temperature=temperature,
                             top_p=args.top_p,
                             max_tokens=min(args.max_new_tokens, 32768-len(ids)),
                             stop=default_lean_stop_sequences(False),
                             include_stop_str_in_output=True, seed=args.seed)
              for ids in tokenized]
    outputs = llm.generate([{"prompt_token_ids": ids} for ids in tokenized], params)
    for path, output in zip(args.files, outputs, strict=True):
        for i, completion in enumerate(output.outputs):
            print(json.dumps({
                "source": str(path), "sample_id": i,
                "response": completion.text,
                "lean_code": parse_generated_lean(completion.text, is_deepseek=False),
                "finish_reason": completion.finish_reason,
            }, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()