"""
Cron-скрипт: liveness probe всех моделей из DB.
Запускается внутри контейнера по расписанию (cron / supervisord).
"""
from __future__ import annotations

import json
import os
import sqlite3
import sys
import time
import urllib.request
from datetime import datetime, timezone

DB_PATH = os.environ.get("DB_PATH", "/data/arena.db")
OMNIROUTE_URL = os.environ.get("OMNIROUTE_URL", "http://192.168.1.101:20128")
OMNIROUTE_KEY = os.environ.get("OMNIROUTE_API_KEY", "")
TIMEOUT = int(os.environ.get("CHECK_TIMEOUT", "30"))
CONCURRENCY = int(os.environ.get("CHECK_CONCURRENCY", "5"))


def check_model(model_id: str) -> tuple[bool, int | None, str | None]:
    t0 = time.time()
    error = None
    try:
        body = json.dumps({
            "model": model_id,
            "messages": [{"role": "user", "content": "2+2=?"}],
            "max_tokens": 5,
        }).encode()
        req = urllib.request.Request(
            f"{OMNIROUTE_URL}/v1/chat/completions",
            data=body,
            headers={
                "Authorization": f"Bearer {OMNIROUTE_KEY}",
                "Content-Type": "application/json",
            },
        )
        with urllib.request.urlopen(req, timeout=TIMEOUT) as r:
            resp = json.loads(r.read())
            content = resp["choices"][0]["message"]["content"]
            ok = bool(content.strip())
    except Exception as e:
        error = f"{type(e).__name__}: {e}"
        print(f"  ✗ {model_id}: {error}", file=sys.stderr)
        ok = False
    latency = int((time.time() - t0) * 1000)
    return ok, latency, error


def run():
    db = sqlite3.connect(DB_PATH, check_same_thread=False)
    models = [r[0] for r in db.execute("SELECT id FROM models").fetchall()]
    db.close()

    print(f"[probe] checking {len(models)} models @ {datetime.now(timezone.utc).isoformat()}")

    from concurrent.futures import ThreadPoolExecutor, as_completed

    def do(mid):
        ok, lat, err = check_model(mid)
        return mid, ok, lat, err

    db = sqlite3.connect(DB_PATH, check_same_thread=False)
    now = datetime.now(timezone.utc).isoformat()

    with ThreadPoolExecutor(max_workers=CONCURRENCY) as ex:
        futs = {ex.submit(do, m): m for m in models}
        done = ok_count = 0
        for f in as_completed(futs):
            mid, ok, lat, err = f.result()
            db.execute(
                "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error) VALUES (?,?,?,?,?)",
                (mid, now, int(ok), lat, err),
            )
            done += 1
            if ok:
                ok_count += 1
            status = "✓" if ok else "✗"
            print(f"  {status} {mid} ({lat}ms)" if lat else f"  {status} {mid}")

    db.commit()
    db.close()
    print(f"[probe] done: {ok_count}/{done} ok")


if __name__ == "__main__":
    run()