"""全局出站发送限流器。

目的：托管账号数量多（如 1000 个）时，避免自动回复/手动发送在同一时刻
大量并发打向抖音接口，占满服务器出站带宽或触发平台级风控。

机制（两层，先到先过）：
  1. 并发上限（Semaphore）：同时在途的发送请求数不超过 KEFU_SEND_MAX_CONCURRENCY；
  2. 令牌桶（QPS）：平均发送速率不超过 KEFU_SEND_RATE_PER_SEC，允许 KEFU_SEND_BURST 的突发。

限流器按事件循环惰性创建（backend 主循环一个实例；独立脚本各自一个），
所有账号共享同一个实例，因此是跨账号的全局闸门。

环境变量：
  KEFU_SEND_MAX_CONCURRENCY  同时在途发送数上限，默认 20，<=0 表示不限
  KEFU_SEND_RATE_PER_SEC     平均每秒发送数上限，默认 10，<=0 表示不限
  KEFU_SEND_BURST            令牌桶容量（突发上限），默认与并发上限相同
"""
from __future__ import annotations

import asyncio
import logging
import os
import time

logger = logging.getLogger("douyin_im.rate_limit")


def _env_float(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, "") or default)
    except ValueError:
        return default


class GlobalSendLimiter:
    """并发上限 + 令牌桶。用法：async with limiter: await 发送。"""

    def __init__(
        self,
        max_concurrency: int = 20,
        rate_per_sec: float = 10.0,
        burst: float | None = None,
    ):
        self.max_concurrency = int(max_concurrency)
        self.rate_per_sec = float(rate_per_sec)
        self.burst = float(burst if burst is not None else max(1, max_concurrency))
        self._sem = (
            asyncio.Semaphore(self.max_concurrency) if self.max_concurrency > 0 else None
        )
        self._tokens = self.burst
        self._last_refill = time.monotonic()
        self._token_lock = asyncio.Lock()
        # 统计：便于日志观察限流是否生效
        self._waited_total = 0.0
        self._acquired_count = 0

    async def _take_token(self) -> None:
        if self.rate_per_sec <= 0:
            return
        async with self._token_lock:
            while True:
                now = time.monotonic()
                self._tokens = min(
                    self.burst, self._tokens + (now - self._last_refill) * self.rate_per_sec
                )
                self._last_refill = now
                if self._tokens >= 1.0:
                    self._tokens -= 1.0
                    return
                await asyncio.sleep((1.0 - self._tokens) / self.rate_per_sec)

    async def __aenter__(self) -> "GlobalSendLimiter":
        start = time.monotonic()
        if self._sem is not None:
            await self._sem.acquire()
        try:
            await self._take_token()
        except BaseException:
            if self._sem is not None:
                self._sem.release()
            raise
        waited = time.monotonic() - start
        self._waited_total += waited
        self._acquired_count += 1
        if waited > 1.0:
            logger.info(
                "send throttled: waited %.1fs (in-flight cap=%s, rate=%s/s)",
                waited, self.max_concurrency or "∞", self.rate_per_sec or "∞",
            )
        return self

    async def __aexit__(self, *exc) -> None:
        if self._sem is not None:
            self._sem.release()


# 每个事件循环一个实例（backend 主循环即全局唯一；脚本自用循环互不影响）
_limiters: dict[int, GlobalSendLimiter] = {}


def get_send_limiter() -> GlobalSendLimiter:
    loop = asyncio.get_running_loop()
    key = id(loop)
    limiter = _limiters.get(key)
    if limiter is None:
        max_conc = int(_env_float("KEFU_SEND_MAX_CONCURRENCY", 20))
        rate = _env_float("KEFU_SEND_RATE_PER_SEC", 10.0)
        burst = _env_float("KEFU_SEND_BURST", max(1, max_conc))
        limiter = GlobalSendLimiter(max_conc, rate, burst)
        _limiters[key] = limiter
        logger.info(
            "global send limiter ready: concurrency=%s, rate=%s/s, burst=%s",
            max_conc if max_conc > 0 else "unlimited",
            rate if rate > 0 else "unlimited",
            burst,
        )
    return limiter