"""
Prepare Omartificial-Intelligence-Space/Arabic_Reasoning_Dataset for reasoning SFT.

Each row (instruction, answer) becomes a ChatML sample where the derivation lives
inside <think>...</think> and the conclusion follows it:

    <|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n<think>\n{reasoning}\n</think>\n{answer}<|im_end|>

Only rows whose derivation ends in an explicit conclusion line ("إذن، ...") are kept.
Expository rows without one have no real "final answer" to place after </think> — their
closing paragraph is a side remark, so training on them would teach the model to think and
then trail off. They are dropped rather than force-split.

Writes data/train.jsonl and data/eval.jsonl.
"""
import json, re, random, unicodedata, collections
from pathlib import Path

import pyarrow.parquet as pq

SRC      = Path("data/data/train-00000-of-00001.parquet")
OUT_DIR  = Path("data")
EVAL_N   = 400
SEED     = 42

# Lines that mark the final conclusion of a derivation.
CONCLUSION = ["إذن،", "إذن ", "لذا،", "لذلك،", "باختصار،", "وبالتالي،", "في النهاية،",
              "الخلاصة", "النتيجة النهائية", "النتيجة:", "الإجابة", "الجواب", "بالتالي،"]
# Chatty sign-offs that are not part of the answer.
FLUFF = ["تذكر", "آمل", "أتمنى", "يرجى", "لا تتردد", "إذا كان لديك أي", "هل لديك",
         "أرجو", "نصيحة:", "ملاحظة:", "إذا كانت لديك"]
# Prompt suffix present on a subset of instructions; stripped so reasoning is unconditional.
SUFFIX_RE = re.compile(r"\s*خذ\s+نفسًا?\s+عميقًا.*$", re.S)

MIN_THINK_CHARS  = 60
MIN_ANSWER_CHARS = 10

MD_NOISE_RE  = re.compile(r"\*\*|__|#{2,}")
# Some source rows bundle several problems; the conclusion of one is followed by the next
# problem's header. Anything from that header on is not part of the answer.
NEXT_PROBLEM_RE = re.compile(r"^\s*(المشكلة|المسألة|السؤال|التمرين|مثال|Problem|Question|Example)\b")
ANSWER_ARTIFACT_RE = re.compile(r"^\s*(\*\*)?Answer:\s*", re.I)


def norm(text: str) -> str:
    text = unicodedata.normalize("NFC", text.replace("‏", "").replace("‎", ""))
    text = re.sub(r"[ \t]+", " ", text)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip()


def clean_instruction(text: str) -> str:
    return norm(SUFFIX_RE.sub("", norm(text))).rstrip(". ").strip() or norm(text)


def is_fluff(line: str) -> bool:
    head = line.lstrip("*#- ").strip()
    return any(head.startswith(f) for f in FLUFF)


def is_conclusion(line: str) -> bool:
    head = line.lstrip("*#- ").strip()
    return any(head.startswith(m) for m in CONCLUSION)


def strip_markup(line: str) -> str:
    """Drop markdown emphasis and the stray "**Answer:" prefix some rows carry."""
    return norm(MD_NOISE_RE.sub("", ANSWER_ARTIFACT_RE.sub("", line)))


def split_answer(answer: str, instruction: str):
    """-> (reasoning, final_answer) or None when the row can't be split cleanly."""
    lines = [strip_markup(l) for l in norm(answer).split("\n")]
    lines = [l for l in lines if l]
    if len(lines) < 2:
        return None

    # Drop trailing chatter first — it belongs to neither part.
    while lines and is_fluff(lines[-1]):
        lines.pop()
    if len(lines) < 2:
        return None

    # A leading restatement of the question adds nothing to the derivation.
    if lines and lines[0][:40] == instruction.strip()[:40]:
        lines.pop(0)

    idx = next((i for i in range(len(lines) - 1, 0, -1) if is_conclusion(lines[i])), None)
    if idx is None:
        return None
    reasoning = "\n".join(lines[:idx])
    tail = lines[idx:]
    cut = next((i for i in range(1, len(tail)) if NEXT_PROBLEM_RE.match(tail[i])), len(tail))
    final = "\n".join(tail[:cut])

    if len(reasoning) < MIN_THINK_CHARS or len(final) < MIN_ANSWER_CHARS:
        return None
    # A "conclusion" longer than the derivation means the split went the wrong way.
    if len(final) > len(reasoning):
        return None
    return reasoning, final


def main():
    table = pq.read_table(SRC).to_pydict()
    rows  = list(zip(table["instruction"], table["answer"]))
    stats = collections.Counter(total=len(rows))

    seen, samples = set(), []
    for raw_ins, raw_ans in rows:
        ins = clean_instruction(raw_ins)
        key = re.sub(r"\W+", "", ins)
        if key in seen:
            stats["dropped_duplicate"] += 1
            continue
        seen.add(key)

        split = split_answer(raw_ans, ins)
        if split is None:
            stats["dropped_no_conclusion"] += 1
            continue
        reasoning, final = split
        stats["kept"] += 1
        samples.append({"instruction": ins, "reasoning": reasoning, "answer": final})

    random.Random(SEED).shuffle(samples)
    eval_set, train_set = samples[:EVAL_N], samples[EVAL_N:]

    OUT_DIR.mkdir(exist_ok=True)
    for name, split in (("train", train_set), ("eval", eval_set)):
        with open(OUT_DIR / f"{name}.jsonl", "w", encoding="utf-8") as fh:
            for s in split:
                fh.write(json.dumps(s, ensure_ascii=False) + "\n")
        print(f"[+] {name}: {len(split)} samples -> {OUT_DIR / f'{name}.jsonl'}")

    print("[*] stats:", dict(stats))


if __name__ == "__main__":
    main()