# deepseek/deepseek_balance.py — DeepSeek 余额监控
# 文件夹化模块：配置（API key / 阈值 / 轮询间隔 / 开关）全部从本目录 setting.json 读取。
# 命令（管理员）: .deepseek help / balance / threshold <int> / polling <int> / alert <on|off>
# 监控: 定时查询余额，低于阈值时向【所有已注册通道】的管理员私聊告警。
import asyncio
import json
import os

import aiohttp

import identity
from channel import push_to_user, send_message

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
SETTING_FILE = os.path.join(BASE_DIR, "setting.json")

# ---------------- 配置（从 setting.json 读取） ----------------
API_KEY = ""
API_BASE = "https://api.deepseek.com"
THRESHOLD = 3
POLLING_MINUTES = 3
ALERT_ENABLED = True
# 启动宽限（秒）：首次余额查询前先等 TG 通道 register / NapCat 反向连接完成，
# 否则首次告警会因通道尚未注册而投递失败（启动竞态）。
MONITOR_STARTUP_DELAY = 30


def _load_setting() -> dict:
    try:
        with open(SETTING_FILE, "r", encoding="utf-8") as f:
            data = json.load(f)
    except (OSError, json.JSONDecodeError):
        data = {}
    return data if isinstance(data, dict) else {}


def _save_setting(data: dict):
    os.makedirs(os.path.dirname(SETTING_FILE), exist_ok=True)
    with open(SETTING_FILE, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def _persist():
    """写回全部配置（保留 api_key / api_base 等字段）"""
    data = _load_setting()
    data.update({
        "api_key": API_KEY,
        "api_base": API_BASE,
        "threshold": THRESHOLD,
        "polling": POLLING_MINUTES,
        "alert_enabled": ALERT_ENABLED,
    })
    _save_setting(data)


# 从持久化文件恢复上次设置（跨重启保留）
_s = _load_setting()
API_KEY = str(_s.get("api_key", ""))
API_BASE = str(_s.get("api_base", API_BASE)).rstrip("/")
THRESHOLD = int(_s.get("threshold", THRESHOLD))
POLLING_MINUTES = int(_s.get("polling", POLLING_MINUTES))
ALERT_ENABLED = bool(_s.get("alert_enabled", ALERT_ENABLED))

BALANCE_URL = f"{API_BASE}/user/balance"

_alerted = False
# 本次低余额事件中已成功投递的目标令牌 {"<channel>:<target_id>"}。
# 逐目标补投的依据：只重发不在其中的目标；余额恢复时清空。
_alert_delivered = set()


async def get_balance():
    if not API_KEY:
        print("[balance] setting.json 缺少 api_key，跳过查询")
        return None
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Accept": "application/json"
    }
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get(BALANCE_URL, headers=headers, timeout=30) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    for info in data.get("balance_infos") or []:
                        if info.get("currency") == "CNY":
                            return float(info["total_balance"])
                    return None
                text = await resp.text()
                print(f"[DEBUG][balance] 查询余额失败，状态码 {resp.status}，响应体: {text[:200]}")
                return None
    except Exception as e:
        print(f"[DEBUG][balance] 查询余额异常: {e}")
        return None


async def check_balance_once():
    global _alerted, _alert_delivered
    if not ALERT_ENABLED:
        return
    balance = await get_balance()
    if balance is None:
        print("[DEBUG][balance] 本轮查询失败，跳过")
        return
    print(f"[DEBUG][balance] 轮询查询成功，当前余额 {balance} 元，阈值 {THRESHOLD} 元")
    if balance < THRESHOLD:
        if ALERT_ENABLED and not _alerted:
            msg = (f"注意到 DeepSeek Platform 余额已低于阈值 （ {THRESHOLD} 元 ），请判别业务影响。\n"
                   f"当前余额 {balance} 元")
            # 逐目标补投：按管理员身份扇出，本轮只对「尚未成功」的目标重发，
            # 已成功的目标绝不重复推送（避免 TG 重复轰炸）。
            pending_total, sent_now = 0, 0
            for uid in identity.admin_uids():
                tokens = {"%s:%s" % t: t for t in identity.push_targets(uid)}
                # 目标被解绑/移除即视为消失：不在 tokens 里的历史令牌不再阻塞结案
                todo = [t for token, t in tokens.items() if token not in _alert_delivered]
                pending_total += len(todo)
                if not todo:
                    continue
                results = await push_to_user(uid, msg, tag="DeepSeek 余额告警", targets=todo)
                for (channel_name, target_id), ok in results.items():
                    token = "%s:%s" % (channel_name, target_id)
                    sent_now += 1 if ok else 0
                    if ok:
                        _alert_delivered.add(token)
                    else:
                        print(f"[DEBUG][balance] 目标 {token} 本轮投递失败，下一轮补投")
            # 结案判定：当前所有管理员的目标都已成功即结案；
            # 目标集合为空（无管理员身份）时也正常结案，避免每轮空转刷「未结案」日志
            remaining = 0
            existing = 0
            for uid in identity.admin_uids():
                for t in identity.push_targets(uid):
                    existing += 1
                    if "%s:%s" % t not in _alert_delivered:
                        remaining += 1
            if remaining == 0:
                _alerted = True
                if existing > 0:
                    print(f"[DEBUG][balance] 余额预警已结案（本事件补投 {sent_now} 个目标，"
                          f"累计成功 {existing} 个），当前余额 {balance} 元，阈值 {THRESHOLD} 元")
                else:
                    print(f"[DEBUG][balance] 余额低于阈值但无管理员推送目标，直接结案，"
                          f"当前余额 {balance} 元，阈值 {THRESHOLD} 元")
            else:
                print(f"[DEBUG][balance] 余额预警未结案：本轮补投 {sent_now} 个成功 / "
                      f"{pending_total} 个待投，剩余 {remaining} 个目标未成功（共 {existing} 个），"
                      f"下一轮只补投未成功目标")
                return
        _alerted = True
    else:
        # 余额恢复：复位告警与逐目标投递记录
        _alerted = False
        _alert_delivered.clear()


async def monitor_loop():
    print(f"[DEBUG][balance] 余额监控已启动，每 {POLLING_MINUTES} 分钟查询一次，阈值 {THRESHOLD} 元")
    # 启动宽限：等 TG 通道 register 与 NapCat 反向连接完成后再做首次查询，
    # 避免重启后的首次真实告警投递到空通道（那也是启动竞态吞告警的根因）。
    await asyncio.sleep(MONITOR_STARTUP_DELAY)
    while True:
        try:
            await check_balance_once()
        except Exception as e:
            print(f"[DEBUG][balance] 轮询异常: {e}")
        await asyncio.sleep(POLLING_MINUTES * 60)


HELP_TEXT = (
    "DeepSeek 余额监控指令\n"
    "• .deepseek help  显示本帮助\n"
    "• .deepseek balance  查询当前余额\n"
    "• .deepseek threshold <int>  设置余额阈值（当前: {threshold} 元）\n"
    "• .deepseek polling <int>  设置查询间隔（当前: {polling} 分钟）\n"
    "• .deepseek alert <on|off>  开关余额监控（当前: {alert}）\n"
)


async def handle_command(channel, msg_type, target_id, raw_text):
    global THRESHOLD, POLLING_MINUTES, ALERT_ENABLED, _alerted, _alert_delivered
    parts = raw_text.split(maxsplit=2)
    subcmd = parts[1].lower() if len(parts) >= 2 else "help"
    if subcmd == "help":
        await send_message(channel, msg_type, target_id, HELP_TEXT.format(
            threshold=THRESHOLD, polling=POLLING_MINUTES,
            alert="开" if ALERT_ENABLED else "关"))
        return
    if subcmd == "balance":
        balance = await get_balance()
        if balance is None:
            await send_message(channel, msg_type, target_id, "余额查询失败，请稍后重试")
        else:
            await send_message(channel, msg_type, target_id, f"DeepSeek 账户当前余额：{balance:.2f} 元")
        return
    if subcmd == "threshold":
        if len(parts) < 3:
            await send_message(channel, msg_type, target_id, "用法: .deepseek threshold <int>，例如 .deepseek threshold 3")
            return
        try:
            value = int(parts[2].strip())
        except ValueError:
            await send_message(channel, msg_type, target_id, "阈值必须为整数，例如 .deepseek threshold 3")
            return
        if value < 0:
            await send_message(channel, msg_type, target_id, "阈值不能为负数")
            return
        THRESHOLD = value
        # 阈值变更构成新的告警判定基准：复位结案标志与投递记录，
        # 否则旧事件的 _alerted/_alert_delivered 会压制新阈值下的告警
        _alerted = False
        _alert_delivered.clear()
        _persist()
        await send_message(channel, msg_type, target_id, f"已设置余额阈值为 {THRESHOLD} 元")
        return
    if subcmd == "polling":
        if len(parts) < 3:
            await send_message(channel, msg_type, target_id, "用法: .deepseek polling <int>，例如 .deepseek polling 3")
            return
        try:
            value = int(parts[2].strip())
        except ValueError:
            await send_message(channel, msg_type, target_id, "查询间隔必须为整数分钟，例如 .deepseek polling 3")
            return
        if value < 1:
            await send_message(channel, msg_type, target_id, "查询间隔至少为 1 分钟")
            return
        POLLING_MINUTES = value
        _persist()
        await send_message(channel, msg_type, target_id, f"已设置查询间隔为 {POLLING_MINUTES} 分钟")
        return
    if subcmd == "alert":
        if len(parts) < 3 or parts[2].strip().lower() not in ("on", "off"):
            await send_message(channel, msg_type, target_id, "用法: .deepseek alert <on|off>")
            return
        ALERT_ENABLED = parts[2].strip().lower() == "on"
        # 开关切换即开始新事件：投递记录一并清空，
        # 否则上次事件残留的成功令牌会让重启后的告警被整体跳过（静默吞告警）
        _alerted = False
        _alert_delivered.clear()
        _persist()
        if ALERT_ENABLED:
            await send_message(channel, msg_type, target_id, "已开启余额监控")
        else:
            await send_message(channel, msg_type, target_id, "已关闭余额监控（暂停轮询）")
        return
    await send_message(channel, msg_type, target_id, "未知子指令，使用 .deepseek help 查看指令范围")