app/usage_store.py

"""Token 用量存储与聚合（SQLite，单文件 logs/usage.db）。

每次 /v1 请求拿到 usage 后记一行（时间戳 + 路径 + 模型 + token 数）。
面板按 24h / 1d / 3d / 7d / 30d 时间窗聚合：分桶时间序列 + 分模型汇总 + 总计。

- 纯标准库 sqlite3，无外部依赖；WAL 模式并发友好，进程内 threading 锁串行写。
- 查询失败/库损坏不阻断主流程（记 usage 是旁路，绝不影响 API 响应）。
"""
from __future__ import annotations

import sqlite3
import threading
import time
from pathlib import Path
from typing import Any

_LOCK = threading.Lock()
_CONN: sqlite3.Connection | None = None
_DB_PATH: Path | None = None

# 时间窗 → 秒数 + 分桶粒度（秒）。粒度决定折线图点数（点数 = 窗口/粒度）。
_WINDOWS: dict[str, tuple[int, int]] = {
    "24h": (24 * 3600, 3600),        # 24 点（每小时）
    "1d": (24 * 3600, 3600),         # 同 24h（别名）
    "3d": (3 * 24 * 3600, 3 * 3600), # 24 点（每 3 小时）
    "7d": (7 * 24 * 3600, 6 * 3600), # 28 点（每 6 小时）
    "30d": (30 * 24 * 3600, 24 * 3600),  # 30 点（每天）
}


def init(db_path: str | Path) -> None:
    """初始化数据库连接与表结构（应用启动时调一次）。"""
    global _CONN, _DB_PATH
    with _LOCK:
        _DB_PATH = Path(db_path)
        _DB_PATH.parent.mkdir(parents=True, exist_ok=True)
        _CONN = sqlite3.connect(str(_DB_PATH), check_same_thread=False)
        _CONN.execute("PRAGMA journal_mode=WAL")
        _CONN.execute("PRAGMA synchronous=NORMAL")
        _CONN.execute(
            """CREATE TABLE IF NOT EXISTS usage (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                ts REAL NOT NULL,
                path TEXT NOT NULL DEFAULT '',
                model TEXT NOT NULL DEFAULT '',
                input_tokens INTEGER NOT NULL DEFAULT 0,
                output_tokens INTEGER NOT NULL DEFAULT 0,
                total_tokens INTEGER NOT NULL DEFAULT 0
            )"""
        )
        _CONN.execute("CREATE INDEX IF NOT EXISTS idx_usage_ts ON usage(ts)")
        _CONN.commit()


def _int(v: Any) -> int:
    try:
        return int(v or 0)
    except (TypeError, ValueError):
        return 0


def record(usage: dict[str, Any] | None, *, path: str = "", model: str = "") -> None:
    """记录一次用量（旁路，异常静默）。usage 支持 OpenAI/Anthropic 两种字段名。"""
    if _CONN is None or not usage:
        return
    inp = _int(usage.get("input_tokens") or usage.get("prompt_tokens"))
    out = _int(usage.get("output_tokens") or usage.get("completion_tokens"))
    tot = _int(usage.get("total_tokens")) or (inp + out)
    if not (inp or out or tot):
        return
    try:
        with _LOCK:
            _CONN.execute(
                "INSERT INTO usage(ts, path, model, input_tokens, output_tokens, total_tokens) "
                "VALUES(?,?,?,?,?,?)",
                (time.time(), path or "", model or "", inp, out, tot),
            )
            _CONN.commit()
    except Exception:  # noqa: BLE001
        pass


def aggregate(window: str = "24h") -> dict[str, Any]:
    """按时间窗聚合：返回分桶时间序列 + 分模型汇总 + 总计。

    返回 ``{window, since, buckets:[{t, input, output, total, requests}],
    by_model:[{model, input, output, total, requests}], total:{...}}``。
    """
    span, gran = _WINDOWS.get(window, _WINDOWS["24h"])
    now = time.time()
    since = now - span
    empty = {
        "window": window, "since": since, "granularity": gran,
        "buckets": [], "by_model": [], "requests": 0,
        "total": {"input": 0, "output": 0, "total": 0, "requests": 0},
    }
    if _CONN is None:
        return empty
    try:
        with _LOCK:
            rows = _CONN.execute(
                "SELECT ts, model, input_tokens, output_tokens, total_tokens "
                "FROM usage WHERE ts >= ? ORDER BY ts",
                (since,),
            ).fetchall()
    except Exception:  # noqa: BLE001
        return empty

    # 时间分桶（对齐到 gran 边界，保证前端 x 轴均匀）
    n_buckets = max(1, int(span // gran))
    start = now - n_buckets * gran
    buckets = [
        {"t": start + i * gran, "input": 0, "output": 0, "total": 0, "requests": 0}
        for i in range(n_buckets)
    ]
    by_model: dict[str, dict[str, int]] = {}
    tot_in = tot_out = tot_tot = tot_req = 0

    for ts, model, inp, out, tot in rows:
        idx = int((ts - start) // gran)
        if 0 <= idx < n_buckets:
            b = buckets[idx]
            b["input"] += inp; b["output"] += out; b["total"] += tot; b["requests"] += 1
        m = by_model.setdefault(model or "(unknown)",
                                {"input": 0, "output": 0, "total": 0, "requests": 0})
        m["input"] += inp; m["output"] += out; m["total"] += tot; m["requests"] += 1
        tot_in += inp; tot_out += out; tot_tot += tot; tot_req += 1

    by_model_list = sorted(
        [{"model": k, **v} for k, v in by_model.items()],
        key=lambda x: x["total"], reverse=True,
    )
    return {
        "window": window, "since": since, "granularity": gran,
        "buckets": buckets, "by_model": by_model_list, "requests": tot_req,
        "total": {"input": tot_in, "output": tot_out, "total": tot_tot, "requests": tot_req},
    }


def purge_older_than(seconds: float) -> int:
    """删除超过保留期的旧记录，返回删除行数（可选维护，默认不调用）。"""
    if _CONN is None:
        return 0
    try:
        with _LOCK:
            cur = _CONN.execute("DELETE FROM usage WHERE ts < ?", (time.time() - seconds,))
            _CONN.commit()
            return cur.rowcount
    except Exception:  # noqa: BLE001
        return 0


def close() -> None:
    """关闭连接（测试/关停用）。"""
    global _CONN
    with _LOCK:
        if _CONN is not None:
            try:
                _CONN.close()
            except Exception:  # noqa: BLE001
                pass
            _CONN = None