"""Cache-aware incremental decoding with per-token workspace routing (forced budget or learned controller).

For every new position t: backbone step with KV cache -> H_t; bank append (K_t, V_t); z_t; fresh S0; actions
(forced or controller-chosen, hard cap r_max); differential readout; logits; sample. The workspace is discarded
after the token (no persistent memory). The explicit budget-0 branch never allocates slots.
"""
from __future__ import annotations

from typing import Callable, Optional

import torch

from .cache import RequestBatchCache
from .model import PrismeModel


@torch.no_grad()
def step_logits(model: PrismeModel, tokens_new: torch.Tensor, cache: RequestBatchCache, max_budget: int,
                cost: Optional[torch.Tensor] = None, forced_actions: Optional[torch.Tensor] = None,
                lambda_cost=None, kappa_risk=None):
    """tokens_new: [B, n] new tokens (prefill n>=1 or decode n=1). Returns logits [B,n,V] and the action trace."""
    cfg = model.cfg
    B, n = tokens_new.shape
    start = cache.length
    H = model.backbone(tokens_new, cache=cache.layers)                           # [B,n,d]
    if not cfg.use_workspace:
        return model.backbone.head(H), None, torch.zeros(B, n, dtype=torch.long, device=H.device)
    if max_budget == 0:
        # explicit bypass: no workspace allocation, no bank read; still extend the bank for later tokens
        hb = model.bank_norm(H)
        cache.bank.append(model.W_K(hb), model.W_V(hb))
        return model.backbone.head(H), None, torch.zeros(B, n, dtype=torch.long, device=H.device)
    hb = model.bank_norm(H)
    cache.bank.append(model.W_K(hb), model.W_V(hb))
    Tb = cache.bank.length
    z = model.P(hb)
    qpos = torch.arange(start, start + n, device=H.device).view(1, n).expand(B, n)
    key_pos = torch.arange(Tb, device=H.device)
    bank_mask = key_pos.view(1, 1, Tb) <= qpos.unsqueeze(-1)
    from .model import EncodedContext
    from .workspace import init_workspace
    entropy = None
    if cfg.ctrl_use_entropy:
        p = torch.softmax(model.backbone.head(H).float(), -1)
        entropy = -(p * torch.log(p.clamp_min(1e-12))).sum(-1)
    ctx = EncodedContext(H=H, K=cache.bank.K, V=cache.bank.V, z=z, bank_mask=bank_mask, S0=init_workspace(model.slot_embed, z),
                         ctrl_in=(hb if cfg.ctrl_input == "h" else z), entropy=entropy)
    S = ctx.S0
    R = min(max_budget, cfg.r_max)
    trace = torch.full((B, n, R), cfg.stop_action, dtype=torch.long, device=H.device)
    n_edits = torch.zeros(B, n, dtype=torch.long, device=H.device)
    stopped = torch.zeros(B, n, dtype=torch.bool, device=H.device)
    for r in range(R):
        step = torch.full((B, n), r, device=H.device, dtype=torch.long)
        remaining = torch.full((B, n), R - r - 1, device=H.device, dtype=torch.long)
        if forced_actions is not None:
            a = forced_actions[:, :, r]
        else:
            preds = model.controller(ctx.ctrl_in, S, step, remaining, ctx.entropy)
            a = model.controller.select(preds, cost, lambda_cost=lambda_cost, kappa_risk=kappa_risk)
        a = torch.where(stopped, torch.full_like(a, cfg.stop_action), a)
        stopped = stopped | (a == cfg.stop_action)
        trace[:, :, r] = a
        S, _ = model.apply_action(S, ctx, a, step, remaining)
        n_edits += (a != cfg.stop_action).long()
        if bool(stopped.all()):
            break
    logits, _ = model.logits_from_state(S, ctx, n_edits > 0)
    return logits, trace, n_edits


@torch.no_grad()
def generate_uncached(model, prompt: torch.Tensor, max_new_tokens: int, max_budget: int = 0, temperature: float = 0.0,
                      eos_id=None, forced_budget_actions=None, generator=None, **_):
    """Fallback for models without the PRISME caches (e.g. the global-recurrence baseline): recompute the full
    sequence at every step with forced actions for all positions. O(T^2) per token; used only for evaluation."""
    cfg = model.cfg
    B, T0 = prompt.shape
    seq = prompt
    done = torch.zeros(B, dtype=torch.bool, device=prompt.device)
    edits = []
    for i in range(max_new_tokens):
        T = seq.shape[1]
        acts = forced_budget_actions(B, T) if (forced_budget_actions is not None and max_budget > 0) else None
        out = model(seq, acts)
        nxt = out["logits"][:, -1]
        if out.get("n_edits") is not None:
            edits.append(out["n_edits"][:, -1])
        if temperature <= 0:
            tok = nxt.argmax(-1, keepdim=True)
        else:
            tok = torch.multinomial(torch.softmax(nxt / temperature, -1), 1, generator=generator)
        if eos_id is not None:
            tok = torch.where(done.unsqueeze(-1), torch.full_like(tok, eos_id), tok)
            done = done | (tok.squeeze(-1) == eos_id)
        seq = torch.cat([seq, tok], 1)
        if bool(done.all()):
            break
    return {"tokens": seq, "n_edits": torch.stack(edits, 1) if edits else None, "cache_bytes": 0}


@torch.no_grad()
def generate(model: PrismeModel, prompt: torch.Tensor, max_new_tokens: int, max_budget: int = 0,
             cost: Optional[torch.Tensor] = None, temperature: float = 0.0, eos_id: Optional[int] = None,
             forced_budget_actions: Optional[Callable[[int, int], torch.Tensor]] = None, generator=None,
             lambda_cost=None, kappa_risk=None):
    """Greedy (temperature 0) or sampled decoding. prompt: [B, T0]. Returns dict(tokens [B,T0+new], n_edits [B,new],
    actions list). Prompt positions are processed with the same per-token routing (prefill)."""
    cfg = model.cfg
    if not hasattr(model, "W_K") and getattr(cfg, "use_workspace", False):
        return generate_uncached(model, prompt, max_new_tokens, max_budget, temperature, eos_id, forced_budget_actions, generator)
    cache = RequestBatchCache(layers=model.backbone.new_cache())
    B, T0 = prompt.shape
    fa = forced_budget_actions(B, T0) if forced_budget_actions is not None else None
    logits, _, _ = step_logits(model, prompt, cache, max_budget, cost, fa, lambda_cost, kappa_risk)
    out = [prompt]
    edits, traces = [], []
    done = torch.zeros(B, dtype=torch.bool, device=prompt.device)
    nxt = logits[:, -1]
    for i in range(max_new_tokens):
        if temperature <= 0:
            tok = nxt.argmax(-1, keepdim=True)
        else:
            probs = torch.softmax(nxt / temperature, -1)
            tok = torch.multinomial(probs, 1, generator=generator)
        if eos_id is not None:
            tok = torch.where(done.unsqueeze(-1), torch.full_like(tok, eos_id), tok)
            done = done | (tok.squeeze(-1) == eos_id)
        out.append(tok)
        if i == max_new_tokens - 1 or bool(done.all()):
            break
        fa = forced_budget_actions(B, 1) if forced_budget_actions is not None else None
        logits, tr, ne = step_logits(model, tok, cache, max_budget, cost, fa, lambda_cost, kappa_risk)
        edits.append(ne[:, 0]); traces.append(tr)
        nxt = logits[:, -1]
    return {"tokens": torch.cat(out, 1), "n_edits": torch.stack(edits, 1) if edits else None, "cache_bytes": cache.nbytes()}