"""One 8-token real text, bounded memory/semantic capacity smoke check."""
import argparse
import json
from pathlib import Path
import resource
import time
import numpy as np
from .model import StreamingLM
from .run import tokenize,write_json


def main():
    p=argparse.ArgumentParser();p.add_argument('--root',required=True,type=Path);a=p.parse_args()
    samples,ids=tokenize(a.root)
    start=time.time()
    lm=StreamingLM(a.root/'weights',json.loads((a.root/'model_config.json').read_text()),seconds=1200)
    def progress(layer):print(json.dumps({'preflight_layer':layer,'seconds':time.time()-start}),flush=True)
    h,_=lm.forward(ids[:1,:8],progress=progress)
    nll=lm.nll(h,ids[:1,:8])
    if not np.isfinite(nll).all():raise FloatingPointError('Nonfinite NLL')
    write_json(a.root/'preflight.json',{'passed':True,'sample_id':samples[0]['id'],'tokens':8,
               'elapsed_seconds':time.time()-start,'peak_rss_kib':resource.getrusage(resource.RUSAGE_SELF).ru_maxrss,
               'mean_nll':float(nll.mean()),'token_nll':nll.tolist(),
               'limitation':'Real-weight capacity check, not an independent correctness oracle or quality benchmark'})


if __name__=='__main__':main()