# -*- coding: utf-8 -*-
"""
API 成本统计 - 记录调用成本，超预算告警
"""
import os
import json
import logging
from pathlib import Path
from datetime import datetime
from typing import Optional

logger = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
COST_FILE = ROOT / "config" / ".api_cost.json"
BUDGET_DAILY = float(os.getenv("API_COST_BUDGET_DAILY", "0"))  # 0=不限制
BUDGET_MONTHLY = float(os.getenv("API_COST_BUDGET_MONTHLY", "0"))

# 单次调用预估成本（元，可按实际调整）
COST_PER_CALL = {
    "gemini_extract": 0.02,
    "gemini_qa": 0.01,
    "gemini_vision": 0.015,
    "openai_chat": 0.02,
    "openai_embedding": 0.001,
    "qwen_enhance": 0.05,
}


def _load() -> dict:
    if not COST_FILE.exists():
        return {"daily": {}, "monthly": {}}
    try:
        with open(COST_FILE, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return {"daily": {}, "monthly": {}}


def _save(data: dict):
    COST_FILE.parent.mkdir(parents=True, exist_ok=True)
    with open(COST_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def record(service: str, count: int = 1, cost: Optional[float] = None):
    """记录一次 API 调用"""
    c = cost if cost is not None else COST_PER_CALL.get(service, 0.01) * count
    today = datetime.now().strftime("%Y-%m-%d")
    month = datetime.now().strftime("%Y-%m")
    data = _load()
    data["daily"][today] = data["daily"].get(today, 0) + c
    data["monthly"][month] = data["monthly"].get(month, 0) + c
    _save(data)
    return c


def get_today_cost() -> float:
    today = datetime.now().strftime("%Y-%m-%d")
    data = _load()
    return data["daily"].get(today, 0)


def get_month_cost() -> float:
    month = datetime.now().strftime("%Y-%m")
    data = _load()
    return data["monthly"].get(month, 0)


async def check_budget_alert():
    """检查是否超预算，超则企微告警"""
    if BUDGET_DAILY <= 0 and BUDGET_MONTHLY <= 0:
        return
    try:
        from config.config import WECHAT_WEBHOOK
        if not WECHAT_WEBHOOK:
            return
        import httpx
        today = get_today_cost()
        month = get_month_cost()
        msg_parts = []
        if BUDGET_DAILY > 0 and today >= BUDGET_DAILY:
            msg_parts.append(f"⚠️ 今日 API 成本 {today:.2f} 元 ≥ 预算 {BUDGET_DAILY} 元")
        if BUDGET_MONTHLY > 0 and month >= BUDGET_MONTHLY:
            msg_parts.append(f"⚠️ 本月 API 成本 {month:.2f} 元 ≥ 预算 {BUDGET_MONTHLY} 元")
        if not msg_parts:
            return
        async with httpx.AsyncClient(timeout=10) as client:
            await client.post(WECHAT_WEBHOOK, json={
                "msgtype": "markdown",
                "markdown": {"content": "\n".join(msg_parts)}
            })
        logger.info(f"[成本] 告警已发送: {msg_parts}")
    except Exception as e:
        logger.warning(f"[成本] 告警失败: {e}")