scripts/seed_pools.py

"""Check which challenge tables have what counts and warn if any are empty.

Usage:
    python scripts/seed_pools.py --check   # just show current state
    python scripts/seed_pools.py           # would seed (placeholder for now)

The actual per-type seed generation happens inside the per-type
generators in app/generators/*.py — they run automatically in the
background once main.py starts. This script is for manual ops:

    * ``--check``  → see how full each pool is
    * ``--watch``  → tail the pool size until each table reaches POOL_TARGET
"""

import argparse
import os
import sys
import time
from pathlib import Path

try:
    from dotenv import load_dotenv
except ImportError:
    sys.exit("Missing python-dotenv. Run: pip install python-dotenv")

import httpx


ROOT = Path(__file__).resolve().parent.parent
ENV_PATH = ROOT / ".env"


POOL_TARGETS = {
    "encryption_challenges": 5,            # crypto (red)
    "code_fixing_challenges": 5,            # code-fixing (blue)
    "log_analysis_challenges": 5,           # log-analysis (blue)
    "vulnerability_hunter_challenges": 15,  # 5 difficulties × 3 vuln types (blue)
}


def _load_env() -> None:
    if ENV_PATH.exists():
        load_dotenv(ENV_PATH)


def _client() -> httpx.Client:
    base = os.environ.get("SUPABASE_URL", "").rstrip("/")
    key = os.environ.get("SUPABASE_ANON_KEY", "")
    return httpx.Client(
        timeout=15,
        base_url=base,
        headers={
            "apikey": key,
            "Authorization": f"Bearer {key}",
            "Content-Type": "application/json",
        },
    )


def _count(client: httpx.Client, table: str, team: str | None = None) -> int:
    q = f"select=id"
    if team:
        q += f"&team_role=eq.{team}"
    r = client.get(f"/rest/v1/{table}?{q}")
    if r.status_code != 200:
        return -1
    return len(r.json())


def _print_table(client: httpx.Client) -> None:
    print(f"{'Table':40s} {'Target':>8s} {'Red':>6s} {'Blue':>6s}  Status")
    for table, target in POOL_TARGETS.items():
        red = _count(client, table, "red")
        blue = _count(client, table, "blue")
        ok = "✓" if (red + blue) >= target else "·"
        print(f"{table:40s} {target:>8d} {red:>6d} {blue:>6d}  {ok}")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--check", action="store_true", help="print current pool sizes")
    parser.add_argument("--watch", action="store_true", help="poll every 10s")
    args = parser.parse_args()

    _load_env()
    if not os.environ.get("SUPABASE_URL"):
        sys.exit("SUPABASE_URL is not set.")

    with _client() as c:
        if args.watch:
            while True:
                _print_table(c)
                time.sleep(10)
        else:
            _print_table(c)

    return 0


if __name__ == "__main__":
    sys.exit(main())