from datetime import date
from models import DailyParseStat, UserGroup
from models import db

class RateLimiter:
    """限流器"""

    @staticmethod
    def check_limit(user_id=None, ip_address=None):
        """检查是否超过限制"""
        today = date.today()

        if user_id:
            # 已登录用户
            stat = DailyParseStat.query.filter_by(
                user_id=user_id,
                date=today
            ).first()

            from models import User
            user = User.query.get(user_id)
            group = UserGroup.query.get(user.group_id)
            limit = group.daily_limit

            current_count = stat.parse_count if stat else 0

            return {
                'allowed': current_count < limit,
                'current': current_count,
                'limit': limit,
                'remaining': max(0, limit - current_count)
            }
        else:
            # 游客
            stat = DailyParseStat.query.filter_by(
                ip_address=ip_address,
                date=today
            ).first()

            from models import SiteConfig
            config = SiteConfig.query.filter_by(config_key='guest_daily_limit').first()
            limit = int(config.config_value) if config else 5

            current_count = stat.parse_count if stat else 0

            return {
                'allowed': current_count < limit,
                'current': current_count,
                'limit': limit,
                'remaining': max(0, limit - current_count)
            }

    @staticmethod
    def increment_count(user_id=None, ip_address=None, success=True):
        """增加计数"""
        today = date.today()

        if user_id:
            stat = DailyParseStat.query.filter_by(
                user_id=user_id,
                date=today
            ).first()

            if not stat:
                stat = DailyParseStat(user_id=user_id, date=today)
                db.session.add(stat)

            stat.parse_count = (stat.parse_count or 0) + 1
            if success:
                stat.success_count = (stat.success_count or 0) + 1
            else:
                stat.fail_count = (stat.fail_count or 0) + 1
        else:
            stat = DailyParseStat.query.filter_by(
                ip_address=ip_address,
                date=today
            ).first()

            if not stat:
                stat = DailyParseStat(ip_address=ip_address, date=today)
                db.session.add(stat)

            stat.parse_count = (stat.parse_count or 0) + 1
            if success:
                stat.success_count = (stat.success_count or 0) + 1
            else:
                stat.fail_count = (stat.fail_count or 0) + 1

        db.session.commit()