"""Shared model/config loading for train_sft.py, train_dpo.py, train_kto.py."""
from __future__ import annotations

import os
from pathlib import Path
from typing import Any

import torch
import yaml
from peft import LoraConfig
from transformers import AutoModelForImageTextToText, AutoProcessor, BitsAndBytesConfig

# All linear layers of the language-model backbone (attention proj + MLP).
# QLoRA (Dettmers et al., 2305.14314) found this — not rank — is the critical
# LoRA hyperparameter: "LoRA on all linear transformer block layers are
# required to match full finetuning performance... the projection dimension r
# does not affect performance." LoRA (Hu et al., 2106.09685) independently
# found spreading a fixed budget across more matrix types beats concentrating
# it in fewer at higher rank. Vision-tower/projector linears deliberately left
# out — adapting only the language backbone is the cheaper, lower-forgetting-
# risk default; revisit if the vision side turns out to be the bottleneck.
DEFAULT_LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj",
]


def load_config(path: str) -> dict[str, Any]:
    with open(path) as f:
        return yaml.safe_load(f)


def build_quantization_config(cfg: dict) -> BitsAndBytesConfig | None:
    if not cfg.get("qlora", {}).get("enabled", False):
        return None
    q = cfg["qlora"]
    return BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type=q.get("quant_type", "nf4"),
        bnb_4bit_compute_dtype=getattr(torch, q.get("compute_dtype", "bfloat16")),
        bnb_4bit_use_double_quant=q.get("double_quant", True),
    )


def build_lora_config(cfg: dict) -> LoraConfig:
    l = cfg["lora"]
    # NOTE: `l.get("target_modules", DEFAULT)` would NOT apply the default when
    # the YAML explicitly sets `target_modules: null` (every configs/*.yaml
    # does) — a present key with value None defeats dict.get's default arg,
    # it only falls back on a MISSING key. `or` catches both missing-and-None.
    # Getting this wrong would silently pass target_modules=None into
    # LoraConfig, which makes PEFT try to auto-detect target modules from a
    # model-type table that has no entry for this checkpoint's "qwen3_5"
    # architecture — almost certainly a hard failure, not a quiet fallback.
    target_modules = l.get("target_modules") or DEFAULT_LORA_TARGET_MODULES
    return LoraConfig(
        r=l.get("r", 16),
        lora_alpha=l.get("alpha", 32),
        lora_dropout=l.get("dropout", 0.05),
        target_modules=target_modules,
        bias="none",
        task_type="CAUSAL_LM",
    )


def load_model_and_processor(cfg: dict):
    model_id = cfg["model"]["base_model"]
    revision = cfg["model"].get("revision")
    quant_config = build_quantization_config(cfg)
    # "sdpa" ships in PyTorch itself — no extra install, no CUDA/torch-version-
    # matched wheel to get wrong. "flash_attention_2" is faster but requires
    # the flash-attn package (NOT in requirements.txt — it needs a build step
    # matched to your exact torch/CUDA version, easy to get wrong on a fresh
    # instance) — opt into it explicitly per-config once flash-attn is
    # confirmed installed, don't make it the default something can silently
    # depend on.
    attn_implementation = cfg["model"].get("attn_implementation", "sdpa")

    processor = AutoProcessor.from_pretrained(
        model_id, revision=revision, trust_remote_code=True,
    )
    model_kwargs = dict(
        revision=revision,
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        quantization_config=quant_config,
        device_map="auto",
        attn_implementation=attn_implementation,
    )
    try:
        model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
    except (ValueError, KeyError) as e:
        # This checkpoint family's HF model card metadata reports
        # transformersInfo.auto_model = "AutoModelForMultimodalLM" — not a
        # standard transformers auto-class name as of when this was written,
        # possibly custom-code-registered via trust_remote_code. If the
        # standard VLM auto-class doesn't recognize the architecture, fall
        # back to the fully generic AutoModel, which resolves purely through
        # the checkpoint's own registered config/model classes and doesn't
        # care what auto-class name shows up in unrelated HF metadata.
        print(
            f"AutoModelForImageTextToText failed ({e!r}) — retrying with generic "
            f"AutoModel (trust_remote_code). See README.md's pre-flight checklist."
        )
        from transformers import AutoModel
        model = AutoModel.from_pretrained(model_id, **model_kwargs)
    return model, processor


def build_vlm_sft_collator(processor):
    """Custom data collator for multimodal SFT.

    Required, not optional: SFTTrainer's default dataset-preparation pipeline
    assumes text-only tokenization/packing and will mishandle (or error on)
    an "images" column — this is why train_sft.py pairs this collator with
    `dataset_kwargs={"skip_prepare_dataset": True}`. DPO/KTO/SimPO's trainers
    are NOT given an equivalent custom collator here — TRL's own VLM examples
    for those trainers consume prompt/chosen/rejected/completion + images
    columns directly via `processing_class` without one; SFT is the one that
    actually needs this. See README.md's pre-flight checklist — this collator
    was written against the general HF VLM processor contract
    (apply_chat_template + images kwarg), not run against this specific
    "qwen3_5" architecture's actual processor.
    """
    tokenizer = getattr(processor, "tokenizer", processor)

    def collate(examples: list[dict]) -> dict:
        texts = []
        images = []
        prompt_lengths = []
        for ex in examples:
            full_text = processor.apply_chat_template(ex["prompt"] + ex["completion"], tokenize=False)
            texts.append(full_text)
            images.append(ex["images"])
            prompt_text = processor.apply_chat_template(
                ex["prompt"], tokenize=False, add_generation_prompt=True
            )
            prompt_lengths.append(len(tokenizer(prompt_text, add_special_tokens=False)["input_ids"]))

        batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
        labels = batch["input_ids"].clone()
        for i, plen in enumerate(prompt_lengths):
            labels[i, :plen] = -100
        if tokenizer.pad_token_id is not None:
            labels[batch["input_ids"] == tokenizer.pad_token_id] = -100
        batch["labels"] = labels
        return batch

    return collate


def init_wandb(cfg: dict, run_name: str) -> None:
    """TRL's *Trainer classes report to wandb automatically when
    report_to=["wandb"] is set in the *Config and WANDB_API_KEY/WANDB_PROJECT
    are in the environment — this just makes the project/run naming explicit
    and fails fast (before an expensive model load) if the key is missing."""
    if "wandb" not in (cfg.get("training", {}).get("report_to") or []):
        return
    if not os.environ.get("WANDB_API_KEY"):
        raise RuntimeError(
            "training.report_to includes 'wandb' but WANDB_API_KEY is not set. "
            "Run `wandb login` or `export WANDB_API_KEY=...` before launching, "
            "or remove 'wandb' from the config to train without it."
        )
    os.environ.setdefault("WANDB_PROJECT", cfg.get("wandb", {}).get("project", "404-coder-training"))
    os.environ.setdefault("WANDB_RUN_GROUP", cfg.get("wandb", {}).get("group", run_name))


def resolve_output_dir(cfg: dict, default: str) -> str:
    out = cfg.get("training", {}).get("output_dir", default)
    Path(out).mkdir(parents=True, exist_ok=True)
    return out


# --- Optional Muon–AdamW hybrid optimizer (opt-in via `optim: "muon_adamw"`) ---
# See muon.py for the what/why and the honest caveat (naive Muon-on-LoRA is a
# heuristic). Kept off the default path — configs ship paged_adamw_8bit — so
# nothing here runs unless a config explicitly asks for it.

def resolve_optim(cfg: dict) -> tuple[str, bool]:
    """Return (hf_optim_string, use_muon). If the config asks for our
    "muon_adamw" extension, we hand HF a harmless valid optim string (its
    optimizer is then replaced by our create_optimizer override — see
    make_muon_trainer) because "muon_adamw" is NOT a valid HF OptimizerNames
    enum and would fail TrainingArguments validation if passed through."""
    o = cfg.get("training", {}).get("optim", "adamw_torch")
    if o == "muon_adamw":
        return "adamw_torch", True
    return o, False


def build_muon_config(cfg: dict) -> dict:
    """The `training.muon` block, with Muon's lr defaulting to the run's base
    learning_rate (LoRA makes Muon robust to a suboptimal lr — arXiv 2605.10468
    — so reusing the Adam lr is the safe default; raise training.muon.lr toward
    Muon's native ~0.01–0.02 once a baseline works)."""
    t = cfg.get("training", {})
    muon = dict(t.get("muon") or {})
    muon.setdefault("lr", t.get("learning_rate", 1e-4))
    return muon


def make_muon_trainer(base_cls):
    """Subclass a TRL/HF trainer so create_optimizer builds the Muon–AdamW
    hybrid from the *already-PEFT-wrapped* self.model. Set `.muon_config` on
    the instance before .train(); if unset, behaves exactly like the base
    trainer. Works uniformly whether PEFT wrapping happened inside the trainer
    (SFT) or before it (KTO with an sft_adapter)."""
    class _MuonTrainer(base_cls):
        muon_config: dict | None = None

        def create_optimizer(self):
            if self.optimizer is None and self.muon_config:
                from muon import build_muon_adamw_optimizer
                self.optimizer = build_muon_adamw_optimizer(self.model, self.muon_config)
            return super().create_optimizer()  # returns self.optimizer once set
    _MuonTrainer.__name__ = f"Muon{base_cls.__name__}"
    return _MuonTrainer