# run_inference.py — drive every sample through a policy -> one JSONL line per sample.
# Works for any policy (SilentPolicy, APIPolicy, RLPolicy later). EVERY question_id must
# appear, even with an empty model_response_list. Robust for paid runs:
#   - per-sample retry (transient Windows Errno 22 / API blips) with empty-record fallback,
#     so ONE bad sample never aborts the whole run;
#   - incremental flushed append, so a crash never loses (or re-pays for) finished samples;
#   - --resume skips question_ids already present in the output.
import glob
import json
import os
import time
import argparse
import threading
from concurrent.futures import ThreadPoolExecutor

from humomni.core.streaming_driver import run_sample


def make_policy_factory(name, config):
    """Return a zero-arg factory that builds a FRESH policy per sample (no state leak)."""
    if name == "silent":
        from humomni.phase1.policies import SilentPolicy
        return lambda: SilentPolicy()
    if name == "api":
        from humomni.phase1.policy_api import APIPolicy           # imported lazily (needs deps + API key)
        return lambda: APIPolicy(config=config)
    raise ValueError(f"unknown policy {name!r}")


def _qid(d):
    with open(os.path.join(d, "question.json"), encoding="utf-8") as f:
        return json.load(f)["question_id"]


def _done_qids(out):
    done = set()
    if os.path.exists(out):
        with open(out, encoding="utf-8") as f:
            for ln in f:
                ln = ln.strip()
                if ln:
                    try:
                        done.add(json.loads(ln)["question_id"])
                    except Exception:
                        pass
    return done


def main(data_dir, out, policy_factory, limit=0, workers=1, resume=False):
    dirs = sorted(d for d in glob.glob(os.path.join(data_dir, "*")) if os.path.isdir(d))
    if limit:
        dirs = dirs[:limit]                              # used by experiment.py for fast partial Phase-2 evals
    done = _done_qids(out) if resume else set()
    if not resume:
        open(out, "w", encoding="utf-8").close()        # truncate (fresh run)
    todo = [d for d in dirs if _qid(d) not in done]
    print(f"{len(dirs)} samples | {len(done)} already done | {len(todo)} to run", flush=True)

    lock = threading.Lock()
    fails = []
    wfile = open(out, "a", encoding="utf-8")

    def process(d):
        for attempt in range(3):
            try:
                return run_sample(d, policy_factory())   # fresh policy state per sample
            except Exception as e:
                if attempt == 2:
                    fails.append((_qid(d), str(e)[:120]))
                    return {"question_id": _qid(d), "model_response_list": []}  # keep id present
                time.sleep(0.5)

    def handle(d):
        rec = process(d)
        with lock:
            wfile.write(json.dumps(rec, ensure_ascii=False) + "\n")
            wfile.flush()
        return rec

    try:
        if workers > 1:
            with ThreadPoolExecutor(max_workers=workers) as ex:
                for n, _ in enumerate(ex.map(handle, todo), 1):
                    if n % 50 == 0:
                        print(f"  ...{n}/{len(todo)}", flush=True)
        else:
            for n, d in enumerate(todo, 1):
                handle(d)
                if n % 50 == 0:
                    print(f"  ...{n}/{len(todo)}", flush=True)
    finally:
        wfile.close()
    print(f"wrote -> {out} (failed samples: {len(fails)})")
    for q, e in fails[:10]:
        print("  FAIL", q, e)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--data_dir", default="data/phase1/data")
    ap.add_argument("--out", default="submission.jsonl")
    ap.add_argument("--policy", default="api", choices=["silent", "api"])
    ap.add_argument("--config", default="config.json", help="orchestrator config (gate/theta/dedup)")
    ap.add_argument("--limit", type=int, default=0, help="process only the first N samples (0=all; experiment.py uses it for fast Phase-2 evals)")
    ap.add_argument("--workers", type=int, default=1, help="parallel samples (independent, still causal)")
    ap.add_argument("--resume", action="store_true", help="skip question_ids already in --out")
    a = ap.parse_args()
    cfg = None
    if a.policy == "api" and os.path.exists(a.config):
        cfg = json.load(open(a.config, encoding="utf-8"))

    if a.policy == "api" and a.workers > 1:
        # Warm BOTH API clients + their lazy pydantic imports in the MAIN thread first;
        # otherwise concurrent first-use across worker threads races and throws
        # "No module named 'pydantic._migration'".
        import humomni.phase1.vlm_client as vlm_client
        try:
            vlm_client._client_().chat.completions.create(
                model=vlm_client.MODEL, max_tokens=1,
                messages=[{"role": "user", "content": "ping"}])
        except Exception:
            pass
        if (cfg or {}).get("dedup", {}).get("use_llm"):
            import humomni.core.judges as judges
            try:
                judges.gemini("ping")
            except Exception:
                pass
        try:
            import humomni.core.emb_cache as emb_cache          # warm MiniLM (cosine dedup prefilter) single-threaded
            emb_cache.embed("ping")
        except Exception:
            pass

    main(a.data_dir, a.out, make_policy_factory(a.policy, cfg),
         limit=a.limit, workers=a.workers, resume=a.resume)