"""KV-cached incremental decoding for the Dory port (dory/modeling.py).

Every (physical attention layer, loop pass) owns its own key/value cache, like
vLLM's ``recurrent_kv_cache_mode="per_loop"``: pass 1 of the recurrent block
attends over pass-1 keys of earlier tokens, pass 2 over pass-2 keys, and so on.

Caches are preallocated once per request:

* full-attention layers: ``[B, kv_heads, capacity, head_dim]`` indexed by position;
* sliding-window layers: a ring buffer of ``window`` slots, slot = position % window.
  The port's window keeps keys with ``0 <= i - j < window``, i.e. the current token
  plus the previous ``window - 1`` tokens, which is exactly what the ring holds.

The decode step always attends over the whole buffer with a boolean validity mask
(slot <= position), so its shapes are static and it can be captured in a CUDA graph.
RoPE is applied to keys before they are cached (as in the port), so no position
bookkeeping is needed at read time. Keys are stored after qk-norm; queries get it too.

Only batch size 1 is exercised, but nothing assumes it except ``sample``.
"""

import math

import torch
import torch.nn.functional as F

from .modeling import DoryForCausalLM, apply_rope_neox, causal_mask, justnorm, rope_cos_sin, slerp_alpha


class KVCache:
    """Per-(layer, loop) caches for one request."""

    def __init__(self, model: DoryForCausalLM, capacity: int, n_loops: int, batch: int = 1, device="cpu",
                 dtype=None):
        cfg = model.config
        p = next(model.parameters())
        dtype = dtype or p.dtype
        self.cfg, self.capacity, self.device = cfg, capacity, torch.device(device)
        self.window = cfg.sliding_window
        self.pos = 0  # number of tokens already cached
        self.k, self.v = {}, {}
        n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers
        for i, lt in enumerate(cfg.layer_types):
            if lt == "mlp":
                continue
            passes = range(n_loops) if n_in <= i < n_in + n_rec else range(1)
            cap = self.window if lt == "sliding_attention" else capacity
            for lp in passes:
                shape = (batch, cfg.num_key_value_heads, cap, cfg.head_dim)
                self.k[i, lp] = torch.zeros(shape, device=device, dtype=dtype)
                self.v[i, lp] = torch.zeros(shape, device=device, dtype=dtype)
        self.slots = torch.arange(max(capacity, self.window), device=device)
        # RoPE tables for every position and both thetas: [capacity, head_dim/2]
        pos = torch.arange(capacity, device=device)[None]
        self.rope = {}
        for lt in set(cfg.layer_types) - {"mlp"}:
            c, s = rope_cos_sin(pos, cfg.head_dim, cfg.rope_parameters[lt]["rope_theta"], device)
            self.rope[lt] = (c[0], s[0])

    def slot_of(self, layer_type, positions):
        return positions % self.window if layer_type == "sliding_attention" else positions


def _qkv(mixer, x, cos, sin):
    b, t, _ = x.shape
    h, kv, d = mixer.num_heads, mixer.num_kv_heads, mixer.head_dim
    q = mixer.q_proj(x).view(b, t, h, d)
    k = mixer.k_proj(x).view(b, t, kv, d)
    v = mixer.v_proj(x).view(b, t, kv, d)
    gate = mixer.g_proj(x)
    q, k = apply_rope_neox(q, cos, sin), apply_rope_neox(k, cos, sin)
    sqk = (mixer.sqk.float() * mixer.sqk_factor).view(h, d)
    q = (sqk.square() * justnorm(q).float()).to(q.dtype)
    k = justnorm(k)
    return q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), gate  # [B, H|kv, T, D]


def _finish(mixer, out, gate):
    b, h, t, d = out.shape
    out = out.transpose(1, 2).reshape(b, t, h * d)
    sg = (mixer.s_gate.float() * mixer.s_gate_factor).square()
    out = (out.float() * torch.sigmoid(gate.float() * sg)).to(out.dtype)
    return mixer.o_proj(justnorm(out))


def _attend(q, k, v, mask, d):
    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=1.0 / math.sqrt(d), enable_gqa=True)


def attn_prefill(mixer, x, cache: KVCache, key, mask):
    """x: [B, T, H] for positions 0..T-1 of an empty cache. Writes the cache, returns the mixer output."""
    lt = mixer.layer_type
    t = x.shape[1]
    positions = torch.arange(t, device=x.device)
    cos, sin = cache.rope[lt]
    q, k, v, gate = _qkv(mixer, x, cos[None, positions], sin[None, positions])
    kc, vc = cache.k[key], cache.v[key]
    if lt == "sliding_attention" and t > cache.window:  # only the newest `window` keys survive
        keep = t - cache.window
        slots = cache.slot_of(lt, positions[keep:])
        kc[:, :, slots], vc[:, :, slots] = k[:, :, keep:], v[:, :, keep:]
    else:
        slots = cache.slot_of(lt, positions)
        kc[:, :, slots], vc[:, :, slots] = k, v
    out = _attend(q, k, v, mask, mixer.head_dim)
    return _finish(mixer, out, gate)


def attn_decode(mixer, x, cache: KVCache, key, pos, valid):
    """x: [B, 1, H] at scalar-tensor position `pos`; `valid[lt]` is the bool key mask [1,1,1,cap]."""
    lt = mixer.layer_type
    cos, sin = cache.rope[lt]
    idx = pos.view(1)  # index_select, not cos[pos]: a 0-d tensor index syncs and breaks graph capture
    q, k, v, gate = _qkv(mixer, x, cos.index_select(0, idx)[None], sin.index_select(0, idx)[None])
    slot = cache.slot_of(lt, pos)
    kc, vc = cache.k[key], cache.v[key]
    kc.index_copy_(2, slot.view(1), k)
    vc.index_copy_(2, slot.view(1), v)
    out = _attend(q, kc, vc, valid[lt], mixer.head_dim)
    return _finish(mixer, out, gate)


def _backbone(model, h, n_loops, run_layer):
    """Shared layer schedule: input layers, n_loops recurrent passes, output layers."""
    cfg = model.config
    layers = model.backbone.layers
    n_in, n_rec = cfg.n_input_layers, cfg.n_recurrent_layers
    for i in range(n_in):
        h = run_layer(i, 0, h)
    input_h = h
    state = input_h
    for lp in range(n_loops):
        h = justnorm(state + input_h)
        for i in range(n_in, n_in + n_rec):
            h = run_layer(i, lp, h)
        state = justnorm(h)
    h = state
    for i in range(n_in + n_rec, len(layers)):
        h = run_layer(i, 0, h)
    return h


def _layer(model, i, lp, h, attn):
    layer = model.backbone.layers[i]
    if layer.is_attention:
        out = attn(layer.mixer, h, (i, lp))
        return slerp_alpha(h, out, layer.attn_alpha, layer.alpha_scale)
    return slerp_alpha(h, layer.mixer(h), layer.mlp_alpha, layer.alpha_scale)


@torch.no_grad()
def prefill(model, input_ids, cache: KVCache, n_loops):
    """Run the prompt (positions 0..T-1), fill the caches, return logits of the last token [B, V]."""
    assert cache.pos == 0
    t = input_ids.shape[1]
    dev = input_ids.device
    masks = {lt: causal_mask(t, model.config.sliding_window if lt == "sliding_attention" else None, dev)
             for lt in set(model.config.layer_types) - {"mlp"}}

    def attn(mixer, h, key):
        return attn_prefill(mixer, h, cache, key, masks[mixer.layer_type])

    h = model.backbone.embeddings(input_ids)
    h = _backbone(model, h, n_loops, lambda i, lp, h: _layer(model, i, lp, h, attn))
    cache.pos = t
    return model.logits(h[:, -1])


@torch.no_grad()
def decode_step(model, token, cache: KVCache, n_loops, pos):
    """One token [B, 1] at position `pos` (0-d long tensor). Returns logits [B, V]."""
    slots = cache.slots
    valid = {"full_attention": (slots[: cache.capacity] <= pos)[None, None, None],
             "sliding_attention": (slots[: cache.window] <= pos)[None, None, None]}

    def attn(mixer, h, key):
        return attn_decode(mixer, h, cache, key, pos, valid)

    h = model.backbone.embeddings(token)
    h = _backbone(model, h, n_loops, lambda i, lp, h: _layer(model, i, lp, h, attn))
    return model.logits(h[:, -1])


def sample(logits, temperature=1.0, top_p=0.95, generator=None):
    """Nucleus sampling on [1, V] float logits; temperature 0 = greedy."""
    if temperature <= 0:
        return logits.argmax(-1, keepdim=True)
    probs = torch.softmax(logits / temperature, dim=-1)
    if top_p < 1.0:
        sp, si = probs.sort(descending=True, dim=-1)
        cum = sp.cumsum(-1)
        sp[cum - sp > top_p] = 0.0  # keep the smallest prefix whose mass >= top_p
        choice = torch.multinomial(sp / sp.sum(-1, keepdim=True), 1, generator=generator)
        return si.gather(-1, choice)
    return torch.multinomial(probs, 1, generator=generator)


class Decoder:
    """Optional CUDA-graph wrapper around decode_step with static input/position buffers."""

    def __init__(self, model, cache, n_loops, use_graph=False):
        self.model, self.cache, self.n_loops = model, cache, n_loops
        dev = cache.device
        self.tok = torch.zeros(1, 1, dtype=torch.long, device=dev)
        self.pos = torch.zeros((), dtype=torch.long, device=dev)
        self.graph = None
        self.use_graph = use_graph and dev.type == "cuda"

    def __call__(self, token, pos):
        self.tok.fill_(int(token))
        self.pos.fill_(int(pos))
        if not self.use_graph:
            return decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos)
        if self.graph is None:
            try:
                self._capture()
            except Exception as e:  # pragma: no cover - depends on the CUDA runtime
                print(f"CUDA graph capture failed ({type(e).__name__}: {e}); decoding eagerly")
                self.use_graph = False
                return decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos)
        self.graph.replay()
        return self.out

    def _capture(self):
        # Warm up on a side stream, then capture. The cache write at `pos` is idempotent, so
        # running the same step several times is harmless.
        s = torch.cuda.Stream()
        s.wait_stream(torch.cuda.current_stream())
        with torch.cuda.stream(s):
            for _ in range(2):
                decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos)
        torch.cuda.current_stream().wait_stream(s)
        graph = torch.cuda.CUDAGraph()
        with torch.cuda.graph(graph):
            self.out = decode_step(self.model, self.tok, self.cache, self.n_loops, self.pos)
        self.graph = graph


@torch.no_grad()
def generate_stream(model, input_ids, max_new_tokens, n_loops=2, temperature=1.0, top_p=0.95,
                    eos_ids=(2, 11), seed=None, use_graph=False, return_logits=False):
    """Yield one generated token id (int) at a time; stops after an EOS id (which is yielded).

    With return_logits=True yields (token, logits[1, V]) pairs instead (for tests)."""
    dev = input_ids.device
    t = input_ids.shape[1]
    cache = KVCache(model, capacity=t + max_new_tokens, n_loops=n_loops, device=dev)
    gen = None
    if seed is not None:
        gen = torch.Generator(device=dev).manual_seed(seed)
    logits = prefill(model, input_ids, cache, n_loops)
    step = Decoder(model, cache, n_loops, use_graph=use_graph)
    for i in range(max_new_tokens):
        nxt = sample(logits, temperature, top_p, gen)
        tok = int(nxt[0, 0])
        yield (tok, logits.clone()) if return_logits else tok
        if tok in eos_ids:
            break
        if i + 1 < max_new_tokens:
            logits = step(tok, t + i)