#!/usr/bin/env python3
"""Sample Ollama Cloud usage into db/ollama_usage.sqlite. Run from cron every minute.

Writes a `samples` row only when something changed; `meta` records every poll so a
gap in `samples` can be told apart from a poll that failed or never ran.
"""

from __future__ import annotations

import json
import sqlite3
import sys
import urllib.error
import urllib.request
from datetime import UTC, datetime
from pathlib import Path

from ollama_usage import API_URL, load_api_key

WORKSPACE = Path(__file__).resolve().parents[3]
DB_PATH = WORKSPACE / "db" / "ollama_usage.sqlite"

SCHEMA = """
CREATE TABLE IF NOT EXISTS samples (
  ts             TEXT PRIMARY KEY,
  session_usage  REAL NOT NULL,
  weekly_usage   REAL NOT NULL,
  session_models TEXT NOT NULL,
  weekly_models  TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS meta (
  id          INTEGER PRIMARY KEY CHECK (id = 1),
  last_ts     TEXT NOT NULL,
  last_status TEXT NOT NULL
);
"""


def fetch_usage() -> dict:
    request = urllib.request.Request(API_URL, headers={
        "Authorization": f"Bearer {load_api_key()}",
        "Accept": "application/json",
    })
    with urllib.request.urlopen(request, timeout=15) as response:
        return json.load(response)


def canonical_models(models: list[dict]) -> str:
    """Stable JSON for change detection — the API returns an unordered list."""
    return json.dumps(
        {m["name"]: m["request_count"] for m in models},
        sort_keys=True,
    )


def connect() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.executescript(SCHEMA)
    return conn


def record_poll(conn: sqlite3.Connection, ts: str, status: str) -> None:
    conn.execute(
        "INSERT INTO meta (id, last_ts, last_status) VALUES (1, ?, ?) "
        "ON CONFLICT(id) DO UPDATE SET last_ts = excluded.last_ts, "
        "last_status = excluded.last_status",
        (ts, status),
    )
    conn.commit()


def latest_sample(conn: sqlite3.Connection) -> tuple | None:
    row = conn.execute(
        "SELECT session_usage, weekly_usage, session_models, weekly_models "
        "FROM samples ORDER BY ts DESC LIMIT 1"
    ).fetchone()
    return row


def main() -> None:
    now = datetime.now(UTC).replace(microsecond=0).isoformat()
    conn = connect()

    try:
        data = fetch_usage()
    except urllib.error.HTTPError as e:
        record_poll(conn, now, f"http_{e.code}")
        print(f"{now} ollama.com API error: HTTP {e.code}", file=sys.stderr)
        return
    except urllib.error.URLError as e:
        record_poll(conn, now, "unreachable")
        print(f"{now} ollama.com unreachable: {e.reason}", file=sys.stderr)
        return

    limits = data.get("limits", {})
    session = limits.get("session", {})
    weekly = limits.get("weekly", {})
    sample = (
        session.get("usage", 0.0),
        weekly.get("usage", 0.0),
        canonical_models(session.get("models", [])),
        canonical_models(weekly.get("models", [])),
    )

    if sample != latest_sample(conn):
        conn.execute("INSERT INTO samples VALUES (?, ?, ?, ?, ?)", (now, *sample))
    record_poll(conn, now, "ok")


if __name__ == "__main__":
    main()