"""
Faz 7 ön-hazırlık — RAG-aware SFT verisi (extractive QA → "bağlamdan cevapla").

Amaç: v1-instruct'a "### Girdi:'deki bağlamdan cevabı çıkar (ezberden değil); bağlamda
yoksa 'bilmiyorum' de" davranışını öğret. A/B testinde 177M'in bağlamı yok sayıp
halüsinasyon ürettiği görüldü → bu, prompt'la değil SFT ile kapanır.

Kaynaklar (extractive, cevap bağlamda span):
  EN = rajpurkar/squad_v2          (answerable + unanswerable=abstain HAZIR)
  TR = ucsahin/TR-Extractive-QA-82K (answer string, context'te span; abstain YOK → sentetik)

Çıktı: {instruction, input, output} JSONL → faz6_sft.py --data'ya verilir.
  instruction = faz7_rag.build_rag_prompt'un talimatıyla BİREBİR (train=inference hizası):
                "<bağlamı kullan, kısa/doğrudan, yoksa bilmiyorum>\nSoru: {q}"
  input       = bağlam (context)
  output      = bağlamdaki cevap span'i  |  "bilmiyorum"/"I don't know" (abstain)

Çalıştırma (Colab/yerel; datasets + sentencepiece kurulu, HF login):
  pip install "datasets>=2.18" sentencepiece
  HF_TOKEN=hf_xxx python faz7_prep_rag_sft.py --out rag_sft.jsonl --n_en 15000 --n_tr 15000
"""
import os, sys, json, re, random, argparse

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


# ───────────── saf-mantık (yerelde gerçek tokenizer'la test edilebilir) ─────────────
# TR tespiti: özel karakter VEYA TR soru-kelimesi (çoğu TR soru "nedir/hangi/neresi" gibi
# kelime taşır ama çğıöşü taşımayabilir — örn "Fotosentez nedir?"). faz7_rag ile AYNI tutulur.
_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ı"}


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 build_instr(question):
    # faz7_rag.build_rag_prompt'un talimat metniyle BİREBİR aynı (dil-duyarlı).
    if is_tr(question):
        return ("Aşağıdaki bağlamı kullanarak soruyu kısa ve doğrudan yanıtla; adım adım düşünme. "
                "Cevap bağlamda yoksa \"bilmiyorum\" de.\nSoru: " + question)
    return ("Answer the question using only the context below, briefly and directly; do not think step by step. "
            "If the answer is not in the context, say \"I don't know\".\nQuestion: " + question)


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 tok_len(sp, instr, out, inp=""):
    """faz6_sft ile aynı tokenizasyon: prompt(+input) + yanıt + eos."""
    return len(sp.encode(build_prompt(instr, inp) + out.strip(), out_type=int)) + 1


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


def make_row(sp, question, context, answer, max_len, sentence=False):
    """{instruction,input,output}; max_len aşılırsa None (KESME YOK — faz6 kalıbı).
    sentence=True → output = span'in geçtiği tam cümle (extractive cümle-düzeyi)."""
    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
    instr = build_instr(q)
    if tok_len(sp, instr, out, ctx) > max_len:
        return None
    return {"instruction": instr, "input": ctx, "output": out}


def synth_tr_abstain(sp, rows, k, max_len, rng):
    """Alakasız soru↔bağlam eşle → 'bilmiyorum' (TR abstain; ucsahin'de doğal abstain yok)."""
    out, n = [], len(rows)
    if n < 2 or k <= 0:
        return out
    for _ in range(k * 4):                       # filtre kaybı için fazladan dene
        if len(out) >= k:
            break
        i, j = rng.randrange(n), rng.randrange(n)
        if i == j:
            continue
        q = rows[i][0]; ctx = rows[j][1]
        if rows[i][2] and rows[i][2] 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 (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_en(sp, args):
    """squad_v2: answerable → span; unanswerable (text=[]) → 'I don't know'."""
    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)
    print(f"[en] {EN_REPO}: answerable={len(pos)} abstain={len(neg)}", flush=True)
    return pos, neg


def gather_tr(sp, args):
    """ucsahin TR-82K: answer span. (abstain doğal yok → sentetik üretilir)"""
    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))
    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["instruction"], r["output"], r.get("input", "")) for r in sample)
    ab = sum(1 for r in rows if r["output"] in ABSTAIN.values())
    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 örnek: abstain_frac kadar negatif (abstain), gerisi pozitif."""
    n_neg = min(len(neg), int(n * abstain_frac))
    n_pos = n - n_neg
    return pos[:n_pos] + neg[:n_neg]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", default="rag_sft.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.2, help="abstain (bilmiyorum) örnek oranı")
    ap.add_argument("--answer_mode", default="span", choices=["span", "sentence"],
                    help="span=kısa cevap | sentence=span'in geçtiği TAM cümle (daha doğal cevap)")
    ap.add_argument("--no_tr_abstain", action="store_true", help="TR sentetik abstain üretme")
    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, 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)
    tr_neg = [] if args.no_tr_abstain else synth_tr_abstain(
        sp, tr_rows, int(args.n_tr * args.abstain_frac), args.max_len, rng)
    tr = take_balanced(tr_pos, tr_neg, args.n_tr, 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()