"""Request-owned device history for captured sampling."""
import torch
from .captured_sampling import CapturedSampler
from .protocol import MUSIC_END, CODEC_OFFSET, CODEC_SIZE, ABC_END, EOD

class DeviceRingSampler(CapturedSampler):
    @torch.inference_mode()
    def __init__(self, logits, sampling, generator, *, phase='semantic'):
        if sampling.penalty_window < 1:
            raise ValueError('Device history requires a positive window')
        self.counter = torch.zeros((),device=logits.device,dtype=torch.long)
        self.expected_step = 0
        super().__init__(logits,sampling,generator,phase=phase)
        # Constructor warmup executes updates; capture itself only records them.
        self.counter.zero_()
        self.recent.fill_(-1)

    def _compute(self):
        self.allow_eos.copy_(self.counter >= self.sampling.min_tokens)
        selected = super()._compute()
        index = (self.counter % self.sampling.penalty_window).reshape(1,1)
        local = torch.where(selected == ABC_END, EOD, selected) if self.phase == 'abc' else selected-MUSIC_END
        self.recent.scatter_(-1,index,local)
        self.counter.add_(1)
        return selected

    @torch.inference_mode()
    def __call__(self,logits,history,step):
        if self.closed:
            raise RuntimeError('DeviceRingSampler is closed')
        if step != self.expected_step or len(history) != step:
            raise ValueError('Device history requires consecutive steps from an empty history')
        if self.phase == 'abc':
            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.graph.replay()
        self.expected_step += 1
        return self.output

    def close(self):
        super().close()
        self.counter = None