"""A text-only Qwen decoder, LoRA and selected output-embedding rows."""

import functools
import inspect
import json
import os
from pathlib import Path

import torch
import torch.nn.functional as F
from peft import LoraConfig, PeftModel, TaskType, get_peft_model
from torch import nn
from transformers import AutoModelForImageTextToText

from .common import atomic_json


def install_device_dispatch():
    """FLA is CUDA-only; keep Transformers' reference implementation usable on CPU."""
    from transformers.models.qwen3_5 import modeling_qwen3_5 as implementation

    for name in ("torch_chunk_gated_delta_rule", "torch_recurrent_gated_delta_rule"):
        original = getattr(implementation, name)
        if getattr(original, "_themis_dispatch", False):
            continue
        reference = inspect.unwrap(original)

        def dispatch(*args, _accelerated=original, _reference=reference, **kwargs):
            return (_reference if args[0].device.type == "cpu" else _accelerated)(*args, **kwargs)

        dispatch = functools.wraps(original)(dispatch)
        dispatch._themis_dispatch = True
        setattr(implementation, name, dispatch)


install_device_dispatch()


def bidirectional_masks(lengths, width, device):
    real = (
        torch.arange(width, device=device)[None, :] < torch.tensor(lengths, device=device)[:, None]
    )
    return {
        "full_attention": real[:, None, None, :].expand(-1, 1, width, -1),
        "linear_attention": None,
    }


def collate(encoded, pad, device):
    lengths = [len(row["input_ids"]) for row in encoded]
    width = max(lengths)
    ids = torch.full((len(encoded), width), pad, dtype=torch.long, device=device)
    rows, positions, targets, sizes = [], [], [], []
    for index, row in enumerate(encoded):
        ids[index, : lengths[index]] = torch.tensor(row["input_ids"], device=device)
        rows += [index] * len(row["positions"])
        positions += row["positions"]
        sizes += [len(keys) for keys in row["option_keys"]]
        targets += row["targets"]
    answer = {
        "input_ids": ids,
        "attention_mask": bidirectional_masks(lengths, width, device),
        "rows": torch.tensor(rows, device=device),
        "positions": torch.tensor(positions, device=device),
        "sizes": sizes,
        "lengths": lengths,
    }
    if targets:
        t = torch.zeros(len(targets), max(sizes), device=device)
        for i, values in enumerate(targets):
            t[i, : len(values)] = torch.tensor(values, device=device)
        answer["target"] = t
    return answer


def soft_ce_sum(logits, target):
    logp = logits.log_softmax(-1).masked_fill(target == 0, 0)
    return -(target * logp).sum()


class DecisionModel(nn.Module):
    def __init__(self, backbone, output_embeddings, code_ids):
        super().__init__()
        self.backbone = backbone
        self.output_embeddings = output_embeddings
        self.output_embeddings.requires_grad_(False)
        self.register_buffer("code_ids", torch.tensor(code_ids, dtype=torch.long), persistent=False)

    def forward(self, batch):
        hidden = self.backbone(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
            use_cache=False,
            return_dict=True,
        ).last_hidden_state
        selected = hidden[batch["rows"], batch["positions"]]
        weights = self.output_embeddings.weight[self.code_ids[: max(batch["sizes"])]]
        logits = F.linear(selected.float(), weights.float())
        valid = (
            torch.arange(logits.shape[1], device=logits.device)[None, :]
            < torch.tensor(batch["sizes"], device=logits.device)[:, None]
        )
        return logits.masked_fill(~valid, -torch.inf)

    def save_adapter(self, folder, codec, cfg):
        from safetensors.torch import save_file

        folder = Path(folder)
        self.backbone.save_pretrained(folder, safe_serialization=True, save_embedding_layers=False)
        codec.tokenizer.save_pretrained(folder / "tokenizer")
        save_file(
            {
                "answer_embedding": self.backbone.get_input_embeddings()
                .weight[codec.marker]
                .detach()
                .cpu()
                .contiguous()
            },
            str(folder / "answer_embedding.safetensors"),
        )
        atomic_json(folder / "themis_config.json", {"config": cfg, "codec": codec.config()})


def load_model(cfg, codec, device="cuda", adapter=None, trainable=True):
    dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
    base = AutoModelForImageTextToText.from_pretrained(
        os.environ.get("THEMIS_BASE_MODEL_PATH", cfg["model"]["repo"]),
        revision=cfg["model"]["revision"],
        dtype=dtype,
        attn_implementation=cfg["model"]["attention"],
    )
    # Mean initialization is deterministic; it must match when adapters are reloaded.
    if base.get_input_embeddings().num_embeddings < len(codec.tokenizer):
        old = base.get_input_embeddings().weight.detach()
        mean = old.float().mean(0).to(dtype)
        old_count = old.shape[0]
        base.resize_token_embeddings(len(codec.tokenizer), mean_resizing=False)
        with torch.no_grad():
            base.get_input_embeddings().weight[old_count:] = mean
            if not base.config.tie_word_embeddings:
                base.get_output_embeddings().weight[old_count:] = mean
    # Qwen pads its embedding matrix beyond the tokenizer vocabulary. The added
    # marker can fit in an existing untrained row without requiring a resize.
    with torch.no_grad():
        embedding = base.get_input_embeddings().weight
        embedding[codec.marker] = embedding[: codec.marker].float().mean(0).to(embedding.dtype)
        if adapter:
            from safetensors.torch import load_file

            stored = load_file(str(Path(adapter) / "answer_embedding.safetensors"))[
                "answer_embedding"
            ]
            if stored.shape != embedding[codec.marker].shape:
                raise ValueError("Saved placeholder embedding shape mismatch")
            embedding[codec.marker] = stored.to(embedding.dtype)
    decoder = base.get_decoder()
    output = base.get_output_embeddings()
    del base  # The vision tower is released and never moved to GPU.
    decoder.requires_grad_(False)
    output.requires_grad_(False)
    decoder.config.use_cache = False
    if adapter:
        saved = json.loads((Path(adapter) / "themis_config.json").read_text())
        if saved["codec"] != codec.config() or saved["config"]["model"] != cfg["model"]:
            raise ValueError("Adapter encoder/model configuration mismatch")
        decoder = PeftModel.from_pretrained(decoder, adapter, is_trainable=trainable)
    else:
        decoder = get_peft_model(
            decoder,
            LoraConfig(
                task_type=TaskType.FEATURE_EXTRACTION,
                target_modules="all-linear",
                r=cfg["model"]["lora_rank"],
                lora_alpha=cfg["model"]["lora_alpha"],
                lora_dropout=cfg["model"]["lora_dropout"],
                bias="none",
            ),
        )
    if trainable and cfg["train"]["checkpointing"]:
        decoder.gradient_checkpointing_enable(
            gradient_checkpointing_kwargs={"use_reentrant": False}
        )
    model = DecisionModel(decoder, output, codec.code_ids).to(device)
    trainable_names = [name for name, p in model.named_parameters() if p.requires_grad]
    if trainable and (not trainable_names or any("lora_" not in name for name in trainable_names)):
        raise ValueError("Only text-decoder LoRA parameters may be trainable")
    return model