"""SFT-train a LoRA adapter on Gemma-3-27B-it using single-token loss on
the answer letter. Reads `sft_data.jsonl` produced by `sft_data_gen.py`.

Loss is masked to a SINGLE token per example — the answer letter token
that immediately follows `<answer>`. Everything else is ignored (-100 in
labels). This installs the preference signal without rewriting the model's
broader behaviour.

Usage on pod:
    /workspace/vllm-venv/bin/python /workspace/code/scripts/sft_train.py \
        --base-model /workspace/models/gemma-3-27b-it \
        --data-dir /workspace/code/logs/sft_data \
        --out /workspace/code/logs/sft_adapter \
        --epochs 1 --batch-size 4 --grad-accum 4 --lr 1e-4
"""
from __future__ import annotations

import argparse
import json
import math
import time
from dataclasses import dataclass
from pathlib import Path

import torch
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader

from transformers import (
    AutoModelForCausalLM, AutoTokenizer,
    get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup,
)
from peft import LoraConfig, TaskType, get_peft_model


LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj",
                "gate_proj", "up_proj", "down_proj"]


@dataclass
class SFTExample:
    kind: str       # "preference" or "chat"
    prompt: str     # for preference: prompt ending in `<answer>`
                    # for chat:       chat-templated prompt ending in `<start_of_turn>model\n`
    target: str     # for preference: "A" or "B"
                    # for chat:       the full response text


class SFTDataset(Dataset):
    def __init__(self, examples: list[SFTExample], tokenizer, letter_ids: dict[str, int],
                 max_len: int = 512, target_suffix: str = ""):
        self.examples = examples
        self.tok = tokenizer
        self.LID = letter_ids
        self.max_len = max_len
        self.suffix_ids = (self.tok(target_suffix, return_tensors="pt", add_special_tokens=False)["input_ids"][0]
                           if target_suffix else torch.empty(0, dtype=torch.long))

    def __len__(self): return len(self.examples)

    def __getitem__(self, i):
        ex = self.examples[i]
        prompt_ids = self.tok(ex.prompt, return_tensors="pt", add_special_tokens=False)["input_ids"][0]

        if ex.kind == "preference":
            # target = letter + suffix; loss on all of them (4 tokens with </answer>)
            target_id = self.LID[ex.target]
            target_ids = torch.cat([
                torch.tensor([target_id], dtype=prompt_ids.dtype),
                self.suffix_ids.to(prompt_ids.dtype),
            ])
        else:  # "chat"
            # target = the response, loss on all of it
            target_ids = self.tok(ex.target, return_tensors="pt", add_special_tokens=False)["input_ids"][0].to(prompt_ids.dtype)

        # Truncate the PROMPT from the left if total too long
        max_prompt_len = self.max_len - len(target_ids)
        if len(prompt_ids) > max_prompt_len:
            prompt_ids = prompt_ids[-max_prompt_len:]
        input_ids = torch.cat([prompt_ids, target_ids])
        # Labels: -100 on prompt, real ids on the target tokens
        labels = torch.full_like(input_ids, -100)
        labels[-len(target_ids):] = target_ids
        return {"input_ids": input_ids, "labels": labels}


def collate_left_pad(batch, pad_id: int):
    """Left-pad to longest in batch. Pad with pad_id; pad positions in labels become -100."""
    L = max(item["input_ids"].shape[0] for item in batch)
    ids_out, lab_out, attn_out = [], [], []
    for item in batch:
        n = item["input_ids"].shape[0]
        pad = L - n
        ids = torch.cat([torch.full((pad,), pad_id, dtype=item["input_ids"].dtype),
                         item["input_ids"]])
        lab = torch.cat([torch.full((pad,), -100, dtype=item["labels"].dtype),
                         item["labels"]])
        attn = torch.cat([torch.zeros(pad, dtype=torch.long),
                          torch.ones(n, dtype=torch.long)])
        ids_out.append(ids); lab_out.append(lab); attn_out.append(attn)
    return {
        "input_ids": torch.stack(ids_out),
        "labels": torch.stack(lab_out),
        "attention_mask": torch.stack(attn_out),
    }


def load_preference_data(data_dir: Path) -> list[SFTExample]:
    out = []
    with (data_dir / "sft_data.jsonl").open() as f:
        for line in f:
            d = json.loads(line)
            out.append(SFTExample(kind="preference", prompt=d["prompt"], target=d["target"]))
    return out


def load_chat_data(path: Path) -> list[SFTExample]:
    """Load self-distill chat examples produced by self_distill_gen.py."""
    out = []
    with path.open() as f:
        for line in f:
            d = json.loads(line)
            # prompt is the FULL chat-templated text ending with <start_of_turn>model\n
            # target is the model's sampled response (no leading newline)
            out.append(SFTExample(
                kind="chat",
                prompt=d.get("chat_prompt_text") or d["prompt"],
                target=d["response"],
            ))
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base-model", required=True)
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--epochs", type=int, default=1)
    ap.add_argument("--batch-size", type=int, default=4)
    ap.add_argument("--grad-accum", type=int, default=4)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--lora-r", type=int, default=32)
    ap.add_argument("--lora-alpha", type=int, default=64)
    ap.add_argument("--warmup-steps", type=int, default=10)
    ap.add_argument("--max-steps", type=int, default=None)
    ap.add_argument("--save-every", type=int, default=0, help="0 = only save at end")
    ap.add_argument("--target-suffix", default="</answer>",
                    help="String appended after the answer letter; included in the loss to teach the model to close cleanly")
    ap.add_argument("--self-distill-data", default=None,
                    help="Optional self-distill jsonl (output of self_distill_gen.py) to mix in")
    ap.add_argument("--lr-schedule", choices=["linear", "cosine"], default="cosine")
    ap.add_argument("--max-len", type=int, default=512)
    args = ap.parse_args()
    out_dir = Path(args.out); out_dir.mkdir(parents=True, exist_ok=True)

    # Reuse the letter-id metadata from data gen to be sure we train on the same token IDs
    meta = json.loads((Path(args.data_dir) / "sft_data_meta.json").read_text())
    LID = meta["letter_ids"]
    print(f"[meta] letter_ids: {LID}", flush=True)

    print(f"[load] {args.base_model}", flush=True); t0 = time.time()
    tok = AutoTokenizer.from_pretrained(args.base_model)
    if tok.pad_token is None: tok.pad_token = tok.eos_token
    model = AutoModelForCausalLM.from_pretrained(
        args.base_model, torch_dtype=torch.bfloat16, device_map="auto",
        attn_implementation="eager",
    )
    print(f"[load] done in {time.time()-t0:.1f}s", flush=True)

    # Wrap in LoRA
    peft_cfg = LoraConfig(
        r=args.lora_r, lora_alpha=args.lora_alpha, lora_dropout=0.0,
        target_modules=LORA_TARGETS, task_type=TaskType.CAUSAL_LM, bias="none",
    )
    model = get_peft_model(model, peft_cfg)
    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    n_total = sum(p.numel() for p in model.parameters())
    print(f"[lora] trainable={n_trainable/1e6:.1f}M / total={n_total/1e9:.1f}B "
          f"({100*n_trainable/n_total:.3f}%)", flush=True)
    model.gradient_checkpointing_enable()
    model.enable_input_require_grads()
    model.train()

    examples = load_preference_data(Path(args.data_dir))
    print(f"[data] {len(examples)} preference examples", flush=True)
    if args.self_distill_data:
        chat_examples = load_chat_data(Path(args.self_distill_data))
        print(f"[data] + {len(chat_examples)} self-distill chat examples", flush=True)
        examples = examples + chat_examples
    ds = SFTDataset(examples, tok, LID, max_len=args.max_len, target_suffix=args.target_suffix)
    n_suffix = len(ds.suffix_ids)
    print(f"[data] {len(examples)} total examples; "
          f"preference target = letter + {args.target_suffix!r} ({1 + n_suffix} loss tokens), "
          f"chat target = full response", flush=True)
    dl = DataLoader(ds, batch_size=args.batch_size, shuffle=True,
                    collate_fn=lambda b: collate_left_pad(b, tok.pad_token_id))

    n_steps_per_epoch = math.ceil(len(dl) / args.grad_accum)
    total_steps = n_steps_per_epoch * args.epochs
    if args.max_steps is not None:
        total_steps = min(total_steps, args.max_steps)
    print(f"[plan] {total_steps} optimizer steps across {args.epochs} epoch(s)", flush=True)

    optim = torch.optim.AdamW(
        [p for p in model.parameters() if p.requires_grad],
        lr=args.lr, betas=(0.9, 0.95), weight_decay=0.0,
    )
    if args.lr_schedule == "cosine":
        sched = get_cosine_schedule_with_warmup(optim, args.warmup_steps, total_steps)
    else:
        sched = get_linear_schedule_with_warmup(optim, args.warmup_steps, total_steps)
    print(f"[plan] LR schedule: {args.lr_schedule}, peak={args.lr}, warmup={args.warmup_steps}, total_steps={total_steps}", flush=True)

    device = next(model.parameters()).device
    step = 0; sample = 0; t0 = time.time(); accum_loss = 0.0
    for epoch in range(args.epochs):
        for batch_idx, batch in enumerate(dl):
            batch = {k: v.to(device) for k, v in batch.items()}
            out = model(**batch, use_cache=False)
            # `out.loss` is mean over non-(-100) positions = mean over the K examples
            # in the batch (one masked-loss token per example). HF's modeling code
            # shifts labels for causal LM, so the label at position t aligns with
            # logits[t-1]. Our last-token label gets compared to the model's
            # next-token prediction at the position BEFORE the answer letter —
            # which is `<answer>`'s last subtoken, the correct position.
            loss = out.loss / args.grad_accum
            loss.backward()
            accum_loss += loss.item() * args.grad_accum
            sample += 1

            if sample % args.grad_accum == 0:
                torch.nn.utils.clip_grad_norm_(
                    [p for p in model.parameters() if p.requires_grad], max_norm=1.0)
                optim.step(); sched.step(); optim.zero_grad()
                step += 1
                avg_loss = accum_loss / args.grad_accum
                accum_loss = 0.0
                elapsed = time.time() - t0
                print(f"  step {step}/{total_steps}  loss={avg_loss:.4f}  "
                      f"lr={sched.get_last_lr()[0]:.2e}  ({elapsed:.0f}s)", flush=True)
                if args.save_every and step % args.save_every == 0:
                    ckpt = out_dir / f"step_{step}"
                    model.save_pretrained(str(ckpt))
                    print(f"  saved checkpoint to {ckpt}", flush=True)
                if args.max_steps and step >= args.max_steps:
                    break
        if args.max_steps and step >= args.max_steps:
            break

    # Final save
    model.save_pretrained(str(out_dir))
    print(f"\nwrote final LoRA adapter to {out_dir}", flush=True)


if __name__ == "__main__":
    main()