"""
Schedule-locked warmup for the Detect-Person chute.

TurboVision challenges fire at blocks that are multiples of `WINDOW_BLOCKS`
(300), ~61 min apart. A cold chute returns HTTP 503 instantly and the
validator's own warmup helper is broken, so a cold window scores 0. This
daemon keeps an instance alive only for the minutes around each window:

  - wakes ~LEAD_SECONDS before the next challenge block
  - pings /predict every PING_EVERY_S until TRAIL_SECONDS past the block
  - sleeps the rest of the hour (instance scales to zero on its own)

Run it wherever it can reach api.chutes.ai and the bittensor chain. It only
needs CHUTES_API_KEY in the environment (source your .env).
"""

import base64
import os
import time
import urllib.request
import json

SLUG = os.getenv("SV_CHUTE_SLUG", "aiyoshida-turbovision-dev0524-shrimp")
WINDOW_BLOCKS = int(os.getenv("SV_WINDOW_BLOCKS", "300"))
# Measured cold-start delay across 16 real windows: min 179s, median 221s, max 508s.
# 420s lead missed the tail (one window came up at 506-508s, after the challenge
# block had already fired). Lead now covers the observed max with real margin.
LEAD_SECONDS = int(os.getenv("SV_LEAD_SECONDS", "720"))     # start warming T-12min
TRAIL_SECONDS = int(os.getenv("SV_TRAIL_SECONDS", "360"))   # keep warm to T+6min
PING_EVERY_S = int(os.getenv("SV_PING_EVERY_S", "30"))
BLOCK_TIME_S = float(os.getenv("SV_BLOCK_TIME_S", "12.3"))
API_KEY = os.environ["CHUTES_API_KEY"]

_TINY_PNG = base64.b64decode(
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+P+/HgAFhAJ/wlseKgAAAABJRU5ErkJggg=="
)
_PING_BODY = json.dumps(
    {"frames": [{"frame_id": 0, "data": base64.b64encode(_TINY_PNG).decode()}],
     "meta": {"batch_size": 1, "n_keypoints": 0}}
).encode()


def log(msg: str) -> None:
    print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {msg}", flush=True)


def current_block() -> int:
    from bittensor import Subtensor

    st = Subtensor(network=os.getenv("BITTENSOR_SUBTENSOR_ENDPOINT", "finney"))
    return int(st.get_current_block())


def ping() -> str:
    req = urllib.request.Request(
        f"https://{SLUG}.chutes.ai/predict",
        data=_PING_BODY,
        headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=120) as r:
            ok = json.loads(r.read()).get("success")
            return f"200 success={ok}"
    except urllib.error.HTTPError as e:
        return f"{e.code} {e.read()[:120].decode(errors='ignore')}"
    except Exception as e:
        return f"{type(e).__name__}: {e}"


def seconds_to_next_window() -> tuple[int, float]:
    b = current_block()
    next_block = ((b // WINDOW_BLOCKS) + 1) * WINDOW_BLOCKS
    return next_block, (next_block - b) * BLOCK_TIME_S


def main() -> None:
    log(f"warmup daemon start: slug={SLUG} lead={LEAD_SECONDS}s trail={TRAIL_SECONDS}s")
    while True:
        next_block, eta = seconds_to_next_window()
        sleep_for = eta - LEAD_SECONDS
        if sleep_for > 0:
            log(f"next challenge block {next_block} in ~{eta/60:.1f} min; sleeping {sleep_for/60:.1f} min")
            time.sleep(min(sleep_for, 1800))
            continue

        deadline = time.time() + (eta if eta > 0 else 0) + TRAIL_SECONDS
        log(f"warming for block {next_block} until T+{TRAIL_SECONDS}s")
        while time.time() < deadline:
            log(f"  ping -> {ping()}")
            time.sleep(PING_EVERY_S)
        log(f"window {next_block} done; releasing instance")
        time.sleep(60)


if __name__ == "__main__":
    main()