"""LLM-assisted formalization of a section's EXERCISES into Lean 4 bench targets.

These are the real lean-bench targets: problems the reader is meant to solve, as
opposed to formalize.py, which turns the book's own stated Theorems/Propositions/
Lemmas/Corollaries into a *candidate library* -- premises meant to be proven once
and then reused, not exam questions. See textbook/README.md for the split.

Exercises often reference the body by number instead of restating the claim (e.g.
"1.6.2. Prove Proposition 1.6.9.", or "1.6.15. Prove Lemma 1.6.26."), so this script
gives the model the exercise text *and* the full body section text together, and
tells it to resolve such references into a self-contained statement using the
body's actual content -- never leave a Lean statement that just says "the thing
Exercise N asks for" without spelling out what that thing is.

Same caveats as formalize.py: never runs Lean itself, the source PDF's math symbols
can be corrupted, and output must go through check_statements.py before trusting it
or spending lean-agent proof attempts on it.
"""
import argparse
import json
import re
from pathlib import Path

try:
    from .llm import ChatClient, load_dotenv
    from .formalize import extract_json_block
except ImportError:
    from llm import ChatClient, load_dotenv
    from formalize import extract_json_block

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

You will receive:
1. The book's own Theorem/Proposition/Lemma/Corollary/Definition statements for \
this section (body text), for reference.
2. The section's numbered exercises (their own independent numbering, unrelated to \
the body's numbering -- Exercise 1.6.5 has nothing to do with body item 1.6.5).
3. A checklist of every exercise number in this section, so you know exactly what \
to cover. EVERY number in the checklist MUST appear in your output as at least one \
entry (an excluded one is fine, e.g. for a pure computation -- see below -- but \
silently leaving a checklist number out entirely is never acceptable and your \
output will be rejected if you do this). Do not invent numbers that are not listed.

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.

Many exercises say things like "Prove Proposition 1.6.9" or "Prove Lemma 1.6.26" \
instead of restating a claim. For those, use the body text to find what that \
referenced item actually says, and write ITS full claim as the Lean statement -- \
never emit something like "the statement of Proposition 1.6.9" as a placeholder; \
the reader of your output has no access to the body text, only the Lean statement.

An exercise with lettered sub-parts (a), (b), (c) that each state an independent \
provable claim should become one entry PER PART (id suffixed "-a", "-b", ... in \
part order), the same way the body pipeline splits multi-part propositions. Keep a \
single entry for one exercise if its parts are really one claim you naturally \
express as a single Lean statement (e.g. an iff with two directions written as \
separate parts).

Produce exactly one JSON object with a single key "exercises" (a JSON array), \
inside one ```json fenced code block, and nothing else outside that block.

Each element of "exercises":
{
  "id": "goodman-1-6-ex-5",
  "source_label": "Exercise 1.6.5",
  "kind": "exercise",
  "statement": "the Lean 4 proposition only, i.e. what goes after `theorem name : `, \
using Mathlib vocabulary wherever possible instead of reinventing definitions",
  "prelude": "Lean import lines (and, rarely, a small needed `def`) 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"],
  "references": ["body labels this exercise restates or depends on, e.g. \
'Proposition 1.6.9', or [] if none"],
  "include_in_bench": true,
  "exclude_reason": null
}

Set "include_in_bench" to false (with a short "exclude_reason") instead of omitting \
an entry when an exercise is not a genuine general claim to formalize and prove -- \
e.g. it asks the reader to write a computer program, perform one specific numeric \
computation with no general statement, answer a yes/no or "discuss" question with \
no single formal claim, 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 "goodman-" + section-prefix-with-hyphens + "-ex-" + the exercise's own \
last number component (and a lettered suffix for sub-parts), e.g. Exercise 1.6.5 -> \
"goodman-1-6-ex-5", part (a) of Exercise 1.6.11 -> "goodman-1-6-ex-11-a". 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, exercises_text, checklist):
    checklist_lines = '\n'.join(f"- Exercise {item['number']}: {item['first_line']}"
                                 for item in checklist)
    return (f"Checklist of exercise numbers in this section:\n{checklist_lines}\n\n"
            f"Body text (for resolving 'Prove Proposition N.N.N'-style references):\n"
            f"{section_text}\n\n"
            f"Exercises text:\n{exercises_text}")


def validate(data, checklist=None):
    if not isinstance(data, dict) or 'exercises' not in data:
        raise ValueError('Response must be a JSON object with an "exercises" array')
    exercises = data['exercises']
    if not isinstance(exercises, list):
        raise ValueError('"exercises" must be an array')
    seen_ids = set()
    covered_numbers = set()
    for item in exercises:
        for field in ('id', 'source_label', 'statement', 'prelude', 'include_in_bench'):
            if field not in item:
                raise ValueError(f'exercise entry missing field {field!r}: {item}')
        if item['id'] in seen_ids:
            raise ValueError(f'Duplicate exercise id: {item["id"]}')
        seen_ids.add(item['id'])
        # "Exercise 1.6.11(a)" or "Exercise 1.6.11" -> "1.6.11"
        m = re.search(r'(\d+\.\d+\.\d+)', item['source_label'])
        if m:
            covered_numbers.add(m.group(1))
    if checklist is not None:
        expected = {c['number'] for c in checklist}
        missing = expected - covered_numbers
        if missing:
            raise ValueError(
                'Model silently dropped checklist items instead of marking them '
                'include_in_bench=false with a reason (never allowed): '
                + ', '.join(sorted(missing, key=lambda s: tuple(int(p) for p in s.split('.'))))
            )
    return exercises


def main():
    parser = argparse.ArgumentParser(description='LLM-formalize a section\'s exercises into Lean 4 bench targets')
    parser.add_argument('section_text', type=Path, help='body text, from extract_section.py')
    parser.add_argument('exercises_text', type=Path, help='from extract_exercises.py')
    parser.add_argument('exercises_index', type=Path, help='from extract_exercises.py')
    parser.add_argument('--out-dataset', type=Path, required=True, help='JSONL of exercises with include_in_bench=true')
    parser.add_argument('--out-excluded', type=Path, help='JSONL of exercise 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.exercises_index.read_text(encoding='utf-8'))
    section_text = args.section_text.read_text(encoding='utf-8')
    exercises_text = args.exercises_text.read_text(encoding='utf-8')
    messages = [{"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt(section_text, exercises_text, checklist)}]
    reply = client.complete(messages)
    data = extract_json_block(reply)
    exercises = validate(data, checklist)

    included = [t for t in exercises if t.get('include_in_bench')]
    excluded = [t for t in exercises 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')
    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)} exercises -> {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}')


if __name__ == '__main__':
    main()