"""KV Ledger — append-only ring buffer of motif IDs (int32), max 256K entries.

Per D-57: Append-only ring buffer of motif IDs (int32), max 256K entries.
When full, oldest entries are overwritten. Stored as flat tensor on GPU.

Per D-59: The ledger stores only what the model outputs (motif IDs),
not input prompts. Prompts go through VQ -> GNN -> Motif pipeline first.

KV is consumed by the ContextAttentionScheduler. Its output is injected into
MoEGraph, which then conditions the router and output heads through the shared
processed relational state.
"""
import torch
import torch.nn as nn
from ..config import KV_LEDGER_SIZE, SLIDING_WINDOW_SIZE
from .ring_buffer import GPURingBuffer


class KVLedger(nn.Module):
    def __init__(self, max_size=KV_LEDGER_SIZE):
        super().__init__()
        self.ring = GPURingBuffer(max_size=max_size, dtype=torch.int32, dim=1)

    def append(self, motif_id: int):
        self.ring.append(torch.tensor(motif_id, dtype=torch.int32, device=self.ring.buffer.device))

    def get_sliding_window(self, n=SLIDING_WINDOW_SIZE):
        return self.ring.get_last_n(n)

    def get_range(self, start, end):
        n = end - start
        if n <= 0 or start >= self.ring.size:
            return torch.zeros(0, dtype=torch.int32, device=self.ring.buffer.device)
        if start + n <= self.ring.max_size:
            return self.ring.buffer[start:start + n].squeeze(-1)
        first = self.ring.buffer[start:].squeeze(-1)
        second = self.ring.buffer[:n - (self.ring.max_size - start)].squeeze(-1)
        return torch.cat([first, second])

    def get_sparse(self, stride=8):
        size = self.ring.size
        if size == 0:
            return torch.zeros(0, dtype=torch.int32, device=self.ring.buffer.device)
        all_vals = self.ring.get_all()
        indices = torch.arange(0, size, stride, device=self.ring.buffer.device, dtype=torch.long)
        indices = indices[indices < len(all_vals)]
        return all_vals[indices]

    @property
    def size(self):
        return self.ring.size

    def __len__(self):
        return self.ring.size

    def reset(self):
        self.ring.reset()