"""
Faz 7 v1.5b — RAG-aware SFT verisi, çok-turlu {messages} chat formatı (kısa-extractive + abstain).

v1 (faz7_prep_rag_sft.py = {instruction,input,output}) KORUNUR; bu YENİ üretici faz6_sft.encode_messages
ile uyumlu {messages} verir. RAG chat'e doğal oturur: system=direktif, user=bağlam+soru, assistant=kısa cevap.
KARAR: RAG-SFT KISA-EXTRACTIVE kalır (sohbet SFT'sinden AYRI; factuality > akıcılık) + abstain ("bilmiyorum").

Kaynaklar (extractive, cevap bağlamda span — şema DOĞRULANDI v1'de):
  EN = rajpurkar/squad_v2           (answerable + unanswerable=abstain HAZIR)
  TR = ucsahin/TR-Extractive-QA-82K (answer span; abstain YOK → sentetik alakasız-eşleme)

Çıktı: {"messages":[system,user,assistant], "lang", "src"} JSONL → faz6_sft.py --data (encode_messages).
Şablon faz6_prep_v15b.render() (### Sistem/### Kullanıcı/### Asistan) — tek-kaynak import.
⚠️ faz7_rag.py inference AYNI SYS/user_msg/is_tr'yi kullanmalı (train=inference hizası) → oradan import edilecek.

Çalıştırma (Colab/yerel; datasets+sentencepiece kurulu, HF login):
  HF_TOKEN=hf_xxx python faz7_prep_rag_v15b.py --out rag_v15b.jsonl --n_en 15000 --n_tr 15000 --answer_mode sentence
"""
import os, sys, json, re, random, argparse
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from faz6_prep_v15b import render, valid, load_tok   # tek-kaynak şablon + tokenizer

EN_REPO = "rajpurkar/squad_v2"
TR_REPO = "ucsahin/TR-Extractive-QA-82K"
ABSTAIN = {"tr": "bilmiyorum", "en": "I don't know"}

# TR tespiti: özel karakter VEYA TR soru-kelimesi (faz7_rag ile AYNI tutulur — inference hizası)
_TR_WORDS = {"nedir", "ne", "hangi", "neresi", "nerede", "nereye", "kim", "kimdir", "neden",
             "niçin", "niye", "mıdır", "midir", "mudur", "müdür", "kaç", "kaçtır", "nasıl",
             "mı", "mi", "mu", "mü", "için", "ile", "kaçıncı"}

# RAG system direktifi (dil-duyarlı) — inference'ta faz7_rag AYNISINI system mesajı yapmalı
SYS = {
    "tr": ("Sen bir RAG asistanısın. SADECE aşağıdaki bağlamı kullanarak soruyu KISA ve doğrudan yanıtla; "
           "ezbere/dışarıdan bilgi ekleme, adım adım düşünme. Cevap bağlamda yoksa \"bilmiyorum\" de."),
    "en": ("You are a RAG assistant. Answer the question using ONLY the context below, briefly and directly; "
           "do not add outside knowledge or think step by step. If the answer is not in the context, say \"I don't know\"."),
}


# ───────────── saf-mantık (yerelde gerçek tokenizer'la test edilebilir) ─────────────
def is_tr(question):
    ql = question.lower()
    if any(ch in ql for ch in "çğıöşü"):
        return True
    words = set(ql.replace("?", " ").replace("'", " ").replace(".", " ").split())
    return bool(words & _TR_WORDS)


def user_msg(question, context):
    q, ctx = question.strip(), context.strip()
    if is_tr(question):
        return f"Bağlam:\n{ctx}\n\nSoru: {q}"
    return f"Context:\n{ctx}\n\nQuestion: {q}"


def sentence_with(context, answer):
    """answer span'inin geçtiği TAM cümleyi döndür (cümle-düzeyi extractive). Bulamazsa span."""
    for s in re.split(r"(?<=[.!?])\s+", context.strip()):
        if answer in s:
            return s.strip()
    return answer


def tok_len(sp, messages):
    return len(sp.encode(render(messages), out_type=int)) + 1


def make_row(sp, question, context, answer, max_len, sentence=False):
    """{messages:[system,user,assistant]}; max_len aşılırsa None (KESME YOK). Cevap KISA (span/cümle/abstain)."""
    q, ctx, a = question.strip(), context.strip(), answer.strip()
    if not (q and ctx and a):
        return None
    out = sentence_with(ctx, a) if sentence else a
    lang = "tr" if is_tr(q) else "en"
    messages = [
        {"role": "system", "content": SYS[lang]},
        {"role": "user", "content": user_msg(q, ctx)},
        {"role": "assistant", "content": out},
    ]
    if not valid(messages) or tok_len(sp, messages) > max_len:
        return None
    return {"messages": messages, "lang": lang, "src": "rag"}


def is_abstain(row):
    return row["messages"][-1]["content"].strip() in ABSTAIN.values()


def synth_tr_abstain(sp, rows, k, max_len, rng):
    """Alakasız soru↔bağlam eşle → 'bilmiyorum' (TR abstain; ucsahin'de doğal yok)."""
    out, n = [], len(rows)
    if n < 2 or k <= 0:
        return out
    for _ in range(k * 4):
        if len(out) >= k:
            break
        i, j = rng.randrange(n), rng.randrange(n)
        if i == j:
            continue
        q, ans_i = rows[i][0], rows[i][2]
        ctx = rows[j][1]
        if ans_i and ans_i in ctx:        # cevap tesadüfen bu bağlamda → mismatch değil, atla
            continue
        r = make_row(sp, q, ctx, ABSTAIN["tr"], max_len)
        if r:
            out.append(r)
    return out[:k]


# ───────────── yükleyiciler ─────────────
def gather_en(sp, args):
    from datasets import load_dataset
    ds = load_dataset(EN_REPO, split="train")
    sent = args.answer_mode == "sentence"
    pos, neg = [], []
    for ex in ds:
        texts = (ex.get("answers") or {}).get("text") or []
        if texts:
            r = make_row(sp, ex.get("question", ""), ex.get("context", ""), texts[0], args.max_len, sentence=sent)
            if r:
                pos.append(r)
        else:
            r = make_row(sp, ex.get("question", ""), ex.get("context", ""), ABSTAIN["en"], args.max_len)
            if r:
                neg.append(r)
        if args.cap_src and len(pos) >= args.cap_src and len(neg) >= args.cap_src:
            break
    print(f"[en] {EN_REPO}: answerable={len(pos)} abstain={len(neg)}", flush=True)
    return pos, neg


def gather_tr(sp, args):
    from datasets import load_dataset
    ds = load_dataset(TR_REPO, split="train")
    sent = args.answer_mode == "sentence"
    pos, rows = [], []
    for ex in ds:
        q = (ex.get("question") or "").strip(); ctx = (ex.get("context") or "").strip()
        a = (ex.get("answer") or "").strip()
        r = make_row(sp, q, ctx, a, args.max_len, sentence=sent)
        if r:
            pos.append(r); rows.append((q, ctx, a))
        if args.cap_src and len(pos) >= args.cap_src:
            break
    print(f"[tr] {TR_REPO}: answerable={len(pos)}", flush=True)
    return pos, rows


def stats(sp, rows, name):
    if not rows:
        print(f"[{name}] 0 örnek", flush=True); return
    sample = rows if len(rows) <= 3000 else random.sample(rows, 3000)
    ls = sorted(tok_len(sp, r["messages"]) for r in sample)
    ab = sum(1 for r in rows if is_abstain(r))
    print(f"[{name}] n={len(rows)} abstain={ab} | token: med={ls[len(ls)//2]} "
          f"p90={ls[int(len(ls)*0.9)]} max={ls[-1]}", flush=True)


def take_balanced(pos, neg, n, abstain_frac):
    n_neg = min(len(neg), int(n * abstain_frac))
    return pos[:n - n_neg] + neg[:n_neg]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="rag_v15b.jsonl")
    ap.add_argument("--max_len", type=int, default=2048)
    ap.add_argument("--n_en", type=int, default=15000)
    ap.add_argument("--n_tr", type=int, default=15000)
    ap.add_argument("--abstain_frac", type=float, default=0.15, help="abstain (bilmiyorum) örnek oranı")
    ap.add_argument("--answer_mode", default="sentence", choices=["span", "sentence"],
                    help="span=kısa | sentence=span'in geçtiği TAM cümle (v1'de sentence 6/6 kazandı)")
    ap.add_argument("--no_tr_abstain", action="store_true")
    ap.add_argument("--cap_src", type=int, default=0, help="smoke test: kaynak başına ~N örnek (0=üretim)")
    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 (squad_v2) ===", flush=True)
    en_pos, en_neg = gather_en(sp, args)
    rng.shuffle(en_pos); rng.shuffle(en_neg)
    en = take_balanced(en_pos, en_neg, args.n_en if not args.cap_src else min(args.n_en, args.cap_src), args.abstain_frac)

    print("=== TR (TR-82K) ===", flush=True)
    tr_pos, tr_rows = gather_tr(sp, args)
    rng.shuffle(tr_pos); rng.shuffle(tr_rows)
    k_neg = int((args.n_tr if not args.cap_src else args.cap_src) * args.abstain_frac)
    tr_neg = [] if args.no_tr_abstain else synth_tr_abstain(sp, tr_rows, k_neg, args.max_len, rng)
    tr = take_balanced(tr_pos, tr_neg, args.n_tr if not args.cap_src else min(args.n_tr, args.cap_src), args.abstain_frac)

    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)} örnek (EN {len(en)} + TR {len(tr)}) -> {args.out}", flush=True)


if __name__ == "__main__":
    main()