"""serve_bon.py — the shippable "smarter version": chat with base + best-of-N + real-exec verifier.

Day-25 proved best-of-N+verifier lifts sealed-holdout solve 11x (greedy 2/60 -> bo32 22/60) with NO weight
changes. This exposes that as a usable thing: an interactive REPL (and a one-shot --prompt mode) that, for a
coding task with stdin/stdout test cases, samples K candidates, runs each against the provided tests, and
returns the FIRST that passes (or the best-partial if none fully pass). For free-form chat (no tests) it
returns the greedy answer. This is base-model + inference search — the "smarter version" used by default.

Usage (pod):
  PYTHONPATH=/workspace/RSI VLLM_USE_FLASHINFER_SAMPLER=0 /venv/main/bin/python3 serve_bon.py            # REPL
  ... serve_bon.py --prompt "problem text" --tests tests.json                                            # one-shot
tests.json = {"inputs": ["..."], "outputs": ["..."]}  (same shape as apps meta). Omit for plain chat.
Env: MODEL, BON_K (default 16), BON_TEMP (0.9), MAXTOK (640), ADAPTER (optional LoRA dir)."""
import os, sys, json, argparse, subprocess, tempfile
sys.path.insert(0,"/workspace/RSI")
from concurrent.futures import ThreadPoolExecutor
from src.utils.external_benchmarks import _extract_code
from src.utils.vllm_backend import VLLMModelLoader

MODEL=os.environ.get("MODEL","/workspace/RSI/expanded_models/qwen3_32b")
BNB={"load_in_4bit":True,"bnb_4bit_compute_dtype":"bfloat16"}
K=int(os.environ.get("BON_K","16")); TEMP=float(os.environ.get("BON_TEMP","0.9"))
MAXTOK=int(os.environ.get("MAXTOK","640")); ADAPTER=os.environ.get("ADAPTER","")
_POOL=ThreadPoolExecutor(max_workers=int(os.environ.get("GRADE_WORKERS","32")))

def _norm(s): return "\n".join(l.rstrip() for l in str(s).strip().splitlines())
def run_stdin(code,inp,t=8):
    d=tempfile.mkdtemp(); p=os.path.join(d,"s.py"); open(p,"w").write(code)
    try:
        try: return subprocess.run([sys.executable,p],input=str(inp),capture_output=True,text=True,timeout=t).stdout
        except Exception: return ""
    finally:
        try: os.remove(p); os.rmdir(d)
        except: pass
def score(code,tests):
    ins=tests.get("inputs") or []; outs=tests.get("outputs") or []
    if not code or not ins: return 0.0
    tc=list(zip(ins,outs)); res=list(_POOL.map(lambda io:_norm(run_stdin(code,io[0]))==_norm(io[1]),tc))
    return sum(res)/max(1,len(res))

def main():
    ap=argparse.ArgumentParser(); ap.add_argument("--prompt"); ap.add_argument("--tests"); ap.add_argument("--k",type=int,default=K)
    a=ap.parse_args()
    loader=VLLMModelLoader(model_path=MODEL,dtype="bfloat16",max_model_len=4096,gpu_memory_utilization=0.85,
        allow_remote_code=True,quantization_config=BNB,max_lora_rank=128,enable_chunked_prefill=False,
        enable_lora=bool(ADAPTER),enforce_eager=True)
    loader.load()
    if ADAPTER and os.path.exists(ADAPTER): loader.set_lora_adapter(ADAPTER); print(f"[serve] adapter {ADAPTER}",flush=True)
    def solve(prompt,tests,k):
        if not tests:  # plain chat -> greedy
            return loader.generate_batch([prompt],max_new_tokens=MAXTOK,temperature=0.0,top_p=1.0)[0], None
        # best-of-K: sample k, return first that fully passes; else best-partial
        cands=list(loader.generate_batch([prompt]*k,max_new_tokens=MAXTOK,temperature=TEMP,top_p=0.95))
        best=None; best_s=-1.0
        for c in cands:
            code=_extract_code(c) or c; s=score(code,tests)
            if s>=0.999: return c, 1.0
            if s>best_s: best,best_s=c,s
        return best, best_s
    if a.prompt is not None:
        tests=json.load(open(a.tests)) if a.tests else None
        ans,s=solve(a.prompt,tests,a.k); print(f"\n[serve] {'SOLVED' if s==1.0 else ('best-partial %.2f'%s if s is not None else 'chat')}\n{ans}")
        return
    print(f"[serve] REPL ready — base+best-of-{K}+verifier. Paste problem; for verified solve add a line '#TESTS <path.json>'. 'exit' quits.")
    while True:
        try: u=input("\nyou> ").strip()
        except (EOFError,KeyboardInterrupt): break
        if u=="exit": break
        if not u: continue
        tests=None
        if "#TESTS" in u:
            parts=u.split("#TESTS"); u=parts[0].strip()
            tp=parts[1].strip()
            if os.path.exists(tp): tests=json.load(open(tp))
        ans,s=solve(u,tests,a.k)
        tag=("SOLVED via best-of-%d"%a.k) if s==1.0 else ("best-partial %.2f"%s if s is not None else "chat")
        print(f"\nmodel[{tag}]> {ans}")

if __name__=="__main__":
    main()