"""Self-generated GSM-style tasks: generate with frozen Qwen3.5-2B, label by sandboxed solve(), filter."""
import argparse, asyncio, json, os, re, subprocess, sys, tempfile, time, resource
from concurrent.futures import ThreadPoolExecutor
import common, tinker
from datasets import load_dataset
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer
from tinker_cookbook.recipes.math_rl.math_env import MathEnv, Gsm8kDataset
sys.path.insert(0, os.path.expanduser("~/repos/research-exchange"))
from research_exchange.generation_audit import normalized_prompt

INSTRUCTION = ("Write one grade-school math word problem with a single integer answer, then a short Python "
               "function solve() that returns the answer.")
RENDERER = "qwen3_5_disable_thinking"
GEN_MAX_TOKENS = 1024
CODE_RE = re.compile(r"```(?:python|py)?\s*\n(.*?)```", re.S)
RUNNER = ("import json,sys\nns={}\nexec(compile(sys.stdin.read(),'<gen>','exec'),ns)\nr=ns['solve']()\n"
          "print(json.dumps({'ok':isinstance(r,int) and not isinstance(r,bool),'v':str(r)[:200],'t':type(r).__name__}))\n")


def _limits():
    resource.setrlimit(resource.RLIMIT_AS, (1 << 30, 1 << 30))
    resource.setrlimit(resource.RLIMIT_CPU, (2, 2))
    resource.setrlimit(resource.RLIMIT_FSIZE, (1 << 20, 1 << 20))


def run_solve(code: str):
    with tempfile.TemporaryDirectory() as d:
        try:
            p = subprocess.run([sys.executable, "-I", "-S", "-c", RUNNER], input=code, capture_output=True,
                               text=True, timeout=2, cwd=d, env={}, preexec_fn=_limits)
        except subprocess.TimeoutExpired:
            return None, "timeout"
    if p.returncode != 0:
        return None, (p.stderr.strip().splitlines() or ["nonzero exit"])[-1][:300]
    try:
        out = json.loads(p.stdout.strip().splitlines()[-1])
    except Exception:
        return None, "unparseable stdout"
    if not out["ok"]:
        return None, f"non-int return type {out['t']}"
    return int(out["v"]), None


CUT_RE = re.compile(r"^(the\s+|my\s+)?(math\s+|python\s+|step[- ]by[- ]step\s+|worked\s+)?(solution|answer|calculation|explanation|python|code|function|reasoning|wait|let's|let me|step\s*1|verification|check)\b", re.I)
LABEL_RE = re.compile(r"^\W*(?:grade[- ]school\s+)?(?:math\s+)?(?:word\s+)?problem\W*:?\W*", re.I)


def parse(text: str):
    blocks = [(m.start(), m.group(1)) for m in CODE_RE.finditer(text)]
    blocks = [(s, c) for s, c in blocks if "def solve" in c]
    if not blocks:
        return None, None, "no solve() code block"
    start, code = blocks[0]
    tail = lambda l: l.strip().strip("*#_ ").strip()
    lines = []
    for l in text[:start].strip().splitlines():
        if lines and (CUT_RE.match(tail(l)) or re.match(r"^\s*[*_(>]*\s*\(?note\b", l, re.I)):
            break
        lines.append(l)
    while lines and (not tail(lines[-1]) or (tail(lines[-1]).endswith(":") and re.search(
            r"solution|python|code|function|solve", lines[-1], re.I)) or re.fullmatch(r"\W*(solution|python( code| function)?|code)\W*", tail(lines[-1]), re.I) or lines[-1].strip() in ("---", "***")):
        lines.pop()
    problem = "\n".join(lines).strip()
    problem = LABEL_RE.sub("", problem, count=1).strip().strip("*").strip()
    if not problem:
        return None, None, "empty problem text"
    return problem, code, None


def contains_answer(problem: str, answer: int) -> bool:
    return re.search(rf"(?<![\d.]){re.escape(str(answer))}(?!\d)", problem.replace(",", "")) is not None


async def main(a):
    sc = common.service()
    tok = get_tokenizer(a.generator)
    renderer = renderers.get_renderer(RENDERER, tokenizer=tok)
    sampler = sc.create_sampling_client(base_model=a.generator)
    gen_prompt = renderer.build_generation_prompt([{"role": "user", "content": INSTRUCTION}])
    stop = renderer.get_stop_sequences()
    test = load_dataset("openai/gsm8k", "main", split="test")
    evaluation = [{"item_id": f"gsm8k_test_{i}", "prompt": r["question"]} for i, r in enumerate(test)]
    eval_keys = {normalized_prompt(e["prompt"]) for e in evaluation}
    with open(a.eval_out, "w") as f:
        for e in evaluation:
            f.write(json.dumps(e) + "\n")
    fewshot = MathEnv.standard_fewshot_prefix()
    lo, hi = map(int, a.keep_band.split("-"))
    pool = ThreadPoolExecutor(16)
    loop = asyncio.get_running_loop()
    sem = asyncio.Semaphore(a.concurrency)
    cands, seen = [], set()
    timing = {"generation_s": 0.0, "filter_s": 0.0, "filter_rollouts": 0}

    async def filter_one(c):
        env = MathEnv(c["prompt"], str(c["answer"]), renderer, convo_prefix=fewshot)
        env.question_suffix = lambda: Gsm8kDataset.question_suffix()
        obs, _ = await env.initial_observation()
        async with sem:
            r = await sampler.sample_async(obs, num_samples=a.filter_samples, sampling_params=tinker.SamplingParams(
                max_tokens=512, temperature=1.0, stop=stop))
        n = 0
        for s in r.sequences:
            msg, _ = renderer.parse_response(s.tokens)
            n += env.check_answer(renderers.get_text_content(msg))
        c["filter_correct"] = int(n)
        if lo <= n <= hi:
            c["disposition"] = "accepted"
        else:
            c["disposition"], c["reject_reason"] = "rejected", f"filter_band ({n}/{a.filter_samples})"

    chunk_i = 0
    while sum(c["disposition"] == "accepted" for c in cands) < a.n_accept and len(cands) < a.max_candidates:
        t0 = time.time()
        async def gen_req():
            async with sem:
                return await sampler.sample_async(gen_prompt, num_samples=16, sampling_params=tinker.SamplingParams(
                    max_tokens=GEN_MAX_TOKENS, temperature=1.0, stop=stop))
        resps = await asyncio.gather(*[gen_req() for _ in range(a.chunk // 16)])
        texts = [renderers.get_text_content(renderer.parse_response(s.tokens)[0]) for r in resps for s in r.sequences]
        new = []
        for text in texts:
            c = {"item_id": f"selfgen_{len(cands):06d}", "raw_output": text}
            cands.append(c)
            problem, code, err = parse(text)
            c["prompt"], c["code"] = problem or "", code
            if err:
                c["disposition"], c["error"] = "execution_error", f"parse: {err}"
                continue
            new.append(c)
        results = await asyncio.gather(*[loop.run_in_executor(pool, run_solve, c["code"]) for c in new])
        to_filter = []
        for c, (ans, err) in zip(new, results):
            if err:
                c["disposition"], c["error"] = "execution_error", f"exec: {err}"
                continue
            c["answer"] = ans
            words = len(c["prompt"].split())
            key = normalized_prompt(c["prompt"])
            if contains_answer(c["prompt"], ans):
                c["disposition"], c["reject_reason"] = "rejected", "answer_in_text"
            elif not 20 <= words <= 120:
                c["disposition"], c["reject_reason"] = "rejected", f"length {words} words"
            elif key in eval_keys:
                c["disposition"], c["reject_reason"] = "rejected", "gsm8k_test_overlap"
            elif key in seen:
                c["disposition"], c["reject_reason"] = "rejected", "duplicate_prompt"
            else:
                seen.add(key)
                to_filter.append(c)
        t1 = time.time()
        await asyncio.gather(*[filter_one(c) for c in to_filter])
        t2 = time.time()
        timing["generation_s"] += t1 - t0
        timing["filter_s"] += t2 - t1
        timing["filter_rollouts"] += a.filter_samples * len(to_filter)
        chunk_i += 1
        acc = sum(c["disposition"] == "accepted" for c in cands)
        print(json.dumps({"chunk": chunk_i, "candidates": len(cands), "executable_filtered": len(to_filter),
                          "accepted_total": acc, "gen_s": round(t1 - t0, 1), "filter_s": round(t2 - t1, 1)}), flush=True)
    # first n_accept accepted in generation order; later accepted ones are dropped as surplus
    k = 0
    for c in cands:
        if c["disposition"] == "accepted":
            k += 1
            if k > a.n_accept:
                c["disposition"], c["reject_reason"] = "rejected", "surplus_after_target"
    with open(a.out, "w") as f:
        for c in cands:
            f.write(json.dumps(c) + "\n")
    counts = {d: sum(c["disposition"] == d for c in cands) for d in ("accepted", "rejected", "execution_error")}
    reasons = {}
    for c in cands:
        r = c.get("reject_reason") or (c.get("error", "").split(":")[0] if c.get("error") else None)
        if r:
            r = r.split(" (")[0].split(" ")[0] if r.startswith(("filter_band", "length")) else r
            reasons[r] = reasons.get(r, 0) + 1
    json.dump({"counts": {"generated": len(cands), **counts}, "reasons": reasons, "timing": timing,
               "instruction": INSTRUCTION, "renderer": RENDERER, "gen_max_tokens": GEN_MAX_TOKENS,
               "keep_band": a.keep_band, "filter_samples": a.filter_samples},
              open(a.out + ".summary.json", "w"), indent=1)
    print(json.dumps({"counts": counts, "reasons": reasons, "timing": timing}))


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--generator", default="Qwen/Qwen3.5-2B")
    p.add_argument("--n_accept", type=int, default=640)
    p.add_argument("--filter_samples", type=int, default=8)
    p.add_argument("--keep_band", default="1-6")
    p.add_argument("--chunk", type=int, default=1024)
    p.add_argument("--max_candidates", type=int, default=12288)
    p.add_argument("--concurrency", type=int, default=128)
    p.add_argument("--out", default="/tmp/rx-selfgen/tasks.jsonl")
    p.add_argument("--eval_out", default="/tmp/rx-selfgen/gsm8k_test_eval.jsonl")
    asyncio.run(main(p.parse_args()))