"""A GPU-free stand-in for viggle-remake-worker that speaks the same queue protocol (viggle_remake_mcp/jobs.py), for
the end-to-end tests. It heartbeats every second, claims $REMAKE_HOME/queue/*.json in sorted-name order, holds a hop
until its prefix exists, checks every job field, and "renders" by desaturating the cond (same size, num_frames
frames). Every claimed job is appended to <home>/fake_<name>.jsonl so a test can read what it got.

    python fake_worker.py --home H --name A [--max480 345] [--max736 345] [--modes swap,edit] [--delay 0]
"""
import argparse
import json
import os
import subprocess
import threading
import time
from pathlib import Path

ap = argparse.ArgumentParser()
ap.add_argument("--home", required=True)
ap.add_argument("--name", default="fake0")
ap.add_argument("--max480", type=int, default=345)
ap.add_argument("--max736", type=int, default=345)
ap.add_argument("--modes", default="swap,edit")
ap.add_argument("--delay", type=float, default=0.0, help="extra seconds per job")
a = ap.parse_args()

home = Path(a.home)
queue, workers = home / "queue", home / "workers"
queue.mkdir(parents=True, exist_ok=True)
workers.mkdir(parents=True, exist_ok=True)
limit = {480: a.max480, 736: a.max736}
state, lock = {"job": None}, threading.Lock()
LENS = {17 * k + 5 for k in range(5, 21)}
KEYS = {"mode", "cond", "audio", "refs", "prompt", "num_frames", "seed", "out", "pin", "tier"}


def beat():
    hb = {"t": time.time(), "host": "fake", "gpu": a.name, "pid": os.getpid(), "job": state["job"],
          "modes": a.modes.split(","), "max_frames": {"480": a.max480, "736": a.max736}}
    with lock:  # the beat thread and the job loop both beat: one shared tmp name
        tmp = workers / f".fake_{a.name}.tmp"
        tmp.write_text(json.dumps(hb))
        tmp.rename(workers / f"fake_{a.name}.json")


def beats():
    while True:
        beat()
        time.sleep(1)


def probe(p):
    """-> (frames, width, height)"""
    r = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0", "-count_packets", "-show_entries",
                        "stream=nb_read_packets,width,height", "-of", "json", str(p)], capture_output=True, text=True, check=True)
    s = json.loads(r.stdout)["streams"][0]
    return int(s["nb_read_packets"]), s["width"], s["height"]


def check(job):
    miss = KEYS - set(job)
    assert not miss, f"missing fields {sorted(miss)}"
    assert job["mode"] in a.modes.split(","), f"mode {job['mode']} not served here"
    nf, tier = job["num_frames"], job["tier"]
    assert nf in LENS, f"num_frames {nf} is not 17k+5 in 90..345"
    assert nf <= limit.get(tier, 0), f"num_frames {nf} > this card's max {limit.get(tier, 0)} at tier {tier}"
    n, w, h = probe(job["cond"])
    assert n == nf, f"cond has {n} frames, num_frames {nf}"
    assert w % 32 == 0 and h % 32 == 0 and w * h <= {480: 832 * 480, 736: 1280 * 736}[tier], f"cond {w}x{h} at tier {tier}"
    assert Path(job["audio"]).is_file(), "no audio"
    assert len(job["refs"]) == len(job["pin"]) and all(Path(r).is_file() for r in job["refs"]), "refs/pin"
    assert all(0 <= p < nf for p in job["pin"]), f"pin {job['pin']} outside [0, {nf})"
    assert job["prompt"].startswith("subject_definitions:"), "prompt does not start with subject_definitions:"
    assert job["refs"] or job.get("prefix"), "no ref and no prefix"
    if job.get("prefix"):
        pn = probe(job["prefix"])[0]
        assert 0 <= job["prefix_at"] and job["prefix_at"] + 17 <= pn, f"prefix_at {job['prefix_at']} of a {pn}-frame prefix"
    return w, h


def render(job):
    out = Path(job["out"])
    tmp = out.parent / f".{out.name}"  # hidden while it is written: kit.takes globs w<i>_*
    subprocess.run(["ffmpeg", "-v", "error", "-y", "-i", job["cond"], "-vf", "hue=s=0", "-frames:v", str(job["num_frames"]),
                    "-an", "-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p", str(tmp)], check=True)
    tmp.rename(job["out"])


threading.Thread(target=beats, daemon=True).start()
log = home / f"fake_{a.name}.jsonl"
while True:
    run = None
    for p in sorted(queue.glob("*.json")):
        try:
            job = json.loads(p.read_text())
        except (OSError, ValueError):
            continue
        if job.get("prefix") and not Path(job["prefix"]).exists():
            continue  # a hop waits for the take it continues
        try:
            os.utime(p)
            p.rename(p.with_suffix(".run"))
        except FileNotFoundError:
            continue
        run = p.with_suffix(".run")
        break
    if run is None:
        time.sleep(0.5)
        continue
    state["job"] = run.stem
    beat()
    t = time.time()
    with open(log, "a") as fh:
        fh.write(json.dumps({"name": run.stem, "job": job}) + "\n")
    try:
        check(job)
        time.sleep(a.delay)
        render(job)
        tmp = run.with_suffix(".tmp")
        tmp.write_text(json.dumps({"out": job["out"], "secs": round(time.time() - t, 1), "host": "fake", "gpu": a.name,
                                   "mode": job["mode"]}))
        tmp.rename(run.with_suffix(".done"))
    except Exception as e:
        run.with_suffix(".err").write_text(f"{type(e).__name__}: {e}")
    run.unlink(missing_ok=True)
    state["job"] = None