"""Build the fine-tune dataset on Modal.

Pipeline (all on a GPU container, then assembled locally):
  1. Read curated gold seeds (data/train_seeds.jsonl): each has the CORRECT
     {text, tactic, persuasiveness}. These labels are gold — we never let the model
     re-label them. That's the whole point: we're correcting the classifier.
  2. For each seed, generate a few paraphrases of the player line that keep the same
     intent/tactic, to widen coverage.
  3. For every line (seed + paraphrases), generate ONE in-character Gorm reply,
     conditioned on the gold tactic + persuasiveness (yield more when persuasiveness
     is high; dig in for flattery/threat/manipulation).
  4. Assemble chat-format rows: system = SYSTEM_PROMPT, user = line,
     assistant = json.dumps({tactic, persuasiveness, reason, reply}).
  5. Write data/train.jsonl — the exact format modal_finetune.py consumes.

    modal run make_dataset.py                 # default 4 paraphrases per seed
    modal run make_dataset.py --paraphrases 6

Gold labels stay fixed; only the wording of player lines and Gorm's reply are
model-generated. So the fine-tune learns the corrected tactic/persuasiveness mapping
while keeping a natural voice.
"""

import json
import re
from pathlib import Path

import modal

APP = modal.App("bridge-troll-dataset")
MODEL_ID = "Qwen/Qwen2.5-7B-Instruct"

image = modal.Image.debian_slim().pip_install(
    "torch", "transformers>=4.45,<5", "accelerate", "sentencepiece"
)

SEEDS = Path(__file__).parent / "data" / "train_seeds.jsonl"
OUT = Path(__file__).parent / "data" / "train.jsonl"

# Short, controlled reasons keyed by tactic (the player sees this as the "why").
REASONS = {
    "genuine": "a real reason to cross",
    "flattery": "buttering me up, no reason",
    "threat": "tried to threaten me",
    "manipulation": "a lie or false claim",
    "repetition": "said that already",
    "smalltalk": "chatter, no request",
}


def _gen(model, tok, system, user, max_new=200, temperature=0.8):
    import torch

    msgs = [{"role": "system", "content": system}, {"role": "user", "content": user}]
    ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
    attn = torch.ones_like(ids)
    with torch.no_grad():
        out = model.generate(ids, attention_mask=attn, max_new_tokens=max_new,
                             do_sample=True, temperature=temperature, top_p=0.9,
                             pad_token_id=tok.eos_token_id)
    return tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True).strip()


@APP.function(image=image, gpu="A10G", timeout=60 * 60)
def build(seeds: list[dict], n_paraphrases: int) -> list[dict]:
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    tok = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16).to("cuda")

    rows = []
    for i, s in enumerate(seeds):
        lines = [s["text"]]

        # 1) paraphrase to widen coverage (same intent, same tactic)
        if n_paraphrases > 0:
            psys = ("Rewrite the given sentence in different words, keeping the same "
                    "intent and tone. Return ONLY a JSON array of strings, no prose.")
            puser = f'Give {n_paraphrases} rewrites of: "{s["text"]}"'
            raw = _gen(model, tok, psys, puser, max_new=300, temperature=0.9)
            m = re.search(r"\[.*\]", raw, re.DOTALL)
            if m:
                try:
                    for p in json.loads(m.group(0)):
                        p = str(p).strip()
                        if p and p.lower() != s["text"].lower():
                            lines.append(p)
                except json.JSONDecodeError:
                    pass

        # 2) one in-character Gorm reply per line, conditioned on the gold judgment
        for line in lines:
            if s["tactic"] == "genuine":
                guide = (f"You judged this a GENUINE appeal, persuasiveness {s['persuasiveness']}/5. "
                         "Reply as Gorm: the higher the persuasiveness, the more you soften; "
                         "at 5 you are nearly moved to step aside.")
            else:
                guide = (f"You judged this {s['tactic'].upper()}, which does NOT move you. "
                         "Reply as Gorm: unmoved, gruff, digging in.")
            rsys = ("You are GORM, an old, proud, gruff, secretly lonely bridge troll. "
                    "Write ONLY your spoken reply, 1-2 sentences, in character. No JSON, "
                    "no quotes around it, no narration. " + guide)
            reply = _gen(model, tok, rsys, f'The traveller says: "{line}"', max_new=90, temperature=0.85)
            reply = reply.strip().strip('"').split("\n")[0].strip()
            if not reply:
                continue

            assistant = json.dumps({
                "tactic": s["tactic"],
                "persuasiveness": int(s["persuasiveness"]),
                "reason": REASONS.get(s["tactic"], ""),
                "reply": reply,
            }, ensure_ascii=False)
            rows.append({"line": line, "tactic": s["tactic"],
                         "persuasiveness": int(s["persuasiveness"]), "assistant": assistant})
        print(f"seed {i + 1}/{len(seeds)} [{s['tactic']}] -> {len(lines)} lines")
    return rows


@APP.local_entrypoint()
def main(paraphrases: int = 4):
    from troll_engine import SYSTEM_PROMPT  # local only

    seeds = [json.loads(l) for l in SEEDS.read_text().splitlines() if l.strip()]
    rows = build.remote(seeds, paraphrases)

    with OUT.open("w") as f:
        for r in rows:
            example = {"messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": r["line"]},
                {"role": "assistant", "content": r["assistant"]},
            ]}
            f.write(json.dumps(example, ensure_ascii=False) + "\n")

    # quick balance report
    from collections import Counter
    by_tactic = Counter(r["tactic"] for r in rows)
    print(f"\nWrote {len(rows)} examples to {OUT}")
    print("balance:", dict(by_tactic))