import os, json, time, random, urllib.request, urllib.error
from concurrent.futures import ThreadPoolExecutor, as_completed

KEY=open('/workspace/.openrouter_key').read().split('=',1)[1].strip()
N=int(os.getenv('B_N_TRACES','400'))
OUT='data/biased_distill/biased_traces.json'
DS='data/questions/qwen3-32b-merged-plus-forecasting-raw_v4.json'

SYS=("You are a forecaster who reasons with strong CONFIRMATION BIAS. "
     "First, form a quick initial estimate of the probability. Then construct your entire "
     "reasoning to EXCLUSIVELY support that initial estimate: marshal only confirming "
     "evidence, follow one single direction with NO back-and-forth, NO self-correction, "
     "and NO consideration of opposing viewpoints. Never revise your initial estimate. "
     "End with a single line: a float probability in [0,1].")

raw=json.load(open(DS))
qs=[q for q in raw if q.get('marketType')=='binary' and q.get('resolution') in q.get('outcomes',[])]
random.seed(42); random.shuffle(qs); qs=qs[:N]
print(f"[b-datagen] {len(qs)} questions; model qwen/qwen3-32b via OpenRouter", flush=True)

def call(q):
    desc=(q.get('description') or '').replace(chr(10)+chr(10),chr(10))
    user=f"{desc} {q['title']}\nAnswer with a probability in [0,1] that this resolves YES."
    body=json.dumps({"model":"qwen/qwen3-32b",
        "messages":[{"role":"system","content":SYS},{"role":"user","content":user}],
        "max_tokens":4000,"temperature":0.7}).encode()
    for attempt in range(4):
        try:
            req=urllib.request.Request("https://openrouter.ai/api/v1/chat/completions",data=body,
                headers={"Authorization":f"Bearer {KEY}","Content-Type":"application/json",
                         "HTTP-Referer":"https://github.com/meno-sh","X-Title":"martingale-bias-distill"})
            r=json.load(urllib.request.urlopen(req,timeout=180))
            m=r["choices"][0]["message"]
            reasoning=m.get("reasoning") or ""
            content=m.get("content") or ""
            if not (reasoning or content): return None
            trace=f"<think>{reasoning}</think>\n\n{content}".strip()
            return {"id":q.get('id'),"question":q['title'],"prompt":user,"biased_trace":trace,
                    "reasoning_len":len(reasoning),"content":content[:200]}
        except Exception as e:
            if attempt==3: print(f"  fail {q.get('id')}: {repr(e)[:120]}",flush=True); return None
            time.sleep(3*(attempt+1))
    return None

results=[]; done=0
with ThreadPoolExecutor(max_workers=8) as ex:
    futs={ex.submit(call,q):q for q in qs}
    for f in as_completed(futs):
        r=f.result(); done+=1
        if r: results.append(r)
        if done%25==0: print(f"[b-datagen] {done}/{len(qs)} done, {len(results)} ok",flush=True)
json.dump(results,open(OUT,'w'),indent=1)
print(f"[b-datagen] DONE: {len(results)}/{len(qs)} traces -> {OUT}",flush=True)