"""SemanticWiki v3 - LoRA SFT of 14B-class students on citation-verified data.

Uses TRL's conversational prompt-completion format: `prompt` = the system+user
messages, `completion` = the assistant wiki page. Loss is computed on the
completion only (TRL default for prompt-completion datasets), so the model is
not trained to regurgitate the code context.

Real run:
  python sw_train.py --base_model Qwen/Qwen2.5-Coder-14B-Instruct \
      --hub_model_id GhostScientist/semanticwiki-coder-14b-v3
Smoke (CPU): add --smoke (tiny subset, 1 step, no push, no trackio).
"""

import argparse
import json
import os

from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoTokenizer, TrainerCallback
from trl import SFTConfig, SFTTrainer

DATASET = "GhostScientist/semanticwiki-data-v3"
MAX_LEN = 24576  # contexts are <=20k tokens; completions add ~1-2k

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


class TrackioCallback(TrainerCallback):
    def on_log(self, args, state, control, logs=None, **kwargs):
        if logs:
            import trackio
            trackio.log({k: v for k, v in logs.items()
                         if isinstance(v, (int, float))},
                        step=state.global_step)


def _norm_content(c):
    """Arrow may infer message content as Json(decode=True), turning some string
    contents into e.g. ["the string"]. Normalize everything back to str."""
    if isinstance(c, str):
        return c
    if isinstance(c, list):
        return "".join(x if isinstance(x, str) else x.get("text", "") for x in c)
    if c is None:
        return ""
    return json.dumps(c)


def to_prompt_completion(ex):
    msgs = ex["messages"]
    if isinstance(msgs, str):  # Arrow can infer the column as string-typed JSON
        msgs = json.loads(msgs)
    msgs = [{"role": m.get("role", "user"), "content": _norm_content(m.get("content"))}
            for m in msgs]
    return {"prompt": msgs[:-1], "completion": [msgs[-1]]}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model", required=True)
    ap.add_argument("--hub_model_id", default=None)
    ap.add_argument("--dataset", default=DATASET)
    ap.add_argument("--epochs", type=float, default=2.0)
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--max_steps", type=int, default=-1)
    ap.add_argument("--max_len", type=int, default=MAX_LEN,
                    help="override sequence length (GPU smoke uses the real 24576)")
    ap.add_argument("--smoke", action="store_true",
                    help="tiny subset, 1 step, no push, no trackio, short seqs")
    args = ap.parse_args()

    tok = AutoTokenizer.from_pretrained(args.base_model)
    ds = load_dataset(args.dataset, split="train")
    assert "messages" in ds.column_names, f"unexpected columns: {ds.column_names}"

    max_len = 2048 if args.smoke else args.max_len
    pc = ds.map(to_prompt_completion, remove_columns=ds.column_names, num_proc=4)
    # fail fast, single process: the chat template must accept the normalized rows
    for ex in pc.select(range(min(2, len(pc)))):
        tok.apply_chat_template(ex["prompt"] + ex["completion"])
    print("[train] chat-template sanity check passed on 2 rows")
    bad = [i for i, ex in enumerate(pc)
           if not (isinstance(ex["prompt"], list) and ex["prompt"]
                   and isinstance(ex["completion"], list) and ex["completion"])]
    print(f"[train] rows with malformed prompt/completion: {len(bad)}")

    def keep(ex):
        ids = tok.apply_chat_template(ex["prompt"] + ex["completion"], tokenize=True)
        if isinstance(ids, dict):  # some transformers versions return a BatchEncoding
            ids = ids["input_ids"]
        return len(ids) <= max_len

    n0 = len(pc)
    pc = pc.filter(keep, num_proc=4)
    print(f"[train] dataset: {n0} -> {len(pc)} rows (length <= {max_len})")
    assert len(pc) > 0, "length filter removed everything"

    if args.smoke:
        pc = pc.select(range(min(4, len(pc))))

    space_id = None if args.smoke else os.environ.get("TRACKIO_SPACE_ID")
    if space_id:
        import trackio
        trackio.init(project="semanticwiki-sft-14b", space_id=space_id)

    # CPU smoke only: bf16 matmul unsupported here; keep bfloat16 on the GPU path.
    dtype = "float32" if args.smoke else "bfloat16"
    cfg = SFTConfig(
        output_dir="/work/out",
        per_device_train_batch_size=1,
        gradient_accumulation_steps=16,
        num_train_epochs=args.epochs,
        max_steps=args.max_steps,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_steps=12,  # ~3% of the ~400-step run (warmup_ratio removed in transformers 5.x)
        logging_steps=5,
        save_strategy="no",
        max_length=max_len,
        packing=False,
        report_to="none",
        bf16=not args.smoke,
        fp16=False,
        model_init_kwargs={"dtype": dtype, "attn_implementation": "sdpa"},
        push_to_hub=bool(args.hub_model_id) and not args.smoke,
        hub_model_id=args.hub_model_id,
        seed=42,
    )
    lora = LoraConfig(
        r=64, lora_alpha=128, lora_dropout=0.05,
        target_modules=TARGET_MODULES, task_type="CAUSAL_LM",
    )

    trainer = SFTTrainer(
        model=args.base_model,
        args=cfg,
        train_dataset=pc,
        peft_config=lora,
        callbacks=[TrackioCallback()] if space_id else [],
    )
    trainer.train()
    trainer.save_model("/work/out/final")
    print("[train] done; final loss history:",
          [round(log.get("loss", -1), 4) for log in trainer.state.log_history
           if "loss" in log][-5:])

    if args.hub_model_id and not args.smoke:
        trainer.push_to_hub()
        print("[train] pushed adapter to " + args.hub_model_id)
    if space_id:
        import trackio
        trackio.finish()


if __name__ == "__main__":
    main()