"""Bounded tensor-backed file state, not a transformer-compiled filesystem.

No SQLite. UTF-8 bytes live in the model session's tensors. The controller uses
explicit tensor indexing; it is deliberately not advertised as learned logic.
"""
from dataclasses import dataclass
import torch
from torch import nn


@dataclass
class MemoryState:
    names: torch.Tensor
    name_length: torch.Tensor
    data: torch.Tensor
    length: torch.Tensor
    version: torch.Tensor
    deleted: torch.Tensor
    clock: torch.Tensor

    def clone(self):
        return MemoryState(**{k:v.clone() for k,v in vars(self).items()})


class TensorFiles(nn.Module):
    def __init__(self, slots=64, versions=8, name_bytes=128, file_bytes=8192):
        super().__init__()
        self.spec = dict(slots=slots,versions=versions,name_bytes=name_bytes,file_bytes=file_bytes)
        self.register_buffer("anchor",torch.zeros(0),persistent=False)

    def empty(self):
        s=self.spec; dev=self.anchor.device
        return MemoryState(
            torch.zeros(s["slots"],s["name_bytes"],dtype=torch.uint8,device=dev),
            torch.zeros(s["slots"],dtype=torch.long,device=dev),
            torch.zeros(s["slots"],s["versions"],s["file_bytes"],dtype=torch.uint8,device=dev),
            torch.zeros(s["slots"],s["versions"],dtype=torch.long,device=dev),
            torch.zeros(s["slots"],s["versions"],dtype=torch.long,device=dev),
            torch.ones(s["slots"],s["versions"],dtype=torch.bool,device=dev),
            torch.zeros((),dtype=torch.long,device=dev))

    def _name(self,name,device):
        b=name.encode("utf-8")
        if not b or len(b)>self.spec["name_bytes"] or b"\0" in b:
            raise ValueError("Invalid UTF-8 filename length or NUL")
        return torch.tensor(list(b.ljust(self.spec["name_bytes"],b"\0")),dtype=torch.uint8,device=device),len(b)

    def _slot(self,state,name,create=False):
        key,n=self._name(name,state.names.device)
        matches=(state.names==key).all(1)&(state.name_length==n)
        if bool(matches.any()): return int(matches.long().argmax())
        if not create: raise KeyError(name)
        free=state.name_length==0
        if not bool(free.any()): raise MemoryError("File directory is full")
        slot=int(free.long().argmax())
        state.names[slot]=key;state.name_length[slot]=n
        return slot

    def _write(self,state,name,content):
        slot=self._slot(state,name,create=content is not None)
        payload=b"" if content is None else content.encode("utf-8")
        if len(payload)>self.spec["file_bytes"]: raise ValueError("File exceeds byte capacity")
        state.clock.add_(1)
        sequence=int((state.version[slot]>0).sum())
        if sequence<self.spec["versions"]:
            column=sequence
        else:
            column=int(state.version[slot].argmin())
        state.data[slot,column].zero_()
        if payload:
            state.data[slot,column,:len(payload)]=torch.tensor(list(payload),dtype=torch.uint8,device=state.data.device)
        state.length[slot,column]=len(payload)
        state.version[slot,column]=state.clock
        state.deleted[slot,column]=content is None
        return int(state.clock)

    def forward(self,state,operation,**kwargs):
        if operation=="put_many":
            files=kwargs["files"]
            if not files or len(files)>self.spec["slots"]: raise ValueError("Invalid transaction size")
            staged=state.clone()
            revisions={}
            for name,content in files.items():
                if not isinstance(name,str) or (content is not None and not isinstance(content,str)):
                    raise TypeError("Names and contents must be strings")
                revisions[name]=self._write(staged,name,content)
            # Commit only after every write passed. No partial externally visible writes.
            for name,value in vars(state).items(): value.copy_(getattr(staged,name))
            return {"saved":list(files),"revisions":revisions}
        if operation=="get":
            name=kwargs["name"];slot=self._slot(state,name)
            revision=kwargs.get("revision")
            if revision is None:
                column=int(state.version[slot].argmax())
            else:
                matches=state.version[slot]==int(revision)
                if not bool(matches.any()): raise KeyError("Revision absent or evicted")
                column=int(matches.long().argmax())
            if bool(state.deleted[slot,column]): raise KeyError(name)
            n=int(state.length[slot,column])
            return {"name":name,"revision":int(state.version[slot,column]),
                    "content":bytes(state.data[slot,column,:n].cpu().tolist()).decode("utf-8")}
        if operation=="list":
            result=[]
            for slot in range(self.spec["slots"]):
                n=int(state.name_length[slot])
                if not n: continue
                column=int(state.version[slot].argmax())
                if bool(state.deleted[slot,column]): continue
                result.append({"name":bytes(state.names[slot,:n].cpu().tolist()).decode("utf-8"),
                               "revision":int(state.version[slot,column])})
            return result
        if operation=="rollback":
            old=self(state,"get",name=kwargs["name"],revision=kwargs["revision"])
            return self(state,"put_many",files={old["name"]:old["content"]})
        raise ValueError(f"Unknown memory operation: {operation}")