"""LoRA supervised fine-tuning for Micro RPG Engine.

Teaches a small base model (1B-4B) to reliably emit the three-block tag protocol
with valid mechanics, using the parser-validated dataset from build_dataset.py.

We use LoRA (PEFT) so it trains on a single consumer/Colab GPU and produces a tiny
adapter (a few MB). Point the engine at it with MICRORPG_ADAPTER to play with your
fine-tuned model.

Quickstart
----------
    pip install -r requirements-train.txt
    python -m finetune.build_dataset --n 1200
    python -m finetune.train \
        --model Qwen/Qwen3-4B-Instruct-2507 \
        --out finetune/out/qwen3-4b-microrpg

Then play with it:
    # PowerShell
    $env:MICRORPG_ADAPTER = "finetune/out/qwen3-4b-microrpg"
    python app.py

Notes
-----
* `--model` accepts any chat model with a chat template (Qwen3-4B, MiniCPM, a Llama
  for the "Llama Champion" quest, etc.). Swap freely — the dataset is model-agnostic.
* For a 4B model on a small GPU, add `--load-4bit` (needs bitsandbytes).
"""

from __future__ import annotations

import argparse
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default=os.environ.get("MICRORPG_MODEL", "Qwen/Qwen3-4B-Instruct-2507"))
    ap.add_argument("--train", default="finetune/data/train.jsonl")
    ap.add_argument("--eval", default="finetune/data/eval.jsonl")
    ap.add_argument("--out", default="finetune/out/microrpg-adapter")
    ap.add_argument("--epochs", type=float, default=3.0)
    ap.add_argument("--lr", type=float, default=2e-4)
    ap.add_argument("--batch", type=int, default=2)
    ap.add_argument("--grad-accum", type=int, default=8)
    ap.add_argument("--max-len", type=int, default=1536)
    ap.add_argument("--lora-r", type=int, default=16)
    ap.add_argument("--lora-alpha", type=int, default=32)
    ap.add_argument("--load-4bit", action="store_true", help="QLoRA via bitsandbytes")
    ap.add_argument("--merge", action="store_true",
                    help="after training, merge the adapter into the base and save full weights")
    args = ap.parse_args()

    # Heavy imports kept inside main so `--help` and import-checks stay light.
    import torch
    from datasets import load_dataset
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import LoraConfig
    from trl import SFTConfig, SFTTrainer

    print(f"Base model : {args.model}")
    print(f"Train file : {args.train}")

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    model_kwargs = {"torch_dtype": torch.bfloat16 if torch.cuda.is_available() else torch.float32}
    if args.load_4bit:
        from transformers import BitsAndBytesConfig
        model_kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    if torch.cuda.is_available():
        model_kwargs["device_map"] = "auto"

    model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs)

    # LoRA adapter on attention + MLP projections — the standard, portable target set.
    peft_config = LoraConfig(
        r=args.lora_r,
        lora_alpha=args.lora_alpha,
        lora_dropout=0.05,
        bias="none",
        task_type="CAUSAL_LM",
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
    )

    data_files = {"train": args.train}
    if os.path.exists(args.eval):
        data_files["eval"] = args.eval
    ds = load_dataset("json", data_files=data_files)

    sft_config = SFTConfig(
        output_dir=args.out,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch,
        gradient_accumulation_steps=args.grad_accum,
        learning_rate=args.lr,
        lr_scheduler_type="cosine",
        warmup_ratio=0.05,
        logging_steps=10,
        save_strategy="epoch",
        eval_strategy="epoch" if "eval" in ds else "no",
        bf16=torch.cuda.is_available(),
        gradient_checkpointing=True,
        max_seq_length=args.max_len,
        packing=False,
        report_to="none",
        # The dataset has a "messages" column → TRL applies the chat template and,
        # by default, masks the prompt so loss is computed only on the assistant turn.
        assistant_only_loss=True,
    )

    trainer = SFTTrainer(
        model=model,
        args=sft_config,
        train_dataset=ds["train"],
        eval_dataset=ds.get("eval"),
        peft_config=peft_config,
        processing_class=tokenizer,
    )

    trainer.train()
    trainer.save_model(args.out)
    tokenizer.save_pretrained(args.out)
    print(f"\nAdapter saved to: {args.out}")

    if args.merge:
        print("Merging adapter into base weights...")
        merged_dir = args.out.rstrip("/\\") + "-merged"
        merged = trainer.model.merge_and_unload()
        merged.save_pretrained(merged_dir)
        tokenizer.save_pretrained(merged_dir)
        print(f"Merged model saved to: {merged_dir}")

    print("\nPlay with it:  set MICRORPG_ADAPTER to the output dir, then run app.py")


if __name__ == "__main__":
    main()