"""LLM-assisted formalization of a textbook section into Lean 4 candidate statements.

Reads a section's extracted text and labeled-statement checklist (from
extract_section.py) and asks an LLM to produce:
  (a) a Lean 4 theorem statement for every genuine Theorem/Proposition/Lemma/Corollary,
      preferring Mathlib's existing vocabulary over reinventing definitions, and
  (b) a manifest of every Definition the theorems depend on, pointing at the matching
      Mathlib name when one exists and only asking for hand-written Lean `def`s when it
      doesn't.

This step never runs Lean. The extracted PDF text can have corrupted math symbols
(missing =/in, dropped superscripts) from the source PDF's font encoding; the model is
told to lean on its own knowledge of standard textbook mathematics plus the section
context, and to flag low-confidence items rather than guess silently. Always run
check_statements.py against a real Mathlib-backed Lean project before trusting the
output, and before spending lean-agent proof attempts on it.
"""
import argparse
import json
import re
from pathlib import Path

try:
    from .llm import ChatClient, load_dotenv  # imported as textbook.formalize
except ImportError:
    from llm import ChatClient, load_dotenv  # run directly: python3 textbook/formalize.py

SYSTEM_PROMPT = """You are formalizing a section of an undergraduate abstract algebra \
textbook (Goodman, "Algebra: Abstract and Concrete") into Lean 4 using Mathlib.

You will receive:
1. The raw text of one book section, extracted from a PDF. WARNING: the source PDF's \
math font encoding is broken in this extraction -- symbols like congruence, set \
membership, comparisons, superscripts and subscripts are frequently dropped or garbled. \
Use your own knowledge of standard undergraduate number theory / algebra to reconstruct \
the intended mathematics from context; do not invent content that is not actually in the \
section.
2. A checklist of every labeled Theorem/Proposition/Lemma/Corollary/Definition in the \
section (by number), so you know exactly what to cover -- do not skip any, and do not \
invent ones that are not listed.

Produce exactly two JSON arrays, "definitions" and "theorems", as a single JSON object \
inside one ```json fenced code block, and nothing else outside that block.

Each element of "definitions" (only for Definition-labeled entries the theorems below \
actually depend on):
{
  "label": "Definition 1.6.3",
  "concept": "short name, e.g. 'prime number'",
  "mathlib_ref": "Mathlib identifier that already means this, e.g. 'Nat.Prime', or null \
if Mathlib has no equivalent",
  "lean_code": "a standalone Lean 4 `def ...` ONLY if mathlib_ref is null, else null",
  "notes": "anything a formalizer should know, e.g. where the book's definition differs \
from Mathlib's convention"
}

Each element of "theorems" (one per Theorem/Proposition/Lemma/Corollary in the \
checklist):
{
  "id": "goodman-1-6-6",
  "source_label": "Theorem 1.6.6",
  "kind": "theorem",
  "statement": "the Lean 4 proposition only, i.e. what goes after `theorem name : `, \
using Mathlib vocabulary (Nat.Prime, Nat.gcd, Int.gcd, etc.) wherever possible instead \
of reinventing definitions",
  "prelude": "Lean import lines this statement needs, usually just \\"import Mathlib\\"",
  "natural_language": "a short original English restatement in your own words (do not \
quote the source text verbatim)",
  "requires": ["list of Mathlib identifiers the statement uses"],
  "include_in_bench": true,
  "exclude_reason": null
}

Set "include_in_bench" to false (with a short "exclude_reason") instead of omitting an \
entry when a checklist item turns out to not be a genuine formalizable claim (e.g. it is \
actually a definition misfiled under a theorem-style label, a purely numerical one-off \
computation with no general statement, or the extracted text is too corrupted to recover \
its meaning confidently) -- still emit the entry so nothing silently disappears, just \
mark it excluded.

id must be the label with dots and spaces replaced by hyphens, lowercased, prefixed \
"goodman-". Do not include a proof, "sorry", or "by" -- statement is only the \
proposition, no theorem keyword or name. Do not wrap the JSON in anything other than a \
single ```json fenced block."""


def user_prompt(section_text, checklist):
    checklist_lines = '\n'.join(f"- {item['kind']} {item['label']}: {item['first_line']}"
                                 for item in checklist)
    return (f"Checklist of labeled items in this section:\n{checklist_lines}\n\n"
            f"Section text:\n{section_text}")


def extract_json_block(text):
    match = re.search(r'```json\s*\n(.*?)```', text, re.S)
    payload = match.group(1) if match else text
    return json.loads(payload)


def validate(data):
    if not isinstance(data, dict) or 'theorems' not in data or 'definitions' not in data:
        raise ValueError('Response must be a JSON object with "theorems" and "definitions" arrays')
    theorems, definitions = data['theorems'], data['definitions']
    if not isinstance(theorems, list) or not isinstance(definitions, list):
        raise ValueError('"theorems" and "definitions" must be arrays')
    seen_ids = set()
    for item in theorems:
        for field in ('id', 'source_label', 'statement', 'prelude', 'include_in_bench'):
            if field not in item:
                raise ValueError(f'theorem entry missing field {field!r}: {item}')
        if item['id'] in seen_ids:
            raise ValueError(f'Duplicate theorem id: {item["id"]}')
        seen_ids.add(item['id'])
    for item in definitions:
        for field in ('label', 'concept', 'mathlib_ref'):
            if field not in item:
                raise ValueError(f'definition entry missing field {field!r}: {item}')
    return theorems, definitions


def main():
    parser = argparse.ArgumentParser(description='LLM-formalize a book section into Lean 4 candidate statements')
    parser.add_argument('section_text', type=Path)
    parser.add_argument('statement_index', type=Path)
    parser.add_argument('--out-dataset', type=Path, required=True, help='JSONL of theorems with include_in_bench=true')
    parser.add_argument('--out-definitions', type=Path, required=True, help='JSON manifest of definitions')
    parser.add_argument('--out-excluded', type=Path, help='JSONL of theorem entries the model excluded, for review')
    parser.add_argument('--model', required=True)
    parser.add_argument('--base-url', default='https://dashscope.aliyuncs.com/compatible-mode/v1')
    parser.add_argument('--api-key-env', default='DASHSCOPE_API_KEY')
    parser.add_argument('--env-file', type=Path, default=Path(__file__).resolve().parents[1] / '.env')
    parser.add_argument('--timeout', type=float, default=300)
    args = parser.parse_args()

    load_dotenv(args.env_file)
    client = ChatClient(args.model, args.base_url, args.api_key_env, args.timeout)
    checklist = json.loads(args.statement_index.read_text(encoding='utf-8'))
    section_text = args.section_text.read_text(encoding='utf-8')
    messages = [{"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt(section_text, checklist)}]
    reply = client.complete(messages)
    data = extract_json_block(reply)
    theorems, definitions = validate(data)

    included = [t for t in theorems if t.get('include_in_bench')]
    excluded = [t for t in theorems if not t.get('include_in_bench')]

    with args.out_dataset.open('w', encoding='utf-8') as stream:
        for item in included:
            stream.write(json.dumps(item, ensure_ascii=False) + '\n')
    args.out_definitions.write_text(json.dumps(definitions, ensure_ascii=False, indent=2) + '\n', encoding='utf-8')
    if args.out_excluded:
        with args.out_excluded.open('w', encoding='utf-8') as stream:
            for item in excluded:
                stream.write(json.dumps(item, ensure_ascii=False) + '\n')

    print(f'{len(included)} theorems -> {args.out_dataset}')
    if excluded:
        note = f' (see {args.out_excluded})' if args.out_excluded else ' (pass --out-excluded to keep them)'
        print(f'{len(excluded)} excluded{note}')
    print(f'{len(definitions)} definitions -> {args.out_definitions}')


if __name__ == '__main__':
    main()