"""Frozen R51 generation path with unused adapter branches removed.

No model hooks or trained modules; see source_manifest.json and parity tests.
"""

import math
import time

import torch
from transformers import TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS


def new_cache(model):
    from transformers.models.granitemoehybrid.modeling_granitemoehybrid import (
        HybridMambaAttentionDynamicCache,
    )

    parameter = next(model.parameters())
    cache = HybridMambaAttentionDynamicCache(
        model.config, batch_size=1, dtype=parameter.dtype, device=parameter.device
    )
    for layer in cache.transformer_layers:
        shape = (
            1,
            model.config.num_key_value_heads,
            0,
            model.config.hidden_size // model.config.num_attention_heads,
        )
        cache.key_cache[layer] = torch.empty(shape, dtype=parameter.dtype, device=parameter.device)
        cache.value_cache[layer] = torch.empty_like(cache.key_cache[layer])
    return cache


def idle_frozen(model):
    if model.training or any(p.requires_grad or p.grad is not None for p in model.parameters()):
        raise ValueError("Runtime requires a frozen evaluation backbone")
    attention = ALL_ATTENTION_FUNCTIONS["sdpa"]
    if (
        getattr(model, "_jev_feedback_active", False)
        or getattr(attention, "_r17_controlled", False)
        or getattr(attention, "_r18_controlled", False)
        or any(m._forward_hooks or m._forward_pre_hooks for m in model.modules())
        or torch.nn.modules.module._global_forward_hooks
        or torch.nn.modules.module._global_forward_pre_hooks
    ):
        raise RuntimeError("Model intervention already active")


def valid_ids(model, ids):
    return (
        isinstance(ids, list)
        and bool(ids)
        and all(type(i) is int and 0 <= i < model.config.vocab_size for i in ids)
    )


def sample_setup(sampling, device):
    if sampling is None:
        return (None, [], None)
    if (
        not isinstance(sampling, dict)
        or set(sampling) != {"seed", "temperature", "top_p"}
        or type(sampling["seed"]) is not int
        or (not 0 <= sampling["seed"] < 2**63)
        or any(
            type(sampling[k]) not in (int, float) or not math.isfinite(sampling[k])
            for k in ("temperature", "top_p")
        )
        or (not 0 < sampling["temperature"] <= 2)
        or (not 0 < sampling["top_p"] <= 1)
    ):
        raise ValueError("Invalid sampling settings")
    settings = dict(sampling, top_k=50)
    generator = torch.Generator(device=device).manual_seed(settings["seed"])
    warpers = [
        TemperatureLogitsWarper(float(settings["temperature"])),
        TopKLogitsWarper(50),
        TopPLogitsWarper(float(settings["top_p"])),
    ]
    return (generator, warpers, settings)


def generate(model, tok, ids, *, limit, eos, context_limit=32768, sampling=None, deadline=None):
    idle_frozen(model)
    if (
        not valid_ids(model, ids)
        or not valid_ids(model, eos)
        or len(eos) != len(set(eos))
        or (type(limit) is not int)
        or (not 1 <= limit <= 2048)
        or (type(context_limit) is not int)
        or (not 1 <= context_limit <= 32768)
        or (context_limit > model.config.max_position_embeddings)
        or (len(ids) + limit > context_limit)
    ):
        raise ValueError("Invalid tokens, generation limit or context")
    device = next(model.parameters()).device
    generator, warpers, settings = sample_setup(sampling, device)
    pending, accepted, past = (list(ids), [], 0)
    cache, traces = (new_cache(model), [])
    tick = time.monotonic()
    forward_pending, attempted_positions = False, 0
    decoding_started = False
    try:
        with torch.no_grad():
            for _ in range(limit):
                if deadline:
                    deadline()
                attempted_positions += len(pending)
                forward_pending = True
                out = model(
                    input_ids=torch.tensor([pending], device=device),
                    attention_mask=torch.ones(
                        (1, past + len(pending)), device=device, dtype=torch.long
                    ),
                    past_key_values=cache,
                    use_cache=True,
                    cache_position=torch.arange(past, past + len(pending), device=device),
                    logits_to_keep=1,
                )
                past += len(pending)
                forward_pending = False
                cache = out.past_key_values
                logits = out.logits[:, -1, :].float()
                if not torch.isfinite(logits).all():
                    raise ValueError("Nonfinite generation logits")
                if generator is None:
                    token = int(logits[0].argmax())
                else:
                    for warper in warpers:
                        logits = warper(None, logits)
                    token = int(
                        torch.multinomial(torch.softmax(logits, dim=-1), 1, generator=generator)[
                            0, 0
                        ]
                    )
                accepted.append(token)
                if token in eos:
                    break
                pending = [token]
        decoding_started = True
        body = accepted[:-1] if accepted[-1] in eos else accepted
        return dict(
            prompt_token_ids=list(ids),
            generated_token_ids=accepted,
            text=tok.decode(body, skip_special_tokens=False, clean_up_tokenization_spaces=False),
            finish_reason="eos" if accepted[-1] in eos else "length",
            processed_tokens=len(ids) + len(accepted) - 1,
            seconds=time.monotonic() - tick,
            events=traces,
            sampling=settings,
        )
    except BaseException as exc:
        body = accepted[:-1] if accepted and accepted[-1] in eos else accepted
        partial = dict(
            prompt_token_ids=list(ids),
            generated_token_ids=list(accepted),
            text=None,
            finish_reason=(
                "time"
                if isinstance(exc, TimeoutError)
                else "error"
                if isinstance(exc, Exception)
                else "interrupted"
            ),
            processed_tokens=past,
            attempted_forward_positions=attempted_positions,
            forward_work_unknown=forward_pending,
            seconds=time.monotonic() - tick,
            events=traces,
            sampling=settings,
        )
        exc.partial_generation = partial
        if not decoding_started:
            try:
                partial["text"] = tok.decode(
                    body, skip_special_tokens=False, clean_up_tokenization_spaces=False
                )
            except BaseException as decoding_error:
                partial["decoding_error_type"] = type(decoding_error).__name__
                if not isinstance(decoding_error, Exception):
                    decoding_error.partial_generation = partial
                    raise
        raise