"""Pinned full-model cold/warm prefill, including explicit hidden-state capture."""
import argparse,hashlib,json,os,subprocess,sys
from pathlib import Path
parser=argparse.ArgumentParser(description=__doc__)
for name in ['checkout','root','native','out']:parser.add_argument(name,type=Path)
parser.add_argument('--hidden',action='store_true')
a=parser.parse_args()
assert subprocess.check_output(['git','-C',str(a.checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
sys.path.insert(0,str(a.checkout))
from mtplx.server.openai import _server_runtime_env_overrides
os.environ.update(_server_runtime_env_overrides(argparse.Namespace(model=str(a.root),generation_mode='mtp' if a.hidden else 'ar'),None))
os.environ['MTPLX_SUSTAINED_PREFILL']='1'
import mlx.core as mx
import numpy as np
from mlx_lm.utils import load_model
from mtplx.models import qwen4_exp
from mtplx.runtime import MTPLXRuntime,_load_tokenizer_resilient
from mtplx.sampling import SamplerConfig
from mtplx.generation import generate_ar,restore_or_prefill_prompt_state
from mtplx.cache_state import snapshot_cache,snapshot_cache_lazy_hybrid
from mtplx.session_bank import SessionBank,_lazy_snapshot_enabled
from mtplx.kernels import moe_sorted_gather_nax as kernel
assert mx.__version__=='0.32.2'
a.out.mkdir(parents=True,exist_ok=True)
model,config=load_model(a.root,lazy=False,strict=True,get_model_classes=lambda **_:(qwen4_exp.Model,qwen4_exp.ModelArgs))
model.post_weight_load(a.root)
tokenizer=_load_tokenizer_resilient(a.root,config)
from mtplx.mtp_patch import MTPContract
if a.hidden: assert model.attach_mtp(a.root)
rt=MTPLXRuntime(model,tokenizer,a.root,a.hidden,MTPContract() if a.hidden else None)
def fingerprint(v):
 values=np.asarray(v.astype(mx.float32)).astype('<f4')
 assert np.isfinite(values).all()
 return dict(shape=list(v.shape),sha256=hashlib.sha256(values.tobytes()).hexdigest())
for chunk in [2048,4096]:
 os.environ['MTPLX_PREFILL_CHUNK_SIZE']=str(chunk)
 bank=SessionBank()
 for phase in ['cold','warm']:
  name=f'{chunk}-{phase}' if a.hidden else f'{chunk}-0-{phase}'
  expected=json.loads((a.native/(name+'.json')).read_text())
  prompt=expected['state']['ids'] if a.hidden else expected['state']['ids'][:expected['metrics']['prompt']]
  if a.hidden:
   result=restore_or_prefill_prompt_state(rt,prompt,mtp_history_policy='cycle',session_bank=bank,session_id='issue112',capture_hidden=True,store_prefix_snapshot=False,prefill_callback=lambda event:print(json.dumps(dict(case='progress',chunk=chunk,phase=phase,event=event),default=str),flush=True))
   cache,logits,hidden,ids,cached=result.trunk_cache,result.logits,result.hidden,prompt,result.cached_tokens
  else:
   result=generate_ar(rt,prompt,max_tokens=1,sampler=SamplerConfig(temperature=0.,top_p=.95,top_k=20),seed=1,stop_token_ids=set(),capture_final_state=True,session_bank=bank,session_id='issue112',store_prefix_snapshot=False,token_callback=lambda ids:print(json.dumps(dict(case='tokens',chunk=chunk,phase=phase,ids=ids)),flush=True))
   final=result.final_state
   assert final.safe_to_commit
   cache,logits,hidden,ids,cached=final.final_trunk_cache,final.final_logits,final.final_hidden,prompt+result.tokens,result.stats.cached_tokens
  states=[[None if v is None else fingerprint(v) for v in state] if state is not None else None for state in (snapshot_cache_lazy_hybrid(cache) if _lazy_snapshot_enabled() else snapshot_cache(cache)).states]
  record=dict(state=dict(ids=ids,logits=fingerprint(logits),states=states),hidden=None if hidden is None else fingerprint(hidden))
  (a.out/(name+'.json')).write_text(json.dumps(record))
  assert record['state']==expected['state'],(chunk,phase,'state differs')
  assert record['hidden']==expected['hidden'],(chunk,phase,'hidden differs')
  assert cached==(expected['cached'] if a.hidden else expected['metrics']['cached'])
  print(json.dumps(dict(case='model_oracle',chunk=chunk,phase=phase,hidden=a.hidden,exact=True,cached=cached,kernel=kernel.stats())),flush=True)
  if phase=='cold':
   assert bank.put(runtime=rt,token_ids=ids,cache=cache,logits=logits,hidden=hidden,hidden_variant='post_norm',session_id='issue112',mtp_history_policy='cycle',snapshot_epoch=len(ids)) is not None