"""
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

from usage_utils import extract_usage

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, dict]:
    """Возвращает (ok, latency_ms, error, usage).

    usage извлекается ДО валидации контента: если парсинг choices падает,
    потраченные токены всё равно сохраняются. Таймаут/ошибка сети без тела
    ответа дают usage=None по всем полям (NULL, а не нули).
    """
    t0 = time.time()
    error = None
    usage = {"input_tokens": None, "output_tokens": None,
             "total_tokens": None, "usage_json": None}
    try:
        body = json.dumps({
            "model": model_id,
            "messages": [{"role": "user", "content": "2+2=?"}],
            "max_tokens": 5,
            "reasoning_effort": "none",
        }).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())
            # Usage — до любой валидации контента.
            usage = extract_usage(resp)
            # HTTP 200 может содержать только reasoning и content=null.
            # Это не успешный текстовый ответ и не ошибка Python.
            choice = resp["choices"][0]
            content = choice["message"].get("content")
            ok = isinstance(content, str) and bool(content.strip())
            if not ok:
                error = f"Empty text response (finish_reason={choice.get('finish_reason')})"
    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, usage


def _ensure_usage_columns(db: sqlite3.Connection):
    """Идемпотентно добавить колонки token usage, если БД ещё со старой схемой.

    Cron стартует раньше uvicorn, поэтому probe не может полагаться на то,
    что init_db из main.py уже прогнал миграцию.
    """
    cols = [r[1] for r in db.execute("PRAGMA table_info(checks)").fetchall()]
    for col, decl in (
        ("input_tokens", "INTEGER"),
        ("output_tokens", "INTEGER"),
        ("total_tokens", "INTEGER"),
        ("usage_json", "TEXT"),
    ):
        if col not in cols:
            db.execute(f"ALTER TABLE checks ADD COLUMN {col} {decl}")
    db.commit()


def run():
    db = sqlite3.connect(DB_PATH, check_same_thread=False)
    _ensure_usage_columns(db)
    model_cols = [r[1] for r in db.execute("PRAGMA table_info(models)")]
    if "active" not in model_cols:
        db.execute("ALTER TABLE models ADD COLUMN active INTEGER NOT NULL DEFAULT 1")
        db.commit()
    models = [r[0] for r in db.execute("SELECT id FROM models WHERE active = 1").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, usage = check_model(mid)
        return mid, ok, lat, err, usage

    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, usage = f.result()
            db.execute(
                "INSERT INTO checks (model_id, checked_at, ok, latency_ms, error, "
                "input_tokens, output_tokens, total_tokens, usage_json) "
                "VALUES (?,?,?,?,?,?,?,?,?)",
                (mid, now, int(ok), lat, err,
                 usage["input_tokens"], usage["output_tokens"],
                 usage["total_tokens"], usage["usage_json"]),
            )
            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()