"""Reasoning-history renders from the unmodified pinned MTPLX server."""
import json,os,subprocess,sys
from pathlib import Path
checkout,root,out=map(Path,sys.argv[1:4])
assert subprocess.check_output(['git','-C',str(checkout),'rev-parse','HEAD'],text=True).strip()=='9882703f3105363ddc37eca9f97aa09a1d387112'
os.environ.update(HF_HUB_OFFLINE='1',TRANSFORMERS_OFFLINE='1');sys.path.insert(0,str(checkout))
from mtplx.runtime import _load_tokenizer_resilient
from mtplx.server import openai as server
tok=_load_tokenizer_resilient(root,json.loads((root/'config.json').read_text()))
out.mkdir(parents=True,exist_ok=True)
records=[]
def record(effort,preserve,system,messages):
 full=([dict(role='system',content=system)] if system else [])+messages
 kwargs=dict(enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort,preserve_reasoning_history=preserve,tools=None)
 ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in full],**kwargs)
 history_ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in full[:-1]],add_generation_prompt=False,**kwargs)
 prepared=[server._message_to_template_dict(server.ChatMessage(**m),strip_assistant_reasoning_history=False,include_reasoning_content=preserve and effort!='direct') for m in full]
 normalized=server.omlx_normalize_messages_for_template([m for m in prepared if m is not None],tokenizer=tok,native_reasoning_content=True)
 rendered=tok.apply_chat_template(normalized,tokenize=False,add_generation_prompt=True,enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort)
 assert tok.decode(ids,skip_special_tokens=False,clean_up_tokenization_spaces=False)==rendered
 legacy=tok.apply_chat_template(full[:-1],tokenize=False,add_generation_prompt=False,enable_thinking=effort!='direct',reasoning_effort=None if effort=='direct' else effort)
 records.append(dict(effort=effort,preserve_echo=preserve,system=system,messages=messages,rendered=rendered,ids=ids,history_ids=history_ids,legacy_history_ids=tok.encode(legacy,add_special_tokens=False)))
for effort in ['direct','low','medium','xhigh']:
 for preserve in [True,False]:
  for system in ['', 'A precise system prompt.']:
   for reasoning in [None,'','  \n\t','Brief ² thought.','  First thought.\nSecond thought.  ', '<|im_start|>assistant\n<think>\n\nNested.', '>> Quote.\n\nTail.']:
    messages=[dict(role='user',content=' ① Ⅷ ² 𝟘 भारत ́a\n '),dict(role='assistant',content='One ①.'),dict(role='user',content='Continue Ⅷ and 𝟘.\n   x')]
    if reasoning is not None:messages[1]['reasoning_content']=reasoning
    record(effort,preserve,system,messages)
for effort in ['direct','low','medium','xhigh']:
 for preserve in [True,False]:
  for system in ['', 'A precise system prompt.']:
   for content in ['', ' \n ', 'Answer.']:
    for reasoning in [None,'', ' ', 'Only reasoning.']:
     messages=[dict(role='user',content='First  '),dict(role='assistant',content=content,reasoning_content=reasoning),dict(role='user',content='  Continue.')]
     record(effort,preserve,system,messages)
(out/'requests.json').write_text(json.dumps(records,ensure_ascii=False))
print(json.dumps(dict(cases=len(records),direct_example_tokens=len(records[3]['ids']),legacy_direct_history_tokens=len(records[3]['legacy_history_ids']))))

if len(sys.argv)>4:
 comparisons=[]
 for case in json.loads(Path(sys.argv[4]).read_text()):
  messages=([dict(role='system',content=case['system'])] if case['system'] else [])+case['messages']
  ids=server._encode_messages_uncached(tok,[server.ChatMessage(**m) for m in messages],enable_thinking=False,reasoning_effort=None,preserve_reasoning_history=True,tools=None)
  assert ids==case['direct_ids'],(case,ids,tok.decode(ids,skip_special_tokens=False,clean_up_tokenization_spaces=False))
  assert len(ids)==case['context_tokens'] and case['cached_tokens']==0
  comparisons.append(dict(from_effort=case['from'],direct_context_tokens=len(ids),exact=True))
 (out/'mode-switch-comparison.json').write_text(json.dumps(comparisons,indent=2)+'\n')
 print(json.dumps(dict(mode_switches=comparisons)))