"""Prepare a pinned German math SFT pilot; optional pyarrow reads MGSM Parquet."""

from __future__ import annotations

import argparse
from collections import Counter, defaultdict
import hashlib
import json
from pathlib import Path
import random
import re
import urllib.request


SOURCES = {
    "mathoctopus.json": {
        "url": "https://huggingface.co/datasets/Mathoctopus/GSM8KInstruct_Parallel/resolve/dd8076b83bacd7c6ade2829317e8a500c6ffbec4/MGSM8KInstruct_Parallel.json",
        "revision": "dd8076b83bacd7c6ade2829317e8a500c6ffbec4", "license": "Apache-2.0",
        "sha256": "57864b498f8ea0bc0f8903e8fd912c079ea1a824fad23b1077a0149b455f5022",
    },
    "gsm8k-train.jsonl": {
        "url": "https://raw.githubusercontent.com/openai/grade-school-math/3101c7d5072418e28b9008a6636bde82a006892c/grade_school_math/data/train.jsonl",
        "revision": "3101c7d5072418e28b9008a6636bde82a006892c", "license": "MIT",
        "sha256": "17f347dc51477c50d4efb83959dbb7c56297aba886e5544ee2aaed3024813465",
    },
    "gsm8k-test.jsonl": {
        "url": "https://raw.githubusercontent.com/openai/grade-school-math/3101c7d5072418e28b9008a6636bde82a006892c/grade_school_math/data/test.jsonl",
        "revision": "3101c7d5072418e28b9008a6636bde82a006892c", "license": "MIT",
        "sha256": "3730d312f6e3440559ace48831e51066acaca737f6eabec99bccb9e4b3c39d14",
    },
    "mgsm-de-test.parquet": {
        "url": "https://huggingface.co/datasets/juletxara/mgsm/resolve/b2f13d426afe3be8d69a7e739b36724db8b66bbc/de/test-00000-of-00001.parquet",
        "revision": "b2f13d426afe3be8d69a7e739b36724db8b66bbc", "license": "CC-BY-SA-4.0",
        "sha256": "8a18ae04bd783a0fe1ed7148a36d5eade77dca1c94c2143216892963f9c95201",
    },
}
PROMPT_SUFFIX = (
    "\n\nLöse die Aufgabe auf Deutsch und erkläre die Rechenschritte. "
    "Schließe mit der Zahl im Format <answer>ZAHL</answer> ab."
)
WRAPPER = re.compile(
    r"\ABelow is an instruction that describes a task\. Write a response that appropriately "
    r"completes the request in (?P<language>\w+)\. Please answer in (?P=language)\."
    r"\s+### Instruction:\s*(?P<question>.*?)\s+### Response:\s*\Z", re.DOTALL
)
NUMBER = re.compile(r"[-+]?\d+(?:[.,]\d+)*")
AUDIT_EXCLUSIONS = {
    "gsm8k/train/00090": "Translation wrongly assigns 30 days to December, January and February.",
    "gsm8k/train/00237": "Translated intermediate egg counts contradict the problem despite matching the final gold.",
}


def digest(path: Path) -> str:
    with path.open("rb") as stream:
        return hashlib.file_digest(stream, "sha256").hexdigest()


def _read(path: Path) -> list[dict]:
    with path.open(encoding="utf-8") as stream:
        return [json.loads(line) for line in stream if line.strip()]


def fetch_sources(directory: Path, *, download: bool = True) -> dict:
    directory.mkdir(parents=True, exist_ok=True)
    report = {}
    for name, source in SOURCES.items():
        path = directory / name
        if not path.exists():
            if not download:
                raise ValueError(f"Missing pinned source: {path}")
            temporary = path.with_suffix(path.suffix + ".partial")
            try:
                with urllib.request.urlopen(source["url"], timeout=60) as response, temporary.open("xb") as stream:
                    while chunk := response.read(1024 * 1024):
                        stream.write(chunk)
                if digest(temporary) != source["sha256"]:
                    raise ValueError(f"Pinned source checksum mismatch: {name}")
                temporary.replace(path)
            finally:
                temporary.unlink(missing_ok=True)
        if digest(path) != source["sha256"]:
            raise ValueError(f"Pinned source checksum mismatch: {name}")
        report[name] = {**source, "path": str(path.resolve()), "bytes": path.stat().st_size}
    return report


def align_german(rows: list[dict], canonical_train: list[dict], canonical_test: list[dict],
                 *, max_chars: int = 4096) -> tuple[list[dict], dict]:
    """Use exact English question+rationale anchors, never a guessed row offset."""
    indices = defaultdict(list)
    for index, row in enumerate(canonical_train):
        indices[row["question"].strip()].append((index, row))
    test_questions = {row["question"].strip() for row in canonical_test}
    anchor = None
    accepted, rejected = [], Counter()
    seen = set()
    german_total = 0
    for row_number, row in enumerate(rows, 1):
        match = WRAPPER.fullmatch(row.get("prompt", ""))
        if match is None:
            anchor = None  # Malformed structure cannot retain a trusted preceding anchor.
            rejected["malformed_wrapper"] += 1
            continue
        language, question = match["language"], match["question"].strip()
        if language == "English":
            anchor = None
            candidates = indices.get(question, [])
            if question not in test_questions and len(candidates) == 1:
                index, original = candidates[0]
                if row.get("chosen", "").strip() == original["answer"].strip():
                    anchor = (index, original)
            continue
        if language != "German":
            continue
        german_total += 1
        if anchor is None:
            rejected["no_exact_unique_english_train_anchor"] += 1
            continue
        index, original = anchor
        source_id = f"gsm8k/train/{index:05d}"
        if source_id in seen:
            rejected["duplicate_german_source_id"] += 1
            continue
        seen.add(source_id)
        if source_id in AUDIT_EXCLUSIONS:
            rejected["audited_incorrect_intermediate_reasoning"] += 1
            continue
        if " ".join(question.split()) == " ".join(original["question"].split()):
            rejected["untranslated_source_question"] += 1
            continue
        gold_match = re.search(r"####\s*(-?\d+)\s*\Z", original["answer"])
        if gold_match is None:
            rejected["unsupported_canonical_answer_format"] += 1
            continue
        gold = str(int(gold_match.group(1)))
        rationale = row.get("chosen", "").strip()
        numbers = NUMBER.findall(rationale)
        if not numbers or not re.fullmatch(r"[-+]?\d+", numbers[-1]):
            rejected["ambiguous_or_missing_final_number"] += 1
            continue
        if int(numbers[-1]) != int(gold):
            rejected["final_number_disagrees_with_canonical_gold"] += 1
            continue
        if any(tag in rationale for tag in ("<think>", "</think>", "<answer>", "</answer>")):
            rejected["preexisting_structural_tags"] += 1
            continue
        prompt = question + PROMPT_SUFFIX
        completion = f"<think>{rationale}</think>\n\n<answer>{gold}</answer>"
        if len(prompt) + len(completion) > max_chars:
            rejected["character_length_limit"] += 1
            continue
        accepted.append({"id": f"mathoctopus-de-{index:05d}", "source_id": source_id,
                         "prompt": prompt, "completion": completion, "solution": gold,
                         "source_row": row_number})
    return accepted, {"source_rows": len(rows), "german_rows": german_total,
                      "accepted": len(accepted), "rejected": dict(rejected),
                      "validation_scope": "Exact source identity and final integer only; intermediate reasoning and translation semantics are not proven correct."}


def mgsm_examples(path: Path) -> list[dict]:
    try:
        import pyarrow.parquet as parquet
    except ImportError as exc:
        raise ImportError("MGSM preparation requires optional pyarrow to read the pinned Parquet file") from exc
    rows = parquet.read_table(path).to_pylist()
    if len(rows) != 250:
        raise ValueError("Expected all 250 pinned German MGSM test examples")
    return [{"id": f"mgsm-de-test-{i:03d}", "prompt": row["question"] + PROMPT_SUFFIX,
             "answer": str(row["answer_number"]), "scorer": "final_numeric",
             "category": "German math", "source_id": f"mgsm/de/test/{i:03d}"}
            for i, row in enumerate(rows)]


def prepare(output_dir: str | Path, train_size: int = 1024, dev_size: int = 128,
            source_dir: str | Path | None = None, download: bool = True,
            max_chars: int = 4096, seed: int = 42) -> dict:
    output = Path(output_dir).resolve()
    for name, value in (("train_size", train_size), ("dev_size", dev_size), ("max_chars", max_chars)):
        if type(value) is not int or value <= 0:
            raise ValueError(f"{name} must be a positive integer")
    if output.exists() and any(output.iterdir()):
        raise ValueError(f"Prepared output must be new or empty: {output}")
    sources = Path(source_dir).resolve() if source_dir else output.parent / "sources"
    provenance = fetch_sources(sources, download=download)
    accepted, coverage = align_german(_read(sources / "mathoctopus.json"),
                                      _read(sources / "gsm8k-train.jsonl"),
                                      _read(sources / "gsm8k-test.jsonl"), max_chars=max_chars)
    if len(accepted) < train_size + dev_size:
        raise ValueError(f"Only {len(accepted)} verified rows available for {train_size + dev_size} requested")
    random.Random(seed).shuffle(accepted)
    train = accepted[:train_size]
    dev = accepted[train_size:train_size + dev_size]
    external = mgsm_examples(sources / "mgsm-de-test.parquet")
    # Matching German text is an additional guard; the source split separation is primary.
    if {r["prompt"] for r in train + dev} & {r["prompt"] for r in external}:
        raise ValueError("German train/development question overlaps MGSM test")
    output.mkdir(parents=True, exist_ok=True)
    files = {}
    for name, rows in (("train.jsonl", train), ("dev.jsonl", dev), ("mgsm-de.jsonl", external)):
        path = output / name
        with path.open("x", encoding="utf-8") as stream:
            for row in rows:
                stream.write(json.dumps(row, ensure_ascii=False) + "\n")
        files[name] = {"path": str(path), "rows": len(rows), "sha256": digest(path),
                       "source_ids": [row["source_id"] for row in rows]}
    report = {"schema_version": 1, "method": "german-math-sft-pilot-v1", "seed": seed,
              "max_chars": max_chars, "sources": provenance, "coverage": coverage,
              "files": files, "source_id_disjoint": True, "mgsm_is_training_data": False,
              "audited_exclusions": AUDIT_EXCLUSIONS,
              "token_length_check": "Required with the selected policy tokenizer before training; character length is only an initial bound.",
              "license_note": "MathOctopus Apache-2.0; canonical GSM8K MIT; MGSM mirror CC-BY-SA-4.0. Preserve attribution; sources/derived data are ignored local run artifacts."}
    (output / "manifest.json").write_text(json.dumps(report, indent=2, ensure_ascii=False) + "\n")
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", type=Path, default=Path("runs/german-qwen/data"))
    parser.add_argument("--source-dir", type=Path)
    parser.add_argument("--train-size", type=int, default=1024)
    parser.add_argument("--dev-size", type=int, default=128)
    parser.add_argument("--max-chars", type=int, default=4096)
    parser.add_argument("--offline", action="store_true")
    args = parser.parse_args()
    report = prepare(args.output, args.train_size, args.dev_size, args.source_dir,
                     not args.offline, args.max_chars)
    print(json.dumps({"coverage": report["coverage"], "files": {
        name: {key: value for key, value in info.items() if key != "source_ids"}
        for name, info in report["files"].items()}}, indent=2))


if __name__ == "__main__":
    main()