"""Keep the fal FLUX inpaint container WARM so customers never hit a cold start.

fal serverless scales the FLUX container to zero when idle; the next request then
pays ~2-3 min to spin it back up (cold). During salon hours this loop sends a small
inpaint request every `interval` seconds so the container stays hot and real
customer requests run warm (seconds-of-spin, not minutes).

COST WARNING (learned the hard way): a ping is NOT free, and NOT cheap-per-pixel.
fal flux-general bills $0.075 per MEGAPIXEL with an effective ~1 MP minimum, so even
a 256px ping costs ~$0.075. At a 45s interval that is ~800 pings/10h ≈ ~$60/day just
to keep warm — far more than paying ~$0.075 per ACTUAL customer photo. Ping keep-warm
is therefore NOT economical for this endpoint; prefer fal dedicated/keep_alive (if its
fixed price beats your volume) or simply accept the occasional cold start + a queue/UX.
This script is now mainly a diagnostic.

CRITICAL TUNING: the interval must be SHORTER than fal's idle timeout, or the
container scales down between pings and customers (and the pings themselves) still
hit cold starts. We saw this live at 150s: pings alternated ~4s (warm) and
100-170s (cold) — i.e. fal's idle window is under 150s. This script now CLASSIFIES
each ping (warm vs COLD) and warns when it's leaking, so you can dial the interval
down empirically (try 90s, then 60s) until COLD stops appearing.

Gated on FAL_KEY (this shell only; never logged/committed). Each ping bills ~$0.075
(see COST WARNING above) — shorter interval = more pings = much more cost.

Stop with Ctrl+C. Pilot Ready: NOT CONFIRMED.
"""
from __future__ import annotations

import argparse
import sys
import time
from io import BytesIO
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from app.services.fal_client import FalUnavailable, fal_inpaint_model, fal_real_enabled, run_fal_inpaint  # noqa: E402


def classify(ok: bool, duration: float, warm_max: float) -> str:
    """Label a ping: 'warm' (fast = container was hot), 'COLD' (slow = it had to
    spin up, i.e. keep-warm is leaking), or 'FAIL' (errored/timed out)."""
    if not ok:
        return "FAIL"
    return "warm" if duration <= warm_max else "COLD"


def should_warn_leak(n: int, label: str) -> bool:
    """True only when a ping signals a real between-ping leak. Ping #1 is the
    expected startup spin-up (container is down at launch) — never a leak; only
    a non-warm ping from #2 on means the container scaled down between pings."""
    return n >= 2 and label != "warm"


def sleep_seconds(interval: float, elapsed: float) -> float:
    """Fixed-cadence sleep: hold a CONSTANT start-to-start interval. Sleeping a flat
    `interval` AFTER each ping let a slow (cold ~28s) ping push the next gap out
    (60+28=88s), which kept tripping fal's idle window and cascaded more cold starts.
    Subtracting the elapsed ping time keeps the real gap == interval."""
    return max(1.0, interval - elapsed)


def is_fatal_error(msg: str) -> bool:
    """A NON-retryable fal failure (billing/auth) — keep-warm must STOP, not keep
    hammering a locked account every cadence with a misleading 'lower the interval'
    hint. Seen live: HTTP 403 'User is locked. Reason: Exhausted balance.'"""
    m = (msg or "").lower()
    return any(s in m for s in (
        "exhausted balance", "user is locked", "http 403", "http 401",
        "unauthorized", "forbidden", "payment", "billing",
    ))


def summary_line(counts: dict) -> str:
    """Rolling tally string. warm-rate = warm / total pings so far."""
    total = sum(counts.values())
    warm = counts.get("warm", 0)
    rate = (100.0 * warm / total) if total else 0.0
    return (f"tally: {warm} warm / {counts.get('COLD', 0)} COLD / "
            f"{counts.get('FAIL', 0)} FAIL  (warm-rate {rate:.0f}%)")


def _tiny_png(size=256):
    from PIL import Image, ImageDraw

    img = Image.new("RGB", (size, size), (128, 128, 128))
    mask = Image.new("L", (size, size), 0)
    ImageDraw.Draw(mask).ellipse([size * 0.3, size * 0.3, size * 0.7, size * 0.7], fill=255)
    a = BytesIO(); img.save(a, "PNG")
    b = BytesIO(); mask.save(b, "PNG")
    return a.getvalue(), b.getvalue()


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Keep the fal FLUX container warm")
    ap.add_argument("--interval", type=int, default=90,
                    help="target seconds START-to-START (fixed cadence); MUST be under "
                         "fal's idle timeout. 60s measured ~55%% warm -> try 45, then 30")
    ap.add_argument("--steps", type=int, default=4, help="cheap ping inference steps")
    ap.add_argument("--size", type=int, default=256, help="tiny ping image size")
    ap.add_argument("--warm-max", type=float, default=20.0,
                    help="ping under this many seconds = warm; above = COLD (spin-up)")
    ap.add_argument("--timeout", type=int, default=240,
                    help="per-ping timeout; must exceed a cold spin-up (~2-3 min)")
    args = ap.parse_args(argv)

    if not fal_real_enabled():
        print("REFUSED: FAL_KEY not set (no pings sent).")
        return 2

    img, mask = _tiny_png(args.size)
    print(f"keep-warm: model={fal_inpaint_model()} every {args.interval}s "
          f"(tiny {args.size}px, {args.steps} steps, warm<= {args.warm_max}s). Ctrl+C to stop.")
    counts = {"warm": 0, "COLD": 0, "FAIL": 0}
    n = 0
    while True:
        n += 1
        t = time.time()
        err_msg = ""
        try:
            run_fal_inpaint(img, mask, "warm", strength=0.2, num_inference_steps=args.steps,
                            timeout_s=args.timeout)
            dur = time.time() - t
            label = classify(True, dur, args.warm_max)
            print(f"  ping #{n}: {label} ({dur:.1f}s)")
        except FalUnavailable as exc:
            label = "FAIL"; err_msg = str(exc)
            print(f"  ping #{n}: FAIL ({err_msg})")
        except Exception as exc:  # noqa: BLE001
            label = "FAIL"; err_msg = type(exc).__name__
            print(f"  ping #{n}: FAIL ({err_msg})")
        counts[label] += 1

        # A billing/auth failure is NOT a warm-up problem — stop instead of hammering
        # a locked account (and stop the misleading "lower the interval" hint).
        if label == "FAIL" and is_fatal_error(err_msg):
            print(f"  {summary_line(counts)}")
            print("  STOP: fatal fal error (billing/auth) — this is NOT a keep-warm "
                  "issue. Fix the account (top up at fal.ai/dashboard/billing) or the "
                  "key, then restart. Halting pings.")
            return 1

        # ping #1 is ALWAYS the startup spin-up (the container is down when this
        # process starts) — expected cold, NOT an interval leak. From ping #2 on,
        # a non-warm ping means the container scaled down *between pings* -> the
        # interval is longer than fal's idle window. Warn only on that bad ping;
        # otherwise show the tally periodically.
        if should_warn_leak(n, label):
            print(f"  {summary_line(counts)}")
            if label == "COLD":
                print(f"  hint: COLD on ping #{n} (between-ping) -> container scaled "
                      f"down; idle timeout < {args.interval}s; lower --interval "
                      f"(try {max(30, args.interval - 15)}).")
            else:
                print(f"  note: ping #{n} FAILED (transient/network) — retrying next cadence.")
        elif n == 1 and label != "warm":
            print("  (ping #1 = startup spin-up, expected cold; not a leak. "
                  "Let it run and watch pings #2+ to judge the interval.)")
        elif label == "warm" and n % 10 == 0:
            print(f"  {summary_line(counts)}")
        time.sleep(sleep_seconds(args.interval, time.time() - t))


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        print("\nkeep-warm: stopped.")