"""
WritingPrompts data prep. Deterministic and seeded: every arm sees the exact
same train split and the exact same 50 held-out eval prompts.

On the SYSTEM PROMPT (rewritten from the prior run)
---------------------------------------------------
The prior run used, with a ~512-token cap and no length guidance:

    "Write a short creative story based on this prompt:\\n\\n{text}"

Result: 100% of eval samples truncated mid-word. The model had no budget to
plan an ending, so it never planned one. Since the judge rubric caps quality at
<=3 for unfinished stories, quality became near-constant within every group ->
zero GRPO advantage -> the quality signal was dead for the whole run.

The fix is a word budget + an explicit ending instruction + a raised token cap.

What the system prompt deliberately does NOT do: it says nothing about being
original, varied, surprising, or avoiding cliche. That would be a confound.
Mode collapse is the dependent variable; instructing the model to diversify
would mask exactly the effect every arm is being measured on. The prompt is
identical and neutral for base, E0, E1, E2, E3 and E4 -- all differences
between arms must come from the training objective, not the prompt.
"""

from __future__ import annotations

import hashlib
import json
import re
from pathlib import Path

from datasets import load_dataset

SYSTEM_PROMPT = (
    "You are a fiction writer. Write a complete short story of 200-500 words "
    "responding to the writing prompt.\n"
    "Write only the story: no title, no preamble, no commentary, no author's note.\n"
    "Finish inside the word budget. The story must reach a real ending, not stop mid-scene."
)

MIN_PROMPT_WORDS = 10
MAX_PROMPT_WORDS = 60
SEED = 42

_TAG = re.compile(r"^\s*\[\s*(WP|EU|CW|TT|RF|IP|PI|PM|MP|OT|SP|FF)\s*\]\s*", re.I)
_WS = re.compile(r"\s+")


def clean_prompt(text: str) -> str:
    """Strip the subreddit tag and normalize the mangled WritingPrompts spacing.

    The euclaise mirror is tokenizer-detokenized: "did n't phase us ." etc.
    Left as-is this leaks a distinctive artifact into every prompt and wastes
    judge tokens, so we repair the obvious cases.
    """
    t = _TAG.sub("", text or "").strip()
    t = t.replace("`` ", '"').replace(" ''", '"').replace("''", '"')
    t = re.sub(r"([“‘(\[])\s+", r"\1", t)      # "“ What" -> "“What"
    t = re.sub(r"\s+([”’)\]])", r"\1", t)      # "cockroach! ”" -> "cockroach!”"
    t = re.sub(r"\s+([,.!?;:])", r"\1", t)
    t = re.sub(r"\bn't\b", "n't", t)
    t = re.sub(r"\s+n't", "n't", t)
    t = re.sub(r"\s+'(s|re|ve|ll|d|m)\b", r"'\1", t)
    t = re.sub(r"<\s*newline\s*>", " ", t, flags=re.I)
    t = _WS.sub(" ", t).strip()
    return t


def _key(t: str) -> str:
    """Dedupe key: lowercase alphanumerics only, so punctuation/spacing
    variants of the same prompt collapse together."""
    return hashlib.sha1(re.sub(r"[^a-z0-9]", "", t.lower()).encode()).hexdigest()


def build_splits(
    n_train: int = 1000,
    n_eval: int = 50,
    out_dir: str | Path = "data",
    seed: int = SEED,
) -> dict:
    ds = load_dataset("euclaise/writingprompts", split="train")
    col = "prompt" if "prompt" in ds.column_names else ds.column_names[0]

    seen: set[str] = set()
    kept: list[str] = []
    for rec in ds:
        t = clean_prompt(rec[col])
        n = len(t.split())
        if not (MIN_PROMPT_WORDS <= n <= MAX_PROMPT_WORDS):
            continue
        k = _key(t)
        if k in seen:
            continue
        seen.add(k)
        kept.append(t)
        if len(kept) >= (n_train + n_eval) * 3:   # oversample, then shuffle
            break

    import random
    random.Random(seed).shuffle(kept)
    need = n_train + n_eval
    if len(kept) < need:
        raise RuntimeError(f"only {len(kept)} usable prompts, need {need}")
    sel = kept[:need]
    eval_prompts, train_prompts = sel[:n_eval], sel[n_eval:need]

    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    for name, rows in (("train", train_prompts), ("eval", eval_prompts)):
        p = out / f"{name}_prompts.jsonl"
        with open(p, "w") as f:
            for i, t in enumerate(rows):
                f.write(json.dumps({"id": f"{name}-{i:04d}", "prompt": t}) + "\n")

    meta = {
        "seed": seed, "n_train": len(train_prompts), "n_eval": len(eval_prompts),
        "total_scanned": len(ds), "usable_after_filter": len(kept),
        "min_words": MIN_PROMPT_WORDS, "max_words": MAX_PROMPT_WORDS,
        "system_prompt": SYSTEM_PROMPT,
    }
    (out / "split_meta.json").write_text(json.dumps(meta, indent=2))
    return meta


def load_prompts(split: str, out_dir: str | Path = "data") -> list[dict]:
    p = Path(out_dir) / f"{split}_prompts.jsonl"
    return [json.loads(l) for l in open(p) if l.strip()]


def chat_messages(prompt: str) -> list[dict]:
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt},
    ]


if __name__ == "__main__":
    m = build_splits()
    print(json.dumps(m, indent=2))
    for split in ("train", "eval"):
        rows = load_prompts(split)
        print(f"\n{split}: {len(rows)}")
        for r in rows[:3]:
            print("   ", r["id"], "|", r["prompt"][:110])