NeteaseUrl/auth_system/ledger.py

# -*- coding: utf-8 -*-
"""预扣账本（开账 / 改名 / 补扣 / 结算 / 对账）。

（自 auth_system.py 拆出，行为与原实现一致；DB 访问已迁移为原生异步。）
"""

import time
from typing import Any

from .db import (
    _q,
    _x,
    write_transaction,
)
from .logs import (
    add_log,
)

# ============ 计费账本（实扣、交付、退款、终态的持久化唯一记录） ============


async def ledger_open(task_id: str, user_id: int, reserved: int) -> None:
    """任务计费开账：记录本任务实际预扣次数（admin/-1 用户为 0）。
    task_id 唯一；重复 task_id 由任务系统在提交侧拒绝，这里幂等保护。"""
    await _x(
        'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) '
        'VALUES (?, ?, ?, \'open\', ?)',
        (str(task_id), int(user_id), max(0, int(reserved)), int(time.time())))


async def reserve_and_open(uid: int, count: int, ledger_key: str) -> int | None:
    """预扣与开账同一事务（原两步之间存在崩溃丢账窗口）。
    足额：扣减+开账并返回剩余；不限/管理员：开账(reserved=0)返回 -1；不足：返回 None（不开账）。"""
    count = max(1, int(count))
    key = str(ledger_key)
    async with write_transaction() as conn:
        row = await (await conn.execute('SELECT role, duration_seconds FROM users WHERE id = ?', (uid,))).fetchone()
        if not row:
            return None
        now = int(time.time())
        if row['role'] == 'admin':
            await conn.execute('UPDATE users SET last_active = ? WHERE id = ?', (now, uid))
            await conn.execute(
                'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) '
                'VALUES (?, ?, 0, \'open\', ?)', (key, int(uid), now))
            return -1
        cur = int(row['duration_seconds'] or 0)
        if cur == -1:
            await conn.execute('UPDATE users SET last_active = ? WHERE id = ?', (now, uid))
            await conn.execute(
                'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) '
                'VALUES (?, ?, 0, \'open\', ?)', (key, int(uid), now))
            return -1
        if cur < count:
            return None
        await conn.execute('UPDATE users SET duration_seconds = ?, last_active = ? WHERE id = ?',
                           (cur - count, now, uid))
        await conn.execute(
            'INSERT OR IGNORE INTO billing_ledger (task_id, user_id, reserved, status, created_at) '
            'VALUES (?, ?, ?, \'open\', ?)', (key, int(uid), count, now))
        return cur - count


async def ledger_rename(old_key: str, new_key: str) -> bool:
    """把请求级临时账本行改名为正式 task_id（任务接手账目）。
    old 不存在/已关闭/new 已存在 → False。"""
    async with write_transaction() as conn:
        dup = await (await conn.execute('SELECT 1 FROM billing_ledger WHERE task_id = ?', (str(new_key),))).fetchone()
        if dup:
            return False
        cur = await conn.execute(
            'UPDATE billing_ledger SET task_id = ? WHERE task_id = ? AND status = \'open\'',
            (str(new_key), str(old_key)))
        return cur.rowcount == 1


async def ledger_supplement(task_id: str, extra: int) -> bool:
    """整台播客等"先扣 1 后按实际数量补扣"场景：
    原子补扣用户余额并同步累加账本 reserved。任一步失败返回 False（调用方必须让任务失败，
    不允许无账本继续交付）。余额不足同样返回 False。"""
    task_id = str(task_id)
    extra = int(extra)
    if extra <= 0:
        return True
    try:
        async with write_transaction() as conn:
            row = await (await conn.execute('SELECT user_id, status FROM billing_ledger WHERE task_id = ?',
                                            (task_id,))).fetchone()
            if not row or row['status'] != 'open':
                return False
            urow = await (await conn.execute('SELECT role, duration_seconds FROM users WHERE id = ?',
                                             (row['user_id'],))).fetchone()
            if not urow or urow['role'] == 'admin' or int(urow['duration_seconds'] or 0) == -1:
                return True  # 不限用户：无需补扣，账本也不加
            cur_credits = int(urow['duration_seconds'] or 0)
            if cur_credits < extra:
                return False   # 余额不足：任务必须失败（不允许扣到 0 硬做）
            await conn.execute('UPDATE users SET duration_seconds = duration_seconds - ? WHERE id = ?',
                               (extra, row['user_id']))
            await conn.execute('UPDATE billing_ledger SET reserved = reserved + ? WHERE task_id = ?',
                               (extra, task_id))
            return True
    except Exception:
        return False


async def ledger_get(task_id: str) -> dict[str, Any] | None:
    rows = await _q('SELECT * FROM billing_ledger WHERE task_id = ?', (str(task_id),))
    return dict(rows[0]) if rows else None


async def ledger_settle(task_id: str, delivered: int, final_status: str) -> int:
    """任务计费结算（幂等）：按账本实扣与实际交付计算应退次数，原子退款并落终态。
    - delivered：本次任务实际交付给用户的数量（单曲/单集 0 或 1；批量为成功入包数）
    - 重复调用/重启后再次调用：账本已非 open 直接返回 0，不重复退款
    - admin/-1 提交（reserved=0）：退款自然为 0（不看结算时的账户状态）
    - 以账本 reserved 为准计算应退；口径更正：-1（不限次）账户不落余额
      （加次数会破坏 -1 哨兵值），账本 refunded 记录应付口径，作废可接受
    返回实际退款次数（落余额的部分；-1/admin 场景为 0）。"""
    task_id = str(task_id)
    delivered = max(0, int(delivered))
    async with write_transaction() as conn:
        row = await (await conn.execute('SELECT * FROM billing_ledger WHERE task_id = ?', (task_id,))).fetchone()
        if not row or row['status'] != 'open':
            return 0
        refund = max(0, int(row['reserved']) - delivered)
        # 退款与账本终态同事务落库：要么都成，要么都不成。
        # 注意 WHERE 必须排除 admin/-1：向 -1（不限次）账户加次数会把 -1 哨兵值破坏成
        # -1+N，向 admin 加无意义——因此任务运行期间账户升级为 -1 时退款不落余额，
        # 账本 refunded 记录的是"按账本应付"口径（该用户已有无限次数，作废可接受）；
        # 原注释"照样退"与实现相反，已更正为此口径。
        if refund > 0:
            await conn.execute('UPDATE users SET duration_seconds = duration_seconds + ? '
                               'WHERE id = ? AND role != \'admin\' AND duration_seconds != -1',
                               (refund, row['user_id']))
        await conn.execute('UPDATE billing_ledger SET delivered = ?, refunded = ?, status = ?, '
                           'settled_at = ? WHERE task_id = ?',
                           (delivered, refund, str(final_status)[:20], int(time.time()), task_id))
        return refund


async def ledger_recover_orphans(max_age_seconds: int = 2 * 3600) -> int:
    """启动对账：任务表在内存，进程重启后账本中仍 open 的行代表
    "已扣款但任务已丢失"——退回全部实扣。仅处理超过宽限期的行，
    避免误伤同库其他存活进程的进行中任务。返回退款笔数。"""
    now = int(time.time())
    recovered = 0
    rows = await _q('SELECT task_id FROM billing_ledger WHERE status = \'open\' AND created_at <= ?',
                    (now - int(max_age_seconds),))
    for r in rows:
        try:
            refunded = await ledger_settle(r['task_id'], 0, 'lost_restart')
            if refunded:
                await add_log('success', 'system', 'system', '重启对账退款',
                              f'任务 {r["task_id"]} 因进程重启丢失，退回 {refunded} 次')
            else:
                await ledger_settle(r['task_id'], 0, 'lost_restart')
            recovered += 1
        except Exception:
            pass
    return recovered


async def ledger_cleanup_settled(max_age_seconds: int = 90 * 86400) -> int:
    """清理已结算的历史账本行（每任务一行、只增不删会让表无限膨胀；
    settled 状态行无对账价值，保留 90 天足够追溯）"""
    try:
        return await _x('DELETE FROM billing_ledger WHERE status != \'open\' AND settled_at <= ?',
                        (int(time.time()) - int(max_age_seconds),))
    except Exception:
        return 0