#!/usr/bin/env python3
"""Build chat SFT JSONL from scriptlib (+ optional TVTropes) for H3 script-LoRA.

Reads numbered screenplays under ../scriptlib/*.txt, splits them into short
scene-ish chunks, and emits two kinds of training rows:

1. **format** — classic slugline/action/dialogue → MiniMax FL2VA scene block
2. **premise** — short premise derived from the chunk → full H3 scene beat

Optionally samples TVTropes titles/tropes as extra premise seeds (no script body).

Output: train_dataset.full.jsonl (append or overwrite).

Example:
  python build_sft_from_scriptlib.py --max-scripts 102 --chunks-per-script 4
"""

from __future__ import annotations

import argparse
import json
import random
import re
from pathlib import Path

ROOT = Path(__file__).resolve().parent
SCRIPTLIB = ROOT.parent / "scriptlib"
TROPES_TV = ROOT.parent / "TVTropesData" / "tv_tropes.csv"
TROPES_MASTER = ROOT.parent / "TVTropesData" / "tropes.csv"
OUT_DEFAULT = ROOT / "train_dataset.full.jsonl"

H3_SYSTEM = (
    "You write ONE MiniMax-H3 FL2VA scene beat for Backlot. "
    "Output fields: ACTION, SHOT, STORYBOARD_PROMPT, H3_MODE, H3_VIDEO_PROMPT "
    "(alignment line + integrated_multimodal_description + overall_soundscape + "
    "non_diegetic_music), LORA, AUDIO, DURATION 5. Prefer a single continuous shot "
    "from Picture 1 to Picture 2. Dialogue only inside <d>[Language] ...</d>."
)

SLUGLINE_RE = re.compile(
    r"^\s*(INT\.|EXT\.|INT/EXT\.|I/E\.|EST\.)\s+.+$",
    re.I | re.M,
)
CHAR_RE = re.compile(r"^\s{10,}([A-Z][A-Z0-9 .'\-]{1,40})\s*(\(.*\))?\s*$")


def split_scenes(text: str, max_chars: int = 1800) -> list[str]:
    """Split screenplay into chunks on sluglines, then size-cap."""
    text = text.replace("\r\n", "\n").replace("\r", "\n")
    # Prefer slugline splits
    parts: list[str] = []
    matches = list(SLUGLINE_RE.finditer(text))
    if matches:
        for i, m in enumerate(matches):
            start = m.start()
            end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
            chunk = text[start:end].strip()
            if len(chunk) > 80:
                parts.append(chunk)
    else:
        # Seinfeld-style parenthetical locations: (Comedy club)
        blocks = re.split(r"\n\s*\([^)\n]{3,80}\)\s*\n", text)
        parts = [b.strip() for b in blocks if len(b.strip()) > 120]

    # Size-cap / merge
    out: list[str] = []
    for p in parts:
        if len(p) <= max_chars:
            out.append(p)
        else:
            # take head of long scene (opening beat)
            out.append(p[:max_chars].rsplit("\n", 1)[0])
    return out


def extract_title(text: str) -> str:
    for line in text.splitlines()[:25]:
        s = line.strip()
        if len(s) > 2 and s.isupper() and not s.startswith("WRITTEN"):
            return s.title()
    return "Untitled"


def extract_dialogue_pairs(chunk: str, limit: int = 2) -> list[tuple[str, str]]:
    lines = chunk.splitlines()
    pairs: list[tuple[str, str]] = []
    i = 0
    while i < len(lines) and len(pairs) < limit:
        m = CHAR_RE.match(lines[i])
        if not m:
            i += 1
            continue
        name = m.group(1).strip().title()
        i += 1
        dial: list[str] = []
        while i < len(lines):
            L = lines[i]
            if CHAR_RE.match(L) or SLUGLINE_RE.match(L):
                break
            if L.strip().startswith("(") and L.strip().endswith(")"):
                i += 1
                continue
            if L.strip():
                dial.append(L.strip())
            elif dial:
                break
            i += 1
        if dial:
            pairs.append((name, " ".join(dial)))
    return pairs


def extract_action_lines(chunk: str, max_sents: int = 2) -> str:
    acts: list[str] = []
    for line in chunk.splitlines():
        s = line.strip()
        if not s or CHAR_RE.match(line) or SLUGLINE_RE.match(line):
            continue
        if s.startswith("(") and s.endswith(")"):
            continue
        # skip all-caps character names
        if s.isupper() and len(s) < 40:
            continue
        if len(s) > 20:
            acts.append(s)
        if len(acts) >= max_sents:
            break
    return " ".join(acts) if acts else "The scene plays out continuously."


def slugline_from_chunk(chunk: str) -> str:
    m = SLUGLINE_RE.search(chunk)
    if m:
        return m.group(0).strip().upper()
    # parenthetical location
    m2 = re.search(r"\(([^)\n]{3,60})\)", chunk)
    if m2:
        return f"INT. {m2.group(1).upper()}"
    return "INT. LOCATION - DAY"


def chunk_to_h3_assistant(chunk: str, title: str) -> str:
    slug = slugline_from_chunk(chunk)
    action = extract_action_lines(chunk)
    dialogue = extract_dialogue_pairs(chunk, limit=1)
    # Compress action for ~5s beat
    action_short = action
    if len(action_short) > 280:
        action_short = action_short[:280].rsplit(" ", 1)[0] + "."

    dial_line = ""
    multimodal_extra = ""
    if dialogue:
        name, line = dialogue[0]
        # keep dialogue short
        if len(line) > 160:
            line = line[:160].rsplit(" ", 1)[0] + "..."
        dial_line = f"DIALOGUE — {name}: {line}\n"
        multimodal_extra = (
            f" {name} (S1) says: <d>[English] {line}</d>"
        )

    storyboard = (
        f"cinematic still from {title}: {action_short[:200]}, "
        f"detailed environment, live-action, film lighting"
    )

    body = (
        f"[Shot 1] Live-action, cinematic, medium shot establishing the scene. "
        f"{action_short} The camera pushes in with small amplitude at slow speed."
        f"{multimodal_extra} "
        f"The framing begins on the composition of Picture 1 and continuously "
        f"evolves until it settles into the composition of Picture 2."
    )

    return (
        f"## SCENE 1 — {slug}\n"
        f"ACTION: {action_short}\n"
        f"{dial_line}"
        f"SHOT: medium shot, push in with small amplitude at slow speed\n"
        f"STORYBOARD_PROMPT: {storyboard}\n"
        f"H3_MODE: FL2VA\n"
        f"H3_VIDEO_PROMPT:\n"
        f"How the reference pictures align with the target video — Picture 1 "
        f"(from Shot 1) aligns with the 0.00-second mark of the target video; "
        f"Picture 2 (from Shot 1) aligns with the 5.00-second mark of the target video.\n\n"
        f"integrated_multimodal_description: {body}\n\n"
        f"overall_soundscape: Room tone and soft environmental ambience matching the location.\n\n"
        f"non_diegetic_music: N/A\n"
        f"LORA: none\n"
        f"AUDIO: none\n"
        f"DURATION: 5"
    )


def row(messages: list[dict]) -> str:
    return json.dumps({"messages": messages}, ensure_ascii=False)


def build_from_scripts(
    script_dir: Path,
    *,
    max_scripts: int,
    chunks_per_script: int,
    rng: random.Random,
) -> list[str]:
    files = sorted(script_dir.glob("*.txt"), key=lambda p: int(p.stem) if p.stem.isdigit() else p.stem)
    if max_scripts > 0:
        files = files[:max_scripts]
    rows: list[str] = []
    for path in files:
        try:
            text = path.read_text(encoding="utf-8", errors="replace")
        except OSError:
            continue
        title = extract_title(text)
        chunks = split_scenes(text)
        if not chunks:
            continue
        rng.shuffle(chunks)
        for chunk in chunks[:chunks_per_script]:
            assistant = chunk_to_h3_assistant(chunk, title)
            # Format transfer: classic excerpt → H3 beat
            rows.append(
                row(
                    [
                        {"role": "system", "content": H3_SYSTEM},
                        {
                            "role": "user",
                            "content": (
                                f"Rewrite this screenplay beat as a single ~5s MiniMax-H3 "
                                f"FL2VA scene for Backlot (first storyboard panel → last panel).\n\n"
                                f"SOURCE TITLE: {title}\n\n"
                                f"SCREENPLAY EXCERPT:\n{chunk[:1600]}"
                            ),
                        },
                        {"role": "assistant", "content": assistant},
                    ]
                )
            )
            # Premise → scene (trope-style)
            premise = extract_action_lines(chunk, max_sents=1)
            rows.append(
                row(
                    [
                        {"role": "system", "content": H3_SYSTEM},
                        {
                            "role": "user",
                            "content": (
                                f"Premise: {premise}\n"
                                f"Setting: derived from {title}\n"
                                f"Tone: cinematic\n"
                                f"Write SCENE 1 now (H3 FL2VA, DURATION 5)."
                            ),
                        },
                        {"role": "assistant", "content": assistant},
                    ]
                )
            )
    return rows


def build_from_tropes(path: Path, *, n: int, rng: random.Random) -> list[str]:
    if not path.exists() or n <= 0:
        return []
    # Stream a sample of lines (file is large)
    import csv

    rows_out: list[str] = []
    with path.open(encoding="utf-8", errors="replace", newline="") as f:
        reader = csv.DictReader(f)
        # sample reservoir
        reservoir: list[dict] = []
        for i, rec in enumerate(reader):
            if i < 5000:
                reservoir.append(rec)
            else:
                j = rng.randint(0, i)
                if j < 5000:
                    reservoir[j] = rec
            if i > 200_000:  # don't scan entire multi-hundred-MB file
                break
    rng.shuffle(reservoir)
    for rec in reservoir[:n]:
        title = (rec.get("Title") or rec.get("title") or "Untitled").strip()
        trope = (rec.get("Trope") or rec.get("trope") or "PlotTwist").strip()
        example = (rec.get("Example") or rec.get("Description") or "").strip()
        if len(example) > 400:
            example = example[:400] + "..."
        premise = f"A scene in {title} illustrating the trope '{trope}'. {example}"
        # Lightweight target (model will learn shape from script rows primarily)
        assistant = (
            f"## SCENE 1 — INT. SETTING - DAY\n"
            f"ACTION: Characters enact a brief beat embodying {trope}.\n"
            f"SHOT: medium shot, static shot\n"
            f"STORYBOARD_PROMPT: cinematic still for {title}, {trope}, live-action\n"
            f"H3_MODE: FL2VA\n"
            f"H3_VIDEO_PROMPT:\n"
            f"How the reference pictures align with the target video — Picture 1 "
            f"(from Shot 1) aligns with the 0.00-second mark of the target video; "
            f"Picture 2 (from Shot 1) aligns with the 5.00-second mark of the target video.\n\n"
            f"integrated_multimodal_description: [Shot 1] Live-action, cinematic, a medium "
            f"shot introduces the situation for {trope}. The camera holds a static shot as "
            f"the beat resolves into the final composition of Picture 2.\n\n"
            f"overall_soundscape: Soft room tone.\n\n"
            f"non_diegetic_music: N/A\n"
            f"LORA: none\n"
            f"AUDIO: none\n"
            f"DURATION: 5"
        )
        rows_out.append(
            row(
                [
                    {"role": "system", "content": H3_SYSTEM},
                    {
                        "role": "user",
                        "content": f"Premise: {premise}\nWrite SCENE 1 now (H3 FL2VA, DURATION 5).",
                    },
                    {"role": "assistant", "content": assistant},
                ]
            )
        )
    return rows_out


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--scriptlib", type=Path, default=SCRIPTLIB)
    ap.add_argument("--out", type=Path, default=OUT_DEFAULT)
    ap.add_argument("--max-scripts", type=int, default=0, help="0 = all")
    ap.add_argument("--chunks-per-script", type=int, default=4)
    ap.add_argument("--tropes", type=int, default=80, help="extra TVTropes premise rows")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--include-seed", action="store_true", help="prepend train_dataset.jsonl")
    args = ap.parse_args()

    rng = random.Random(args.seed)
    if not args.scriptlib.is_dir():
        raise SystemExit(f"scriptlib not found: {args.scriptlib}")

    rows = build_from_scripts(
        args.scriptlib,
        max_scripts=args.max_scripts,
        chunks_per_script=args.chunks_per_script,
        rng=rng,
    )
    rows += build_from_tropes(TROPES_TV, n=args.tropes, rng=rng)

    seed_path = ROOT / "train_dataset.jsonl"
    if args.include_seed and seed_path.exists():
        seed_rows = [ln for ln in seed_path.read_text().splitlines() if ln.strip()]
        rows = seed_rows + rows

    rng.shuffle(rows)
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text("\n".join(rows) + "\n", encoding="utf-8")
    print(f"wrote {len(rows)} rows → {args.out}")


if __name__ == "__main__":
    main()