"""Token-counted retrieval and repeat-prefix acceptance, not a broad quality eval."""
import argparse
import json
import random
import time
import urllib.request
from pathlib import Path

p=argparse.ArgumentParser()
p.add_argument('--tokens',type=int,default=28000)
p.add_argument('--output',type=Path,required=True)
p.add_argument('--url',default='http://127.0.0.1:8000')
a=p.parse_args()
key=(Path(__file__).resolve().parents[1]/'secrets/api-key').read_text().strip()
def request(path,body):
    return urllib.request.urlopen(urllib.request.Request(a.url+path,
        data=json.dumps(body).encode(),headers={'Content-Type':'application/json',
        'Authorization':'Bearer '+key}),timeout=900)

rng=random.Random(6000)
lines=['Internal inventory. Find the exact special verification code when asked.\n',
       'The special verification code is VIOLET-74219.\n']
lines += [f'Record {i:06d}: batch {rng.randrange(10**8):08d}, location shelf {rng.randrange(1000):03d}, status archived.\n' for i in range(a.tokens//8+1000)]
question='\nWhat is the special verification code? Reply with only that code.'
def messages(n):
    return [{'role':'user','content':''.join(lines[:n])+question}]
lo,hi=2,len(lines)
while lo < hi:
    mid=(lo+hi+1)//2
    with request('/tokenize',dict(model='qwen3.8-flash-next',messages=messages(mid),
                 chat_template_kwargs={'reasoning_effort':'low'})) as resp:
        count=json.load(resp)['count']
    if count<=a.tokens: lo=mid
    else: hi=mid-1
a.output.parent.mkdir(parents=True,exist_ok=True)
with a.output.open('x') as out:
    for repetition in range(2):
        start=time.monotonic(); first=None; content=''; usage={}; finish=None
        with request('/v1/chat/completions',dict(model='qwen3.8-flash-next',messages=messages(lo),
                     temperature=0,reasoning_effort='low',max_tokens=2048,stream=True,
                     stream_options={'include_usage':True})) as resp:
            for raw in resp:
                if not raw.startswith(b'data: '): continue
                data=raw[6:].strip()
                if data==b'[DONE]': break
                item=json.loads(data)
                usage=item.get('usage') or usage
                for c in item.get('choices',[]):
                    d=c.get('delta',{})
                    if first is None and any(d.get(x) for x in ['content','reasoning','reasoning_content']): first=time.monotonic()
                    content+=d.get('content') or ''
                    finish=c.get('finish_reason') or finish
        row=dict(repetition=repetition,content=content,usage=usage,finish_reason=finish,
                 ttft_s=None if first is None else first-start,elapsed_s=time.monotonic()-start,
                 passed=content.strip()=='VIOLET-74219' and finish=='stop')
        line=json.dumps(row); out.write(line+'\n'); out.flush(); print(line,flush=True)
        if not row['passed']: raise SystemExit('Long-context retrieval failed')