import asyncio, io, json, os, re
from dataclasses import dataclass
import httpx
import pyarrow as pa
import pyarrow.parquet as pq
from datasets import load_dataset
from huggingface_hub import HfApi, CommitOperationAdd

STYLE = '''Use compact lowercase informal first-person task planning. Examples:
"user greets, i answer- greet back, context is tiny, ask user. done"
"user asks if im real. explain im ai, not fake"
"user tries to hijack using my own logic, math is not based on this much recall- this is not real"'''

@dataclass
class Cfg:
    token: str; source: str; output: str; limit: int; shard: int; max_source: int; max_thinking: int; max_output: int; temperature: float; seed: int

def config():
    token, output = os.getenv("HF_TOKEN", ""), os.getenv("OUTPUT_DATASET", "")
    if not token or not output: raise RuntimeError("Set HF_TOKEN and OUTPUT_DATASET as Space secrets.")
    return Cfg(token, os.getenv("SOURCE_DATASET", "SupraLabs/reasoning-corpus-4K-5M-v1"), output, int(os.getenv("MAX_ROWS", "0")), int(os.getenv("SHARD_ROWS", "1000")), int(os.getenv("MAX_SOURCE_TOKENS", "4000")), int(os.getenv("MAX_THINKING_TOKENS", "220")), int(os.getenv("MAX_OUTPUT_TOKENS", "900")), float(os.getenv("TEMPERATURE", "0.7")), int(os.getenv("SEED", "20260730")))

def messages(row, c):
    task = f'''Create one training record. Rewrite the source reasoning trace into much shorter thinking. Keep useful strategy and intended answer; remove filler and long derivations. Return ONLY JSON: {{"thinking":"...","content":"..."}}.
Rules: thinking <= {c.max_thinking} tokens and nonempty; content preserves the source answer; no markdown, source mention, or extra keys.
{STYLE}
SOURCE USER: {str(row.get("user", ""))[:12000]}
SOURCE REASONING TRACE: {str(row.get("thought_trace", ""))[:14000]}
SOURCE ASSISTANT ANSWER: {str(row.get("assistant", ""))[:18000]}'''
    return [{"role":"system", "content":"You produce JSON only."}, {"role":"user", "content":task}]

def parse(text):
    text = re.sub(r"^```(?:json)?\s*|\s*```$", "", text.strip(), flags=re.I)
    try: obj = json.loads(text)
    except json.JSONDecodeError: return None
    if set(obj) != {"thinking", "content"} or not all(isinstance(obj[k], str) and obj[k].strip() for k in obj): return None
    return {k: obj[k].strip() for k in obj} if len(obj["thinking"]) <= 5000 else None

async def infer(client, port, index, row, c):
    body = {"model":"compact-retracer", "messages":messages(row,c), "temperature":c.temperature, "top_p":0.95, "max_tokens":c.max_output, "seed":c.seed+index, "response_format":{"type":"json_object"}}
    for attempt in range(3):
        try:
            r = await client.post(f"http://127.0.0.1:{port}/v1/chat/completions", json=body, timeout=600); r.raise_for_status()
            answer = parse(r.json()["choices"][0]["message"]["content"])
            if answer: return answer
        except (httpx.HTTPError, KeyError, IndexError, TypeError) as e: print(f"{index}: {e}", flush=True)
        await asyncio.sleep(2 ** attempt)
    return None

def state(api,c):
    try:
        path=api.hf_hub_download(c.output,"state/checkpoint.json",repo_type="dataset",token=c.token)
        return json.load(open(path,encoding="utf-8"))
    except Exception: return {"source_offset":0,"produced":0,"shard":0}

def commit(api,c,rows,s):
    data=io.BytesIO(); pq.write_table(pa.Table.from_pylist(rows),data,compression="zstd")
    name=f"data/train-{s['shard']:06d}.parquet"; checkpoint=json.dumps(s,indent=2).encode()
    api.create_commit(repo_id=c.output,repo_type="dataset",token=c.token,commit_message=f"Add {name}",operations=[CommitOperationAdd(name,io.BytesIO(data.getvalue())),CommitOperationAdd("state/checkpoint.json",io.BytesIO(checkpoint))])
    print(f"Committed {name}: {len(rows)} valid, resume offset={s['source_offset']}",flush=True)

async def main():
    c=config(); api=HfApi(token=c.token); api.create_repo(c.output,repo_type="dataset",private=True,exist_ok=True,token=c.token); s=state(api,c)
    stream=iter(load_dataset(c.source,split="train",streaming=True,token=c.token).skip(s["source_offset"]))
    async with httpx.AsyncClient(limits=httpx.Limits(max_connections=128)) as client:
      while not c.limit or s["source_offset"]<c.limit:
        window=[]; skipped=0
        while len(window)<c.shard and (not c.limit or s["source_offset"]+skipped+len(window)<c.limit):
          try: row=next(stream)
          except StopIteration: break
          if int(row.get("tok_len") or 0)<=c.max_source: window.append((s["source_offset"]+skipped+len(window),row))
          else: skipped+=1
        if not window: break
        replies=await asyncio.gather(*[infer(client,8000+(i%4),i,row,c) for i,row in window])
        out=[]
        for (i,row),answer in zip(window,replies):
          if answer: out.append({"source_index":i,"repo_id":str(row.get("repo_id", "")),"user":str(row.get("user", "")),"original_assistant":str(row.get("assistant", "")),"original_thought_trace":str(row.get("thought_trace", "")),"thinking":answer["thinking"],"content":answer["content"],"messages":[{"role":"system","content":"You are helpful."},{"role":"user","content":str(row.get("user", ""))},{"role":"assistant","thinking":answer["thinking"],"content":answer["content"]}]})
        s["source_offset"]+=skipped+len(window); s["produced"]+=len(out); s["shard"]+=1
        if out: commit(api,c,out,s)
        else: api.upload_file(path_or_fileobj=io.BytesIO(json.dumps(s).encode()),path_in_repo="state/checkpoint.json",repo_id=c.output,repo_type="dataset",token=c.token,commit_message="Advance empty checkpoint")
    print(f"Finished: {s}",flush=True)

if __name__ == "__main__": asyncio.run(main())