"""
V5-SFT Asama 2: ChatML jsonl → packed token bin + loss mask.

Format (her ornek):
    ### Sistem:
    {system}
    ### Kullanici:
    {user}
    ### Asistan:
    {assistant}<|endoftext|>

Loss mask:
    - Prompt kismi (Sistem + Kullanici + "### Asistan:\n" dahil) → -100 (ignore)
    - Asistan cevabi + <|endoftext|> → gercek token id (loss hesaplanir)

Cikti:
    data/sft/sft_train_tokens.bin     uint16, packed
    data/sft/sft_train_mask.bin       uint8, 1=loss var, 0=ignore
    data/sft/sft_val_tokens.bin
    data/sft/sft_val_mask.bin
    data/sft/sft_meta.json            istatistikler + format ornegi
"""

import argparse
import json
import random
import sys
from pathlib import Path

import numpy as np

try:
    from tokenizers import Tokenizer
    from tqdm import tqdm
except ImportError:
    print("! pip install tokenizers tqdm")
    sys.exit(1)


# Format delimiterlari
SYS_HEADER = "### Sistem:\n"
USER_HEADER = "### Kullanici:\n"
ASST_HEADER = "### Asistan:\n"
EOT_ID = 0  # <|endoftext|>
N_EOT_TRAILING = 4  # response sonuna kac EOT — model "STOP"u guclu ogrenir


def build_prompt_and_response(rec: dict) -> tuple[str, str]:
    """ChatML kayittan (messages list) prompt (mask=0) ve response (mask=1) ayir."""
    msgs = rec["messages"]
    sys_msg = ""
    user_msg = ""
    asst_msg = ""
    for m in msgs:
        role = m.get("role", "")
        content = m.get("content", "")
        if role == "system":
            sys_msg = content
        elif role == "user":
            user_msg = content
        elif role == "assistant":
            asst_msg = content

    if not user_msg or not asst_msg:
        return "", ""

    # Prompt — masked portion
    parts = []
    if sys_msg:
        parts.append(SYS_HEADER + sys_msg + "\n")
    parts.append(USER_HEADER + user_msg + "\n")
    parts.append(ASST_HEADER)
    prompt = "".join(parts)

    # Response — loss applied
    response = asst_msg  # endoftext ayrica eklenir

    return prompt, response


def tokenize_records(records: list, tok: Tokenizer, max_len: int,
                      pad_id: int = 0, show: bool = True):
    """Her kaydi tokenize et, prompt+response ayir, mask olustur.

    Return: (tokens_list, mask_list) — her ikisi list of np.ndarray (uint16 / uint8)
    """
    all_tokens = []
    all_masks = []
    dropped = 0
    long_truncated = 0

    iterator = tqdm(records, desc="  tokenize", disable=not show)
    for rec in iterator:
        prompt, response = build_prompt_and_response(rec)
        if not prompt or not response:
            dropped += 1
            continue

        prompt_ids = tok.encode(prompt).ids
        # Response sonuna N_EOT_TRAILING adet <|endoftext|> ekle
        # (model "STOP"u net ogrensin diye — tek EOT %0.17 ile cok zayif sinyal)
        response_ids = tok.encode(response).ids + [EOT_ID] * N_EOT_TRAILING

        total_len = len(prompt_ids) + len(response_ids)
        if total_len > max_len:
            # Cok uzun — response'tan kes (prompt korunur, en azindan kisa cevap olur)
            avail = max_len - len(prompt_ids)
            if avail < 32:
                # Prompt zaten cok uzun — at
                dropped += 1
                continue
            # Response'tan kes ama N_EOT_TRAILING EOT kalsin
            response_ids = response_ids[:avail - N_EOT_TRAILING] + [EOT_ID] * N_EOT_TRAILING
            long_truncated += 1

        tokens = np.array(prompt_ids + response_ids, dtype=np.uint16)
        mask = np.zeros(len(tokens), dtype=np.uint8)
        # response_ids icin loss aktif
        mask[len(prompt_ids):] = 1

        all_tokens.append(tokens)
        all_masks.append(mask)

    return all_tokens, all_masks, dropped, long_truncated


def pack_sequences(tokens_list: list, masks_list: list, block_size: int):
    """Liste of uint16 array'i back-to-back paketle.
    Document separator olarak <|endoftext|> her zaten son token, ek separator yok.

    Return: packed uint16 tokens + uint8 mask, ayni uzunlukta.
    """
    total_len = sum(len(t) for t in tokens_list)
    print(f"  Toplam token: {total_len:,}")

    packed_tokens = np.empty(total_len, dtype=np.uint16)
    packed_mask = np.empty(total_len, dtype=np.uint8)
    offset = 0
    for t, m in zip(tokens_list, masks_list):
        n = len(t)
        packed_tokens[offset:offset+n] = t
        packed_mask[offset:offset+n] = m
        offset += n
    assert offset == total_len
    return packed_tokens, packed_mask


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=str,
                        default="data/sft/01_collected.jsonl")
    parser.add_argument("--tokenizer", type=str,
                        default="data/tokenizer-tr-v5.json")
    parser.add_argument("--out-dir", type=str, default="data/sft")
    parser.add_argument("--val-frac", type=float, default=0.01,
                        help="Validation orani (default %1 = ~1.3K)")
    parser.add_argument("--max-len", type=int, default=2048,
                        help="Tek ornek max token (model block_size ile ayni)")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    print(f"Tokenizer: {args.tokenizer}")
    tok = Tokenizer.from_file(args.tokenizer)
    print(f"  Vocab: {tok.get_vocab_size():,}")
    print(f"  EOT (id={EOT_ID}): {tok.id_to_token(EOT_ID)!r}")

    # Yukle
    print(f"\nInput: {args.input}")
    records = []
    with open(args.input, encoding="utf-8") as f:
        for line in f:
            records.append(json.loads(line))
    print(f"  Yuklendi: {len(records):,} ornek")

    # Train/val split (kategori-stratified DEGIL — sadece random)
    random.seed(args.seed)
    random.shuffle(records)
    n_val = int(len(records) * args.val_frac)
    val_records = records[:n_val]
    train_records = records[n_val:]
    print(f"\nSplit: train {len(train_records):,}  val {len(val_records):,}")

    # Tokenize
    print(f"\nTRAIN tokenize ediliyor...")
    train_tokens, train_masks, td, tt = tokenize_records(
        train_records, tok, args.max_len
    )
    print(f"  Drop: {td}  | Truncated: {tt}")

    print(f"\nVAL tokenize ediliyor...")
    val_tokens, val_masks, vd, vt = tokenize_records(
        val_records, tok, args.max_len
    )
    print(f"  Drop: {vd}  | Truncated: {vt}")

    # Paketle
    print(f"\nTRAIN paketleniyor...")
    train_t, train_m = pack_sequences(train_tokens, train_masks, args.max_len)
    print(f"\nVAL paketleniyor...")
    val_t, val_m = pack_sequences(val_tokens, val_masks, args.max_len)

    # Stats
    train_loss_tokens = int(train_m.sum())
    val_loss_tokens = int(val_m.sum())
    print(f"\nİstatistik:")
    print(f"  Train total tokens: {len(train_t):,}")
    print(f"  Train loss tokens:  {train_loss_tokens:,} "
          f"({100*train_loss_tokens/len(train_t):.1f}%)")
    print(f"  Val total tokens:   {len(val_t):,}")
    print(f"  Val loss tokens:    {val_loss_tokens:,} "
          f"({100*val_loss_tokens/len(val_t):.1f}%)")

    # Yaz
    train_t.tofile(out_dir / "sft_train_tokens.bin")
    train_m.tofile(out_dir / "sft_train_mask.bin")
    val_t.tofile(out_dir / "sft_val_tokens.bin")
    val_m.tofile(out_dir / "sft_val_mask.bin")
    print(f"\n[OK] Yazildi: {out_dir}/sft_{{train,val}}_{{tokens,mask}}.bin")

    # Meta
    meta = {
        "tokenizer": args.tokenizer,
        "vocab_size": tok.get_vocab_size(),
        "eot_id": EOT_ID,
        "max_len": args.max_len,
        "train_examples": len(train_records),
        "val_examples": len(val_records),
        "train_total_tokens": len(train_t),
        "train_loss_tokens": train_loss_tokens,
        "val_total_tokens": len(val_t),
        "val_loss_tokens": val_loss_tokens,
        "train_truncated": tt,
        "val_truncated": vt,
        "train_dropped": td,
        "val_dropped": vd,
        "format": {
            "sys_header": SYS_HEADER,
            "user_header": USER_HEADER,
            "asst_header": ASST_HEADER,
            "endoftext_id": EOT_ID,
        },
        "example": (
            f"{SYS_HEADER}Sen yardimci bir asistansin.\n"
            f"{USER_HEADER}2+2 kac?\n"
            f"{ASST_HEADER}4'tur.<|endoftext|>"
        ),
    }
    with open(out_dir / "sft_meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, indent=2, ensure_ascii=False)
    print(f"  Meta: {out_dir}/sft_meta.json")

    # Ornek decode — ilk train sample'i goster
    if train_tokens:
        sample_tokens = train_tokens[0]
        sample_mask = train_masks[0]
        decoded = tok.decode(sample_tokens.tolist())
        n_mask = int(sample_mask.sum())
        print(f"\n--- Ornek decode (ilk train sample) ---")
        print(f"Token uzunluk: {len(sample_tokens)} (loss: {n_mask})")
        print(decoded[:800] + ("..." if len(decoded) > 800 else ""))


if __name__ == "__main__":
    main()