"""
用量追踪中间件：每日对话次数限制与强制执行
"""
import logging
from datetime import date
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from sqlalchemy.orm import Session

from app.core.database import SessionLocal
from app.core.security import decode_access_token
from app.models.user import User

logger = logging.getLogger(__name__)

# 套餐每日限额
TIER_DAILY_LIMITS = {
    "free": 50,
    "pro": 500,
    "enterprise": 0,  # 0 = 无限
}

# 需要计数的路径前缀
COUNTED_PATH_PREFIXES = [
    "/api/v1/agent-chat/",
]

# 不计数路径（导出、会话管理等非对话操作）
EXCLUDED_PATH_PATTERNS = [
    "/sessions/",
    "/search-messages",
    "/export",
]


class UsageMiddleware(BaseHTTPMiddleware):
    """用量追踪中间件：检查每日限额 → 放行 → 递增计数"""

    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        should_count = (
            request.method == "POST"
            and any(path.startswith(p) for p in COUNTED_PATH_PREFIXES)
            and not any(p in path for p in EXCLUDED_PATH_PATTERNS)
        )

        if not should_count:
            return await call_next(request)

        # 从 Authorization header 解码 user_id
        user_id = None
        auth_header = request.headers.get("Authorization", "")
        if auth_header.startswith("Bearer "):
            token = auth_header[7:]
            try:
                payload = decode_access_token(token)
                user_id = payload.get("sub")
            except Exception:
                pass

        if not user_id:
            return await call_next(request)

        db: Session = SessionLocal()
        try:
            user = db.query(User).filter(User.id == user_id).first()
            if not user:
                return await call_next(request)

            today = date.today()

            # 日期重置
            if user.daily_usage_date and user.daily_usage_date < today:
                user.daily_usage_count = 0
                user.daily_usage_date = today
            if not user.daily_usage_date:
                user.daily_usage_date = today

            tier = user.subscription_tier or "free"
            limit = TIER_DAILY_LIMITS.get(tier, 50)
            current = user.daily_usage_count or 0

            # 检查是否已达限额
            if limit > 0 and current >= limit:
                db.close()
                logger.warning("用户 %s 已达每日限额 %s/%s，拒绝请求", user_id, current, limit)
                return JSONResponse(
                    status_code=429,
                    content={
                        "detail": f"今日对话次数已用完（{current}/{limit}）。升级 Pro 版享受每日 {TIER_DAILY_LIMITS.get('pro', 500)} 次对话。",
                        "daily_used": current,
                        "daily_limit": limit,
                        "upgrade_required": True,
                    },
                    headers={"X-Upgrade-Required": "true"},
                )

            # 未达限额：先放行，再递增
            response = await call_next(request)

            if response.status_code < 400:
                # 递增计数（需要重新查询，因为 request 可能耗时较长）
                db_user = db.query(User).filter(User.id == user_id).first()
                if db_user:
                    db_user.daily_usage_count = (db_user.daily_usage_count or 0) + 1
                    db.commit()

                    new_count = db_user.daily_usage_count
                    if limit > 0 and new_count >= limit:
                        logger.info("用户 %s 已达每日限额 %s/%s", user_id, new_count, limit)

            return response
        except Exception as e:
            db.rollback()
            logger.error("用量追踪失败: %s", e)
            return await call_next(request)
        finally:
            db.close()