"""速率限制中间件 — Redis 按天配额 + 内存滑动窗口突发保护"""

import asyncio
import logging
import time
from collections import defaultdict

from fastapi import Request
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.config import settings
from app.core.tenant_context import tenant_ctx

# 每秒最大请求数（突发保护）
BURST_MAX_PER_SECOND = 30
# 滑动窗口时长（秒）
BURST_WINDOW = 1

logger = logging.getLogger(__name__)


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Redis 原子计数器 + 5分钟缓存 plan 查询 + 内存滑动窗口突发保护

    每日配额依赖 Redis（重启后仍准确）；Redis 不可用时仅保留突发保护。
    """

    def __init__(self, app):
        super().__init__(app)
        self._redis = None
        self._quota_cache: dict[str, tuple[int, float]] = {}
        self._lock = asyncio.Lock()
        # 突发保护：{tenant_id: [(timestamp, ...)]}
        self._burst_windows: dict[str, list[float]] = defaultdict(list)

    async def _get_redis(self):
        if self._redis is not None:
            return self._redis
        try:
            from redis.asyncio import Redis
            r = Redis.from_url(settings.REDIS_URL, decode_responses=True, socket_connect_timeout=1)
            await r.ping()
            self._redis = r
        except Exception:
            self._redis = False
        return self._redis if self._redis else None

    async def _get_quota(self, tenant_id: str) -> int:
        now = time.time()
        if tenant_id in self._quota_cache:
            q, exp = self._quota_cache[tenant_id]
            if now < exp:
                return q
        try:
            from sqlalchemy import select

            from app.core.plans import PLANS
            from app.db import async_session
            from app.models.user import Tenant
            async with async_session() as db:
                r = await db.execute(select(Tenant.plan_type).where(Tenant.id == tenant_id))
                plan_type = r.scalar() or "free"
            q = PLANS.get(plan_type, PLANS.get("free", {})).get("api_quota_per_day", 100)
        except Exception:
            logger.exception("plan_quota_fetch_failed")
            q = 100
        self._quota_cache[tenant_id] = (q, now + 300)
        return q

    def _check_burst(self, tid: str) -> bool:
        """检查每秒突发限制，返回 True 表示允许通过"""
        now = time.time()
        window = self._burst_windows[tid]
        # 移除过期时间戳
        cutoff = now - BURST_WINDOW
        i = 0
        while i < len(window) and window[i] < cutoff:
            i += 1
        if i > 0:
            del window[:i]
        if len(window) >= BURST_MAX_PER_SECOND:
            return False
        window.append(now)
        return True

    async def dispatch(self, request: Request, call_next):
        if not request.url.path.startswith("/api/v1/"):
            return await call_next(request)

        tid = tenant_ctx.get()
        if not tid:
            auth_header = request.headers.get("Authorization", "")
            if auth_header.startswith("Bearer "):
                try:
                    from jose import jwt as jose_jwt
                    token = auth_header[7:]
                    payload = jose_jwt.decode(token, settings.JWT_SECRET, algorithms=[settings.JWT_ALGORITHM])
                    tid = payload.get("tid")
                except Exception:
                    pass

        if not tid:
            path = request.url.path
            # webhooks 和 approvals 维持豁免（有独立认证）
            if any(path.startswith(p) for p in ("/api/v1/webhooks/", "/api/v1/approvals/")):
                return await call_next(request)

            client_ip = request.client.host if request.client else "unknown"

            # auth 路由跳过所有限速（登录失败锁已在 auth.py 独立处理）
            if path.startswith("/api/v1/auth/"):
                return await call_next(request)

            # 公开路由 + 匿名：Redis 每 IP 限速（60 req/min）+ 内存突发保护
            r = await self._get_redis()
            if r:
                ip_key = f"ratelimit:anon:{client_ip}"
                count = await r.incr(ip_key)
                if count == 1:
                    await r.expire(ip_key, 60)
                if count > 60:
                    return JSONResponse(
                        status_code=429,
                        content={"success": False, "error": {"code": "RATE_LIMITED", "message": "请求过于频繁，请稍后再试"}},
                        headers={"Retry-After": "60"},
                    )

            if not self._check_burst(f"anon:{client_ip}"):
                return JSONResponse(
                    status_code=429,
                    content={"success": False, "error": {"code": "RATE_LIMITED", "message": "Too many requests."}},
                    headers={"Retry-After": "1"},
                )
            return await call_next(request)

        # 突发保护
        if not self._check_burst(tid):
            return JSONResponse(
                status_code=429,
                content={"success": False, "error": {"code": "RATE_LIMITED", "message": "Too many requests per second."}},
                headers={"Retry-After": "1"},
            )

        today = time.strftime("%Y-%m-%d")
        limit = await self._get_quota(tid)
        if limit <= 0:
            return await call_next(request)

        # Redis 原子计数器
        r = await self._get_redis()
        if r:
            key = f"ratelimit:{tid}:{today}"
            count = await r.incr(key)
            if count == 1:
                await r.expire(key, 86400)
            if count > limit:
                await r.decr(key)
                return self._rate_limited_response(limit, 0)
            response = await call_next(request)
            response.headers["X-RateLimit-Remaining"] = str(max(0, limit - count))
            return response

        # Redis 不可用：跳过每日配额，仅突发保护生效
        logger.warning("Redis unavailable — daily rate limit skipped for tenant %s", tid)
        response = await call_next(request)
        response.headers["X-RateLimit-Remaining"] = "0"
        return response

    def _rate_limited_response(self, limit: int, remaining: int):
        return JSONResponse(
            status_code=429,
            content={"success": False, "error": {"code": "RATE_LIMITED", "message": f"API quota exceeded. Limit: {limit}/day."}},
            headers={"X-RateLimit-Limit": str(limit), "X-RateLimit-Remaining": "0", "Retry-After": "86400"},
        )
