#!/usr/bin/env python3
"""Flow Auto-Pilot local queue.

Source of work:  queue.jsonl   — one task per line: {"prompt": "...", "kind": "video"|"image", "model": "Veo 3.1 - Fast"}
                                  (id optional; auto-assigned by line number)
Output:          results.jsonl — one record per finished task: {id, status, assets:[urls], error, ts}

Behaviour:
  * /next-task  -> hands out the next PENDING task, video-first (rinse the pricey credits first).
  * /task-completed -> appends to results.jsonl, marks done.
  * On startup, tasks already in results.jsonl are skipped -> safe to restart during a month-long run.

Harvested `assets` are googleusercontent URLs; they need your logged-in session to download,
so this just records them. ponytail: fetch+save the blobs from the extension later if you want local files.
"""
import json, os, time
from http.server import HTTPServer, BaseHTTPRequestHandler

HERE = os.path.dirname(os.path.abspath(__file__))
QUEUE_FILE = os.path.join(HERE, "queue.jsonl")
RESULTS_FILE = os.path.join(HERE, "results.jsonl")

# Fallback demo queue if queue.jsonl is absent (so it runs out of the box).
DEMO = [
    {"prompt": "Cinematic synthwave keyboard warrior fighting digital demons, 8k", "kind": "video", "model": "Veo 3.1 - Fast"},
    {"prompt": "Op-shop record crate, warm 90s film grain, product hero shot", "kind": "image", "model": "Nano Banana 2"},
]

def load_queue():
    rows = []
    if os.path.exists(QUEUE_FILE):
        with open(QUEUE_FILE) as f:
            for line in f:
                line = line.strip()
                if line:
                    rows.append(json.loads(line))
    else:
        rows = list(DEMO)
    # assign stable ids by position; video first so the expensive credits burn first
    for i, r in enumerate(rows):
        r.setdefault("id", f"task_{i:05d}")
        r.setdefault("kind", "image")
    rows.sort(key=lambda r: 0 if r.get("kind") == "video" else 1)
    return rows

def done_ids():
    ids = set()
    if os.path.exists(RESULTS_FILE):
        with open(RESULTS_FILE) as f:
            for line in f:
                line = line.strip()
                if line:
                    ids.add(json.loads(line).get("id"))
    return ids

class Handler(BaseHTTPRequestHandler):
    def _send(self, code=200, body=None):
        self.send_response(code)
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")
        self.send_header("Content-Type", "application/json")
        self.end_headers()
        if body is not None:
            self.wfile.write(json.dumps(body).encode())

    def do_OPTIONS(self):
        self._send(200, {})

    def do_GET(self):
        if self.path != "/next-task":
            return self._send(404, {"error": "not found"})
        done = done_ids()
        for t in load_queue():
            if t["id"] not in done and t["id"] not in INFLIGHT:
                INFLIGHT[t["id"]] = time.time()
                remaining = sum(1 for x in load_queue() if x["id"] not in done)
                print(f"[SERVER] -> {t['id']} ({t.get('kind')}) | {remaining} left")
                return self._send(200, t)
        return self._send(204)  # queue drained

    def do_POST(self):
        if self.path != "/task-completed":
            return self._send(404, {"error": "not found"})
        n = int(self.headers.get("Content-Length", 0))
        rec = json.loads(self.rfile.read(n).decode() or "{}")
        rec["ts"] = time.time()
        INFLIGHT.pop(rec.get("id"), None)
        with open(RESULTS_FILE, "a") as f:
            f.write(json.dumps(rec) + "\n")
        assets = rec.get("assets") or []
        print(f"[SERVER] <- {rec.get('id')} {str(rec.get('status')).upper()} | {len(assets)} asset(s)")
        self._send(200, {"ok": True})

    def log_message(self, *a):  # quiet the default noise
        pass

INFLIGHT = {}

def selftest():
    # video-first ordering + resume-skip
    global load_queue
    _orig = load_queue
    load_queue = lambda: [  # noqa: E731
        {"id": "a", "kind": "image"}, {"id": "b", "kind": "video"}]
    assert [t["id"] for t in load_queue()] == ["a", "b"]  # loader keeps as-is here
    q = sorted(load_queue(), key=lambda r: 0 if r["kind"] == "video" else 1)
    assert q[0]["id"] == "b", "video must be handed out first"
    load_queue = _orig
    print("selftest ok")

def run(port=8017):
    print(f"--- Flow queue on http://localhost:{port} ---")
    print(f"queue:   {QUEUE_FILE if os.path.exists(QUEUE_FILE) else '(demo — no queue.jsonl)'}")
    print(f"pending: {sum(1 for t in load_queue() if t['id'] not in done_ids())}")
    HTTPServer(("", port), Handler).serve_forever()

if __name__ == "__main__":
    import sys
    if "--selftest" in sys.argv:
        selftest()
    else:
        try:
            run()
        except KeyboardInterrupt:
            print("\nstopped")