"""QLoRA fine-tune of google/gemma-4-31B-it with Unsloth.

Recipe notes (hard-won — see docs/eval-roadmap.md):
- Formats with the tokenizer's NATIVE Gemma-4 chat template (<|turn> markers), the
  same one llama-server --jinja serves. Never override it; a train/serve template
  mismatch degraded four straight fine-tunes.
- Masks loss to the assistant response only; LR 5e-5.

Run on an A100/H100 80GB (Colab Pro+/RunPod/Lambda). For a lighter/cheaper run,
swap BASE_MODEL to the Gemma 4 E4B edge variant — see PLAN.md fallback.

  pip install -r training/requirements-train.txt
  python training/make_dataset.py
  python training/train_qlora.py
"""
from __future__ import annotations

import inspect
import os

# Unsloth must be imported before trl/transformers/peft so its patches apply.
from unsloth import FastLanguageModel

from datasets import load_dataset  # noqa: E402
from trl import SFTConfig, SFTTrainer  # noqa: E402


def _filter_kwargs(cls, kwargs: dict) -> dict:
    """Keep only kwargs the class' __init__ accepts.

    trl renames fields across releases (max_seq_length -> max_length, the
    SFTTrainer `tokenizer` -> `processing_class`); filtering by signature lets
    this script run across trl versions without pinning an exact one.
    """
    valid = set(inspect.signature(cls.__init__).parameters)
    return {k: v for k, v in kwargs.items() if k in valid}

BASE_MODEL = os.environ.get("BASE_MODEL", "google/gemma-4-31B-it")
MAX_SEQ_LEN = int(os.environ.get("MAX_SEQ_LEN", "4096"))
NUM_EPOCHS = float(os.environ.get("NUM_EPOCHS", "2"))
OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "training/outputs/gemma4-cal-lora")


def main():
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=BASE_MODEL,
        max_seq_length=MAX_SEQ_LEN,
        load_in_4bit=True,
    )
    model = FastLanguageModel.get_peft_model(
        model,
        r=16,
        lora_alpha=16,
        lora_dropout=0.0,
        target_modules=[
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj",
        ],
        use_gradient_checkpointing="unsloth",
    )
    # Use the tokenizer's NATIVE chat template. Gemma-4 ships chat_template.jinja
    # with `<|turn>role ... <turn|>` markers — the SAME template convert_hf_to_gguf
    # embeds in the GGUF and llama-server --jinja serves at inference. Overriding it
    # with unsloth's legacy "gemma" template (<start_of_turn>) trained a turn format
    # inference never uses, and every fine-tune degraded monotonically with more
    # steps (docs/eval-roadmap.md Steps 3-6). Train format MUST equal serve format.
    probe = tokenizer.apply_chat_template(
        [{"role": "user", "content": "ping"}], tokenize=False, add_generation_prompt=True
    )
    assert "<|turn>" in probe, (
        "tokenizer's native chat template doesn't look like Gemma-4 (no <|turn> marker); "
        f"refusing to train on a mismatched format. Rendered: {probe[:200]!r}"
    )
    print(f"native chat template OK: {probe[:120]!r}")

    # Train on the hand-authored set + (if present) the git-ignored SMCalFlow import
    # (training/import_smcalflow.py). Both are in the same chat-messages format.
    # The hand-authored rows are thread-style (matching the app/eval domain) but are
    # outnumbered ~16:1 by SMCalFlow's command-style utterances — upsample them so the
    # model doesn't drift toward command-only phrasing (HAND_UPSAMPLE=1 disables).
    from datasets import concatenate_datasets

    upsample = int(os.environ.get("HAND_UPSAMPLE", "4"))
    hand = load_dataset("json", data_files="training/data/dataset.jsonl", split="train")
    parts = [hand] * upsample
    if os.path.exists("training/data/smcalflow_train.jsonl"):
        parts.append(load_dataset(
            "json", data_files="training/data/smcalflow_train.jsonl", split="train"))
    ds = concatenate_datasets(parts).shuffle(seed=42)
    print(f"training on {len(ds)} examples ({len(hand)} hand-authored x{upsample} + smcalflow)")

    def fmt(batch):
        texts = [
            tokenizer.apply_chat_template(m, tokenize=False, add_generation_prompt=False)
            for m in batch["messages"]
        ]
        return {"text": texts}

    ds = ds.map(fmt, batched=True)

    # Build SFTConfig version-tolerantly: pass both the old and new field names for
    # the sequence-length cap and drop whichever the installed trl doesn't accept.
    args = SFTConfig(**_filter_kwargs(SFTConfig, dict(
        output_dir=OUTPUT_DIR,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        num_train_epochs=NUM_EPOCHS,
        learning_rate=5e-5,  # was 2e-4; gentler on a near-ceiling base (see eval roadmap)
        warmup_ratio=0.03,
        logging_steps=10,
        save_strategy="epoch",
        bf16=True,
        dataset_text_field="text",
        max_seq_length=MAX_SEQ_LEN,  # older trl
        max_length=MAX_SEQ_LEN,      # newer trl (renamed)
    )))

    # SFTTrainer renamed `tokenizer` -> `processing_class` in trl 0.13.
    trainer_kwargs = dict(model=model, train_dataset=ds, args=args)
    tr_params = set(inspect.signature(SFTTrainer.__init__).parameters)
    if "processing_class" in tr_params:
        trainer_kwargs["processing_class"] = tokenizer
    else:
        trainer_kwargs["tokenizer"] = tokenizer
    trainer = SFTTrainer(**trainer_kwargs)

    # Mask loss to the assistant turn only (the ActionPlan JSON). Without this, the
    # long SYSTEM prompt dominates the token count and most gradient signal is spent
    # re-learning the prompt instead of the answer. Markers are Gemma-4's native ones.
    try:
        from unsloth.chat_templates import train_on_responses_only

        trainer = train_on_responses_only(
            trainer,
            instruction_part="<|turn>user\n",
            response_part="<|turn>model\n",
        )
        print("loss masked to assistant responses only")
    except Exception as e:  # noqa: BLE001
        print(f"WARNING: response-only masking unavailable ({e}); training on full text")

    trainer.train()

    # Merge LoRA into fp16 weights, ready for GGUF conversion (see export_gguf.sh).
    model.save_pretrained_merged(
        OUTPUT_DIR + "-merged", tokenizer, save_method="merged_16bit"
    )
    print(f"merged model saved to {OUTPUT_DIR}-merged")


if __name__ == "__main__":
    main()