"""Convert Agnes with an exact SwiGLU branch merge into the Qwen3.5 GGUF graph.

Run with the same arguments as llama.cpp/convert_hf_to_gguf.py.
The HF checkpoint and its native architecture remain unchanged.
"""
import copy
import json
import runpy
import sys
from pathlib import Path

import torch

ROOT = Path(__file__).resolve().parent
LLAMA = ROOT / "llama.cpp"
sys.path.insert(0, str(LLAMA))
sys.path.insert(0, str(LLAMA / "gguf-py"))

import conversion
from conversion.base import ModelBase, gguf
from conversion.qwen import Qwen3_5TextModel
from conversion.qwen3vl import Qwen3VLVisionModel


@ModelBase.register("AgnesForConditionalGeneration")
class AgnesTextGGUF(Qwen3_5TextModel):
    model_arch = gguf.MODEL_ARCH.QWEN35

    def get_vocab_base(self):
        from transformers import PreTrainedTokenizerFast
        tokenizer = PreTrainedTokenizerFast.from_pretrained(self.dir_model, fix_mistral_regex=False)
        pre = self.get_vocab_base_pre(tokenizer)
        vocab = tokenizer.get_vocab()
        size = self.hparams["vocab_size"]
        if max(vocab.values()) >= size:
            raise ValueError("Tokenizer exceeds embedding vocabulary")
        reverse = {idx: token for token, idx in vocab.items()}
        added = tokenizer.added_tokens_decoder
        tokens, types = [], []
        for idx in range(size):
            token = reverse.get(idx, f"[PAD{idx}]")
            kind = gguf.TokenType.NORMAL
            if idx not in reverse:
                kind = gguf.TokenType.UNUSED
            elif idx in added:
                if not added[idx].normalized:
                    token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False))
                kind = gguf.TokenType.CONTROL if added[idx].special or self.does_token_look_special(token) else gguf.TokenType.USER_DEFINED
            tokens.append(token)
            types.append(kind)
        return tokens, types, pre

    def __init__(self, dir_model, *args, **kwargs):
        config = copy.deepcopy(kwargs.get("hparams") or json.loads((dir_model / "config.json").read_text()))
        text = config.get("text_config", config)
        text["intermediate_size"] += text["parallel_ffn_intermediate_size"]
        text["layer_types"] = [
            {"agnes_delta_attention": "linear_attention", "agnes_global_attention": "full_attention"}[k]
            for k in text["layer_types"]
        ]
        text["full_attention_interval"] = text.get("global_attention_interval", 4)
        text["mtp_num_hidden_layers"] = 0
        kwargs["hparams"] = config
        type(self).no_mtp = True
        super().__init__(dir_model, *args, **kwargs)

    @classmethod
    def filter_tensors(cls, item):
        name, gen = item
        if name.startswith(("mtp.", "model.mtp.")):
            return None
        name = name.replace(".delta_attn.", ".linear_attn.").replace(".global_attn.", ".self_attn.")
        return super().filter_tensors((name, gen))

    def index_tensors(self, remote_hf_model_id=None):
        tensors = super().index_tensors(remote_hf_model_id)
        branches = [name for name in tensors if ".mlp.parallel_ffn." in name]
        if len(branches) != 3 * self.hparams["text_config"]["num_hidden_layers"]:
            raise ValueError("Expected exactly three parallel MLP matrices per decoder layer")
        for branch in branches:
            main = branch.replace(".parallel_ffn", "")
            if main not in tensors:
                raise ValueError(f"Missing main branch: {main}")
            a, b = tensors[main], tensors.pop(branch)
            dim = 1 if main.endswith("down_proj.weight") else 0
            tensors[main] = lambda a=a, b=b, dim=dim: torch.cat((a(), b()), dim=dim)
        return tensors

    def get_vocab_base_pre(self, tokenizer):
        backend = json.loads(tokenizer.backend_tokenizer.to_str())
        expected = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?[\p{L}\p{M}]+|\p{N}| ?[^\s\p{L}\p{M}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"
        sequence = backend["pre_tokenizer"]
        if sequence["type"] != "Sequence":
            raise ValueError("Unexpected pre-tokenizer structure")
        steps = sequence["pretokenizers"]
        if len(steps) != 2 or steps[0].get("pattern", {}).get("Regex") != expected:
            raise ValueError("Agnes pre-tokenizer differs from verified Qwen3.5 regex")
        if steps[1]["type"] != "ByteLevel" or steps[1].get("add_prefix_space") or steps[1].get("use_regex"):
            raise ValueError("Unexpected ByteLevel settings")
        return "qwen35"


@ModelBase.register("AgnesForConditionalGeneration")
class AgnesVisionGGUF(Qwen3VLVisionModel):
    model_arch = Qwen3VLVisionModel.model_arch


conversion.TEXT_MODEL_MAP["AgnesForConditionalGeneration"] = "qwen"
conversion.MMPROJ_MODEL_MAP["AgnesForConditionalGeneration"] = "qwen3vl"

if __name__ == "__main__":
    runpy.run_path(str(LLAMA / "convert_hf_to_gguf.py"), run_name="__main__")