app/services/usage_store.py

"""用量统计 + 审计日志服务。"""
from __future__ import annotations

import json
import time
from typing import Any, Optional

from ..database import get_conn


def _now_ts() -> int:
    return int(time.time())


# ===== 用量统计 =====

def usage_summary(*, hours: int = 24, access_key: Optional[str] = None) -> dict:
    """统计最近 N 小时的用量。"""
    since = _now_ts() - hours * 3600
    sql = "SELECT COUNT(*) AS c, COALESCE(SUM(prompt_tokens),0) AS p, COALESCE(SUM(completion_tokens),0) AS q, COALESCE(SUM(total_tokens),0) AS t, COALESCE(SUM(CASE WHEN ok=1 THEN 1 ELSE 0 END),0) AS ok FROM usage_log WHERE ts >= ?"
    args: list = [since]
    if access_key:
        sql += " AND access_key = ?"
        args.append(access_key)
    with get_conn() as conn:
        row = conn.execute(sql, args).fetchone()
        # 按 provider 分组（含 prompt/completion tokens）
        sql2 = ("SELECT provider, COUNT(*) AS c, "
                "COALESCE(SUM(prompt_tokens),0) AS p, "
                "COALESCE(SUM(completion_tokens),0) AS q, "
                "COALESCE(SUM(total_tokens),0) AS t, "
                "COALESCE(SUM(CASE WHEN ok=1 THEN 1 ELSE 0 END),0) AS ok "
                "FROM usage_log WHERE ts >= ?")
        args2: list = [since]
        if access_key:
            sql2 += " AND access_key = ?"
            args2.append(access_key)
        sql2 += " GROUP BY provider ORDER BY c DESC"
        by_provider = conn.execute(sql2, args2).fetchall()
        # 按 model 分组（top 10）
        sql3 = ("SELECT model, COUNT(*) AS c, "
                "COALESCE(SUM(prompt_tokens),0) AS p, "
                "COALESCE(SUM(completion_tokens),0) AS q, "
                "COALESCE(SUM(total_tokens),0) AS t, "
                "COALESCE(SUM(CASE WHEN ok=1 THEN 1 ELSE 0 END),0) AS ok "
                "FROM usage_log WHERE ts >= ?")
        args3: list = [since]
        if access_key:
            sql3 += " AND access_key = ?"
            args3.append(access_key)
        sql3 += " GROUP BY model ORDER BY c DESC LIMIT 10"
        by_model = conn.execute(sql3, args3).fetchall()
        # 错误码分布
        sql4 = "SELECT error_code, COUNT(*) AS c FROM usage_log WHERE ts >= ? AND ok = 0"
        args4: list = [since]
        if access_key:
            sql4 += " AND access_key = ?"
            args4.append(access_key)
        sql4 += " GROUP BY error_code ORDER BY c DESC"
        by_error = conn.execute(sql4, args4).fetchall()
    return {
        "hours": hours,
        "total_requests": int(row["c"]) if row else 0,
        "success": int(row["ok"]) if row else 0,
        "failed": (int(row["c"]) - int(row["ok"])) if row else 0,
        "prompt_tokens": int(row["p"]) if row else 0,
        "completion_tokens": int(row["q"]) if row else 0,
        "total_tokens": int(row["t"]) if row else 0,
        "by_provider": [{
            "provider": r["provider"],
            "requests": int(r["c"]),
            "success": int(r["ok"]),
            "failed": int(r["c"]) - int(r["ok"]),
            "prompt_tokens": int(r["p"]),
            "completion_tokens": int(r["q"]),
            "tokens": int(r["t"]),
        } for r in by_provider],
        "by_model": [{
            "model": r["model"],
            "requests": int(r["c"]),
            "success": int(r["ok"]),
            "failed": int(r["c"]) - int(r["ok"]),
            "prompt_tokens": int(r["p"]),
            "completion_tokens": int(r["q"]),
            "tokens": int(r["t"]),
        } for r in by_model],
        "by_error": [{"code": r["error_code"], "count": int(r["c"])} for r in by_error],
    }


def usage_timeline(*, hours: int = 24, granularity: str = "hour", access_key: Optional[str] = None) -> list[dict]:
    """时间线（按小时/天分组）。

    返回每个时间桶的：requests / success / failed / tokens / prompt_tokens / completion_tokens。
    """
    since = _now_ts() - hours * 3600
    if granularity == "day":
        group_expr = "(ts / 86400)"
    else:
        group_expr = "(ts / 3600)"
    sql = (
        f"SELECT {group_expr} AS bucket, COUNT(*) AS c, "
        f"COALESCE(SUM(total_tokens),0) AS t, "
        f"COALESCE(SUM(prompt_tokens),0) AS p, "
        f"COALESCE(SUM(completion_tokens),0) AS q, "
        f"COALESCE(SUM(CASE WHEN ok=1 THEN 1 ELSE 0 END),0) AS ok "
        f"FROM usage_log WHERE ts >= ?"
    )
    args: list = [since]
    if access_key:
        sql += " AND access_key = ?"
        args.append(access_key)
    sql += f" GROUP BY bucket ORDER BY bucket ASC"
    with get_conn() as conn:
        rows = conn.execute(sql, args).fetchall()
    return [
        {
            "bucket": int(r["bucket"]),
            "bucket_start": int(r["bucket"]) * (86400 if granularity == "day" else 3600),
            "requests": int(r["c"]),
            "success": int(r["ok"]),
            "failed": int(r["c"]) - int(r["ok"]),
            "tokens": int(r["t"]),
            "prompt_tokens": int(r["p"]),
            "completion_tokens": int(r["q"]),
        }
        for r in rows
    ]


# ===== 审计日志 =====

def audit(
    *,
    action: str,
    actor: Optional[str] = None,
    target: Optional[str] = None,
    detail: Any = None,
) -> None:
    try:
        detail_str = json.dumps(detail, ensure_ascii=False) if detail is not None else None
        with get_conn() as conn:
            conn.execute(
                "INSERT INTO audit_log(ts, action, actor, target, detail) VALUES(?,?,?,?,?)",
                (_now_ts(), action, actor, target, detail_str),
            )
    except Exception as e:
        print(f"[audit] record failed: {e}")


def list_audit(*, hours: int = 24, limit: int = 100, offset: int = 0, action: Optional[str] = None) -> list[dict]:
    since = _now_ts() - hours * 3600
    sql = "SELECT id, ts, action, actor, target, detail FROM audit_log WHERE ts >= ?"
    args: list = [since]
    if action:
        sql += " AND action = ?"
        args.append(action)
    sql += " ORDER BY ts DESC LIMIT ? OFFSET ?"
    args += [limit, offset]
    with get_conn() as conn:
        rows = conn.execute(sql, args).fetchall()
    out = []
    for r in rows:
        item = dict(r)
        if r["detail"]:
            try:
                item["detail"] = json.loads(r["detail"])
            except Exception:
                pass
        out.append(item)
    return out