"""Request-owned CUDA graph for ABC and semantic top-k/top-p sampling."""
import torch

from .protocol import MUSIC_END, CODEC_OFFSET, CODEC_SIZE, ABC_END, EOD


class CapturedSampler:
    """Return one vocabulary ID in shared output storage, valid until next call.

    The caller owns the RNG and already validates the generated history. Capture
    and warmup restore its state; only real replays consume random numbers.
    """

    @torch.inference_mode()
    def __init__(self, logits, sampling, generator, *, phase="semantic"):
        if phase not in {"abc", "semantic"}:
            raise ValueError("Captured sampling phase must be abc or semantic")
        if logits.device.type != "cuda" or logits.ndim != 2 or logits.shape[0] != 1:
            raise ValueError("Captured sampling requires one CUDA logit row after CFG")
        required = ABC_END + 1 if phase == "abc" else CODEC_OFFSET + CODEC_SIZE
        if logits.shape[-1] < required or sampling.temperature <= 0:
            raise ValueError("Captured sampling requires the phase vocabulary and positive temperature")
        # Torch 2.10 can report a CUDA generator without its device index.
        if (generator.device.type != "cuda" or
                generator.device.index is not None and generator.device.index != logits.device.index):
            raise ValueError("Sampling generator and logits must share a CUDA device")
        self.sampling = sampling
        self.phase = phase
        self.offset = 0 if phase == "abc" else MUSIC_END
        self.eos_index = EOD if phase == "abc" else 0
        self.generator = generator
        self.graph = self.output = None
        self.closed = False
        size = EOD + 1 if phase == "abc" else CODEC_OFFSET + CODEC_SIZE - MUSIC_END
        self.logits = torch.zeros((1, size),
                                  device=logits.device, dtype=torch.float32)
        self.recent = torch.full((1, sampling.penalty_window), -1,
                                 device=logits.device, dtype=torch.long)
        self.allow_eos = torch.ones((), device=logits.device, dtype=torch.bool)
        saved = generator.get_state()
        try:
            with torch.cuda.device(logits.device):
                current = torch.cuda.current_stream()
                warmup = torch.cuda.Stream()
                warmup.wait_stream(current)
                with torch.cuda.stream(warmup):
                    for _ in range(3):
                        self._compute()
                current.wait_stream(warmup)
                torch.cuda.synchronize()
                self.graph = torch.cuda.CUDAGraph()
                self.graph.register_generator_state(generator)
                with torch.cuda.graph(self.graph):
                    self.output = self._compute()
        finally:
            generator.set_state(saved)

    def _compute(self):
        sampling = self.sampling
        scores = self.logits.clone()
        scores[..., self.eos_index] = torch.where(self.allow_eos, scores[..., self.eos_index], -torch.inf)
        if sampling.repetition_penalty != 1:
            frequency = torch.zeros_like(scores)
            frequency.scatter_add_(-1, self.recent.clamp_min(0), (self.recent >= 0).float())
            alpha = sampling.repetition_penalty ** frequency
            scores = torch.where(scores < 0, scores * alpha, scores / alpha)
        if sampling.temperature != 1:
            scores /= sampling.temperature
        values, indices = scores.topk(min(sampling.top_k, scores.shape[-1]))
        if sampling.top_p < 1:
            probabilities = values.softmax(-1)
            removed = probabilities.cumsum(-1) - probabilities > sampling.top_p
            removed[..., :1] = False
            values = values.masked_fill(removed, -torch.inf)
        choice = torch.multinomial(values.softmax(-1), 1, generator=self.generator)
        selected = indices.gather(-1, choice)
        if self.phase == "abc":
            return torch.where(selected == EOD, ABC_END, selected)
        return selected + MUSIC_END

    @torch.inference_mode()
    def __call__(self, logits, history, step):
        if self.closed:
            raise RuntimeError("CapturedSampler is closed")
        recent = [token - self.offset for token in history[-self.sampling.penalty_window:]]
        recent += [-1] * (self.sampling.penalty_window - len(recent))
        self.recent.copy_(torch.tensor([recent], dtype=torch.long, device="cpu"))
        if self.phase == "abc":
            # Ordinary IDs stay unchanged; only the separate EOS is appended.
            self.logits[..., :EOD].copy_(logits[..., :EOD])
            self.logits[..., EOD].copy_(logits[..., ABC_END])
        else:
            self.logits.copy_(logits[..., MUSIC_END:CODEC_OFFSET + CODEC_SIZE])
        self.allow_eos.fill_(step >= self.sampling.min_tokens)
        self.graph.replay()
        return self.output

    def close(self):
        self.graph = self.output = self.generator = None
        self.logits = self.recent = self.allow_eos = None
        self.closed = True