"""Small shared offline runtime, output and wall-clock controls for Labs 17–20."""
import contextlib, hashlib, importlib.metadata, json, os, platform, signal, time
from datetime import datetime, timezone
from pathlib import Path
from inspect_residual_stream import artifacts, REPOSITORY, REVISION, VERSIONS


def safe_path(path):
    path=Path(path).absolute()
    if any(p.is_symlink() for p in (path,*path.parents)):
        raise ValueError('Path contains a symlink')
    return path


class Output:
    def __init__(self,path,limit=None):
        self.path=safe_path(path)
        self.path.mkdir(mode=0o700,parents=True,exist_ok=False)
        self.limit=limit
    def check(self,extra=0):
        size=sum(p.stat().st_size for p in self.path.rglob('*') if p.is_file())
        if self.limit is not None and size+extra>self.limit:
            raise ValueError('Saved-data byte limit exceeded')
    def json(self,name,value):
        data=(json.dumps(value,indent=2,allow_nan=False)+'\n').encode()
        target=self.path/name
        if target.exists() or target.is_symlink():raise ValueError('Evidence file already exists')
        self.check(len(data))
        target.write_bytes(data)
    def hash(self,name):return hashlib.sha256((self.path/name).read_bytes()).hexdigest()


class Budget:
    def __init__(self,seconds,forward_limit=None):
        self.started=time.perf_counter();self.deadline=self.started+seconds
        self.seconds=seconds;self.forwards=0;self.forward_limit=forward_limit
    def check(self):
        if time.perf_counter()>=self.deadline:raise TimeoutError('Cumulative wall-time limit')
    def forward(self):
        self.check()
        if self.forward_limit is not None and self.forwards>=self.forward_limit:
            raise ValueError('Forward limit')
        self.forwards+=1
    @contextlib.contextmanager
    def execution(self):
        if not hasattr(signal,'setitimer'):raise RuntimeError('POSIX wall timer required')
        def expired(signum,frame):raise TimeoutError('Cumulative wall-time limit')
        old=signal.signal(signal.SIGALRM,expired)
        signal.setitimer(signal.ITIMER_REAL,max(0.001,self.deadline-time.perf_counter()))
        try:yield
        finally:
            signal.setitimer(signal.ITIMER_REAL,0);signal.signal(signal.SIGALRM,old)


def runtime():
    if platform.python_version()!='3.13.13':raise ValueError('Use CPython 3.13.13 reference runtime')
    for package,expected in VERSIONS.items():
        if importlib.metadata.version(package).split('+')[0]!=expected:
            raise ValueError('Pinned runtime differs: '+package)
    # Offline controls precede importing Transformers. No acquisition is performed here.
    os.environ['HF_HUB_OFFLINE']='1';os.environ['TRANSFORMERS_OFFLINE']='1'
    import torch
    torch.set_num_threads(1);torch.set_num_interop_threads(1)
    torch.manual_seed(0)
    return torch


def host_record():
    host=os.uname()
    return {'os':' '.join(host),'cpu':host.machine,'architecture':host.machine,'cpu_count':os.cpu_count()}


def environment(torch):
    return {'python':platform.python_version(),'packages':{k:importlib.metadata.version(k) for k in VERSIONS},
            **host_record(),
            'device':'cpu','dtype':'float32','intra_op':torch.get_num_threads(),
            'inter_op':torch.get_num_interop_threads(),'timestamp_utc':datetime.now(timezone.utc).isoformat()}


def load_specimen(directory,torch):
    directory=safe_path(directory)
    if not directory.is_dir():raise ValueError('Verified existing artifact directory required')
    hashes=artifacts(directory,offline=True)
    from transformers import AutoTokenizer,GPTNeoXForCausalLM
    model=GPTNeoXForCausalLM.from_pretrained(str(directory),local_files_only=True,
        trust_remote_code=False,use_safetensors=True,dtype=torch.float32,attn_implementation='eager').to('cpu').eval()
    tokenizer=AutoTokenizer.from_pretrained(str(directory),local_files_only=True,trust_remote_code=False,use_fast=True)
    if (model.config.num_hidden_layers,model.config.hidden_size,model.config.vocab_size,
        model.config.use_parallel_residual)!=(6,128,50304,True):raise ValueError('Specimen architecture differs')
    if not tokenizer.is_fast or any(p.device.type!='cpu' or p.dtype!=torch.float32 for p in model.parameters()):
        raise ValueError('CPU float32 fast-tokenizer contract failed')
    if model.gpt_neox.embed_in.weight.data_ptr()==model.embed_out.weight.data_ptr():
        raise ValueError('Expected untied embeddings')
    return model,tokenizer,{'repository':REPOSITORY,'revision':REVISION,'artifacts':hashes}


def source_hashes(names):
    return {name:hashlib.sha256(Path(__file__).with_name(name).read_bytes()).hexdigest() for name in names}