#!/usr/bin/env python
# infer_simple.py — the minimal causal inference loop, for reading and reference.
#
# Pick a policy (api or rl), stream each sample's frames one at a time in time order, and write
# one JSONL line per sample. A failing sample is skipped (empty record, run continues) so the
# output always has every id. The production runner (run_inference.py) adds the rest for a real
# paid 500-sample run — 3x retry, --resume, parallel workers, client warm-up — which this file
# deliberately omits.
import argparse
import glob
import json
import os

from humomni.core.streaming_driver import _emissions   # SAME normalization as run_inference -> identical output


def make_policy(name):
    """A fresh policy per sample (no state leak). Imported lazily — each pulls heavy deps."""
    if name == "api":
        from humomni.phase1.policy_api import APIPolicy
        config = json.load(open("config.json", encoding="utf-8"))
        return lambda: APIPolicy(config=config)
    raise SystemExit(f"unknown policy: {name!r} (use 'api')")


def _video_length(sample_dir):
    """The latest frame timestamp present — the 'video_length' the loop steps up to (frames sit on a
    0.5 s grid: 0.5.jpg, 1.0.jpg, …)."""
    ts = []
    for p in glob.glob(os.path.join(sample_dir, "*.jpg")):
        try:
            ts.append(float(os.path.basename(p)[:-4]))
        except ValueError:
            pass                                             # skip non-frame files
    return max(ts) if ts else 0.0


def infer_one(sample_dir, policy):
    """One sample, matching the competition inference pseudocode:

        input(question)
        for current_time in range(0.5, video_length, 0.5):
            input(f"{current_time}.png")          # frames at t > current_time are PROHIBITED
            if respond_now: output(... @ current_time)

    The policy gets one frame per tick, in ascending time, and never sees t > current_time.
    policy.step() may return None / a str / a list of strs — one frame can emit several answers
    (the v5+v3 ensemble); we append each AT current_time, exactly like streaming_driver.drive().
    A transient failure returns an empty-but-present record (the run never aborts and every id
    stays); a causality AssertionError is not caught. No retry/resume/parallelism — that's
    run_inference.py.
    """
    with open(os.path.join(sample_dir, "question.json"), encoding="utf-8") as f:
        q = json.load(f)
    question, qid = q["question"], q["question_id"]
    try:
        policy.reset(question=question)
        responses, last_t = [], -1.0
        video_length = _video_length(sample_dir)
        for step in range(1, int(round(video_length / 0.5)) + 1):       # current_time = 0.5 .. video_length
            current_time = step * 0.5
            frame_path = os.path.join(sample_dir, f"{current_time:.1f}.jpg")
            if not os.path.exists(frame_path):                          # missing tick on the grid — skip
                continue
            assert current_time > last_t, "frames must arrive in order (no future frames)"
            last_t = current_time
            reply = policy.step(current_time, frame_path)
            for content in _emissions(reply):                           # 0..N answers, all at current_time
                responses.append({"time": current_time, "content": content})
        return {"question_id": qid, "model_response_list": responses}
    except AssertionError:
        raise
    except Exception as e:
        print(f"  ! {qid} failed ({type(e).__name__}: {str(e)[:80]}) — empty record")
        return {"question_id": qid, "model_response_list": []}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--policy", choices=["api"], default="api")
    ap.add_argument("--data_dir", default="data/phase1/data")
    ap.add_argument("--out", default="submission.jsonl")
    a = ap.parse_args()

    new_policy = make_policy(a.policy)
    dirs = [d for d in sorted(glob.glob(os.path.join(a.data_dir, "*"))) if os.path.isdir(d)]
    total = len(dirs)
    with open(a.out, "w", encoding="utf-8") as w:
        for i, sample_dir in enumerate(dirs, 1):
            rec = infer_one(sample_dir, new_policy())    # fresh policy per sample
            w.write(json.dumps(rec, ensure_ascii=False) + "\n")
            print(f"[{i}/{total}] {rec['question_id']} -> {len(rec['model_response_list'])} responses",
                  flush=True)                            # flush so progress shows live
    print(f"done {total} samples -> {a.out}")


if __name__ == "__main__":
    main()