"""LoRA / QLoRA fine-tuning on Sinhala MCQs.

Training data = jsonl in the same format as the dev set (question, options,
answer, level). The prompt is built by the SAME code as inference, and the
loss is only on the answer token - exactly what predict.py scores.

    python train.py --config configs/train_qlora_qwen25_7b.yaml

After training, predict with the adapter:
    model.adapter_path: outputs/<name>/adapter   (in a predict config)

Rules to remember:
  * never train on test questions (rule 04)
  * LoRA parameters count towards the 8B limit (base + adapter)
  * data generated with closed APIs (GPT/Claude/Gemini) must be disclosed
"""

from __future__ import annotations

import argparse
import json
import math
import shutil
import sys
from pathlib import Path

import torch

from sinmmlu.config import load_config
from sinmmlu.data import has_gold, load_questions
from sinmmlu.metrics import accuracy_by_level, format_report
from sinmmlu.modeling import describe_budget, load_model, load_tokenizer, pick_dtype, prepare_for_kbit_training, set_determinism
from sinmmlu.prompts import build_messages, render_prompt
from sinmmlu.runinfo import append_result, write_run_info
from sinmmlu.scoring import Scorer, label_token_id


def build_examples(records: list[dict], tok, cfg: dict) -> list[dict]:
    """One example = prompt tokens + the correct label token.

    Options are rotated (deterministically) so the right answer appears in
    every position equally often - otherwise the model can learn "the answer
    is usually 2" instead of learning the content.
    """
    p = cfg["prompt"]
    max_len = cfg["inference"]["max_length"]
    rotate = cfg["train"].get("rotate_options", True)
    examples, skipped = [], 0
    for i, r in enumerate(records):
        n = len(r["options"])
        shift = (i % n) if rotate else 0
        options = r["options"][shift:] + r["options"][:shift]
        gold_pos = (r["gold"] - shift) % n  # where the correct option ends up
        msgs = build_messages(r, p["template"], None, options)
        text, prefix = render_prompt(tok, msgs, p["use_chat_template"], p["chat_template_kwargs"])
        ids = tok(text, add_special_tokens=False)["input_ids"]
        if len(ids) + 1 > max_len:
            skipped += 1
            continue
        target = label_token_id(tok, text, prefix + str(gold_pos + 1))
        examples.append({"input_ids": ids + [target], "labels": [-100] * len(ids) + [target]})
    if skipped:
        print(f"[train] skipped {skipped} examples longer than {max_len} tokens")
    return examples


class Collator:
    def __init__(self, pad_id: int):
        self.pad_id = pad_id

    def __call__(self, batch: list[dict]) -> dict:
        width = max(len(b["input_ids"]) for b in batch)
        ids, labels, mask = [], [], []
        for b in batch:
            pad = width - len(b["input_ids"])
            ids.append([self.pad_id] * pad + b["input_ids"])
            labels.append([-100] * pad + b["labels"])
            mask.append([0] * pad + [1] * len(b["input_ids"]))
        return {
            "input_ids": torch.tensor(ids),
            "labels": torch.tensor(labels),
            "attention_mask": torch.tensor(mask),
        }


def main(argv=None) -> int:
    from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
    from transformers import Trainer, TrainingArguments

    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--config", required=True)
    ap.add_argument("--who", default="")
    ap.add_argument("--max-steps", type=int, default=-1, help="stop early (debugging)")
    args = ap.parse_args(argv)

    cfg = load_config(args.config)
    tc = cfg["train"]
    set_determinism(cfg["seed"])
    if not tc["train_file"]:
        sys.exit("train.train_file is required")

    train_records = load_questions(tc["train_file"], cfg["data"]["answer_base"])
    if not has_gold(train_records):
        sys.exit("training file has no gold answers")
    print(f"[train] {len(train_records)} training questions from {tc['train_file']}")

    tok = load_tokenizer(cfg["model"])
    model = load_model(cfg["model"], for_training=True)
    if cfg["model"]["load_in_4bit"] and cfg["model"].get("text_only"):
        # peft's version would cast Gemma 4's 2.8 B-parameter embedding table to float32 (10.5 GiB -> OOM on a T4)
        model = prepare_for_kbit_training(model, use_gradient_checkpointing=tc["gradient_checkpointing"])
    elif cfg["model"]["load_in_4bit"]:
        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=tc["gradient_checkpointing"])
    elif tc["gradient_checkpointing"]:
        model.gradient_checkpointing_enable()
        model.enable_input_require_grads()

    lora = LoraConfig(
        r=tc["lora_r"],
        lora_alpha=tc["lora_alpha"],
        lora_dropout=tc["lora_dropout"],
        target_modules=tc["target_modules"],
        task_type="CAUSAL_LM",
    )
    model = get_peft_model(model, lora)
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    print(f"[lora] adapter parameters: {trainable:,}")

    examples = build_examples(train_records, tok, cfg)
    steps_per_epoch = math.ceil(len(examples) / (tc["batch_size"] * tc["grad_accum"]))
    total_steps = args.max_steps if args.max_steps > 0 else steps_per_epoch * tc["epochs"]
    out_dir = Path(tc["output_dir"])
    dtype = pick_dtype(cfg["model"]["dtype"])
    on_gpu = torch.cuda.is_available()
    targs = TrainingArguments(
        output_dir=str(out_dir / "checkpoints"),
        num_train_epochs=tc["epochs"],
        max_steps=args.max_steps,
        learning_rate=tc["learning_rate"],
        per_device_train_batch_size=tc["batch_size"],
        gradient_accumulation_steps=tc["grad_accum"],
        lr_scheduler_type="cosine",
        warmup_steps=math.ceil(0.03 * total_steps),  # 3% warm-up
        logging_steps=tc["logging_steps"],
        save_strategy="no",
        report_to=[],
        seed=cfg["seed"],
        data_seed=cfg["seed"],
        bf16=on_gpu and dtype == torch.bfloat16,
        fp16=on_gpu and dtype == torch.float16,
        remove_unused_columns=False,
        dataloader_num_workers=0,
    )
    trainer = Trainer(model=model, args=targs, train_dataset=examples, data_collator=Collator(tok.pad_token_id))
    trainer.train()

    adapter_dir = out_dir / "adapter"
    model.save_pretrained(adapter_dir)
    tok.save_pretrained(adapter_dir)
    shutil.copy(args.config, out_dir / "train_config.yaml")
    total = sum(p.numel() for p in model.parameters())
    print(f"[save] adapter -> {adapter_dir}")
    print(f"[budget] base + adapter as loaded: {describe_budget(total)}")

    extra = {"train_file": tc["train_file"], "n_train": len(examples), "adapter_params": trainable}
    if tc["eval_file"]:
        model.eval()
        eval_records = load_questions(tc["eval_file"], cfg["data"]["answer_base"])
        preds = Scorer(model, tok, cfg).predict(eval_records)
        result = accuracy_by_level([r["level"] for r in eval_records], preds, [r["gold"] for r in eval_records])
        print(f"[eval] {tc['eval_file']}\n" + format_report(result))
        cfg_logged = json.loads(json.dumps(cfg))
        cfg_logged["model"]["adapter_path"] = str(adapter_dir)
        append_result(cfg_logged, tc["eval_file"], result, args.who, "after training")
        extra["eval_result"] = result
    write_run_info(out_dir / "train.run.json", cfg, extra)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())