"""Gemma 3 (text) for the ExecuTorch WebGPU playground: model, export and CPU validation.

ExecuTorch 1.5.1's export_llm has no Gemma 3 model class, and its llama_transformer lacks the
pieces Gemma 3 needs (post-attention/post-FFN norms, embedding scaling, GELU-tanh MLP, per-layer
RoPE base, sliding-window layers). This module re-implements the Gemma 3 decoder with the ops the
public WebGPU runtime registers and lowers it through the same LLMEdgeManager pipeline export_llm
uses (same quantization transforms, edge config, to_executorch config), with WebGPUPartitioner.

Program contract (same as the other playground LLMs):
  forward(tokens int64 [1, 1..prefill_max], input_pos int64 [1]) -> logits float32 [1, vocab]
  (last position), KV caches inside the program as fp32 mutable buffers.

Attention:
  * global layers: llama.sdpa_with_kv_cache (writes K/V at input_pos, causal over [0, pos+S)).
  * sliding-window layers: the public runtime has no masked/windowed KV-cache op and its
    stand-alone update_cache does not accept a dynamic position. So each local layer keeps a
    position-indexed cache with PAD leading rows (position x lives at row x + PAD) and
      1. writes K/V with sdpa_with_kv_cache(k, k, v, cache, pos + PAD) (1-head dummy query;
         its output is discarded, it only performs the cache write),
      2. reads a fixed-length window of L = window + prefill_max - 1 rows ending at the last
         token of the chunk, and attends with bmm -> + mask -> softmax -> bmm.
    The additive mask is band[prefill_max - S : prefill_max] (relative positions, constant) plus
    padbias[pos + S - 1 : ... + L] (masks rows that hold no position yet). This reproduces
    transformers' sliding mask exactly: key visible iff 0 <= q_pos - k_pos < sliding_window.
"""

import contextlib
import json
import re
import sys
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import nn

NEG = -1.0e30  # finite: WGSL does not guarantee IEEE infinities; NEG + NEG stays finite


def _load_custom_ops():
    import executorch.extension.llm.custom_ops.custom_ops  # noqa: F401  (llama::sdpa_with_kv_cache)


class RMSNorm(nn.Module):
    """Gemma RMSNorm with the (1 + w) offset folded into `weight` (fp32), in the op order the
    Vulkan rms_norm pattern matches: x * rsqrt(mean(x^2) + eps) * weight."""

    def __init__(self, dim, eps):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim), requires_grad=False)

    def forward(self, x):
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x * self.weight


def _rope(q, k, cos, sin):
    # HF rotate_half RoPE through the exact module the Vulkan/WebGPU hf_rope pattern matches
    # (fused into et_vk.apply_rotary_emb_hf); in fp32 it computes q*cos + rotate_half(q)*sin.
    from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingPattern

    return HfRotaryEmbeddingPattern()(q, k, cos, sin)


class Layer(nn.Module):
    def __init__(self, cfg, sliding, ctx, pad):
        super().__init__()
        d, hd = cfg["hidden_size"], cfg["head_dim"]
        self.nh, self.nkv, self.hd = (
            cfg["num_attention_heads"],
            cfg["num_key_value_heads"],
            hd,
        )
        self.sliding = sliding
        eps = cfg["rms_norm_eps"]
        self.wq = nn.Linear(d, self.nh * hd, bias=False)
        self.wk = nn.Linear(d, self.nkv * hd, bias=False)
        self.wv = nn.Linear(d, self.nkv * hd, bias=False)
        self.wo = nn.Linear(self.nh * hd, d, bias=False)
        self.q_norm = RMSNorm(hd, eps)
        self.k_norm = RMSNorm(hd, eps)
        self.attention_norm = RMSNorm(d, eps)
        self.post_attention_norm = RMSNorm(d, eps)
        self.ffn_norm = RMSNorm(d, eps)
        self.post_ffn_norm = RMSNorm(d, eps)
        self.w1 = nn.Linear(d, cfg["intermediate_size"], bias=False)  # gate_proj
        self.w3 = nn.Linear(d, cfg["intermediate_size"], bias=False)  # up_proj
        self.w2 = nn.Linear(cfg["intermediate_size"], d, bias=False)  # down_proj
        self.rows = ctx + (pad if sliding else 0)
        self.make_buffers()

    def make_buffers(self):
        shape = (1, self.rows, self.nkv, self.hd)
        with torch.device("cpu"):
            self.register_buffer("k_cache", torch.zeros(shape), persistent=False)
            self.register_buffer("v_cache", torch.zeros(shape), persistent=False)

    def attention(self, x, p, S, cos, sin, mask, pad, win):
        q = self.wq(x).view(1, S, self.nh, self.hd)
        k = self.wk(x).view(1, S, self.nkv, self.hd)
        v = self.wv(x).view(1, S, self.nkv, self.hd)
        q, k = _rope(self.q_norm(q), self.k_norm(k), cos, sin)
        if not self.sliding:
            out = torch.ops.llama.sdpa_with_kv_cache(
                q, k, v, self.k_cache, self.v_cache, p, S, None, 0.0, True, None
            )
            return out.view(1, S, self.nh * self.hd)
        # cache write only (output unused); query = k keeps it to a single head
        torch.ops.llama.sdpa_with_kv_cache(
            k, k, v, self.k_cache, self.v_cache, p + pad, S, None, 0.0, True, None
        )
        start = p + S - 1
        kw = self.k_cache.narrow(1, start, win).view(1, win, self.hd).transpose(1, 2)
        vw = self.v_cache.narrow(1, start, win).view(1, win, self.hd)
        qf = q.reshape(1, S * self.nh, self.hd)  # row = s * nh + h (single KV head)
        scores = torch.bmm(qf, kw) * (self.hd**-0.5)
        scores = scores.view(1, S, self.nh, win) + mask
        probs = torch.softmax(scores, dim=-1).view(1, S * self.nh, win)
        return torch.bmm(probs, vw).view(1, S, self.nh * self.hd)

    def forward(self, x, p, S, cos, sin, mask, pad, win):
        h = x + self.post_attention_norm(
            self.wo(self.attention(self.attention_norm(x), p, S, cos, sin, mask, pad, win))
        )
        x = self.ffn_norm(h)
        ff = self.w2(F.gelu(self.w1(x), approximate="tanh") * self.w3(x))
        return h + self.post_ffn_norm(ff)


class Gemma3Text(nn.Module):
    def __init__(self, cfg, max_context, prefill_max):
        super().__init__()
        assert cfg["num_key_value_heads"] == 1, (
            "window attention folds heads; needs 1 KV head"
        )
        assert cfg["query_pre_attn_scalar"] == cfg["head_dim"], (
            "sdpa uses 1/sqrt(head_dim)"
        )
        assert cfg.get("attn_logit_softcapping") is None
        assert cfg.get("final_logit_softcapping") is None
        assert cfg["hidden_activation"] == "gelu_pytorch_tanh"
        self.cfg = cfg
        self.ctx, self.prefill_max = max_context, prefill_max
        self.window = cfg["sliding_window"]
        self.win = self.window + prefill_max - 1  # rows read per local layer
        self.pad = self.win - 1
        self.layer_types = layer_types(cfg)
        d, V = cfg["hidden_size"], cfg["vocab_size"]
        self.tok_embeddings = nn.Embedding(V, d)
        # HF: embed(ids) * tensor(hidden_size ** 0.5) in the weight dtype (fp32 here)
        self.embed_scale = float(torch.tensor(d**0.5, dtype=torch.float32, device="cpu"))
        self.layers = nn.ModuleList(
            Layer(cfg, t == "sliding_attention", max_context, self.pad)
            for t in self.layer_types
        )
        self.norm = RMSNorm(d, cfg["rms_norm_eps"])
        self.output = nn.Linear(d, V, bias=False)
        self.make_buffers()

    def make_buffers(self):
        """Constant/state buffers (rope tables, window masks, KV caches), always built on CPU."""
        with torch.device("cpu"):
            cos_g, sin_g, cos_l, sin_l = rope_tables(self.cfg, self.ctx)
            self.register_buffer("cos_global", cos_g, persistent=False)
            self.register_buffer("sin_global", sin_g, persistent=False)
            self.register_buffer("cos_local", cos_l, persistent=False)
            self.register_buffer("sin_local", sin_l, persistent=False)
            r = torch.arange(self.prefill_max)[:, None]
            j = torch.arange(self.win)[None, :]
            band = torch.where((j >= r) & (j <= r + self.window - 1), 0.0, NEG)
            self.register_buffer("band", band.float(), persistent=False)
            padbias = torch.zeros(self.pad + self.ctx)
            padbias[: self.pad] = NEG
            self.register_buffer("padbias", padbias, persistent=False)
            for layer in self.layers:
                layer.make_buffers()

    def forward(self, tokens, input_pos):
        S = tokens.shape[1]
        p = input_pos[0].item()
        torch._check_is_size(p)
        torch._check(p + S <= self.ctx)
        h = self.tok_embeddings(tokens) * self.embed_scale
        cg, sg = self.cos_global.narrow(0, p, S), self.sin_global.narrow(0, p, S)
        cl, sl = self.cos_local.narrow(0, p, S), self.sin_local.narrow(0, p, S)
        mask = self.band.narrow(0, self.prefill_max - S, S) + self.padbias.narrow(
            0, p + S - 1, self.win
        )
        mask = mask.view(1, S, 1, self.win)
        for layer in self.layers:
            if layer.sliding:
                h = layer(h, p, S, cl, sl, mask, self.pad, self.win)
            else:
                h = layer(h, p, S, cg, sg, None, self.pad, self.win)
        h = self.norm(h[:, -1, :])
        return self.output(h)


def layer_types(cfg):
    if cfg.get("layer_types"):
        return list(cfg["layer_types"])
    pat = cfg.get("sliding_window_pattern", 6)
    return [
        "sliding_attention" if (i + 1) % pat else "full_attention"
        for i in range(cfg["num_hidden_layers"])
    ]


def rope_tables(cfg, ctx):
    """cos/sin tables computed by transformers' own Gemma3RotaryEmbedding (bit-identical to HF)."""
    from transformers import Gemma3TextConfig
    from transformers.models.gemma3.modeling_gemma3 import Gemma3RotaryEmbedding

    hf_cfg = Gemma3TextConfig(**{k: v for k, v in cfg.items() if not k.startswith("_")})
    rope = Gemma3RotaryEmbedding(hf_cfg)
    pos = torch.arange(ctx)[None]
    x = torch.zeros(1, dtype=torch.float32)
    cg, sg = rope(x, pos, "full_attention")
    cl, sl = rope(x, pos, "sliding_attention")
    return (
        cg[0].contiguous(),
        sg[0].contiguous(),
        cl[0].contiguous(),
        sl[0].contiguous(),
    )


HF_TO_ET = {
    "input_layernorm": "attention_norm",
    "post_attention_layernorm": "post_attention_norm",
    "pre_feedforward_layernorm": "ffn_norm",
    "post_feedforward_layernorm": "post_ffn_norm",
    "self_attn.q_proj": "wq",
    "self_attn.k_proj": "wk",
    "self_attn.v_proj": "wv",
    "self_attn.o_proj": "wo",
    "self_attn.q_norm": "q_norm",
    "self_attn.k_norm": "k_norm",
    "mlp.gate_proj": "w1",
    "mlp.up_proj": "w3",
    "mlp.down_proj": "w2",
}


def build(snapshot_dir, max_context, prefill_max, linear_dtype=torch.bfloat16):
    """Gemma3Text with checkpoint weights. Linear/embedding weights stay in `linear_dtype`
    (the checkpoint's bf16, so quantization sees the stored values, as export_llm does); norms
    are (1 + w) in fp32 exactly as transformers computes them; output is tied to the embedding."""
    from safetensors.torch import load_file

    snap = Path(snapshot_dir)
    cfg = json.loads((snap / "config.json").read_text())
    sd = load_file(str(snap / "model.safetensors"))
    with torch.device("meta"):
        model = Gemma3Text(cfg, max_context, prefill_max)
    model = model.to_empty(device="cpu")
    model.make_buffers()
    used = set()

    def take(key):
        used.add(key)
        return sd[key]

    def norm(dst, key):
        dst.weight = nn.Parameter(1.0 + take(key).float(), requires_grad=False)

    def lin(dst, key):
        dst.weight = nn.Parameter(take(key).to(linear_dtype), requires_grad=False)

    pre = "model." if "model.embed_tokens.weight" in sd else ""
    emb = take(f"{pre}embed_tokens.weight").to(linear_dtype)
    model.tok_embeddings.weight = nn.Parameter(emb, requires_grad=False)
    model.output.weight = nn.Parameter(emb.clone(), requires_grad=False)  # tied head
    norm(model.norm, f"{pre}norm.weight")
    for i, layer in enumerate(model.layers):
        for hf, et in HF_TO_ET.items():
            key = f"{pre}layers.{i}.{hf}.weight"
            dst = layer.get_submodule(et)
            (norm if isinstance(dst, RMSNorm) else lin)(dst, key)
    unused = set(sd) - used - {"lm_head.weight"}
    assert not unused, f"unused checkpoint tensors: {sorted(unused)[:5]}"
    return model.eval(), cfg


def reset_caches(model):
    for layer in model.layers:
        layer.k_cache.zero_()
        layer.v_cache.zero_()


def quantize(model, group_size, embedding_group_size):
    """Same transforms and order as export_llm (embedding first, then linears), quantizing the
    checkpoint-dtype weights, then casting everything else to fp32."""
    from executorch.examples.models.llama.source_transformation.quantize import (
        get_quant_embedding_transform,
        get_quant_weight_transform,
    )
    from executorch.extension.llm.export.builder import DType

    model = get_quant_embedding_transform(
        f"4,{embedding_group_size}", False, quantize_with_hqq=True
    )(model)
    model = get_quant_weight_transform(
        quantization_mode="4w",
        group_size=group_size,
        computation_dtype=DType.fp32,
        checkpoint_dtype=DType.bf16,
        quantize_with_hqq=True,
    )(model)
    return model.to(dtype=torch.float32)


def edge_manager(model, prefill_max, max_context, cfg, metadata_extra=None):
    from executorch.extension.llm.export.builder import DType, LLMEdgeManager

    metadata = {
        "get_max_seq_len": prefill_max + 1,
        "get_max_context_len": max_context,
        "get_n_layers": cfg["num_hidden_layers"],
        "get_vocab_size": cfg["vocab_size"],
        "use_kv_cache": True,
        "use_sdpa_with_kv_cache": True,
        "enable_dynamic_shape": True,
        "get_bos_id": cfg["bos_token_id"],
        "get_eos_ids": list(cfg["eos_token_id"]),
        **(metadata_extra or {}),
    }
    return LLMEdgeManager(
        model=model,
        modelname="gemma3",
        max_seq_len=prefill_max + 1,
        dtype=DType.fp32,
        use_kv_cache=True,
        example_inputs=(
            torch.tensor([[2, 3, 4]], dtype=torch.long),
            torch.tensor([0], dtype=torch.long),
        ),
        dynamic_shapes=({1: torch.export.Dim("token_dim", max=prefill_max)}, None),
        enable_dynamic_shape=True,
        metadata=metadata,
    )


def webgpu_partitioner():
    from executorch.backends.webgpu.partitioner import WebGPUPartitioner

    # Same compile options export_llm passes to VulkanPartitioner for LLMs (dynamic shapes, fp32).
    return WebGPUPartitioner(
        compile_options={"require_dynamic_shapes": True, "force_fp16": False}
    )


_KV_CACHE = re.compile(r"(^|\.)(k|v)_cache$")


@contextlib.contextmanager
def live_kv_cache_reads():
    """Lowering-time fix so sliding-window layers read the live KV cache.

    The Vulkan/WebGPU lowering serializes mutable buffers (KV caches) as delegate constants.
    llama.sdpa_with_kv_cache handles its own prepacking, so it binds the cache tensor itself and
    its in-place writes persist across calls. Every other reader (here: the slice_copy that cuts
    the attention window) gets an et_vk.prepack node, i.e. a separate GPU copy made once at
    load, so it would only ever see the initial zeros. This wraps the backend's
    insert_prepack_nodes pass and, for buffers named *.k_cache / *.v_cache only, removes those
    prepack nodes and points their users at the cache tensor itself, so the slices bind the
    same GPU buffer sdpa_with_kv_cache writes. The runtime is unchanged.
    """
    import executorch.backends.vulkan.vulkan_preprocess as vp
    from executorch.exir.dialects._ops import ops as exir_ops

    orig = vp.insert_prepack_nodes
    rewired = []

    def patched(program):
        program = orig(program)
        buffers = program.graph_signature.inputs_to_buffers
        graph = program.graph_module.graph
        for node in list(graph.nodes):
            if node.target != exir_ops.edge.et_vk.prepack.default:
                continue
            src = node.args[0]
            if src.op == "placeholder" and _KV_CACHE.search(buffers.get(src.name, "")):
                node.replace_all_uses_with(src)
                graph.erase_node(node)
                rewired.append(buffers[src.name])
        program.graph_module.recompile()
        return program

    vp.insert_prepack_nodes = patched
    try:
        yield rewired
    finally:
        vp.insert_prepack_nodes = orig


def lower(model, prefill_max, max_context, cfg):
    _load_custom_ops()
    em = edge_manager(model, prefill_max, max_context, cfg).export()
    em.run_canonical_optimizations()
    em = em.export_to_edge()
    with live_kv_cache_reads() as rewired:
        em = em.to_backend([webgpu_partitioner()])
    n_local = sum(t == "sliding_attention" for t in layer_types(cfg)[: len(model.layers)])
    assert len(rewired) == 2 * n_local, f"rewired {len(rewired)} cache reads, expected {2 * n_local}: {rewired}"
    em = em.to_executorch()
    em.rewired_kv_caches = sorted(rewired)
    return em