from __future__ import annotations

import asyncio
import time
from typing import Any

from app.enums.private_portrait import PRIVATE_PORTRAIT_ACTION_QPS_LIMITS, PrivatePortraitEventSource, PrivatePortraitEventStatus, PrivatePortraitEventType
from app.services.operation_log_service import log_operation_event
from app.utils.redis import get_redis

DOMAIN = "private_portrait"


class PrivatePortraitRateLimitExceeded(RuntimeError):
    pass


async def acquire_private_portrait_action_token(
    *,
    action: str,
    wait_timeout_seconds: float = 2.0,
    for_celery: bool = False,
) -> bool:
    """Redis 分布式 QPS 限制。Redis 不可用时降级放行，避免影响主功能。"""
    limit = int(PRIVATE_PORTRAIT_ACTION_QPS_LIMITS.get(action, 1))
    if limit <= 0:
        return True
    deadline = time.monotonic() + max(0.0, wait_timeout_seconds)
    while True:
        ok = await _try_take(action, limit)
        if ok:
            return True
        if time.monotonic() >= deadline:
            event_type = PrivatePortraitEventType.ARK_API_RATE_LIMIT_WAIT.value if for_celery else PrivatePortraitEventType.ARK_API_RATE_LIMIT_REJECT.value
            log_operation_event(
                domain=DOMAIN,
                event_type=event_type,
                event_status=PrivatePortraitEventStatus.SKIPPED.value if for_celery else PrivatePortraitEventStatus.FAILED.value,
                source=PrivatePortraitEventSource.CELERY.value if for_celery else PrivatePortraitEventSource.API.value,
                remote_action=action,
                detail={"limit": limit, "wait_timeout_seconds": wait_timeout_seconds},
                message="火山私域真人素材 API 触发本地 QPS 限制",
            )
            if for_celery:
                return False
            raise PrivatePortraitRateLimitExceeded("请求过于频繁，请稍后再试")
        await asyncio.sleep(0.05)


async def _try_take(action: str, limit: int) -> bool:
    client = get_redis()
    if client is None:
        return True
    key = f"private_portrait:qps:{action}:{int(time.time())}"
    try:
        count = await client.incr(key)
        if count == 1:
            await client.expire(key, 2)
        return int(count) <= limit
    except Exception:
        return True