#!/usr/bin/env python3
"""Train the TenaOS single task-tagged Gemma 4 LoRA adapter with Unsloth."""

from __future__ import annotations

import argparse
import json
import math
import traceback
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


DEFAULT_BASE_MODEL = "unsloth/gemma-4-E4B-it"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--train-jsonl", type=Path, required=True)
    parser.add_argument("--validation-jsonl", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL)
    parser.add_argument("--merge-output-dir", type=Path, help="Optional directory for merged weights.")
    parser.add_argument(
        "--max-seq-length",
        type=int,
        default=32768,
        help=(
            "Real multi-turn conversations run much longer than the old flattened single-completion "
            "rows. Measured char-length distribution across the reconstructed corpus (rough proxy, "
            "~3.5 chars/token): form median ~22.4K chars / p95 ~50.5K chars / max ~125.5K chars; "
            "report median ~9.8K / p95 ~40.9K / max ~66K; cds/patient_education/scribe all comfortably "
            "under 16K chars. 32768 tokens covers the large majority without the extreme few-outlier "
            "tail; those are truncated by SFTConfig's max_length rather than raising the default so high "
            "that most batches waste compute on padding."
        ),
    )
    parser.add_argument("--epochs", type=float, default=3.0)
    parser.add_argument("--max-steps", type=int, default=-1, help="Optional cap for smoke tests or short runs.")
    parser.add_argument("--per-device-train-batch-size", type=int, default=1)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    parser.add_argument("--warmup-ratio", type=float, default=0.03)
    parser.add_argument("--logging-steps", type=int, default=10)
    parser.add_argument("--eval-steps", type=int, default=100)
    parser.add_argument("--save-steps", type=int, default=250)
    parser.add_argument("--seed", type=int, default=3407)
    parser.add_argument("--load-in-4bit", action="store_true", help="Use QLoRA fallback if BF16 LoRA does not fit.")
    parser.add_argument("--dataset-num-proc", type=int, default=4)
    parser.add_argument("--packing", action="store_true", help="Enable TRL packing for short examples.")
    parser.add_argument(
        "--chat-template",
        default=None,
        help=(
            "Optional Unsloth canned chat_template name (see unsloth.chat_templates.CHAT_TEMPLATES). "
            "Defaults to the base model's own tokenizer chat_template (recommended: this is the same "
            "template baked into the merged GGUF that llama.cpp renders with --jinja at serve time, "
            "so training and inference stay byte-identical without depending on Unsloth shipping a "
            "matching canned template for this specific base model)."
        ),
    )
    parser.add_argument(
        "--instruction-part",
        default="<|turn>user\n",
        help=(
            "Turn-start marker for non-assistant turns, used by train_on_responses_only masking. "
            "NOTE: this Gemma 4 base model's chat template uses '<|turn>role ... <turn|>' markers, "
            "NOT the classic Gemma 2/3 '<start_of_turn>role ... <end_of_turn>' tokens — verified by "
            "rendering a sample conversation with the tokenizer's own apply_chat_template. Using the "
            "wrong markers silently zeroes out every example's loss mask (train_on_responses_only "
            "finds no matches), which crashes the RandomSampler with num_samples=0."
        ),
    )
    parser.add_argument(
        "--response-part",
        default="<|turn>model\n",
        help="Turn-start marker for assistant turns, used by train_on_responses_only masking.",
    )
    parser.add_argument(
        "--no-response-only-masking",
        action="store_true",
        help="Disable response-only loss masking (NOT recommended — trains loss over tool/user tokens too).",
    )
    parser.add_argument(
        "--use-flex-attention",
        action="store_true",
        help=(
            "Force torch's flex_attention backend instead of Unsloth's default sdpa. NOT "
            "recommended: under Unsloth's custom gradient checkpointing the layer recompute runs "
            "outside torch.compile, so flex_attention's backward dispatches to its eager dense "
            "decomposition (torch/_higher_order_ops/flex_attention.py::sdpa_dense_backward), which "
            "materializes full quadratic (B, H, N, N) score/grad tensors — ~9 GiB each at "
            "N=24576 with 8 heads — and reliably OOMs an 80 GB A100 on the first backward pass. "
            "The original rationale (head_dim 512 exceeding fused-kernel limits) was wrong: this "
            "model's config has head_dim=256, which sdpa's fused flash/mem-efficient kernels fully "
            "support."
        ),
    )
    parser.add_argument(
        "--no-length-filter",
        action="store_true",
        help=(
            "Disable dropping examples whose rendered token length exceeds --max-seq-length. By default "
            "such outliers are dropped entirely rather than silently truncated by SFTConfig's max_length, "
            "since truncation can cut a conversation off mid tool-call/JSON payload and teach the model a "
            "corrupted, incomplete turn. Measured on the current corpus, only ~0.1% of examples (the "
            "extreme long tail of the 'form' task) exceed 24576 tokens, so dropping them is cheap."
        ),
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    args.output_dir.mkdir(parents=True, exist_ok=True)

    try:
        from unsloth import FastVisionModel
        from unsloth.chat_templates import CHAT_TEMPLATES, get_chat_template, train_on_responses_only
        import torch
        from datasets import load_dataset
        from trl import SFTConfig, SFTTrainer
    except ImportError as exc:
        raise SystemExit(
            "Missing training dependencies. Install Unsloth, TRL, datasets, transformers, peft, and torch "
            "in the training environment before running this script."
        ) from exc

    model, tokenizer = FastVisionModel.from_pretrained(
        args.base_model,
        load_in_4bit=args.load_in_4bit,
        use_gradient_checkpointing="unsloth",
        max_seq_length=args.max_seq_length,
    )

    # Keep Unsloth's default sdpa, but force the repeat-KV path instead of enable_gqa.
    #
    # Verified by intercepting the actual SDPA calls during a 24576-token forward:
    #   - sliding-window layers (35/42): q/k (1, 8, N, 256) + boolean mask
    #     -> fused memory-efficient kernel, linear memory. Fine as-is.
    #   - global full-attention layers (7/42): q (1, 8, N, 512), k/v (1, 2, N, 512)
    #     with enable_gqa=True and is_causal.
    # No fused SDPA kernel supports head_dim=512 *with* enable_gqa (flash caps at
    # head_dim 256; the memory-efficient kernel rejects enable_gqa), so those 7 layers
    # silently dispatch to the quadratic-memory math kernel, whose backward materializes
    # an (8, N, N) fp32 tensor -- an 18 GiB single allocation at N=24576, which OOMs the
    # 80 GB A100. With K/V repeated to the full 8 heads the memory-efficient kernel
    # accepts head_dim=512 (measured: 2.49 GiB peak for fwd+bwd of one such layer at
    # N=24576, vs OOM before). repeat_kv is mathematically identical to enable_gqa,
    # so this changes memory/speed only, never outputs.
    #
    # (A previous revision instead forced flex_attention here, which OOMed even harder:
    # under Unsloth's gradient checkpointing the recompute runs uncompiled, and eager
    # flex_attention's backward decomposes to dense sdpa_dense_backward with full
    # (B, H, N, N) score tensors on *every* layer.)
    if args.use_flex_attention:
        model.set_attn_implementation("flex_attention")
        print(f"Attention implementation forced to: {model.config._attn_implementation}")
    else:
        import transformers.integrations.sdpa_attention as _sdpa_integration

        _sdpa_integration.use_gqa_in_sdpa = lambda attention_mask, key: False
        print(
            "SDPA enable_gqa disabled (forcing repeat_kv) so the 512-dim global-attention "
            "layers use the fused memory-efficient kernel instead of the quadratic math fallback."
        )

    if args.chat_template:
        if args.chat_template not in CHAT_TEMPLATES:
            raise SystemExit(
                f"--chat-template={args.chat_template!r} is not a known Unsloth canned template "
                f"(available: {sorted(CHAT_TEMPLATES.keys())}). Omit --chat-template to use the base "
                "model's own tokenizer chat_template instead (recommended)."
            )
        tokenizer = get_chat_template(tokenizer, chat_template=args.chat_template)
    elif not getattr(tokenizer, "chat_template", None):
        raise SystemExit(
            f"Base model {args.base_model!r} has no tokenizer chat_template and no --chat-template "
            "override was given. A previous run hardcoded chat_template='gemma-4-thinking', which does "
            "not exist in this Unsloth version's CHAT_TEMPLATES registry and would crash here — pass an "
            "explicit --chat-template only if you have verified it renders tool-role turns the same way "
            "llama.cpp's --jinja rendering of the merged GGUF does (see the template-check step)."
        )
    print(f"Using chat_template (first 200 chars): {str(tokenizer.chat_template)[:200]!r}")

    model = FastVisionModel.get_peft_model(
        model,
        finetune_vision_layers=False,
        finetune_language_layers=True,
        finetune_attention_modules=True,
        finetune_mlp_modules=True,
        r=16,
        lora_alpha=32,
        # NOTE: the installed peft version's MoE ParamWrapper path (used for
        # `experts.gate_up_proj` / `experts.down_proj`) raises
        # `ValueError: lora.ParamWrapper does not work with lora_dropout != 0`
        # for this model's grouped-expert parameters, so dropout must stay 0.
        lora_dropout=0.0,
        bias="none",
        random_state=args.seed,
        use_rslora=False,
        loftq_config=None,
        target_modules="all-linear",
    )

    dataset = load_dataset(
        "json",
        data_files={
            "train": str(args.train_jsonl),
            "validation": str(args.validation_jsonl),
        },
    )

    def formatting_prompts_func(examples: dict[str, Any]) -> dict[str, list[str]]:
        texts = [
            tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False).removeprefix("<bos>")
            for convo in examples["conversations"]
        ]
        return {"text": texts}

    dataset = dataset.map(
        formatting_prompts_func,
        batched=True,
        num_proc=args.dataset_num_proc,
        remove_columns=dataset["train"].column_names,
    )

    if not args.no_length_filter:
        def _add_token_length(examples: dict[str, Any]) -> dict[str, list[int]]:
            # Gemma4's tokenizer is a full multimodal processor; it requires `text=`
            # passed as an explicit keyword (positional args get swallowed as `images`
            # by Unsloth's patched __call__ wrapper, leaving text=None internally).
            # padding=False is required too: batched calls otherwise pad every row in
            # the (default size-1000) map batch out to that batch's single longest
            # sequence, which silently inflates the reported length of every short
            # example that happens to share a batch with one long outlier.
            return {
                "_n_tokens": [
                    len(ids)
                    for ids in tokenizer(
                        text=examples["text"], add_special_tokens=False, padding=False
                    )["input_ids"]
                ]
            }

        for split in ("train", "validation"):
            before = len(dataset[split])
            dataset[split] = dataset[split].map(_add_token_length, batched=True, num_proc=args.dataset_num_proc)
            dataset[split] = dataset[split].filter(
                lambda ex: ex["_n_tokens"] <= args.max_seq_length,
                num_proc=args.dataset_num_proc,
            )
            dataset[split] = dataset[split].remove_columns(["_n_tokens"])
            after = len(dataset[split])
            print(
                f"Length filter ({split}): dropped {before - after}/{before} examples exceeding "
                f"--max-seq-length={args.max_seq_length} tokens (kept {after})."
            )
    else:
        print(
            "WARNING: --no-length-filter set; over-length examples will be silently truncated by "
            "SFTConfig's max_length instead of being dropped."
        )

    train_size = len(dataset["train"])
    effective_batch = max(1, args.per_device_train_batch_size * args.gradient_accumulation_steps)
    estimated_steps = max(1, math.ceil(train_size / effective_batch) * max(1, math.ceil(args.epochs)))
    warmup_steps = max(1, int(estimated_steps * args.warmup_ratio))

    train_args = SFTConfig(
        output_dir=str(args.output_dir),
        num_train_epochs=args.epochs,
        max_steps=args.max_steps,
        per_device_train_batch_size=args.per_device_train_batch_size,
        gradient_accumulation_steps=args.gradient_accumulation_steps,
        learning_rate=args.learning_rate,
        warmup_steps=warmup_steps,
        logging_steps=args.logging_steps,
        # In-loop evaluation is disabled: Unsloth's patched unsloth_prediction_step
        # (unsloth/models/rl.py) unconditionally sets UNSLOTH_RETURN_LOGITS=1 and
        # calls compute_loss(..., return_outputs=True) during eval, ignoring
        # prediction_loss_only entirely. That materializes the full
        # (batch, seq_len, ~262K vocab) logits tensor and casts it to fp32 inside
        # the forward pass itself -- verified to OOM (41.6 GiB single allocation)
        # on this model's vocab size regardless of TrainingArguments settings.
        # prediction_loss_only only affects what the Trainer keeps afterward, by
        # which point the OOM has already happened. Real validation/quality
        # checks happen post-training via evaluate_lora_ab.py against the held-out
        # test split instead; save_steps below still gives periodic checkpoints.
        eval_strategy="no",
        save_steps=args.save_steps,
        save_total_limit=3,
        optim="adamw_8bit",
        weight_decay=0.01,
        lr_scheduler_type="linear",
        seed=args.seed,
        bf16=torch.cuda.is_available() and torch.cuda.is_bf16_supported(),
        fp16=torch.cuda.is_available() and not torch.cuda.is_bf16_supported(),
        dataset_num_proc=args.dataset_num_proc,
        dataset_text_field="text",
        max_length=args.max_seq_length,
        packing=args.packing,
        report_to="none",
    )

    trainer = SFTTrainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
        args=train_args,
    )

    if not args.no_response_only_masking:
        # Every training example is now a real multi-turn conversation
        # (system/user/assistant tool-calls/tool results) reconstructed from
        # production traces, not a single flattened completion. Without this,
        # SFT loss would run over system/user/tool tokens too — including
        # tool results the model never generates itself — which both wastes
        # signal and actively trains the model to try to "predict" search
        # results and prior context verbatim. Mask to assistant-authored
        # tokens only, matched via the chat template's own turn markers.
        trainer = train_on_responses_only(
            trainer,
            instruction_part=args.instruction_part,
            response_part=args.response_part,
        )
        print(
            f"Response-only masking enabled (instruction_part={args.instruction_part!r}, "
            f"response_part={args.response_part!r})."
        )
    else:
        print("WARNING: response-only masking disabled via --no-response-only-masking.")

    train_result = trainer.train()
    trainer.save_model(str(args.output_dir))
    tokenizer.save_pretrained(str(args.output_dir))

    metadata = {
        "schema_version": "tenaos_lora_training_run_v1",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "base_model": args.base_model,
        "train_jsonl": str(args.train_jsonl),
        "validation_jsonl": str(args.validation_jsonl),
        "output_dir": str(args.output_dir),
        "load_in_4bit": args.load_in_4bit,
        "chat_template_name": args.chat_template or "<tokenizer_native>",
        "chat_template_prefix": str(tokenizer.chat_template)[:200],
        "response_only_masking": not args.no_response_only_masking,
        "instruction_part": args.instruction_part,
        "response_part": args.response_part,
        "hyperparameters": {
            "lora_rank": 16,
            "lora_alpha": 32,
            "lora_dropout": 0.0,
            "epochs": args.epochs,
            "max_steps": args.max_steps,
            "max_seq_length": args.max_seq_length,
            "per_device_train_batch_size": args.per_device_train_batch_size,
            "gradient_accumulation_steps": args.gradient_accumulation_steps,
            "learning_rate": args.learning_rate,
            "warmup_steps": warmup_steps,
            "seed": args.seed,
            "target_modules": "all-linear",
            "finetune_language_layers": True,
            "finetune_attention_modules": True,
            "finetune_mlp_modules": True,
            "finetune_vision_layers": False,
        },
        "train_metrics": train_result.metrics,
    }
    (args.output_dir / "training_metadata.json").write_text(
        json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True) + "\n",
        encoding="utf-8",
    )

    if args.merge_output_dir:
        args.merge_output_dir.mkdir(parents=True, exist_ok=True)
        model.save_pretrained_merged(str(args.merge_output_dir), tokenizer, save_method="merged_16bit")
        metadata["merge_output_dir"] = str(args.merge_output_dir)
        (args.output_dir / "training_metadata.json").write_text(
            json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True) + "\n",
            encoding="utf-8",
        )

    print(json.dumps(metadata, indent=2, ensure_ascii=False, sort_keys=True))


if __name__ == "__main__":
    try:
        main()
    except BaseException:
        traceback.print_exc()
        raise