"""
Faz 6 — SFT (instruction fine-tuning) iskeleti — SmartCore V1.

Base modeli (kdirgul/smartcore-v1 son ckpt) instruction verisiyle fine-tune eder.
Veri: JSONL. İKİ format desteklenir (encode_any tespit eder):
  • tek-tur (v1): {"instruction": str, "input": str(ops), "output": str}  → ### Talimat/### Yanıt
  • çok-turlu (v1.5b): {"messages": [{"role","content"},...]}                → ### Sistem/### Kullanıcı/### Asistan
KRİTİK: loss SADECE asistan/yanıt tokenlerinde (+eos) — prompt/kullanıcı -100 ile maskelenir.
NOT: --max_len ≤ pretraining seq_len (2048) tutulmalı (GQA-RoPE extrapolation'ı önle).

Ortam: Colab GPU + resmî mamba-ssm (wheel-cache). Model tanımı gömülü (import bağımlılığı yok).
Yerelde test edilebilir (mamba-ssm'siz): build_prompt / encode_example / encode_messages / collate / lr_at.

Kullanım:
  HF_TOKEN=hf_xxx python faz6_sft.py --data sft.jsonl --epochs 3 --lr 2e-5 \
      --micro_batch 8 --grad_accum 4 --max_len 1024 --save_repo kdirgul/smartcore-v1
"""
import os, sys, json, math, time, random, argparse
import torch, torch.nn as nn, torch.nn.functional as F
from functools import partial

# Fork: Colab'da var; yerelde yoksa None → saf-python fonksiyonlar yine test edilebilir.
try:
    from mamba_ssm.modules.block import Block
    from mamba_ssm.modules.mamba3 import Mamba3
    from mamba_ssm.modules.mlp import GatedMLP
    from mamba_ssm.ops.triton.layer_norm import RMSNorm
    FORK = True
except Exception:
    Block = Mamba3 = GatedMLP = RMSNorm = None
    FORK = False


# ───────────── model (faz3_train.py ile birebir; forward -> logits) ─────────────
def _rms(x, w, eps=1e-5):
    return (x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)) * w


def _rot_half(x):
    a, b = x.chunk(2, -1)
    return torch.cat((-b, a), -1)


class GQAMixer(nn.Module):
    def __init__(self, dim, n_heads=12, n_kv=3, base=10000.0, layer_idx=None, device=None, dtype=None):
        super().__init__()
        self.nh, self.nkv, self.hd = n_heads, n_kv, dim // n_heads
        self.rep = n_heads // n_kv
        fk = {"device": device, "dtype": dtype}
        self.q_proj = nn.Linear(dim, n_heads * self.hd, bias=False, **fk)
        self.k_proj = nn.Linear(dim, n_kv * self.hd, bias=False, **fk)
        self.v_proj = nn.Linear(dim, n_kv * self.hd, bias=False, **fk)
        self.out_proj = nn.Linear(n_heads * self.hd, dim, bias=False, **fk)
        self.qn = nn.Parameter(torch.ones(self.hd, **fk))
        self.kn = nn.Parameter(torch.ones(self.hd, **fk))
        self.register_buffer(
            "inv", 1.0 / (base ** (torch.arange(0, self.hd, 2, device=device).float() / self.hd)),
            persistent=False)

    def _rope(self, x, T):
        f = torch.outer(torch.arange(T, device=x.device, dtype=torch.float32), self.inv)
        e = torch.cat((f, f), -1)
        return (x * e.cos()[None, None] + _rot_half(x) * e.sin()[None, None]).to(x.dtype)

    def forward(self, x, **kw):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.nh, self.hd).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.nkv, self.hd).transpose(1, 2)
        q = _rms(q.float(), self.qn.float()).to(x.dtype)
        k = _rms(k.float(), self.kn.float()).to(x.dtype)
        q, k = self._rope(q, T), self._rope(k, T)
        k = k.repeat_interleave(self.rep, 1)
        v = v.repeat_interleave(self.rep, 1)
        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out_proj(y.transpose(1, 2).contiguous().view(B, T, -1))


class HybridLM(nn.Module):
    def __init__(self, cfg, device=None, dtype=None):
        super().__init__()
        self.cfg = cfg
        self.vocab = cfg["vocab_size"]
        self.scaled_embed = cfg.get("scaled_embed", False)
        d = cfg["d_model"]
        self.embedding = nn.Embedding(self.vocab, d, device=device, dtype=dtype)
        self.layers = nn.ModuleList()
        for i in range(cfg["n_layers"]):
            is_attn = ((i + 1) % cfg["attn_every"] == 0) and i != 0 and i != cfg["n_layers"] - 1
            fk = {"device": device, "dtype": dtype}
            if is_attn:
                mixer_cls = partial(GQAMixer, n_heads=cfg["n_heads"], n_kv=cfg["n_kv_heads"], layer_idx=i, **fk)
            else:
                ssm = dict(d_state=cfg["d_state"], expand=cfg["expand"], headdim=cfg["head_dim"],
                           ngroups=cfg["ngroups"], rope_fraction=cfg["rope_fraction"],
                           is_outproj_norm=False, is_mimo=cfg["is_mimo"], mimo_rank=cfg["mimo_rank"],
                           chunk_size=cfg["chunk_size"])
                mixer_cls = partial(Mamba3, layer_idx=i, **ssm, **fk)
            blk = Block(d, mixer_cls,
                        partial(GatedMLP, hidden_features=cfg["d_intermediate"], out_features=d, **fk),
                        norm_cls=partial(RMSNorm, eps=1e-5, **fk), fused_add_norm=True, residual_in_fp32=True)
            blk.layer_idx = i
            self.layers.append(blk)
        self.norm_f = RMSNorm(d, eps=1e-5, device=device, dtype=dtype)
        self.lm_head = nn.Linear(d, self.vocab, bias=False, device=device, dtype=dtype)
        self.lm_head.weight = self.embedding.weight

    def forward(self, ids):
        h = self.embedding(ids)
        if self.scaled_embed:
            h = h * (self.cfg["d_model"] ** 0.5)
        res = None
        for l in self.layers:
            h, res = l(h, res)
        h = self.norm_f((h + res) if res is not None else h)
        return self.lm_head(h.to(self.lm_head.weight.dtype))


# ───────────── veri (saf-python: yerelde test edilebilir) ─────────────
def build_prompt(ex):
    """instruction (+ opsiyonel input) → prompt metni (### Yanıt:\\n ile biter)."""
    instr = ex["instruction"].strip()
    inp = (ex.get("input") 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 encode_example(sp, ex, max_len, eos_id):
    """(input_ids, labels) — labels: prompt -100 (maskeli), yanıt+eos öğrenilir.
    Çok uzunsa SOLDAN kırpılır (yanıt korunur)."""
    p_ids = sp.encode(build_prompt(ex), out_type=int)
    r_ids = sp.encode(ex["output"].strip(), out_type=int) + [eos_id]
    ids = p_ids + r_ids
    labels = [-100] * len(p_ids) + r_ids
    if len(ids) > max_len:                 # yanıtı koru: son max_len token
        ids, labels = ids[-max_len:], labels[-max_len:]
    return ids, labels


# çok-turlu şablon — faz6_prep_v15b.py render() ile BİREBİR AYNI olmalı (SP özel-token YOK)
SYS_PFX, USER_PFX, ASST_PFX = "### Sistem:\n", "### Kullanıcı:\n", "### Asistan:\n"
_PFX = {"system": SYS_PFX, "user": USER_PFX, "assistant": ASST_PFX}


def encode_messages(sp, ex, max_len, eos_id):
    """Çok-turlu {messages} → (ids, labels). SADECE asistan içeriği + her turun eos'u öğrenilir;
    sistem/kullanıcı turları + tüm prefix'ler + ayraçlar -100 (maskeli). Uzunsa SOLDAN kırpılır
    (son yanıt korunur). Şablon faz6_prep_v15b.render() ile aynı (segment-segment tokenize)."""
    msgs = ex["messages"]
    ids, labels = [], []
    for i, m in enumerate(msgs):
        role = m["role"]; content = (m.get("content") or "").strip()
        prefix = _PFX[role]
        sep = "\n\n" if i < len(msgs) - 1 else ""     # turlar \n\n ile birleşir (render ile aynı)
        if role == "assistant":
            p = sp.encode(prefix, out_type=int)                    # "### Asistan:\n" → maskeli (inference'ta biz veririz)
            c = sp.encode(content, out_type=int) + [eos_id]        # içerik + eos → ÖĞRENİLİR
            s = sp.encode(sep, out_type=int) if sep else []
            ids += p + c + s
            labels += [-100] * len(p) + c + [-100] * len(s)
        else:                                                      # system/user → tümü maskeli
            seg = sp.encode(prefix + content + sep, out_type=int)
            ids += seg
            labels += [-100] * len(seg)
    if len(ids) > max_len:
        ids, labels = ids[-max_len:], labels[-max_len:]           # son yanıtı koru
    return ids, labels


def encode_any(sp, ex, max_len, eos_id):
    """Format tespiti: {messages}→çok-turlu; {instruction,output}→tek-tur (v1 uyumlu)."""
    if ex.get("messages"):
        return encode_messages(sp, ex, max_len, eos_id)
    return encode_example(sp, ex, max_len, eos_id)


def collate(batch, pad_id):
    """Sağdan pad. input_ids pad_id ile, labels -100 ile doldurulur."""
    maxlen = max(len(ids) for ids, _ in batch)
    B = len(batch)
    input_ids = torch.full((B, maxlen), pad_id, dtype=torch.long)
    labels = torch.full((B, maxlen), -100, dtype=torch.long)
    for i, (ids, lab) in enumerate(batch):
        input_ids[i, :len(ids)] = torch.tensor(ids, dtype=torch.long)
        labels[i, :len(lab)] = torch.tensor(lab, dtype=torch.long)
    return input_ids, labels


def lr_at(step, total, peak, warmup, floor_ratio=0.1):
    if step < warmup:
        return peak * (step + 1) / max(1, warmup)
    prog = (step - warmup) / max(1, total - warmup)
    return floor_ratio * peak + 0.5 * (1 - floor_ratio) * peak * (1 + math.cos(math.pi * prog))


def masked_loss(logits, labels):
    """Kaydırmalı (next-token) CE; -100 maskeli. logits[:, :-1] -> labels[:, 1:]."""
    V = logits.size(-1)
    return F.cross_entropy(logits[:, :-1].reshape(-1, V).float(),
                           labels[:, 1:].reshape(-1), ignore_index=-100)


# ───────────── yükleme yardımcıları (Colab) ─────────────
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 base_ckpt(spec, token, subdir="checkpoints_350m"):
    if spec and spec != "latest_hf":
        if os.path.exists(spec):
            return spec
        from huggingface_hub import hf_hub_download   # HF yolu (örn sft/epoch_2/ckpt.pt) → devam SFT
        print(f"[base] HF: {spec}", flush=True)
        return hf_hub_download("kdirgul/smartcore-v1", spec, repo_type="model", token=token)
    from huggingface_hub import HfApi, hf_hub_download   # latest_hf → subdir'deki son pretrain ckpt (v1.5b=checkpoints_350m)
    api = HfApi(token=token)
    fs = [f for f in api.list_repo_files("kdirgul/smartcore-v1", repo_type="model")
          if f.startswith(f"{subdir}/step_") and f.endswith("ckpt.pt")]
    if not fs:
        sys.exit(f"[hata] '{subdir}/' altında pretrain ckpt yok. (v1.0 için --base_subdir checkpoints)")
    latest = max(fs); print(f"[base] {latest}", flush=True)
    return hf_hub_download("kdirgul/smartcore-v1", latest, repo_type="model", token=token)


def load_examples(path):
    rows = []
    with open(path, encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            ex = json.loads(line)
            if ex.get("messages") or (ex.get("instruction") and ex.get("output")):
                rows.append(ex)
    return rows


def load_replay(token, n, max_len, seed):
    """Unutma önleme (replay): pretraining shard'larından n adet düz-LM örneği çek.
    labels = ids (MASKE YOK → tüm tokenler öğrenilir, normal LM loss). Karışım için
    her kaynaktan 1 shard indirir (en_fineweb_edu / tr_fineweb2_hq / code / math)."""
    if n <= 0:
        return []
    import glob
    import pyarrow.parquet as pq
    from huggingface_hub import snapshot_download
    root = snapshot_download("kdirgul/smartcore-v1-data", repo_type="dataset", token=token,
                             allow_patterns=["*/shard_00000.parquet"])
    files = sorted(glob.glob(os.path.join(root, "*", "shard_00000.parquet")))
    if not files:
        print("[replay] UYARI: shard bulunamadı, replay atlandı", flush=True)
        return []
    rng = random.Random(seed); per = max(1, n // len(files)); out = []
    for f in files:
        col = pq.read_table(f, columns=["input_ids"]).column("input_ids")
        for i in rng.sample(range(len(col)), min(per, len(col))):
            ids = list(col[i].as_py())[:max_len]
            out.append((ids, list(ids)))     # maskesiz: labels = ids
    rng.shuffle(out)
    return out[:n]


# ───────────── eğitim ─────────────
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--data", required=True, help="instruction JSONL yolu")
    ap.add_argument("--base", default="latest_hf", help="latest_hf | HF yolu (sft/epoch_2/ckpt.pt) | yerel .pt")
    ap.add_argument("--base_subdir", default="checkpoints_350m",
                    help="latest_hf namespace (v1.5b=checkpoints_350m, v1.0=checkpoints)")
    ap.add_argument("--epochs", type=int, default=3)
    ap.add_argument("--lr", type=float, default=2e-5)
    ap.add_argument("--micro_batch", type=int, default=8)
    ap.add_argument("--grad_accum", type=int, default=4)
    ap.add_argument("--max_len", type=int, default=1024)
    ap.add_argument("--warmup_ratio", type=float, default=0.03)
    ap.add_argument("--val_frac", type=float, default=0.02)
    ap.add_argument("--replay_frac", type=float, default=0.2,
                    help="unutmayı önlemek için SFT'ye karışacak pretraining (düz-LM) örnek oranı (0=kapalı)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--save_repo", default="kdirgul/smartcore-v1")
    ap.add_argument("--save_dir", default="/content/sft")
    ap.add_argument("--save_subdir", default="sft",
                    help="HF çıktı alt-dizini; RAG-SFT için 'sft_rag' ver → genel SFT epoch'larını EZME")
    ap.add_argument("--log_every", type=int, default=10)
    args = ap.parse_args()

    assert FORK, "mamba-og fork yok — önce wheel kurulum hücresini çalıştır."
    assert torch.cuda.is_available(), "CUDA yok (Colab GPU gerekir)."
    dev = "cuda"
    torch.manual_seed(args.seed); random.seed(args.seed)
    torch.set_float32_matmul_precision("high")
    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); eos_id = sp.eos_id(); pad_id = max(sp.pad_id(), 0)
    st = torch.load(base_ckpt(args.base, token, args.base_subdir), map_location="cpu", weights_only=False)
    cfg = st["cfg"]
    model = HybridLM(cfg, device=dev, dtype=torch.bfloat16)
    model.load_state_dict(st["model"], strict=False); model.train()
    print(f"[model] base step={st.get('step','?')} | {'MIMO' if cfg.get('is_mimo') else 'SISO'}", flush=True)

    rows = load_examples(args.data)
    random.shuffle(rows)
    n_val = max(1, int(len(rows) * args.val_frac)) if args.val_frac > 0 else 0
    val_rows, train_rows = rows[:n_val], rows[n_val:]
    enc = [encode_any(sp, ex, args.max_len, eos_id) for ex in train_rows]
    val_enc = [encode_any(sp, ex, args.max_len, eos_id) for ex in val_rows]
    n_multi = sum(1 for ex in train_rows if ex.get("messages"))
    print(f"[format] çok-turlu(messages)={n_multi} | tek-tur(instruction)={len(train_rows)-n_multi}", flush=True)
    n_sft = len(enc)
    if args.replay_frac > 0:
        n_replay = int(n_sft * args.replay_frac / (1 - args.replay_frac))
        rep = load_replay(token, n_replay, args.max_len, args.seed)
        enc = enc + rep
        random.shuffle(enc)
        print(f"[replay] {len(rep)} pretraining örneği karıştırıldı (hedef ~%{int(100*args.replay_frac)}) — unutma önleme", flush=True)
    print(f"[veri] SFT={n_sft} | toplam(train)={len(enc)} | val={len(val_enc)}", flush=True)

    decay = [p for p in model.parameters() if p.ndim >= 2]
    nod = [p for p in model.parameters() if p.ndim < 2]
    opt = torch.optim.AdamW([{"params": decay, "weight_decay": 0.0},
                             {"params": nod, "weight_decay": 0.0}],
                            lr=args.lr, betas=(0.9, 0.95), eps=1e-8, fused=True)

    eff = args.micro_batch * args.grad_accum
    steps_per_epoch = max(1, len(enc) // eff)
    total_steps = steps_per_epoch * args.epochs
    warmup = max(1, int(total_steps * args.warmup_ratio))
    print(f"[plan] {len(enc)} örnek | eff_batch={eff} | {steps_per_epoch} step/epoch | "
          f"{total_steps} step | warmup {warmup} | lr {args.lr}", flush=True)

    def val_loss():
        if not val_enc:
            return float("nan")
        model.eval(); tot = 0.0; nb = 0
        with torch.no_grad():
            for j in range(0, len(val_enc), args.micro_batch):
                ii, ll = collate(val_enc[j:j + args.micro_batch], pad_id)
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    tot += masked_loss(model(ii.to(dev)), ll.to(dev)).item(); nb += 1
        model.train(); return tot / max(1, nb)

    os.makedirs(args.save_dir, exist_ok=True)
    gstep = 0; t0 = time.perf_counter()
    for epoch in range(args.epochs):
        random.shuffle(enc)
        for s in range(steps_per_epoch):
            opt.zero_grad(set_to_none=True)
            lr = lr_at(gstep, total_steps, args.lr, warmup)
            for g in opt.param_groups:
                g["lr"] = lr
            loss_acc = 0.0
            base = s * eff
            for a in range(args.grad_accum):
                chunk = enc[base + a * args.micro_batch: base + (a + 1) * args.micro_batch]
                if not chunk:
                    continue
                ii, ll = collate(chunk, pad_id)
                with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
                    loss = masked_loss(model(ii.to(dev)), ll.to(dev))
                (loss / args.grad_accum).backward()
                loss_acc += loss.item() / args.grad_accum
            gn = torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step(); gstep += 1
            if gstep % args.log_every == 0:
                tps = (gstep * eff) / (time.perf_counter() - t0)
                print(f"e{epoch} step {gstep}/{total_steps} | loss {loss_acc:.4f} | gnorm {gn:5.2f} | "
                      f"lr {lr:.2e} | {tps:.1f} ex/s", flush=True)
        vl = val_loss()
        print(f"[epoch {epoch} bitti] val_loss={vl:.4f}", flush=True)
        # her epoch sonu kaydet + HF push
        d = os.path.join(args.save_dir, f"epoch_{epoch}")
        os.makedirs(d, exist_ok=True)
        torch.save({"model": model.state_dict(), "cfg": cfg, "epoch": epoch, "sft": True,
                    "val_loss": vl}, os.path.join(d, "ckpt.pt"))
        if token and args.save_repo:
            try:
                from huggingface_hub import HfApi
                HfApi(token=token).upload_folder(folder_path=d, repo_id=args.save_repo, repo_type="model",
                                                 path_in_repo=f"{args.save_subdir}/epoch_{epoch}",
                                                 commit_message=f"{args.save_subdir} epoch {epoch} val={vl:.3f}")
                print(f"[ckpt] HF push OK {args.save_subdir}/epoch_{epoch}", flush=True)
            except Exception as e:
                print(f"[ckpt] push HATA: {repr(e)[:160]}", flush=True)
    print("[bitti] SFT tamamlandı.", flush=True)


if __name__ == "__main__":
    main()