import logging
import argparse
import json
from pathlib import Path

def main() -> int:
	ap = argparse.ArgumentParser()
	root = Path(__file__).resolve().parents[1]
	ap.add_argument("--bench", default=str(root / "data" / "benchmarks" / "kpi.json"))
	ap.add_argument("--wm", default=str(root / "models" / "wm_mlp" / "metrics.json"))
	ap.add_argument("--verifier", default=str(root / "models" / "verifier_calib" / "metrics.json"))
	ap.add_argument("--devloop", default=str(root / "data" / "dashboards" / "devloop.json"))
	ap.add_argument("--coder-rl", default=str(root / "models" / "coder_rl" / "metrics.json"))
	ap.add_argument("--coder-nearmiss", default=str(root / "models" / "coder_nearmiss" / "model.json"))
	ap.add_argument("--planner", default=str(root / "data" / "planner_prefs" / "metrics.json"))
	ap.add_argument("--offpolicy", default=str(root / "models" / "wm_offpolicy" / "metrics.json"))
	ap.add_argument("--shaping", default=str(root / "data" / "planner" / "reward_shaping.json"))
	ap.add_argument("--out", default=str(root / "data" / "dashboards" / "summary.json"))
	ap.add_argument("--use-budgeted", action="store_true", help="Prefer budgeted success rates if available")
	ap.add_argument("--traces", default=str(root / "data" / "traces" / "loop_run.jsonl"), help="Optional traces JSONL to compute planner self-eval summary")
	ap.add_argument("--practice", default=str(root / "data" / "logs" / "practice_results.jsonl"), help="Optional coding practice results JSONL for summary")
	ap.add_argument("--scheduler", default=str(root / "data" / "logs" / "scheduler_runs.jsonl"), help="Optional scheduler run logs JSONL to compute dev KPIs")
	ap.add_argument("--matrix", default=str(root / "data" / "ci" / "matrix_results.json"), help="Optional CI matrix results JSON to include")
	ap.add_argument("--unsafe-logs", default=str(root / "data" / "logs" / "scheduler_runs.jsonl"), help="Optional logs path to compute unsafe edit counts")
	ap.add_argument("--scheduler", default=str(root / "data" / "logs" / "scheduler_runs.jsonl"), help="Optional scheduler run logs JSONL to compute dev KPIs")
	ap.add_argument("--registry", default=str(root / "data" / "registry" / "registry.json"), help="Optional dataset/artifact registry snapshot to include")
	ap.add_argument("--batch-audit", default=str(root / "data" / "benchmarks" / "batch_audit.json"), help="Optional batch audit JSON to include")
	ap.add_argument("--external", default=str(root / "data" / "benchmarks" / "external_results.jsonl"), help="Optional external benchmark results JSONL for summary")
	ap.add_argument("--devtools-metrics", default=str(root / "data" / "devtools" / "metrics.json"), help="Optional devtools metrics JSON to include")
	args = ap.parse_args()

	bench = {}
	wm = {}
	ver = {}
	devloop = {}
	coder_rl = {}
	coder_nearmiss = {}
	planner = {}
	try:
		bench = json.loads(Path(args.bench).read_text(encoding="utf-8"))
	except Exception:
		bench = {}
	try:
		wm = json.loads(Path(args.wm).read_text(encoding="utf-8"))
	except Exception:
		wm = {}
	try:
		ver = json.loads(Path(args.verifier).read_text(encoding="utf-8"))
	except Exception:
		ver = {}
	try:
		dlp = Path(args.devloop)
		if dlp.exists():
			devloop = json.loads(dlp.read_text(encoding="utf-8"))
	except Exception:
		devloop = {}
	try:
		rl = Path(args.coder_rl)
		if rl.exists():
			coder_rl = json.loads(rl.read_text(encoding="utf-8"))
	except Exception:
		coder_rl = {}
	try:
		nm = Path(args.coder_nearmiss)
		if nm.exists():
			coder_nearmiss = json.loads(nm.read_text(encoding="utf-8"))
	except Exception:
		coder_nearmiss = {}
	try:
		planner = json.loads(Path(args.planner).read_text(encoding="utf-8"))
	except Exception:
		planner = {}
	# Optional offpolicy metrics
	offpolicy = {}
	try:
		opp = Path(args.offpolicy)
		if opp.exists():
			offpolicy = json.loads(opp.read_text(encoding="utf-8"))
	except Exception:
		offpolicy = {}
	# Optional reward shaping summary
	reward_shaping = {}
	try:
		shp = Path(args.shaping)
		if shp.exists():
			reward_shaping = json.loads(shp.read_text(encoding="utf-8"))
	except Exception:
		reward_shaping = {}
	registry = {}
	try:
		rp = Path(args.registry)
		if rp.exists():
			registry = json.loads(rp.read_text(encoding="utf-8"))
	except Exception:
		registry = {}
	# Batch/Cache audit
	batch_audit = {}
	try:
		ba = Path(args.batch_audit)
		if ba.exists():
			batch_audit = json.loads(ba.read_text(encoding="utf-8"))
	except Exception:
		batch_audit = {}
	# External summary
	external_summary = {}
	try:
		pe = Path(args.external)
		if pe.exists():
			total = 0
			ok = 0
			by_domain = {}
			with pe.open("r", encoding="utf-8") as f:
				for line in f:
					line = line.strip()
					if not line:
						continue
					try:
						rec = json.loads(line)
					except Exception:
						continue
					dom = str(rec.get("domain", ""))
					res = rec.get("result", {}) if isinstance(rec, dict) else {}
					status = str(res.get("status", "")).lower()
					total += 1
					if dom:
						dd = by_domain.get(dom, {"total": 0, "ok": 0})
						dd["total"] += 1
						by_domain[dom] = dd
					if status == "ok":
						ok += 1
						if dom:
							by_domain[dom]["ok"] += 1
			external_summary = {
			"total": int(total),
			"ok": int(ok),
			"success_rate": float(round((float(ok) / max(1.0, float(total))), 3)),
			"by_domain": {k: {"total": v["total"], "ok": v["ok"], "success_rate": float(round((float(v["ok"]) / max(1.0, float(v["total"]))), 3)) } for k, v in by_domain.items()},
		}
	except Exception:
		external_summary = {}

	# Optionally compute a view with budgeted success rates for convenience
	bench_view = dict(bench)
	if bool(args.use_budgeted):
		try:
			cli_sum = dict(bench_view.get("cli_summary", {}))
			dom_sum = dict(bench_view.get("dom_summary", {}))
			if "success_rate_budgeted" in cli_sum:
				cli_sum["success_rate_effective"] = float(round(float(cli_sum.get("success_rate_budgeted", cli_sum.get("success_rate", 0.0))), 3))
			else:
				cli_sum["success_rate_effective"] = float(round(float(cli_sum.get("success_rate", 0.0)), 3))
			if "success_rate_budgeted" in dom_sum:
				dom_sum["success_rate_effective"] = float(round(float(dom_sum.get("success_rate_budgeted", dom_sum.get("success_rate", 0.0))), 3))
			else:
				dom_sum["success_rate_effective"] = float(round(float(dom_sum.get("success_rate", 0.0)), 3))
			bench_view["cli_summary"] = cli_sum
			bench_view["dom_summary"] = dom_sum
		except Exception:
			pass

	# Optional planner self-eval aggregation from traces
	self_eval_summary = {}
	try:
		p = Path(args.traces)
		if p.exists():
			cnt = 0
			sel = []
			with p.open("r", encoding="utf-8") as f:
				for line in f:
					line = line.strip()
					if not line:
						continue
					try:
						rec = json.loads(line)
					except Exception:
						continue
					pc = rec.get("planner_candidates") if isinstance(rec, dict) else None
					if isinstance(pc, dict):
						vals = pc.get("self_eval") or []
						for v in vals:
							try:
								fv = float(v)
								sel.append(fv)
								cnt += 1
							except Exception:
								pass
			if cnt > 0:
				avg = sum(sel) / float(cnt)
				self_eval_summary = {"count": int(cnt), "avg": float(round(avg, 3)), "max": float(round(max(sel), 3)), "min": float(round(min(sel), 3))}
	except Exception:
		self_eval_summary = {}

	# Optional coding micro-suite/practice results aggregation
	practice_summary = {}
	try:
		pp = Path(args.practice)
		if pp.exists():
			n = 0
			ok = 0
			rc_0 = 0
			repos = {}
			with pp.open("r", encoding="utf-8") as f:
				for line in f:
					line = line.strip()
					if not line:
						continue
					try:
						rec = json.loads(line)
					except Exception:
						continue
					n += 1
					rc = int(rec.get("returncode", 0))
					rc_0 += 1 if rc == 0 else 0
					ok += 1 if rc == 0 else 0
					name = str(rec.get("repo", ""))
					if name:
						d = repos.get(name, {"n": 0, "ok": 0})
						d["n"] += 1
						d["ok"] += 1 if rc == 0 else 0
						repos[name] = d
			per_repo = {k: {"n": v["n"], "ok": v["ok"], "ok_rate": float(round((float(v["ok"]) / max(1.0, float(v["n"]))), 3)) } for k, v in repos.items()}
			practice_summary = {"runs": int(n), "repos": int(len(repos)), "ok": int(ok), "ok_rate": float(round((float(ok) / max(1.0, float(n))), 3)), "rc0": int(rc_0), "per_repo": per_repo}
	except Exception:
		practice_summary = {}

	# Optional dev KPIs from scheduler runs (time-to-fix, revert rate proxy)
	dev_kpis = {}
	try:
		sp = Path(args.scheduler)
		if sp.exists():
			n = 0
			ok = 0
			elapsed = []
			with sp.open("r", encoding="utf-8") as f:
				for line in f:
					line = line.strip()
					if not line:
						continue
					try:
						rec = json.loads(line)
					except Exception:
						continue
					n += 1
					if int(rec.get("returncode", 1)) == 0:
						ok += 1
					if "elapsed_sec" in rec:
						try:
							elapsed.append(float(rec.get("elapsed_sec", 0.0)))
						except Exception:
							pass
			ok_rate = float(round((float(ok) / max(1.0, float(n))), 3))
			dev_kpis = {
				"runs": int(n),
				"ok": int(ok),
				"ok_rate": ok_rate,
				"time_to_fix_p50": float(round(sorted(elapsed)[len(elapsed)//2], 3)) if elapsed else 0.0,
				"time_to_fix_p90": float(round(sorted(elapsed)[int(len(elapsed)*0.9)], 3)) if elapsed else 0.0,
				"revert_rate_proxy": float(round(1.0 - ok_rate, 3)),
			}
	except Exception:
		dev_kpis = {}

	# Optional: include CI matrix summary
	matrix = {}
	try:
		mp = Path(args.matrix)
		if mp.exists():
			matrix = json.loads(mp.read_text(encoding="utf-8"))
	except Exception:
		matrix = {}

	# Optional: compute unsafe edit counts from logs
	unsafe_edits = {}
	try:
		lp = Path(args.unsafe_logs)
		if lp.exists():
			unsafe = 0
			total = 0
			with lp.open("r", encoding="utf-8") as f:
				for line in f:
					line = line.strip()
					if not line:
						continue
					try:
						obj = json.loads(line)
					except Exception:
						continue
					stderr_tail = str(obj.get("stderr_tail", ""))
					stdout_tail = str(obj.get("stdout_tail", ""))
					total += 1
					if ("diff_validation_failed" in stderr_tail) or ("diff_validation_failed" in stdout_tail):
						unsafe += 1
			unsafe_edits = {"total": int(total), "unsafe": int(unsafe)}
	except Exception:
		unsafe_edits = {}

	# Optional compact interpretation to help quickly assess status
	interpretation = {}
	try:
		cli_sr = float(bench_view.get("cli_summary", {}).get("success_rate", 0.0))
		dom_sr = float(bench_view.get("dom_summary", {}).get("success_rate", 0.0))
		dom_lat = float(bench_view.get("dom_summary", {}).get("avg_latency_sec", 0.0))
		ver_ece = float(ver.get("cal_ece", ver.get("base_ece", 0.0))) if isinstance(ver, dict) else 0.0
		wm_auc = float(wm.get("wm_auc", 0.0)) if isinstance(wm, dict) else 0.0
		interpretation = {
			"cli_ok": (cli_sr >= 0.8),
			"dom_ok": (dom_sr >= 0.8),
			"dom_latency_ok": (dom_lat <= 10.0),
			"verifier_calib_ok": (ver_ece <= 0.15),
			"wm_auc_ok": (wm_auc >= 0.7),
		}
	except Exception:
		interpretation = {}

	# Devtools metrics
	devtools = {}
	try:
		dp = Path(args.devtools_metrics)
		if dp.exists():
			devtools = json.loads(dp.read_text(encoding="utf-8"))
	except Exception:
		devtools = {}

	summary = {"bench": bench_view, "wm": wm, "verifier": ver, "devloop": devloop, "coder_rl": coder_rl, "coder_nearmiss": coder_nearmiss, "planner": planner, "offpolicy": offpolicy, "reward_shaping": reward_shaping, "planner_self_eval": self_eval_summary, "practice": practice_summary, "dev_kpis": dev_kpis, "devtools": devtools, "matrix": matrix, "unsafe_edits": unsafe_edits, "registry": registry, "external": external_summary, "batch_audit": batch_audit, "interpretation": interpretation}
	outp = Path(args.out)
	outp.parent.mkdir(parents=True, exist_ok=True)
	outp.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(outp)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict, List

def callers_of(index_graph: Dict[str, Any], symbol: str) -> List[str]:
    res: List[str] = []
    calls = (index_graph.get("calls") or {}) if isinstance(index_graph, dict) else {}
    for f, call_list in calls.items():
        try:
            if any((c.get("name") == symbol) for c in (call_list or [])):
                res.append(f)
        except Exception:
            continue
    return res

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    idx_path = root / "data" / "sandbox" / "tmp" / "index.json"
    out_path = root / "data" / "sandbox" / "tmp" / "query_result.json"
    try:
        idx = json.loads(idx_path.read_text(encoding="utf-8")) if idx_path.exists() else {}
    except Exception:
        idx = {}
    q = ""  # symbol to query; kept simple for scaffold
    try:
        import os
        q = str(os.environ.get("AGI_QUERY_SYMBOL", ""))
    except Exception:
        q = ""
    res = callers_of((idx.get("graph") or {}), q) if q else []
    out = {"ok": True, "symbol": q, "callers": res[:200]}
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"ok": True, "out": str(out_path), "n": len(res)}))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict, List

def callers_of(index_graph: Dict[str, Any], symbol: str) -> List[str]:
    res: List[str] = []
    calls = (index_graph.get("calls") or {}) if isinstance(index_graph, dict) else {}
    for f, call_list in calls.items():
        try:
            if any((c.get("name") == symbol) for c in (call_list or [])):
                res.append(f)
        except Exception:
            continue
    return res

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    idx_path = root / "data" / "sandbox" / "tmp" / "index.json"
    out_path = root / "data" / "sandbox" / "tmp" / "query_result.json"
    try:
        idx = json.loads(idx_path.read_text(encoding="utf-8")) if idx_path.exists() else {}
    except Exception:
        idx = {}
    q = ""  # symbol to query; kept simple for scaffold
    try:
        import os
        q = str(os.environ.get("AGI_QUERY_SYMBOL", ""))
    except Exception:
        q = ""
    res = callers_of((idx.get("graph") or {}), q) if q else []
    out = {"ok": True, "symbol": q, "callers": res[:200]}
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")


main

from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict, List

def callers_of(index_graph: Dict[str, Any], symbol: str) -> List[str]:
    res: List[str] = []
    calls = (index_graph.get("calls") or {}) if isinstance(index_graph, dict) else {}
    for f, call_list in calls.items():
        try:
            if any((c.get("name") == symbol) for c in (call_list or [])):
                res.append(f)
        except Exception:
            continue
    return res

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    idx_path = root / "data" / "sandbox" / "tmp" / "index.json"
    out_path = root / "data" / "sandbox" / "tmp" / "query_result.json"
    try:
        idx = json.loads(idx_path.read_text(encoding="utf-8")) if idx_path.exists() else {}
    except Exception:
        idx = {}
    q = ""  # symbol to query; kept simple for scaffold
    try:
        import os
        q = str(os.environ.get("AGI_QUERY_SYMBOL", ""))
    except Exception:
        q = ""
    res = callers_of((idx.get("graph") or {}), q) if q else []
    out = {"ok": True, "symbol": q, "callers": res[:200]}
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"ok": True, "out": str(out_path), "n": len(res)}))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Dict, Any, List

def shard_by_top_level(root: Path, index_graph: Dict[str, Any]) -> Dict[str, List[str]]:
    shards: Dict[str, List[str]] = {}
    files = list((index_graph.get("functions") or {}).keys()) if isinstance(index_graph, dict) else []
    for f in files:
        try:
            p = Path(f)
            top = (p.parts or [""])[0]
        except Exception:
            top = ""
        shards.setdefault(top, []).append(f)
    return shards

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    idx_path = root / "data" / "sandbox" / "tmp" / "index.json"
    out_path = root / "data" / "sandbox" / "tmp" / "shards.json"
    try:
        idx = json.loads(idx_path.read_text(encoding="utf-8")) if idx_path.exists() else {}
    except Exception:
        idx = {}
    shards = shard_by_top_level(root, (idx.get("graph") or {})) if isinstance(idx, dict) else {}
    out = {"ok": True, "strategy": "top_level", "shards": {k: v[:200] for k, v in shards.items()}}
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"ok": True, "out": str(out_path), "n": len(shards)}))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


import logging

import json
from pathlib import Path
from typing import Dict, Any, List

def shard_by_top_level(root: Path, index_graph: Dict[str, Any]) -> Dict[str, List[str]]:
    shards: Dict[str, List[str]] = {}
    files = list((index_graph.get("functions") or {}).keys()) if isinstance(index_graph, dict) else []
    for f in files:
        try:
            p = Path(f)
            top = (p.parts or [""])[0]
        except Exception:
            top = ""
        shards.setdefault(top, []).append(f)
    return shards

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    idx_path = root / "data" / "sandbox" / "tmp" / "index.json"
    out_path = root / "data" / "sandbox" / "tmp" / "shards.json"
    try:
        idx = json.loads(idx_path.read_text(encoding="utf-8")) if idx_path.exists() else {}
    except Exception:
        idx = {}
    shards = shard_by_top_level(root, (idx.get("graph") or {})) if isinstance(idx, dict) else {}
    out = {"ok": True, "strategy": "top_level", "shards": {k: v[:200] for k, v in shards.items()}}
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps({"ok": True, "out": str(out_path), "n": len(shards)}))
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import ast
import sys
from pathlib import Path

TEMPLATE = """
from agi_dw.flags.sdk import FLAGS as _FLAGS

def {name}_dual_route(*args, **kwargs):
    if _FLAGS.is_enabled('{flag_name}', default=True):
        return {primary}(*args, **kwargs)
    return {secondary}(*args, **kwargs)
"""

def inject_dual_route(path: Path, symbol_primary: str, symbol_secondary: str, flag_name: str) -> bool:
    src = path.read_text(encoding="utf-8")
    if f"is_enabled('{flag_name}'" in src or f"is_enabled(\"{flag_name}\"" in src:
        return False
    # Append function template at end of file
    name = f"{symbol_primary}_route"
    snippet = TEMPLATE.format(name=name, flag_name=flag_name, primary=symbol_primary, secondary=symbol_secondary)
    path.write_text(src + ("\n\n" if not src.endswith("\n") else "\n") + snippet, encoding="utf-8")
    return True

def main() -> int:
    if len(sys.argv) < 5:
        print("Usage: dual_route_inject.py <file.py> <primary_symbol> <secondary_symbol> <FLAG_NAME>")
        return 2
    target = Path(sys.argv[1])
    primary = sys.argv[2]
    secondary = sys.argv[3]
    flag_name = sys.argv[4]
    if not target.exists():
        print(f"No such file: {target}")
        return 2
    changed = inject_dual_route(target, primary, secondary, flag_name)
    print({"ok": True, "changed": bool(changed), "file": str(target)})
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


from agi_dw.flags.sdk import FLAGS as _FLAGS

def {name}_dual_route(*args, **kwargs):
    if _FLAGS.is_enabled('{flag_name}', default=True):
        return {primary}(*args, **kwargs)
    return {secondary}(*args, **kwargs)
"""

def inject_dual_route(path: Path, symbol_primary: str, symbol_secondary: str, flag_name: str) -> bool:
    src = path.read_text(encoding="utf-8")
    if f"is_enabled('{flag_name}'" in src or f"is_enabled(\"{flag_name}\"" in src:
        return False
    # Append function template at end of file
    name = f"{symbol_primary}_route"
    snippet = TEMPLATE.format(name=name, flag_name=flag_name, primary=symbol_primary, secondary=symbol_secondary)
    path.write_text(src + ("\n\n" if not src.endswith("\n") else "\n") + snippet, encoding="utf-8")
    return True

def main() -> int:
    if len(sys.argv) < 5:
        print("Usage: dual_route_inject.py <file.py> <primary_symbol> <secondary_symbol> <FLAG_NAME>")
        return 2
    target = Path(sys.argv[1])
    primary = sys.argv[2]
    secondary = sys.argv[3]
    flag_name = sys.argv[4]
    if not target.exists():
        print(f"No such file: {target}")
        return 2
    changed = inject_dual_route(target, primary, secondary, flag_name)
    print({"ok": True, "changed": bool(changed), "file": str(target)})
    return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Any, List

def inject_logging(text: str) -> str:
	if "import logging" in text:
		return text
	# Insert at top after shebang/encoding if present
	lines = text.splitlines()
	insert_at = 0
	if lines and lines[0].startswith("#!/"):
		insert_at = 1
	lines.insert(insert_at, "import logging")
	return "\n".join(lines)

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	report_path = root / "data" / "sandbox" / "tmp" / "obs_inject_report.json"
	changes: List[str] = []
	for p in root.rglob("*.py"):
		try:
			text = p.read_text(encoding="utf-8", errors="ignore")
		except Exception:
			continue
		if "data/sandbox" in str(p) or "/.venv/" in str(p):
			continue
		new_text = inject_logging(text)
		if new_text != text:
			try:
				p.write_text(new_text, encoding="utf-8")
				changes.append(str(p.resolve()))
			except Exception:
				continue
	out = {"ok": True, "touched": len(changes), "files": changes[:200]}
	report_path.parent.mkdir(parents=True, exist_ok=True)
	report_path.write_text(json.dumps(out, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(report_path)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations

import json
from pathlib import Path
from typing import Dict, Any, List

def inject_logging(text: str) -> str:
	if "import logging" in text:
		return text
	# Insert at top after shebang/encoding if present
	lines = text.splitlines()
	insert_at = 0
	if lines and lines[0].startswith("#!/"):
		insert_at = 1
	lines.insert(insert_at, "import logging")
	return "\n".join(lines)

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	report_path = root / "data" / "sandbox" / "tmp" / "obs_inject_report.json"
	changes: List[str] = []
	for p in root.rglob("*.py"):
		try:
			text = p.read_text(encoding="utf-8", errors="ignore")
		except Exception:
			continue
		if "data/sandbox" in str(p) or "/.venv/" in str(p):
			continue
		new_text = inject_logging(text)
		if new_text != text:
			try:
				p.write_text(new_text, encoding="utf-8")
				changes.append(str(p.resolve()))
			except Exception:
				continue


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, Any, List, Tuple

def run_gates(repo_dir: Path) -> Tuple[bool, Dict[str, Any]]:
    """Run a minimal gate suite: flake/mypy/pytest if present; summarize exit codes."""
    summary: Dict[str, Any] = {}
    ok = True
    cmds = [
        ("pytest -q", "tests"),
        ("flake8 .", "."),
        ("mypy .", "."),
    ]
    for cmd, _ in cmds:
        try:
            rc = subprocess.run(cmd, shell=True, cwd=str(repo_dir), capture_output=True, text=True).returncode
        except Exception:
            rc = 0
        summary[cmd] = rc
        ok = ok and (rc == 0)
    return ok, summary

def apply_codemod(repo_dir: Path, codemod_script: Path, args: List[str]) -> Tuple[bool, str]:
    try:
        proc = subprocess.run(["python", str(codemod_script), *args], cwd=str(repo_dir), capture_output=True, text=True)
        return (proc.returncode == 0), proc.stdout.strip()
    except Exception as e:
        return False, str(e)

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    target_repo = root  # local repo by default
    # Shadow workspace
    with tempfile.TemporaryDirectory() as td:
        shadow = Path(td) / "repo"
        shutil.copytree(str(target_repo), str(shadow))
        # Example codemod: no-op (can be replaced by real scripts)
        # Run gates
        ok_before, gates_before = run_gates(shadow)
        # For safety in scaffold, do not actually mutate files; persist reports
        ok_after, gates_after = run_gates(shadow)
        report = {
            "ok": bool(ok_before and ok_after),
            "gates_before": gates_before,
            "gates_after": gates_after,
        }
        out = root / "data" / "sandbox" / "tmp" / "codemod_report.json"
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
        print(json.dumps({"ok": bool(report.get("ok", True)), "out": str(out)}))
        return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, Any, List, Tuple

def run_gates(repo_dir: Path) -> Tuple[bool, Dict[str, Any]]:
    """Run a minimal gate suite: flake/mypy/pytest if present; summarize exit codes."""
    summary: Dict[str, Any] = {}
    ok = True
    cmds = [
        ("pytest -q", "tests"),
        ("flake8 .", "."),
        ("mypy .", "."),
    ]
    for cmd, _ in cmds:
        try:
            rc = subprocess.run(cmd, shell=True, cwd=str(repo_dir), capture_output=True, text=True).returncode
        except Exception:
            rc = 0
        summary[cmd] = rc
        ok = ok and (rc == 0)
    return ok, summary

def apply_codemod(repo_dir: Path, codemod_script: Path, args: List[str]) -> Tuple[bool, str]:
    try:
        proc = subprocess.run(["python", str(codemod_script), *args], cwd=str(repo_dir), capture_output=True, text=True)
        return (proc.returncode == 0), proc.stdout.strip()
    except Exception as e:
        return False, str(e)

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    target_repo = root  # local repo by default
    # Shadow workspace
    with tempfile.TemporaryDirectory() as td:
        shadow = Path(td) / "repo"
        shutil.copytree(str(target_repo), str(shadow))
        # Example codemod: no-op (can be replaced by real scripts)
        # Run gates
        ok_before, gates_before = run_gates(shadow)


def run_gates(repo_dir: Path) -> Tuple[bool, Dict[str, Any]]:
    """Run a minimal gate suite: flake/mypy/pytest if present; summarize exit codes."""
    summary: Dict[str, Any] = {}
    ok = True
    cmds = [
        ("pytest -q", "tests"),
        ("flake8 .", "."),
        ("mypy .", "."),
    ]
    for cmd, _ in cmds:
        try:
            rc = subprocess.run(cmd, shell=True, cwd=str(repo_dir), capture_output=True, text=True).returncode
        except Exception:
            rc = 0
        summary[cmd] = rc
        ok = ok and (rc == 0)
    return ok, summary

def apply_codemod(repo_dir: Path, codemod_script: Path, args: List[str]) -> Tuple[bool, str]:
    try:
        proc = subprocess.run(["python", str(codemod_script), *args], cwd=str(repo_dir), capture_output=True, text=True)
        return (proc.returncode == 0), proc.stdout.strip()
    except Exception as e:
        return False, str(e)

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    target_repo = root  # local repo by default
    # Shadow workspace
    with tempfile.TemporaryDirectory() as td:
        shadow = Path(td) / "repo"
        shutil.copytree(str(target_repo), str(shadow))
        # Example codemod: no-op (can be replaced by real scripts)
        # Run gates
        ok_before, gates_before = run_gates(shadow)
        # For safety in scaffold, do not actually mutate files; persist reports
        ok_after, gates_after = run_gates(shadow)
        report = {
            "ok": bool(ok_before and ok_after),
            "gates_before": gates_before,
            "gates_after": gates_after,
        }
        out = root / "data" / "sandbox" / "tmp" / "codemod_report.json"


        ("mypy .", "."),
    ]
    for cmd, _ in cmds:
        try:
            rc = subprocess.run(cmd, shell=True, cwd=str(repo_dir), capture_output=True, text=True).returncode
        except Exception:
            rc = 0
        summary[cmd] = rc
        ok = ok and (rc == 0)
    return ok, summary

def apply_codemod(repo_dir: Path, codemod_script: Path, args: List[str]) -> Tuple[bool, str]:
    try:
        proc = subprocess.run(["python", str(codemod_script), *args], cwd=str(repo_dir), capture_output=True, text=True)
        return (proc.returncode == 0), proc.stdout.strip()
    except Exception as e:
        return False, str(e)

def main() -> int:
    root = Path(__file__).resolve().parents[2]
    target_repo = root  # local repo by default
    # Shadow workspace
    with tempfile.TemporaryDirectory() as td:
        shadow = Path(td) / "repo"
        shutil.copytree(str(target_repo), str(shadow))
        # Example codemod: no-op (can be replaced by real scripts)
        # Run gates
        ok_before, gates_before = run_gates(shadow)
        # For safety in scaffold, do not actually mutate files; persist reports
        ok_after, gates_after = run_gates(shadow)
        report = {
            "ok": bool(ok_before and ok_after),
            "gates_before": gates_before,
            "gates_after": gates_after,
        }
        out = root / "data" / "sandbox" / "tmp" / "codemod_report.json"
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
        print(json.dumps({"ok": bool(report.get("ok", True)), "out": str(out)}))
        return 0

if __name__ == "__main__":
    raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict

def _load_json(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		return {}

def _safe_load_yaml(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		import yaml  # type: ignore
		data = yaml.safe_load(path.read_text(encoding="utf-8"))
		return data if isinstance(data, dict) else {}
	except Exception:
		return {}

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	policies_path = root / "data" / "sandbox" / "config" / "policies.yaml"
	plan_risk_path = root / "data" / "sandbox" / "tmp" / "plan_risk.json"
	out_path = root / "release" / "plan.yaml"

	policies = _safe_load_yaml(policies_path)
	risk = _load_json(plan_risk_path)
	max_canary = 0.1 if float(risk.get("score", 0.0) or 0.0) > 1.0 else 0.25
	plan = {
		"canary_percent": max_canary,
		"health_checks": ["tests_pass", "error_rate<1%", "p95_ms<budget"],
		"rollback": {"strategy": "immediate", "triggers": ["test_fail", "error_spike", "p95_regress"]},
		"policy_caps": risk.get("caps", {}),
	}
	out_path.parent.mkdir(parents=True, exist_ok=True)
	# Write as YAML if available, else JSON
	try:
		import yaml  # type: ignore
		out_path.write_text(yaml.safe_dump(plan, sort_keys=False), encoding="utf-8")
	except Exception:
		out_path.write_text(json.dumps(plan, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(out_path)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict

def _load_json(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		return {}

def _safe_load_yaml(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		import yaml  # type: ignore
		data = yaml.safe_load(path.read_text(encoding="utf-8"))
		return data if isinstance(data, dict) else {}
	except Exception:
		return {}

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	policies_path = root / "data" / "sandbox" / "config" / "policies.yaml"
	plan_risk_path = root / "data" / "sandbox" / "tmp" / "plan_risk.json"
	out_path = root / "release" / "plan.yaml"

	policies = _safe_load_yaml(policies_path)


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from pathlib import Path
from typing import Any, Dict

def _load_json(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		return {}

def _safe_load_yaml(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		import yaml  # type: ignore
		data = yaml.safe_load(path.read_text(encoding="utf-8"))
		return data if isinstance(data, dict) else {}
	except Exception:
		return {}

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	policies_path = root / "data" / "sandbox" / "config" / "policies.yaml"
	plan_risk_path = root / "data" / "sandbox" / "tmp" / "plan_risk.json"
	out_path = root / "release" / "plan.yaml"

	policies = _safe_load_yaml(policies_path)
	risk = _load_json(plan_risk_path)
	max_canary = 0.1 if float(risk.get("score", 0.0) or 0.0) > 1.0 else 0.25
	plan = {
		"canary_percent": max_canary,
		"health_checks": ["tests_pass", "error_rate<1%", "p95_ms<budget"],
		"rollback": {"strategy": "immediate", "triggers": ["test_fail", "error_spike", "p95_regress"]},
		"policy_caps": risk.get("caps", {}),
	}
	out_path.parent.mkdir(parents=True, exist_ok=True)
	# Write as YAML if available, else JSON
	try:


def _load_json(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		return json.loads(path.read_text(encoding="utf-8"))
	except Exception:
		return {}

def _safe_load_yaml(path: Path) -> Dict[str, Any]:
	if not path.exists():
		return {}
	try:
		import yaml  # type: ignore
		data = yaml.safe_load(path.read_text(encoding="utf-8"))
		return data if isinstance(data, dict) else {}
	except Exception:
		return {}

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	policies_path = root / "data" / "sandbox" / "config" / "policies.yaml"
	plan_risk_path = root / "data" / "sandbox" / "tmp" / "plan_risk.json"
	out_path = root / "release" / "plan.yaml"

	policies = _safe_load_yaml(policies_path)
	risk = _load_json(plan_risk_path)
	max_canary = 0.1 if float(risk.get("score", 0.0) or 0.0) > 1.0 else 0.25
	plan = {
		"canary_percent": max_canary,
		"health_checks": ["tests_pass", "error_rate<1%", "p95_ms<budget"],
		"rollback": {"strategy": "immediate", "triggers": ["test_fail", "error_spike", "p95_regress"]},
		"policy_caps": risk.get("caps", {}),
	}
	out_path.parent.mkdir(parents=True, exist_ok=True)
	# Write as YAML if available, else JSON
	try:
		import yaml  # type: ignore
		out_path.write_text(yaml.safe_dump(plan, sort_keys=False), encoding="utf-8")
	except Exception:
		out_path.write_text(json.dumps(plan, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(out_path)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	plan_path = root / "release" / "plan.yaml"
	status_path = root / "data" / "sandbox" / "tmp" / f"release_{datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}.json"
	# Simulate a publish according to plan
	plan_text = plan_path.read_text(encoding="utf-8") if plan_path.exists() else "{}"
	status: Dict[str, Any] = {
		"ok": True,
		"plan": str(plan_path),
		"timestamp": datetime.utcnow().isoformat() + "Z",
		"result": {"canary": "started", "notes": "simulated publish"},
	}
	status_path.parent.mkdir(parents=True, exist_ok=True)
	status_path.write_text(json.dumps(status, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(status_path)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


#!/usr/bin/env python3
from __future__ import annotations
import logging

import json
from datetime import datetime
from pathlib import Path
from typing import Any, Dict

def main() -> int:
	root = Path(__file__).resolve().parents[2]
	status_path = root / "data" / "sandbox" / "tmp" / f"release_rollback_{datetime.utcnow().strftime('%Y%m%dT%H%M%SZ')}.json"
	status: Dict[str, Any] = {
		"ok": True,
		"timestamp": datetime.utcnow().isoformat() + "Z",
		"result": {"rollback": "completed", "notes": "simulated rollback"},
	}
	status_path.parent.mkdir(parents=True, exist_ok=True)
	status_path.write_text(json.dumps(status, ensure_ascii=False, indent=2), encoding="utf-8")
	print(json.dumps({"ok": True, "out": str(status_path)}))
	return 0

if __name__ == "__main__":
	raise SystemExit(main())


from typing import List, Dict, Any, Tuple
import sys, io, re, json, math, statistics, itertools, functools, copy, random
import torch
from collections import Counter
from agi_dw.core.llm.hf_client import HFClient  # type: ignore
from agi_dw.core.utils.prompt_logger import get_prompt_logger  # type: ignore
from transformers import AutoModelForCausalLM, AutoTokenizer  # type: ignore
from transformers.generation.logits_process import (  # type: ignore
    LogitsProcessor,
    LogitsProcessorList,
)
from agi_dw.core.llm.adapter_cache import AdapterCache  # type: ignore

import os
import subprocess
import shlex
from pathlib import Path
from agi_dw.scripts.selfplay.modules.paths import get_data_root, data_path  # type: ignore
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError
from agi_dw.scripts.selfplay.modules.lora import _loraify_linear, _inject_lora, _forward_with_lora_delta, _set_lora_enabled, _patch_forwards
from agi_dw.scripts.selfplay.modules.model_io import load_seed, save_adapters
from agi_dw.scripts.selfplay.modules.tasks import _PRIMS, _PRIMS_EASY, sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import SanitizeLogits, sample, generate_text
from agi_dw.scripts.selfplay.modules.healing import heal_code, _strip_triple_quoted, is_disallowed
from agi_dw.scripts.selfplay.modules.sandbox import _safe_builtins, _worker_run_one, _worker_run_one_subproc, _values_equal, run_and_test
from agi_dw.scripts.selfplay.modules.train import sft_step, contrastive_step, _autocast_cuda
from agi_dw.scripts.selfplay.modules.evolution import mutate_prompt, heuristic_reward, evolve_population, evolution_loop
from agi_dw.scripts.selfplay.modules.devrepo import _dev_episode, _ensure_local_repo, _is_local_repo_usable
from agi_dw.scripts.selfplay.modules.wm import _resolve_best_wm_model_path
from agi_dw.scripts.selfplay.modules.curriculum import Memory, Curriculum, _maybe_gate_curriculum
from agi_dw.scripts.selfplay.modules.reward import _compute_reward
from agi_dw.scripts.selfplay.modules.episode import episode_loop

if __name__ == "__main__":
    # Allow environment overrides for key knobs so Make can scale without code edits
    def _env_int(name: str, default: int) -> int:
        try:
            v = os.environ.get(name)
            return int(v) if v is not None and str(v).strip() != "" else int(default)
        except Exception:
            return int(default)

    def _env_float(name: str, default: float) -> float:
        try:
            v = os.environ.get(name)
            return float(v) if v is not None and str(v).strip() != "" else float(default)
        except Exception:
            return float(default)

    def _env_str(name: str, default: str) -> str:
        v = os.environ.get(name)
        return str(v) if v is not None and str(v).strip() != "" else str(default)

    def _env_bool(name: str, default: bool) -> bool:
        v = os.environ.get(name)
        if v is None:
            return bool(default)
        s = str(v).strip().lower()
        return s in ("1", "true", "yes", "y", "on")

    cfg = {
        "model_name": _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "n_samples": _env_int("SELFPLAY_N_SAMPLES", 4),
        "max_tokens": _env_int("SELFPLAY_MAX_TOKENS", 512),
        "lr": _env_float("SELFPLAY_LR", 1e-4),
        "train_steps": _env_int("SELFPLAY_TRAIN_STEPS", 10),
        "neg_steps": _env_int("SELFPLAY_NEG_STEPS", 3),
        "ckpt_every": _env_int("SELFPLAY_CKPT_EVERY", 200),
        "ckpt_path": _env_str("SELFPLAY_CKPT_PATH", "adapters.pt"),
        "adapter_r": _env_int("SELFPLAY_ADAPTER_R", 8),
        "adapter_alpha": _env_int("SELFPLAY_ADAPTER_ALPHA", 16),
        # Optional dev-loop mode: default to arena (no dev_repo unless provided via env)
        "dev_repo": (_env_str("SELFPLAY_DEV_REPO", "").strip() or None),
        "dev_args": _env_str("SELFPLAY_DEV_ARGS", ""),
        # Evolution mode and parameters
        "evolution_mode": _env_bool("SELFPLAY_EVOLUTION_MODE", False),
        "evo_num_agents": _env_int("SELFPLAY_EVO_NUM_AGENTS", 5),
        "evo_num_generations": _env_int("SELFPLAY_EVO_NUM_GENERATIONS", 10),
        "evo_top_k_survivors": _env_int("SELFPLAY_EVO_TOP_K", 2),
        "evo_max_prompt_len": _env_int("SELFPLAY_EVO_MAX_PROMPT_LEN", 80),
        "evo_max_new_tokens": _env_int("SELFPLAY_EVO_MAX_NEW_TOKENS", 100),
        "evo_train_survivors": _env_bool("SELFPLAY_EVO_TRAIN_SURVIVORS", True),
        "evo_train_steps": _env_int("SELFPLAY_EVO_TRAIN_STEPS", 3),
        # HITL curriculum approval
        "require_curriculum_approval": _env_bool("SELFPLAY_REQUIRE_CURRICULUM_APPROVAL", True),
        "approval_timeout": _env_int("SELFPLAY_APPROVAL_TIMEOUT", 60),
        # HITL import approval (default off for arena primitives/stdlib)
        "require_import_approval": _env_bool("SELFPLAY_REQUIRE_IMPORT_APPROVAL", False),
        # Logging/memory
        "mem_path": _env_str(
            "SELFPLAY_MEM_PATH",
            str(data_path("selfplay", "progress.jsonl")),
        ),
        "log_prompts": _env_bool("SELFPLAY_LOG_PROMPTS", True),
        # Advanced
        "apply_meta": _env_bool("SELFPLAY_APPLY_META", True),
        "enable_metaopt": _env_bool("SELFPLAY_ENABLE_METAOPT", True),
        "force_lora_inject": _env_bool("SELFPLAY_FORCE_LORA", False),
        "auto_lora_if_frozen": _env_bool("SELFPLAY_AUTO_LORA_IF_FROZEN", False),
        "adapter_dir": _env_str("SELFPLAY_ADAPTER_DIR", ""),
        "direct_model": _env_bool("SELFPLAY_DIRECT_MODEL", False),
        "log_console": _env_bool("SELFPLAY_LOG_CONSOLE", True),
        # Decode policy knobs
        "beam_verify": _env_bool("SELFPLAY_BEAM_VERIFY", False),
        "num_beams": _env_int("SELFPLAY_NUM_BEAMS", 4),
        # Speculative decoding draft model (used when SELFPLAY_SPECULATIVE=1)
        "draft_model": _env_str("SELFPLAY_DRAFT_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
    }
    episode_loop(cfg)


from agi_dw.scripts.selfplay.modules.paths import get_data_root, data_path  # type: ignore
import multiprocessing as mp
from concurrent.futures import ProcessPoolExecutor, TimeoutError
from agi_dw.scripts.selfplay.modules.lora import _loraify_linear, _inject_lora, _forward_with_lora_delta, _set_lora_enabled, _patch_forwards
from agi_dw.scripts.selfplay.modules.model_io import load_seed, save_adapters
from agi_dw.scripts.selfplay.modules.tasks import _PRIMS, _PRIMS_EASY, sample_task, render_prompt
from agi_dw.scripts.selfplay.modules.generation import SanitizeLogits, sample, generate_text
from agi_dw.scripts.selfplay.modules.healing import heal_code, _strip_triple_quoted, is_disallowed
from agi_dw.scripts.selfplay.modules.sandbox import _safe_builtins, _worker_run_one, _worker_run_one_subproc, _values_equal, run_and_test
from agi_dw.scripts.selfplay.modules.train import sft_step, contrastive_step, _autocast_cuda
from agi_dw.scripts.selfplay.modules.evolution import mutate_prompt, heuristic_reward, evolve_population, evolution_loop
from agi_dw.scripts.selfplay.modules.devrepo import _dev_episode, _ensure_local_repo, _is_local_repo_usable
from agi_dw.scripts.selfplay.modules.wm import _resolve_best_wm_model_path
from agi_dw.scripts.selfplay.modules.curriculum import Memory, Curriculum, _maybe_gate_curriculum
from agi_dw.scripts.selfplay.modules.reward import _compute_reward
from agi_dw.scripts.selfplay.modules.episode import episode_loop

if __name__ == "__main__":
    # Allow environment overrides for key knobs so Make can scale without code edits
    def _env_int(name: str, default: int) -> int:
        try:
            v = os.environ.get(name)
            return int(v) if v is not None and str(v).strip() != "" else int(default)
        except Exception:
            return int(default)

    def _env_float(name: str, default: float) -> float:
        try:
            v = os.environ.get(name)
            return float(v) if v is not None and str(v).strip() != "" else float(default)
        except Exception:
            return float(default)

    def _env_str(name: str, default: str) -> str:
        v = os.environ.get(name)
        return str(v) if v is not None and str(v).strip() != "" else str(default)

    def _env_bool(name: str, default: bool) -> bool:
        v = os.environ.get(name)
        if v is None:
            return bool(default)
        s = str(v).strip().lower()
        return s in ("1", "true", "yes", "y", "on")

    cfg = {


from agi_dw.scripts.selfplay.modules.healing import heal_code, _strip_triple_quoted, is_disallowed
from agi_dw.scripts.selfplay.modules.sandbox import _safe_builtins, _worker_run_one, _worker_run_one_subproc, _values_equal, run_and_test
from agi_dw.scripts.selfplay.modules.train import sft_step, contrastive_step, _autocast_cuda
from agi_dw.scripts.selfplay.modules.evolution import mutate_prompt, heuristic_reward, evolve_population, evolution_loop
from agi_dw.scripts.selfplay.modules.devrepo import _dev_episode, _ensure_local_repo, _is_local_repo_usable
from agi_dw.scripts.selfplay.modules.wm import _resolve_best_wm_model_path
from agi_dw.scripts.selfplay.modules.curriculum import Memory, Curriculum, _maybe_gate_curriculum
from agi_dw.scripts.selfplay.modules.reward import _compute_reward
from agi_dw.scripts.selfplay.modules.episode import episode_loop

if __name__ == "__main__":
    # Allow environment overrides for key knobs so Make can scale without code edits
    def _env_int(name: str, default: int) -> int:
        try:
            v = os.environ.get(name)
            return int(v) if v is not None and str(v).strip() != "" else int(default)
        except Exception:
            return int(default)

    def _env_float(name: str, default: float) -> float:
        try:
            v = os.environ.get(name)
            return float(v) if v is not None and str(v).strip() != "" else float(default)
        except Exception:
            return float(default)

    def _env_str(name: str, default: str) -> str:
        v = os.environ.get(name)
        return str(v) if v is not None and str(v).strip() != "" else str(default)

    def _env_bool(name: str, default: bool) -> bool:
        v = os.environ.get(name)
        if v is None:
            return bool(default)
        s = str(v).strip().lower()
        return s in ("1", "true", "yes", "y", "on")

    cfg = {
        "model_name": _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "n_samples": _env_int("SELFPLAY_N_SAMPLES", 4),
        "max_tokens": _env_int("SELFPLAY_MAX_TOKENS", 512),
        "lr": _env_float("SELFPLAY_LR", 1e-4),
        "train_steps": _env_int("SELFPLAY_TRAIN_STEPS", 10),
        "neg_steps": _env_int("SELFPLAY_NEG_STEPS", 3),
        "ckpt_every": _env_int("SELFPLAY_CKPT_EVERY", 200),


from agi_dw.scripts.selfplay.modules.reward import _compute_reward
from agi_dw.scripts.selfplay.modules.episode import episode_loop

if __name__ == "__main__":
    # Allow environment overrides for key knobs so Make can scale without code edits
    def _env_int(name: str, default: int) -> int:
        try:
            v = os.environ.get(name)
            return int(v) if v is not None and str(v).strip() != "" else int(default)
        except Exception:
            return int(default)

    def _env_float(name: str, default: float) -> float:
        try:
            v = os.environ.get(name)
            return float(v) if v is not None and str(v).strip() != "" else float(default)
        except Exception:
            return float(default)

    def _env_str(name: str, default: str) -> str:
        v = os.environ.get(name)
        return str(v) if v is not None and str(v).strip() != "" else str(default)

    def _env_bool(name: str, default: bool) -> bool:
        v = os.environ.get(name)
        if v is None:
            return bool(default)
        s = str(v).strip().lower()
        return s in ("1", "true", "yes", "y", "on")

    cfg = {
        "model_name": _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "n_samples": _env_int("SELFPLAY_N_SAMPLES", 4),
        "max_tokens": _env_int("SELFPLAY_MAX_TOKENS", 512),
        "lr": _env_float("SELFPLAY_LR", 1e-4),
        "train_steps": _env_int("SELFPLAY_TRAIN_STEPS", 10),
        "neg_steps": _env_int("SELFPLAY_NEG_STEPS", 3),
        "ckpt_every": _env_int("SELFPLAY_CKPT_EVERY", 200),
        "ckpt_path": _env_str("SELFPLAY_CKPT_PATH", "adapters.pt"),
        "adapter_r": _env_int("SELFPLAY_ADAPTER_R", 8),
        "adapter_alpha": _env_int("SELFPLAY_ADAPTER_ALPHA", 16),
        # Optional dev-loop mode: default to arena (no dev_repo unless provided via env)


if __name__ == "__main__":
    # Allow environment overrides for key knobs so Make can scale without code edits
    def _env_int(name: str, default: int) -> int:
        try:
            v = os.environ.get(name)
            return int(v) if v is not None and str(v).strip() != "" else int(default)
        except Exception:
            return int(default)

    def _env_float(name: str, default: float) -> float:
        try:
            v = os.environ.get(name)
            return float(v) if v is not None and str(v).strip() != "" else float(default)
        except Exception:
            return float(default)

    def _env_str(name: str, default: str) -> str:
        v = os.environ.get(name)
        return str(v) if v is not None and str(v).strip() != "" else str(default)

    def _env_bool(name: str, default: bool) -> bool:
        v = os.environ.get(name)
        if v is None:
            return bool(default)
        s = str(v).strip().lower()
        return s in ("1", "true", "yes", "y", "on")

    cfg = {
        "model_name": _env_str("SELFPLAY_MODEL", "meta-llama/Llama-3.1-8B-Instruct"),
        "n_samples": _env_int("SELFPLAY_N_SAMPLES", 4),
        "max_tokens": _env_int("SELFPLAY_MAX_TOKENS", 512),
        "lr": _env_float("SELFPLAY_LR", 1e-4),
        "train_steps": _env_int("SELFPLAY_TRAIN_STEPS", 10),
        "neg_steps": _env_int("SELFPLAY_NEG_STEPS", 3),
        "ckpt_every": _env_int("SELFPLAY_CKPT_EVERY", 200),
        "ckpt_path": _env_str("SELFPLAY_CKPT_PATH", "adapters.pt"),
        "adapter_r": _env_int("SELFPLAY_ADAPTER_R", 8),
        "adapter_alpha": _env_int("SELFPLAY_ADAPTER_ALPHA", 16),
        # Optional dev-loop mode: default to arena (no dev_repo unless provided via env)
        "dev_repo": (_env_str("SELFPLAY_DEV_REPO", "").strip() or None),
        "dev_args": _env_str("SELFPLAY_DEV_ARGS", ""),
        # Evolution mode and parameters
        "evolution_mode": _env_bool("SELFPLAY_EVOLUTION_MODE", False),
        "evo_num_agents": _env_int("SELFPLAY_EVO_NUM_AGENTS", 5),
        "evo_num_generations": _env_int("SELFPLAY_EVO_NUM_GENERATIONS", 10),
        "evo_top_k_survivors": _env_int("SELFPLAY_EVO_TOP_K", 2),


import os
import types
import torch

from agi_dw.scripts.selfplay.modules import tasks, generation, sandbox, healing

class FakeTokenizer:
    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt

def test_generation_sample_extracts_fenced_block():
    code_text = (
        "```python\n"
        "def solve(n:int)->int:\n"
        "    return n*2\n"
        "```"
    )
    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")
    assert "return n*2" in outs[0]

def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"
        "    return n*2\n"
    )
    tok = FakeTokenizer([raw])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")

def test_heal_code_improves_failure(monkeypatch):
    spec = make_spec_int_unary("add_one", lambda n: n + 1)
    failing = (
        "def solve(n:int)->int:\n"
        "    return n+2\n"
    )

    def fake_generate_text(tok, model, prompt: str, max_new_tokens: int):
        return (
            "def solve(n:int)->int:\n"
            "    return n+1\n"
        )

    # monkeypatch the generate_text used inside healing
    monkeypatch.setattr(healing, "generate_text", fake_generate_text)
    healed = healing.heal_code(None, None, "prompt", failing, {"first_failure": "wrong answer"}, max_new_tokens=64)
    ok, rep = sandbox.run_and_test(healed, spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")


FakeTokenizer

import os
import types
import torch

from agi_dw.scripts.selfplay.modules import tasks, generation, sandbox, healing

class FakeTokenizer:
    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out


FakeModel

class FakeTokenizer:
    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt


class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt

def test_generation_sample_extracts_fenced_block():
    code_text = (
        "```python\n"
        "def solve(n:int)->int:\n"
        "    return n*2\n"
        "```"
    )
    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1


        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt

def test_generation_sample_extracts_fenced_block():
    code_text = (
        "```python\n"
        "def solve(n:int)->int:\n"
        "    return n*2\n"
        "```"
    )
    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")
    assert "return n*2" in outs[0]

def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"


def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt

def test_generation_sample_extracts_fenced_block():
    code_text = (
        "```python\n"
        "def solve(n:int)->int:\n"
        "    return n*2\n"
        "```"
    )
    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")
    assert "return n*2" in outs[0]

def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"
        "    return n*2\n"
    )
    tok = FakeTokenizer([raw])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")


    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt

def test_generation_sample_extracts_fenced_block():
    code_text = (
        "```python\n"
        "def solve(n:int)->int:\n"
        "    return n*2\n"
        "```"
    )
    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")
    assert "return n*2" in outs[0]

def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"
        "    return n*2\n"
    )
    tok = FakeTokenizer([raw])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")

def test_heal_code_improves_failure(monkeypatch):
    spec = make_spec_int_unary("add_one", lambda n: n + 1)
    failing = (
        "def solve(n:int)->int:\n"
        "    return n+2\n"
    )

    def fake_generate_text(tok, model, prompt: str, max_new_tokens: int):
        return (
            "def solve(n:int)->int:\n"
            "    return n+1\n"


    tok = FakeTokenizer([code_text])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")
    assert "return n*2" in outs[0]

def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"
        "    return n*2\n"
    )
    tok = FakeTokenizer([raw])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")

def test_heal_code_improves_failure(monkeypatch):
    spec = make_spec_int_unary("add_one", lambda n: n + 1)
    failing = (
        "def solve(n:int)->int:\n"
        "    return n+2\n"
    )

    def fake_generate_text(tok, model, prompt: str, max_new_tokens: int):
        return (
            "def solve(n:int)->int:\n"
            "    return n+1\n"
        )

    # monkeypatch the generate_text used inside healing
    monkeypatch.setattr(healing, "generate_text", fake_generate_text)
    healed = healing.heal_code(None, None, "prompt", failing, {"first_failure": "wrong answer"}, max_new_tokens=64)
    ok, rep = sandbox.run_and_test(healed, spec)
    assert ok is True


def test_generation_sample_accepts_raw_def_solve():
    raw = (
        "def solve(n:int)->int:\n"
        "    return n*2\n"
    )
    tok = FakeTokenizer([raw])
    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")

def test_heal_code_improves_failure(monkeypatch):
    spec = make_spec_int_unary("add_one", lambda n: n + 1)
    failing = (
        "def solve(n:int)->int:\n"
        "    return n+2\n"
    )

    def fake_generate_text(tok, model, prompt: str, max_new_tokens: int):
        return (
            "def solve(n:int)->int:\n"
            "    return n+1\n"
        )

    # monkeypatch the generate_text used inside healing
    monkeypatch.setattr(healing, "generate_text", fake_generate_text)
    healed = healing.heal_code(None, None, "prompt", failing, {"first_failure": "wrong answer"}, max_new_tokens=64)
    ok, rep = sandbox.run_and_test(healed, spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")


    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [


import os
import types
import torch

from agi_dw.scripts.selfplay.modules import tasks, generation, sandbox, healing

class FakeTokenizer:
    def __init__(self, texts):
        self._texts = list(texts)
        self.eos_token_id = 0
        self.pad_token_id = 0

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3


generate

    def __call__(self, prompt: str, return_tensors: str = "pt"):
        # minimal encoding with fixed length 3
        return {
            "input_ids": torch.tensor([[1, 2, 3]]),
            "attention_mask": torch.tensor([[1, 1, 1]]),
        }

    def decode(self, tokens, skip_special_tokens: bool = True):
        # return next prepared text regardless of tokens
        if self._texts:
            return self._texts.pop(0)
        return ""

class FakeModel(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # a single parameter to provide device
        self._p = torch.nn.Parameter(torch.tensor(0.0))

    def generate(self, **kwargs):
        # produce a tensor with the right shape: (num_return_sequences, in_len + k)
        num = int(kwargs.get("num_return_sequences", 1))
        input_ids = kwargs["input_ids"]
        in_len = int(input_ids.shape[1])
        k = 3
        out = torch.zeros((num, in_len + k), dtype=torch.long)
        # prefix with dummy input ids
        out[:, :in_len] = input_ids[0]
        # suffix arbitrary ids
        out[:, in_len:] = 9
        return out

def make_spec_int_unary(name: str, ref):
    return {
        "name": name,
        "signature": "def solve(n:int)->int:",
        "tests": [
            {"inp": [0], "out": ref(0)},
            {"inp": [3], "out": ref(3)},
            {"inp": [-4], "out": ref(-4)},
        ],
    }

def test_render_prompt_contains_signature_and_fence():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    prompt = tasks.render_prompt(spec)
    assert spec["signature"] in prompt
    assert "```" in prompt


    model = FakeModel()
    outs = generation.sample(tok, model, prompt="ignored", n=1, max_new_tokens=16)
    assert len(outs) >= 1
    assert outs[0].strip().startswith("def solve(")

def test_sandbox_run_and_test_wraps_body_only():
    spec = make_spec_int_unary("double", lambda n: n * 2)
    ok, rep = sandbox.run_and_test("return n*2", spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")

def test_heal_code_improves_failure(monkeypatch):
    spec = make_spec_int_unary("add_one", lambda n: n + 1)
    failing = (
        "def solve(n:int)->int:\n"
        "    return n+2\n"
    )

    def fake_generate_text(tok, model, prompt: str, max_new_tokens: int):
        return (
            "def solve(n:int)->int:\n"
            "    return n+1\n"
        )

    # monkeypatch the generate_text used inside healing
    monkeypatch.setattr(healing, "generate_text", fake_generate_text)
    healed = healing.heal_code(None, None, "prompt", failing, {"first_failure": "wrong answer"}, max_new_tokens=64)
    ok, rep = sandbox.run_and_test(healed, spec)
    assert ok is True
    assert rep.get("passed") == rep.get("total")


from .common_imports import *
from .text_utils import extract_first_fenced_block
import ast

def _safe_builtins() -> Dict[str, Any]:
    # Build a minimal allowed builtins dict and expose it both as __builtins__ and globals
    allowed: Dict[str, Any] = {}
    for k in [
        "abs","all","any","enumerate","len","range","min","max","sum","map","filter",
        # add numeric and sequence helpers commonly needed
        "zip","sorted","reversed","list","tuple","dict","set", "bool", "int", "float", "str", "print",
        # allow stdlib imports (episode gate handles policy)
        "__import__"
    ]:
        try:
            allowed[k] = getattr(__builtins__, k)
        except Exception:
            pass
    b: Dict[str, Any] = {"__builtins__": allowed}
    b.update(allowed)
    return b

def _worker_run_one(code_str: str, sig: str, args: list) -> Tuple[bool, Any]:
    """Top-level worker: exec the candidate in a fresh sandbox and run one test.
    Returns (ok, value_or_error_string)."""
    # Use a single shared namespace so imports are visible to solve() globals
    env: Dict[str, Any] = _safe_builtins()
    # Build full source robustly to avoid nested def solve
    try:
        src = _compose_src(code_str, sig)
    except Exception:
        src = (code_str if re.match(r"^def\s+solve\s*\(", code_str)
               else sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in code_str.splitlines()))
    try:
        exec(src, env, env)
        solve = env.get("solve")
        if not callable(solve):
            return (False, "no_solve")
        return (True, solve(*args))
    except Exception as e:
        return (False, str(e))

def _compose_src(code_str: str, sig: str) -> str:
    """Compose executable source. Preserve top-level imports + solve() when present.

    Fallbacks:
    - If only a solve signature is present, extract its indented body and attach to provided signature.
    - If body-only code is provided, wrap it under the provided signature.
    """
    s = code_str.lstrip("\ufeff")
    # Try to preserve imports + solve block using AST
    try:
        tree = ast.parse(s)
        lines = s.splitlines()
        # collect import spans
        import_spans: List[Tuple[int, int]] = []
        for node in getattr(tree, "body", []):
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                a = max(1, getattr(node, "lineno", 1))
                b = max(a, getattr(node, "end_lineno", a))
                import_spans.append((a, b))
        # find solve()
        solve_node = None
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef) and getattr(node, "name", "") == "solve":
                solve_node = node
                break
        if solve_node is not None:
            s_a = max(1, getattr(solve_node, "lineno", 1))
            s_b = max(s_a, getattr(solve_node, "end_lineno", s_a))
            chunks: List[str] = []
            seen: set[str] = set()
            for a, b in import_spans:
                block = "\n".join(lines[a - 1:b]).strip()
                if block and block not in seen:
                    seen.add(block)
                    chunks.append(block)
            solve_block = "\n".join(lines[s_a - 1:s_b])
            if chunks:
                chunks.append("")
            chunks.append(solve_block)
            return "\n".join(chunks)
    except Exception:
        pass
    # If text already starts with a full def solve, accept it as-is
    s_strip = s.lstrip()
    if s_strip.startswith("def solve"):
        return s_strip
    # If a solve signature exists somewhere, extract its indented body and attach to provided signature
    m = re.search(r"^\s*def\s+solve\s*\(.*\):\s*$", s, flags=re.MULTILINE)
    if m:
        lines = s.splitlines()
        idx = 0
        for i, ln in enumerate(lines):
            if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
                idx = i
                break
        body_lines: List[str] = []
        for ln in lines[idx + 1:]:
            if not ln.strip():
                body_lines.append(ln)
                continue
            if re.match(r"^\S", ln):
                break
            body_lines.append(ln)
        return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in body_lines)
    # Otherwise treat input as body-only and wrap
    return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in s.splitlines())

def _worker_run_one_subproc(code_str: str, sig: str, args: list, timeout_sec: float) -> Tuple[bool, Any]:
    """Execute one test in an isolated tiny subprocess to avoid heavy imports/pickling."""
    try:
        import base64 as _b64  # type: ignore
        import tempfile as _tmp  # type: ignore
        # Pre-compose full source in parent to avoid quoting issues in child
        try:
            src_full = _compose_src(code_str, sig)
        except Exception:
            src_full = code_str
        payload_env = {
            "SRC": _b64.b64encode(src_full.encode("utf-8")).decode("ascii"),
            "A": _b64.b64encode(json.dumps(args).encode("utf-8")).decode("ascii"),
        }
        script_content = (
            "import os, json, base64, builtins\n"
            "src = base64.b64decode(os.environ['SRC']).decode('utf-8')\n"
            "args = json.loads(base64.b64decode(os.environ['A']).decode('utf-8'))\n"
            # Build a minimal, explicit builtins dict and expose it
            "allowed = {}\n"
            "for k in ['abs','all','any','enumerate','len','range','min','max','sum','map','filter','zip','sorted','reversed','list','tuple','dict','set','bool','int','float','str','print','__import__']:\n"
            "    try: allowed[k] = getattr(builtins, k)\n"
            "    except Exception: pass\n"
            "env = {'__builtins__': allowed}\n"
            "env.update(allowed)\n"
            "try:\n"
            "    exec(src, env, env)\n"
            "    solve = env.get('solve')\n"
            "    if not callable(solve):\n"
            "        print(json.dumps({'ok': False, 'err': 'no_solve'})); raise SystemExit(0)\n"
            "    out = solve(*args)\n"
            "    print(json.dumps({'ok': True, 'out': out}))\n"
            "except Exception as e:\n"
            "    print(json.dumps({'ok': False, 'err': str(e)}))\n"
        )
        tmp = None
        try:
            tmp = _tmp.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8")
            tmp.write(script_content)
            tmp.flush()
            tmp_path = tmp.name
        finally:
            try:
                if tmp is not None:
                    tmp.close()
            except Exception:
                pass
        try:
            p = subprocess.run([sys.executable, tmp_path], env={**os.environ, **payload_env}, capture_output=True, text=True, timeout=float(timeout_sec))
        finally:
            try:
                os.unlink(tmp_path)
            except Exception:
                pass
        txt = (p.stdout or "").strip()
        if not txt:
            err = f"no_output rc={p.returncode} stderr={(p.stderr or '').strip()}"
            return (False, err)
        try:
            obj = json.loads(txt)
        except Exception:
            err = f"bad_json rc={p.returncode} stderr={(p.stderr or '').strip()} out={txt}"
            return (False, err)
        if bool(obj.get("ok")):
            return True, obj.get("out")
        return (False, str(obj.get('err')))
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except Exception as e:
        return False, str(e)

def _values_equal(a: Any, b: Any) -> bool:
    # Numeric tolerance for floats
    try:
        import numbers as _nums  # type: ignore
        if isinstance(a, _nums.Number) and isinstance(b, _nums.Number):
            # Use isclose for floats; ints will compare exactly with isclose as well
            return math.isclose(float(a), float(b), rel_tol=1e-9, abs_tol=1e-9)
    except Exception:
        pass
    # Sequences
    if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
        if len(a) != len(b):
            return False
        return all(_values_equal(x, y) for x, y in zip(a, b))
    # Dicts
    if isinstance(a, dict) and isinstance(b, dict):
        if set(a.keys()) != set(b.keys()):
            return False
        return all(_values_equal(a[k], b[k]) for k in a.keys())
    return a == b

def run_and_test(code: str, spec: Dict[str, Any], ex: Any | None = None) -> Tuple[bool, Dict[str, Any]]:
    g: Dict[str, Any] = _safe_builtins()
    loc: Dict[str, Any] = {}
    try:
        # Normalize incoming code: accept either full def solve(...) or body-only
        txt = code.strip()
        # Strip accidental code fences via shared helper
        blk = extract_first_fenced_block(txt)
        if blk is not None:
            txt = blk.strip()
        # Per-test timeout using tiny subprocess per test case
        try:
            test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
        except Exception:
            test_timeout = 2.5
        passed, total = 0, len(spec["tests"])
        timeouts = 0
        first_failure: Any | None = None
        first_mismatch: Dict[str, Any] | None = None
        for tcase in spec["tests"]:
            ok_case, out_val = _worker_run_one_subproc(txt, spec["signature"], tcase["inp"], timeout_sec=test_timeout)
            if (not ok_case) and out_val == "timeout":
                timeouts += 1
            if ok_case and _values_equal(out_val, tcase["out"]):
                passed += 1
            else:
                # Record first exception or first mismatch for healing context
                if (not ok_case) and first_failure is None:
                    first_failure = out_val
                if ok_case and first_mismatch is None:
                    try:
                        first_mismatch = {"inp": tcase.get("inp"), "want": tcase.get("out"), "got": out_val}
                    except Exception:
                        first_mismatch = None
        ok = passed == total
        rep: Dict[str, Any] = {"passed": passed, "total": total, "timeouts": timeouts}
        if first_failure is not None:
            rep["first_failure"] = first_failure
        if first_mismatch is not None:
            rep["first_mismatch"] = first_mismatch
        return ok, rep
    except Exception as e:
        return False, {"err": str(e)}

def run_subset_tests(code: str, spec: Dict[str, Any], k: int = 3) -> Tuple[float, Dict[str, Any]]:
    """Run a quick subset of tests (first k) to cheaply screen candidates.

    Returns (pass_ratio_on_subset, rep_dict_like_full_runner).
    """
    g: Dict[str, Any] = _safe_builtins()
    loc: Dict[str, Any] = {}
    try:
        # Normalize incoming code similar to run_and_test
        txt = code.strip()
        blk2 = extract_first_fenced_block(txt)
        if blk2 is not None:
            txt = blk2.strip()
        try:
            test_timeout = float(os.environ.get("SELFPLAY_TEST_TIMEOUT_SEC", "2.5") or 2.5)
        except Exception:
            test_timeout = 2.5
        try:
            subset_k = max(1, int(k))
        except Exception:
            subset_k = max(1, int(os.environ.get("SELFPLAY_PRECHECK_K", "3") or 3))
        tests = spec.get("tests", [])
        # Deterministic prefix subset for stability
        sub = tests[: min(len(tests), subset_k)]
        passed, total = 0, len(sub)
        timeouts = 0
        first_failure: Any | None = None
        for tcase in sub:
            ok_case, out_val = _worker_run_one_subproc(txt, spec["signature"], tcase["inp"], timeout_sec=test_timeout)
            if (not ok_case) and out_val == "timeout":
                timeouts += 1
            if ok_case and _values_equal(out_val, tcase["out"]):
                passed += 1
            elif (not ok_case) and first_failure is None:
                first_failure = out_val
        rep: Dict[str, Any] = {"passed": passed, "total": total, "timeouts": timeouts}
        if first_failure is not None:
            rep["first_failure"] = first_failure
        pr = (float(passed) / float(max(1, total))) if total > 0 else 0.0
        return pr, rep
    except Exception as e:
        return 0.0, {"err": str(e), "passed": 0, "total": 0, "timeouts": 0}


from .common_imports import *
from .text_utils import extract_first_fenced_block
import ast

def _safe_builtins() -> Dict[str, Any]:
    # Build a minimal allowed builtins dict and expose it both as __builtins__ and globals
    allowed: Dict[str, Any] = {}
    for k in [
        "abs","all","any","enumerate","len","range","min","max","sum","map","filter",
        # add numeric and sequence helpers commonly needed
        "zip","sorted","reversed","list","tuple","dict","set", "bool", "int", "float", "str", "print",
        # allow stdlib imports (episode gate handles policy)
        "__import__"
    ]:
        try:
            allowed[k] = getattr(__builtins__, k)
        except Exception:
            pass
    b: Dict[str, Any] = {"__builtins__": allowed}
    b.update(allowed)
    return b

def _worker_run_one(code_str: str, sig: str, args: list) -> Tuple[bool, Any]:
    """Top-level worker: exec the candidate in a fresh sandbox and run one test.
    Returns (ok, value_or_error_string)."""
    # Use a single shared namespace so imports are visible to solve() globals
    env: Dict[str, Any] = _safe_builtins()
    # Build full source robustly to avoid nested def solve
    try:
        src = _compose_src(code_str, sig)
    except Exception:
        src = (code_str if re.match(r"^def\s+solve\s*\(", code_str)
               else sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in code_str.splitlines()))
    try:
        exec(src, env, env)
        solve = env.get("solve")
        if not callable(solve):
            return (False, "no_solve")
        return (True, solve(*args))
    except Exception as e:
        return (False, str(e))


import ast

def _safe_builtins() -> Dict[str, Any]:
    # Build a minimal allowed builtins dict and expose it both as __builtins__ and globals
    allowed: Dict[str, Any] = {}
    for k in [
        "abs","all","any","enumerate","len","range","min","max","sum","map","filter",
        # add numeric and sequence helpers commonly needed
        "zip","sorted","reversed","list","tuple","dict","set", "bool", "int", "float", "str", "print",
        # allow stdlib imports (episode gate handles policy)
        "__import__"
    ]:
        try:
            allowed[k] = getattr(__builtins__, k)
        except Exception:
            pass
    b: Dict[str, Any] = {"__builtins__": allowed}
    b.update(allowed)
    return b

def _worker_run_one(code_str: str, sig: str, args: list) -> Tuple[bool, Any]:
    """Top-level worker: exec the candidate in a fresh sandbox and run one test.
    Returns (ok, value_or_error_string)."""
    # Use a single shared namespace so imports are visible to solve() globals
    env: Dict[str, Any] = _safe_builtins()
    # Build full source robustly to avoid nested def solve
    try:
        src = _compose_src(code_str, sig)
    except Exception:
        src = (code_str if re.match(r"^def\s+solve\s*\(", code_str)
               else sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in code_str.splitlines()))
    try:
        exec(src, env, env)
        solve = env.get("solve")
        if not callable(solve):
            return (False, "no_solve")
        return (True, solve(*args))
    except Exception as e:
        return (False, str(e))

def _compose_src(code_str: str, sig: str) -> str:
    """Compose executable source. Preserve top-level imports + solve() when present.

    Fallbacks:
    - If only a solve signature is present, extract its indented body and attach to provided signature.
    - If body-only code is provided, wrap it under the provided signature.
    """
    s = code_str.lstrip("\ufeff")
    # Try to preserve imports + solve block using AST
    try:
        tree = ast.parse(s)
        lines = s.splitlines()
        # collect import spans
        import_spans: List[Tuple[int, int]] = []
        for node in getattr(tree, "body", []):
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                a = max(1, getattr(node, "lineno", 1))
                b = max(a, getattr(node, "end_lineno", a))
                import_spans.append((a, b))


def _worker_run_one(code_str: str, sig: str, args: list) -> Tuple[bool, Any]:
    """Top-level worker: exec the candidate in a fresh sandbox and run one test.
    Returns (ok, value_or_error_string)."""
    # Use a single shared namespace so imports are visible to solve() globals
    env: Dict[str, Any] = _safe_builtins()
    # Build full source robustly to avoid nested def solve
    try:
        src = _compose_src(code_str, sig)
    except Exception:
        src = (code_str if re.match(r"^def\s+solve\s*\(", code_str)
               else sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in code_str.splitlines()))
    try:
        exec(src, env, env)
        solve = env.get("solve")
        if not callable(solve):
            return (False, "no_solve")
        return (True, solve(*args))
    except Exception as e:
        return (False, str(e))

def _compose_src(code_str: str, sig: str) -> str:
    """Compose executable source. Preserve top-level imports + solve() when present.

    Fallbacks:
    - If only a solve signature is present, extract its indented body and attach to provided signature.
    - If body-only code is provided, wrap it under the provided signature.
    """
    s = code_str.lstrip("\ufeff")
    # Try to preserve imports + solve block using AST
    try:
        tree = ast.parse(s)
        lines = s.splitlines()
        # collect import spans
        import_spans: List[Tuple[int, int]] = []
        for node in getattr(tree, "body", []):
            if isinstance(node, (ast.Import, ast.ImportFrom)):
                a = max(1, getattr(node, "lineno", 1))
                b = max(a, getattr(node, "end_lineno", a))
                import_spans.append((a, b))
        # find solve()
        solve_node = None
        for node in ast.walk(tree):
            if isinstance(node, ast.FunctionDef) and getattr(node, "name", "") == "solve":
                solve_node = node
                break
        if solve_node is not None:
            s_a = max(1, getattr(solve_node, "lineno", 1))
            s_b = max(s_a, getattr(solve_node, "end_lineno", s_a))
            chunks: List[str] = []
            seen: set[str] = set()
            for a, b in import_spans:
                block = "\n".join(lines[a - 1:b]).strip()
                if block and block not in seen:
                    seen.add(block)
                    chunks.append(block)
            solve_block = "\n".join(lines[s_a - 1:s_b])
            if chunks:
                chunks.append("")
            chunks.append(solve_block)
            return "\n".join(chunks)
    except Exception:
        pass
    # If text already starts with a full def solve, accept it as-is
    s_strip = s.lstrip()
    if s_strip.startswith("def solve"):
        return s_strip
    # If a solve signature exists somewhere, extract its indented body and attach to provided signature
    m = re.search(r"^\s*def\s+solve\s*\(.*\):\s*$", s, flags=re.MULTILINE)
    if m:
        lines = s.splitlines()
        idx = 0
        for i, ln in enumerate(lines):
            if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
                idx = i
                break
        body_lines: List[str] = []
        for ln in lines[idx + 1:]:
            if not ln.strip():
                body_lines.append(ln)
                continue
            if re.match(r"^\S", ln):
                break
            body_lines.append(ln)
        return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in body_lines)
    # Otherwise treat input as body-only and wrap
    return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in s.splitlines())

def _worker_run_one_subproc(code_str: str, sig: str, args: list, timeout_sec: float) -> Tuple[bool, Any]:
    """Execute one test in an isolated tiny subprocess to avoid heavy imports/pickling."""
    try:
        import base64 as _b64  # type: ignore
        import tempfile as _tmp  # type: ignore
        # Pre-compose full source in parent to avoid quoting issues in child
        try:
            src_full = _compose_src(code_str, sig)
        except Exception:
            src_full = code_str
        payload_env = {
            "SRC": _b64.b64encode(src_full.encode("utf-8")).decode("ascii"),
            "A": _b64.b64encode(json.dumps(args).encode("utf-8")).decode("ascii"),
        }
        script_content = (
            "import os, json, base64, builtins\n"
            "src = base64.b64decode(os.environ['SRC']).decode('utf-8')\n"
            "args = json.loads(base64.b64decode(os.environ['A']).decode('utf-8'))\n"


    if m:
        lines = s.splitlines()
        idx = 0
        for i, ln in enumerate(lines):
            if re.match(r"^\s*def\s+solve\s*\(.*\):\s*$", ln):
                idx = i
                break
        body_lines: List[str] = []
        for ln in lines[idx + 1:]:
            if not ln.strip():
                body_lines.append(ln)
                continue
            if re.match(r"^\S", ln):
                break
            body_lines.append(ln)
        return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in body_lines)
    # Otherwise treat input as body-only and wrap
    return sig + "\n" + "\n".join((("    " + ln) if ln.strip() else ln) for ln in s.splitlines())

def _worker_run_one_subproc(code_str: str, sig: str, args: list, timeout_sec: float) -> Tuple[bool, Any]:
    """Execute one test in an isolated tiny subprocess to avoid heavy imports/pickling."""
    try:
        import base64 as _b64  # type: ignore
        import tempfile as _tmp  # type: ignore
        # Pre-compose full source in parent to avoid quoting issues in child
        try:
            src_full = _compose_src(code_str, sig)
        except Exception:
            src_full = code_str
        payload_env = {
            "SRC": _b64.b64encode(src_full.encode("utf-8")).decode("ascii"),
            "A": _b64.b64encode(json.dumps(args).encode("utf-8")).decode("ascii"),
        }
        script_content = (
            "import os, json, base64, builtins\n"
            "src = base64.b64decode(os.environ['SRC']).decode('utf-8')\n"
            "args = json.loads(base64.b64decode(os.environ['A']).decode('utf-8'))\n"
            # Build a minimal, explicit builtins dict and expose it
            "allowed = {}\n"
            "for k in ['abs','all','any','enumerate','len','range','min','max','sum','map','filter','zip','sorted','reversed','list','tuple','dict','set','bool','int','float','str','print','__import__']:\n"
            "    try: allowed[k] = getattr(builtins, k)\n"
            "    except Exception: pass\n"
            "env = {'__builtins__': allowed}\n"
            "env.update(allowed)\n"
            "try:\n"
            "    exec(src, env, env)\n"
            "    solve = env.get('solve')\n"
            "    if not callable(solve):\n"
            "        print(json.dumps({'ok': False, 'err': 'no_solve'})); raise SystemExit(0)\n"
            "    out = solve(*args)\n"
            "    print(json.dumps({'ok': True, 'out': out}))\n"
            "except Exception as e:\n"
            "    print(json.dumps({'ok': False, 'err': str(e)}))\n"
        )
        tmp = None
        try:
            tmp = _tmp.NamedTemporaryFile(mode="w", suffix=".py", delete=False, encoding="utf-8")
            tmp.write(script_content)
            tmp.flush()
            tmp_path = tmp.name
        finally:
            try:
                if tmp is not None:
                    tmp.close()
            except Exception:
                pass
        try:
            p = subprocess.run([sys.executable, tmp_path], env={**os.environ, **payload_env}, capture_output=True, text=True, timeout=float(timeout_sec))
        finally:
            try:
                os.unlink(tmp_path)
            except Exception:
                pass
        txt = (p.stdout or "").strip()
        if not txt:
            err = f"no_output rc={p.returncode} stderr={(p.stderr or '').strip()}"
            return (False, err)
        try:
            obj = json.loads(txt)
        except Exception:
            err = f"bad_json rc={p.returncode} stderr={(p.stderr or '').strip()} out={txt}"
            return (False, err)
        if bool(obj.get("ok")):
            return True, obj.get("out")
        return (False, str(obj.get('err')))
    except subprocess.TimeoutExpired:
        return False, "timeout"
    except Exception as e:
        return False, str(e)

def _values_equal(a: Any, b: Any) -> bool:
    # Numeric tolerance for floats
    try:
        import numbers as _nums  # type: ignore
        if isinstance(a, _nums.Number) and isinstance(b, _nums.Number):
            # Use isclose for floats; ints will compare exactly with isclose as well
            return math.isclose(float(a), float(b), rel_tol=1e-9, abs_tol=1e-9)
    except Exception:
        pass
    # Sequences
    if isinstance(a, (list, tuple)) and isinstance(b, (list, tuple)):
        if len(a) != len(b):
            return False
        return all(_values_equal(x, y) for x, y in zip(a, b))
    # Dicts
    if isinstance(a, dict) and isinstance(b, dict):
        if set(a.keys()) != set(b.keys()):
            return False
        return all(_values_equal(a[k], b[k]) for k in a.keys())