magicciv/tools/sprite-generation/orchestrate_starter.py

#!/usr/bin/env python3
"""Orchestrate + monitor quality for the Game-1 starter sprite set.

Registers the curated manifest, generates variants, scores through the ranker
pipeline, and prints a live quality funnel until sprites reach review or the
attempt cap is hit.

Usage:
    python3 orchestrate_starter.py register
    python3 orchestrate_starter.py run --backend grok --variants 2
    python3 orchestrate_starter.py status
    python3 orchestrate_starter.py monitor
"""
from __future__ import annotations

import argparse
import asyncio
import json
import sys
import time
from collections import defaultdict
from datetime import datetime
from pathlib import Path

TOOL_DIR = Path(__file__).resolve().parent
sys.path.insert(0, str(TOOL_DIR))

from engine.factory import create_generator, with_backend
from engine.registry import SpriteRegistry
from engine.starter import load_manifest, register_starter_set, starter_sprite_ids

DB_PATH = TOOL_DIR / "spritegen.db"
RAW_DIR = TOOL_DIR / "raw"
CONFIG_PATH = TOOL_DIR / "sprite-config.json"


def _registry() -> SpriteRegistry:
    return SpriteRegistry(DB_PATH)


def _load_config(backend: str | None) -> dict:
    config = json.loads(CONFIG_PATH.read_text())
    if backend:
        return with_backend(config, backend)
    return config


def cmd_register(_args: argparse.Namespace) -> None:
    reg = _registry()
    report = register_starter_set(reg, reset_status=True)
    print(f"Starter set: {len(report.registered)} sprites registered/reset")
    if report.missing_data:
        print("Missing data:", ", ".join(report.missing_data))
    for sid in report.registered:
        print(f"  {sid}")


def _quality_snapshot(reg: SpriteRegistry, sprite_ids: list[str]) -> dict:
    placeholders = ",".join("?" * len(sprite_ids))
    rows = reg.conn.execute(
        f"""
        SELECT s.id, s.status,
               COUNT(v.id) AS variants,
               SUM(CASE WHEN v.job_status='completed' THEN 1 ELSE 0 END) AS completed,
               SUM(CASE WHEN COALESCE(v.rating, -1) > 0 THEN 1 ELSE 0 END) AS good,
               MAX(COALESCE(v.review_tier, 0)) AS max_tier,
               MAX(COALESCE(v.rating, 0)) AS best_rating
        FROM sprites s
        LEFT JOIN variants v ON v.sprite_id = s.id
        WHERE s.id IN ({placeholders})
        GROUP BY s.id, s.status
        ORDER BY s.id
        """,
        sprite_ids,
    ).fetchall()
    return [dict(r) for r in rows]


def _print_funnel(snapshot: list[dict]) -> None:
    by_status: dict[str, int] = defaultdict(int)
    for row in snapshot:
        by_status[row["status"]] += 1

    print(f"\n{'Sprite':<32} {'status':<8} {'var':>4} {'good':>4} {'tier':>4} {'best':>4}")
    print("-" * 60)
    for row in snapshot:
        print(
            f"{row['id']:<32} {row['status']:<8} "
            f"{row['variants']:>4} {row['good']:>4} "
            f"{row['max_tier']:>4} {row['best_rating']:>4}"
        )
    print("-" * 60)
    print("Status totals:", dict(by_status))


def _top_failing_gates(reg: SpriteRegistry, sprite_ids: list[str], limit: int = 8) -> None:
    placeholders = ",".join("?" * len(sprite_ids))
    rows = reg.conn.execute(
        f"""
        SELECT v.sprite_id, v.notes
        FROM variants v
        WHERE v.sprite_id IN ({placeholders})
          AND v.job_status='completed'
          AND v.notes IS NOT NULL
        ORDER BY v.id DESC
        LIMIT 200
        """,
        sprite_ids,
    ).fetchall()

    fails: dict[str, int] = defaultdict(int)
    for row in rows:
        try:
            notes = json.loads(row["notes"])
        except json.JSONDecodeError:
            continue
        gates = notes.get("gates") or {}
        for gate, passed in gates.items():
            if passed is False:
                fails[gate] += 1

    if not fails:
        return
    print("\nTop failing gates (recent variants):")
    for gate, count in sorted(fails.items(), key=lambda x: -x[1])[:limit]:
        print(f"  {gate}: {count}")


def cmd_status(_args: argparse.Namespace) -> None:
    ids = starter_sprite_ids()
    reg = _registry()
    snap = _quality_snapshot(reg, ids)
    print(f"Starter set ({len(ids)} sprites)")
    _print_funnel(snap)
    _top_failing_gates(reg, ids)


def cmd_monitor(_args: argparse.Namespace) -> None:
    ids = starter_sprite_ids()
    reg = _registry()
    print("Monitoring starter set quality (Ctrl+C to stop)...")
    try:
        while True:
            snap = _quality_snapshot(reg, ids)
            ready = sum(1 for r in snap if r["status"] == "review" and r["good"] >= 1)
            needed = sum(1 for r in snap if r["status"] == "needed")
            print(f"\n[{datetime.now():%H:%M:%S}] review-ready={ready} needed={needed}")
            _print_funnel(snap)
            _top_failing_gates(reg, ids)
            sys.stdout.flush()
            time.sleep(30)
    except KeyboardInterrupt:
        print("\nMonitor stopped.")


async def _run_loop(args: argparse.Namespace) -> None:
    from engine.ranker import SpriteRanker

    config = _load_config(getattr(args, "backend", None))
    reg = _registry()
    report = register_starter_set(reg, reset_status=args.reset)
    sprite_ids = starter_sprite_ids()
    print(f"Starter orchestration — {len(sprite_ids)} sprites ({config['backend']})")
    if report.missing_data:
        print("WARN missing data:", report.missing_data)

    gen = create_generator(config=config, registry=reg, raw_dir=RAW_DIR)
    ranker = SpriteRanker(registry=reg, raw_dir=RAW_DIR)

    regen_counts: dict[str, int] = {}
    max_attempts = args.max_attempts

    async def _on_complete(variant_id: int, sprite_id: str) -> None:
        await ranker.advance_sprite(sprite_id)

    for loop in range(1, max_attempts + 1):
        needed_rows = reg.conn.execute(
            f"""
            SELECT s.id FROM sprites s
            LEFT JOIN variants v ON v.sprite_id = s.id
            WHERE s.id IN ({",".join("?" * len(sprite_ids))})
              AND s.status = 'needed'
            GROUP BY s.id
            ORDER BY COUNT(v.id) ASC, s.id ASC
            """,
            sprite_ids,
        ).fetchall()
        needed = [r[0] for r in needed_rows]
        if not needed:
            print(f"\n[loop {loop}] All starter sprites moved past 'needed'.")
            break

        batch = needed[: args.batch_size]
        print(f"\n[loop {loop}] Generating {len(batch)} sprites x {args.variants} variants...")
        # Grok completes synchronously; score once per sprite after the batch.
        on_complete = _on_complete if config["backend"] == "model-boss" else None
        submitted = await gen.submit_batch(
            sprite_ids=batch,
            variants_per=args.variants,
            priority="high",
            on_complete=on_complete,
        )
        if config["backend"] == "model-boss" and submitted:
            await gen.collect_pending(on_complete=_on_complete)

        for sid in batch:
            result = await ranker.rank_and_filter(sid)
            good = result["good_count"]
            total = len(result["ranked"])
            if result["needs_regen"]:
                regen_counts[sid] = regen_counts.get(sid, 0) + 1
                if regen_counts[sid] < max_attempts:
                    reg.update_sprite_status(sid, "needed")
                    print(f"  {sid}: {good}/{total} good — re-queue ({regen_counts[sid]}/{max_attempts})")
                else:
                    reg.update_sprite_status(sid, "review")
                    print(f"  {sid}: cap hit — moving to review")
            else:
                reg.update_sprite_status(sid, "review")
                print(f"  {sid}: {good}/{total} good — review")

        snap = _quality_snapshot(reg, sprite_ids)
        _print_funnel(snap)
        _top_failing_gates(reg, sprite_ids)

        if args.once:
            break

    print("\nDone. Open Theater for human picks:")
    print(f"  python3 cli.py start --port 5850")
    print(f"  http://localhost:5850/?spriteTheater=true")


def cmd_run(args: argparse.Namespace) -> None:
    asyncio.run(_run_loop(args))


def main() -> None:
    parser = argparse.ArgumentParser(description="Starter sprite orchestration")
    from engine.factory import BACKENDS

    sub = parser.add_subparsers(dest="command", required=True)

    sub.add_parser("register", help="Register starter manifest into spritegen.db").set_defaults(func=cmd_register)
    sub.add_parser("status", help="Quality funnel for starter set").set_defaults(func=cmd_status)
    sub.add_parser("monitor", help="Live quality monitor").set_defaults(func=cmd_monitor)

    p_run = sub.add_parser("run", help="Generate + score starter set")
    p_run.add_argument("--backend", choices=BACKENDS, default="grok")
    p_run.add_argument("--variants", type=int, default=2)
    p_run.add_argument("--batch-size", type=int, default=4, help="Sprites per loop")
    p_run.add_argument("--max-attempts", type=int, default=5)
    p_run.add_argument("--reset", action="store_true", help="Force all starter sprites back to needed")
    p_run.add_argument("--once", action="store_true", help="Single loop then stop")
    p_run.set_defaults(func=cmd_run)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()