"""logits-processor: the generation last-mile, loadable through `kernels`.

Two fused ops that every serving stack runs on the hot path:

- ``apply_token_bitmask`` - structured / guided decoding. Set the logits of
  disallowed tokens to -inf from a packed allow-mask (XGrammar convention:
  bit 1 = allowed), so the next softmax cannot pick them. This is the kernel
  behind constrained JSON and tool-calling output.
- ``sample`` - temperature, then the intersection of top-k / top-p (nucleus) /
  min-p filtering, then a multinomial draw, fused into one logits -> token op
  with a counter-based RNG that is deterministic in ``(seed, row, offset)``.
  ``temperature <= 0`` is greedy (argmax).

Both take ``[B, V]`` logits in float / half / bf16. Sampling parameters are
per-row: pass a scalar to broadcast, or a length-B tensor for a heterogeneous
batch (continuous batching, per-request grammars).
"""
from typing import Union

import torch

from ._ops import ops

__all__ = ["apply_token_bitmask", "sample"]

Param = Union[float, int, torch.Tensor]


def _row(x: Param, B: int, device, dtype) -> torch.Tensor:
    if isinstance(x, torch.Tensor):
        return x.to(device=device, dtype=dtype).contiguous()
    return torch.full((B,), x, device=device, dtype=dtype)


def apply_token_bitmask(logits: torch.Tensor, bitmask: torch.Tensor) -> torch.Tensor:
    """In place: set disallowed tokens' logits to -inf.

    logits: ``[B, V]`` (float/half/bf16). bitmask: ``[B, ceil(V/32)]`` int32,
    bit 1 = token allowed. Returns ``logits`` for chaining.
    """
    ops.apply_token_bitmask_inplace(logits, bitmask.to(torch.int32).contiguous())
    return logits


def sample(logits: torch.Tensor, temperature: Param = 1.0, top_k: Param = 0,
           top_p: Param = 1.0, min_p: Param = 0.0, seed: int = 0,
           offset: int = 0) -> torch.Tensor:
    """Fused temperature + top-k ∩ top-p ∩ min-p + multinomial sample.

    top_k <= 0, top_p >= 1, and min_p <= 0 each disable that filter.
    temperature <= 0 selects the argmax. Returns int64 token ids ``[B]``.
    Same ``(seed, offset)`` and logits give the same tokens.
    """
    assert logits.dim() == 2, "logits must be [B, V]"
    B, dev = logits.size(0), logits.device
    return ops.sample(
        logits.contiguous(),
        _row(temperature, B, dev, torch.float32),
        _row(top_p, B, dev, torch.float32),
        _row(top_k, B, dev, torch.int32),
        _row(min_p, B, dev, torch.float32),
        int(seed), int(offset),
    )