"""
Faz 8 ön-hazırlık — DPO tercih verisi (chosen/rejected çiftleri).

Ufuk 1 / Adım 2: v1-instruct-rag'i (sft_rag3/epoch_2) tercih hizalama ile parlat.
⚠️ 177M'de DPO kazancı marjinal + MC benchmark'ı değiştirmez; üretim tonu/formatını iyileştirir.

Kaynaklar (ikisi de {prompt, chosen[msg], rejected[msg]}):
  EN = HuggingFaceH4/ultrafeedback_binarized  (split train_prefs)
  TR = selimc/orpo-dpo-mix-TR-20k             (split train)

Çıktı: {instruction, chosen, rejected} JSONL → faz8_dpo.py --data.
  instruction = prompt (SFT şablonuyla sarılır), chosen/rejected = assistant yanıtları.

Çalıştırma (Colab/yerel; datasets + sentencepiece + HF login):
  HF_TOKEN=hf_xxx python faz8_prep_dpo.py --out dpo.jsonl --n_en 8000 --n_tr 8000
"""
import os, sys, json, random, argparse

EN_REPO = "HuggingFaceH4/ultrafeedback_binarized"
TR_REPO = "selimc/orpo-dpo-mix-TR-20k"


# ───────────── saf-mantık (yerelde gerçek tokenizer'la test edilebilir) ─────────────
def build_prompt(instr, inp=""):
    instr = instr.strip(); inp = (inp or "").strip()
    if inp:
        return f"### Talimat:\n{instr}\n\n### Girdi:\n{inp}\n\n### Yanıt:\n"
    return f"### Talimat:\n{instr}\n\n### Yanıt:\n"


def last_assistant(messages):
    """chosen/rejected mesaj listesinden son assistant içeriğini çıkar (str gelirse aynen)."""
    if isinstance(messages, str):
        return messages.strip()
    for m in reversed(messages or []):
        if isinstance(m, dict) and m.get("role") == "assistant":
            return (m.get("content") or "").strip()
    return ""


def tok_len(sp, instr, resp):
    """faz8_dpo ile aynı: prompt + yanıt + eos."""
    return len(sp.encode(build_prompt(instr) + resp.strip(), out_type=int)) + 1


def make_pair(sp, prompt, chosen, rejected, max_len):
    """{instruction,chosen,rejected}; ikisi de max_len'e sığmalı, chosen≠rejected."""
    p, c, r = prompt.strip(), chosen.strip(), rejected.strip()
    if not (p and c and r) or c == r:
        return None
    if tok_len(sp, p, c) > max_len or tok_len(sp, p, r) > max_len:
        return None
    return {"instruction": p, "chosen": c, "rejected": r}


# ───────────── yükleyiciler (datasets gerekir) ─────────────
def load_tok(token):
    import sentencepiece as spm
    from huggingface_hub import hf_hub_download
    p = hf_hub_download("kdirgul/smartcore-v1", "tokenizer/tokenizer.model", repo_type="model", token=token)
    return spm.SentencePieceProcessor(model_file=p)


def gather(sp, repo, split, max_len, want):
    from datasets import load_dataset
    ds = load_dataset(repo, split=split)
    out = []
    for ex in ds:
        prompt = (ex.get("prompt") or ex.get("question") or "").strip()
        row = make_pair(sp, prompt, last_assistant(ex.get("chosen")),
                        last_assistant(ex.get("rejected")), max_len)
        if row:
            out.append(row)
        if len(out) >= want:
            break
    print(f"[{repo.split('/')[-1]}] tutuldu {len(out)}", flush=True)
    return out


def stats(sp, rows, name):
    if not rows:
        print(f"[{name}] 0 örnek", flush=True); return
    sample = rows if len(rows) <= 2000 else random.sample(rows, 2000)
    lc = sorted(tok_len(sp, r["instruction"], r["chosen"]) for r in sample)
    lr = sorted(tok_len(sp, r["instruction"], r["rejected"]) for r in sample)
    print(f"[{name}] n={len(rows)} | chosen tok med={lc[len(lc)//2]} p90={lc[int(len(lc)*0.9)]} "
          f"| rejected med={lr[len(lr)//2]}", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="dpo.jsonl")
    ap.add_argument("--max_len", type=int, default=1024, help="prompt+yanıt token tavanı (DPO çift forward → kısa tut)")
    ap.add_argument("--n_en", type=int, default=8000)
    ap.add_argument("--n_tr", type=int, default=8000)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    token = os.environ.get("HF_TOKEN")
    try:
        from huggingface_hub import get_token
        token = token or get_token()
    except Exception:
        pass
    sp = load_tok(token)
    rng = random.Random(args.seed)

    print("=== EN (ultrafeedback) ===", flush=True)
    en = gather(sp, EN_REPO, "train_prefs", args.max_len, args.n_en)
    print("=== TR (orpo-dpo-mix-TR) ===", flush=True)
    tr = gather(sp, TR_REPO, "train", args.max_len, args.n_tr)

    stats(sp, en, "EN"); stats(sp, tr, "TR")
    data = en + tr; rng.shuffle(data)
    with open(args.out, "w", encoding="utf-8") as f:
        for r in data:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    stats(sp, data, "TOPLAM")
    print(f"\n[bitti] {len(data)} çift (EN {len(en)} + TR {len(tr)}) -> {args.out}", flush=True)


if __name__ == "__main__":
    main()