# middlewares/throttling.py

from __future__ import annotations

import asyncio
import logging
import time
from dataclasses import dataclass
from typing import (
    Any,
    Awaitable,
    Callable,
    Dict,
    Final,
    Optional,
)

from aiogram import BaseMiddleware
from aiogram.exceptions import TelegramBadRequest
from aiogram.types import (
    CallbackQuery,
    Message,
    TelegramObject,
)

try:
    from core.config import settings
except Exception:
    # Middleware مستقل می‌ماند و اگر config هنگام import
    # در دسترس نبود، ربات به خاطر throttling از کار نمی‌افتد.
    settings = None


logger = logging.getLogger(__name__)


# ============================================================
# Types / Constants
# ============================================================

EVENT_MESSAGE: Final[str] = "message"
EVENT_CALLBACK: Final[str] = "callback"
EVENT_OTHER: Final[str] = "other"

CacheKey = tuple[int, str]


@dataclass(slots=True)
class ThrottleDecision:
    allowed: bool
    retry_after: float = 0.0
    rate_limit: float = 0.0


# ============================================================
# Middleware
# ============================================================

class ThrottlingMiddleware(BaseMiddleware):
    """
    Throttling Middleware برای Aiogram 3.

    ویژگی‌ها:
        - Rate Limit مستقل برای Message و CallbackQuery
        - جلوگیری از Spam و Double-click
        - Cache محدود با Cleanup خودکار
        - asyncio.Lock برای جلوگیری از Race Condition
        - Whitelist دستی
        - Whitelist خودکار ADMIN_IDS
        - Warning cooldown
        - Reset per-user / per-event
        - Statistics
        - API سازگار با نسخه قبلی

    نکته:
        محدودسازی داخل حافظه همین Process است.
        اگر ربات با چند Worker/Process اجرا شود، برای Rate Limit
        سراسری باید Redis یا Storage مشترک استفاده شود.
    """

    def __init__(
        self,
        rate_limit: float = 1.0,
        callback_rate_limit: Optional[float] = None,
        cleanup_interval: float = 60.0,
        cache_ttl: float = 300.0,
        max_cache_size: int = 10_000,
        warning_enabled: bool = True,
        warning_cooldown: float = 5.0,
        whitelist: Optional[set[int]] = None,
        *,
        whitelist_admins: bool = True,
        log_throttled_requests: bool = True,
    ) -> None:
        super().__init__()

        self.rate_limit = self._validate_non_negative_float(
            "rate_limit",
            rate_limit,
        )

        self.callback_rate_limit = (
            self._validate_non_negative_float(
                "callback_rate_limit",
                callback_rate_limit,
            )
            if callback_rate_limit is not None
            else self.rate_limit
        )

        self.cleanup_interval = self._validate_positive_float(
            "cleanup_interval",
            cleanup_interval,
        )

        self.cache_ttl = self._validate_positive_float(
            "cache_ttl",
            cache_ttl,
        )

        try:
            max_cache_size = int(max_cache_size)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                "max_cache_size must be an integer."
            ) from exc

        if max_cache_size <= 0:
            raise ValueError(
                "max_cache_size must be > 0."
            )

        self.max_cache_size = max_cache_size

        self.warning_enabled = bool(
            warning_enabled
        )

        self.warning_cooldown = (
            self._validate_non_negative_float(
                "warning_cooldown",
                warning_cooldown,
            )
        )

        self.log_throttled_requests = bool(
            log_throttled_requests
        )

        self.whitelist: set[int] = set()

        for user_id in whitelist or set():
            self.add_to_whitelist(
                user_id,
                log=False,
            )

        if whitelist_admins:
            for admin_id in self._load_admin_ids():
                self.add_to_whitelist(
                    admin_id,
                    log=False,
                )

        # ----------------------------------------------------
        # Cache
        # ----------------------------------------------------
        #
        # مهم:
        # کلید شامل event_type است.
        #
        # قبلاً:
        #   users_cache[user_id]
        #
        # بود و Message و Callback همدیگر را throttle می‌کردند.
        #
        # حالا:
        #   activity_cache[(user_id, "message")]
        #   activity_cache[(user_id, "callback")]
        #
        # مستقل هستند.
        # ----------------------------------------------------

        self.activity_cache: Dict[
            CacheKey,
            float,
        ] = {}

        # سازگاری اسمی با نسخه قبلی.
        self.users_cache = self.activity_cache

        # Warning بر اساس user + event type.
        self.warning_cache: Dict[
            CacheKey,
            float,
        ] = {}

        self._lock = asyncio.Lock()

        self._last_cleanup = (
            time.monotonic()
        )

        # Statistics
        self._accepted_count = 0
        self._throttled_count = 0
        self._whitelisted_count = 0

        logger.info(
            "ThrottlingMiddleware initialized | "
            "message_rate=%.2fs callback_rate=%.2fs "
            "cache_ttl=%.0fs max_cache=%d whitelist=%d",
            self.rate_limit,
            self.callback_rate_limit,
            self.cache_ttl,
            self.max_cache_size,
            len(self.whitelist),
        )

    # ========================================================
    # Validation
    # ========================================================

    @staticmethod
    def _validate_non_negative_float(
        name: str,
        value: Any,
    ) -> float:
        try:
            result = float(value)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                f"{name} must be a number."
            ) from exc

        if result < 0:
            raise ValueError(
                f"{name} must be >= 0."
            )

        return result

    @staticmethod
    def _validate_positive_float(
        name: str,
        value: Any,
    ) -> float:
        try:
            result = float(value)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                f"{name} must be a number."
            ) from exc

        if result <= 0:
            raise ValueError(
                f"{name} must be > 0."
            )

        return result

    @staticmethod
    def _normalize_user_id(
        user_id: Any,
    ) -> int:
        try:
            value = int(user_id)
        except (TypeError, ValueError) as exc:
            raise ValueError(
                "user_id must be an integer."
            ) from exc

        if value <= 0:
            raise ValueError(
                "user_id must be greater than zero."
            )

        return value

    # ========================================================
    # Admin Whitelist
    # ========================================================

    @staticmethod
    def _load_admin_ids() -> set[int]:
        """
        ADMIN_IDS پروژه را بدون Hardcode بارگذاری می‌کند.
        """
        if settings is None:
            return set()

        method = getattr(
            settings,
            "get_admin_list",
            None,
        )

        if callable(method):
            try:
                return {
                    int(user_id)
                    for user_id in method()
                    if int(user_id) > 0
                }
            except Exception:
                logger.exception(
                    "Could not load admin whitelist "
                    "from settings.get_admin_list()."
                )

        raw = str(
            getattr(
                settings,
                "ADMIN_IDS",
                "",
            )
            or ""
        )

        result: set[int] = set()

        for item in raw.split(","):
            item = item.strip()

            if not item:
                continue

            try:
                user_id = int(item)

                if user_id > 0:
                    result.add(user_id)

            except ValueError:
                logger.warning(
                    "Invalid ADMIN_IDS entry ignored "
                    "by throttling middleware: %r",
                    item,
                )

        return result

    # ========================================================
    # Event Helpers
    # ========================================================

    @staticmethod
    def _get_user(
        event: TelegramObject,
    ):
        if isinstance(
            event,
            Message,
        ):
            return event.from_user

        if isinstance(
            event,
            CallbackQuery,
        ):
            return event.from_user

        return None

    @staticmethod
    def _get_event_type(
        event: TelegramObject,
    ) -> str:
        if isinstance(
            event,
            CallbackQuery,
        ):
            return EVENT_CALLBACK

        if isinstance(
            event,
            Message,
        ):
            return EVENT_MESSAGE

        return EVENT_OTHER

    def _get_rate_limit(
        self,
        event: TelegramObject,
    ) -> float:
        if isinstance(
            event,
            CallbackQuery,
        ):
            return (
                self.callback_rate_limit
            )

        if isinstance(
            event,
            Message,
        ):
            return self.rate_limit

        return 0.0

    def _cache_key(
        self,
        user_id: int,
        event: TelegramObject,
    ) -> CacheKey:
        return (
            int(user_id),
            self._get_event_type(
                event
            ),
        )

    # ========================================================
    # Whitelist
    # ========================================================

    def add_to_whitelist(
        self,
        user_id: int,
        *,
        log: bool = True,
    ) -> None:
        user_id = self._normalize_user_id(
            user_id
        )

        self.whitelist.add(
            user_id
        )

        if log:
            logger.info(
                "User %s added to throttling whitelist.",
                user_id,
            )

    def remove_from_whitelist(
        self,
        user_id: int,
    ) -> None:
        user_id = self._normalize_user_id(
            user_id
        )

        self.whitelist.discard(
            user_id
        )

        logger.info(
            "User %s removed from throttling whitelist.",
            user_id,
        )

    def is_whitelisted(
        self,
        user_id: int,
    ) -> bool:
        try:
            user_id = int(user_id)
        except (TypeError, ValueError):
            return False

        return (
            user_id in self.whitelist
        )

    def refresh_admin_whitelist(
        self,
    ) -> int:
        """
        ADMIN_IDS جدید را بدون حذف Whitelist دستی اضافه می‌کند.
        """
        before = len(
            self.whitelist
        )

        self.whitelist.update(
            self._load_admin_ids()
        )

        return (
            len(self.whitelist)
            - before
        )

    # ========================================================
    # Warning
    # ========================================================

    async def _reserve_warning_slot(
        self,
        key: CacheKey,
        now: float,
    ) -> bool:
        """
        Warning cooldown هم داخل Lock کنترل می‌شود تا دو Task
        همزمان دو Alert جداگانه نفرستند.
        """
        if (
            not self.warning_enabled
            or self.warning_cooldown < 0
        ):
            return False

        async with self._lock:
            last_warning = (
                self.warning_cache.get(
                    key
                )
            )

            if (
                last_warning is not None
                and (
                    now - last_warning
                    < self.warning_cooldown
                )
            ):
                return False

            self.warning_cache[
                key
            ] = now

            return True

    async def _send_throttle_warning(
        self,
        event: TelegramObject,
        retry_after: float,
    ) -> None:
        """
        برای CallbackQuery Alert نمایش داده می‌شود.

        برای Message پیام جدا ارسال نمی‌کنیم تا Rate Limiter
        خودش باعث Spam نشود.
        """
        if not self.warning_enabled:
            return

        user = self._get_user(
            event
        )

        if not user:
            return

        user_id = int(
            user.id
        )

        key = self._cache_key(
            user_id,
            event,
        )

        now = time.monotonic()

        should_send = (
            await self._reserve_warning_slot(
                key,
                now,
            )
        )

        if not should_send:
            # Callback باید Ack شود تا spinner کاربر گیر نکند.
            if isinstance(
                event,
                CallbackQuery,
            ):
                try:
                    await event.answer()
                except Exception:
                    pass

            return

        retry_after = max(
            0.0,
            float(retry_after),
        )

        try:
            if isinstance(
                event,
                CallbackQuery,
            ):
                await event.answer(
                    (
                        "⚠️ لطفاً کمی آهسته‌تر!\n\n"
                        f"حدود {retry_after:.1f} ثانیه "
                        "دیگر دوباره تلاش کنید."
                    ),
                    show_alert=True,
                )

            elif isinstance(
                event,
                Message,
            ):
                # عمداً هیچ پیام جدیدی نمی‌فرستیم.
                return

        except TelegramBadRequest:
            logger.debug(
                "Unable to send throttle warning "
                "to user %s.",
                user_id,
                exc_info=True,
            )

        except Exception:
            logger.exception(
                "Unexpected error sending throttle "
                "warning to user %s.",
                user_id,
            )

    # ========================================================
    # Cache Cleanup
    # ========================================================

    @staticmethod
    def _remove_oldest(
        cache: Dict[
            CacheKey,
            float,
        ],
        max_size: int,
    ) -> int:
        overflow = (
            len(cache)
            - max_size
        )

        if overflow <= 0:
            return 0

        oldest = sorted(
            cache.items(),
            key=lambda item: item[1],
        )[:overflow]

        for key, _ in oldest:
            cache.pop(
                key,
                None,
            )

        return len(oldest)

    def _cleanup_cache_locked(
        self,
        now: float,
    ) -> None:
        """
        فقط زمانی صدا زده شود که self._lock گرفته شده باشد.
        """
        if (
            now - self._last_cleanup
            < self.cleanup_interval
        ):
            return

        self._last_cleanup = now

        expired_activity = [
            key
            for key, timestamp
            in self.activity_cache.items()
            if (
                now - timestamp
                > self.cache_ttl
            )
        ]

        for key in expired_activity:
            self.activity_cache.pop(
                key,
                None,
            )

        expired_warnings = [
            key
            for key, timestamp
            in self.warning_cache.items()
            if (
                now - timestamp
                > self.cache_ttl
            )
        ]

        for key in expired_warnings:
            self.warning_cache.pop(
                key,
                None,
            )

        trimmed_activity = (
            self._remove_oldest(
                self.activity_cache,
                self.max_cache_size,
            )
        )

        trimmed_warnings = (
            self._remove_oldest(
                self.warning_cache,
                self.max_cache_size,
            )
        )

        if (
            expired_activity
            or expired_warnings
            or trimmed_activity
            or trimmed_warnings
        ):
            logger.debug(
                "Throttling cache cleanup | "
                "expired_activity=%d "
                "expired_warnings=%d "
                "trimmed_activity=%d "
                "trimmed_warnings=%d",
                len(expired_activity),
                len(expired_warnings),
                trimmed_activity,
                trimmed_warnings,
            )

    async def cleanup_cache(
        self,
        *,
        force: bool = False,
    ) -> None:
        """
        Cleanup عمومی برای تست/مانیتورینگ.

        force=True:
            cleanup را بدون انتظار cleanup_interval اجرا می‌کند.
        """
        now = time.monotonic()

        async with self._lock:
            if force:
                self._last_cleanup = (
                    now
                    - self.cleanup_interval
                    - 1
                )

            self._cleanup_cache_locked(
                now
            )

    # ========================================================
    # Check Rate Limit
    # ========================================================

    async def _check_rate_limit(
        self,
        user_id: int,
        event: TelegramObject,
    ) -> tuple[
        bool,
        float,
    ]:
        """
        API سازگار با نسخه قبلی.

        Returns:
            (allowed, retry_after)
        """
        decision = (
            await self.check_rate_limit(
                user_id,
                event,
            )
        )

        return (
            decision.allowed,
            decision.retry_after,
        )

    async def check_rate_limit(
        self,
        user_id: int,
        event: TelegramObject,
    ) -> ThrottleDecision:
        user_id = (
            self._normalize_user_id(
                user_id
            )
        )

        rate_limit = (
            self._get_rate_limit(
                event
            )
        )

        if rate_limit <= 0:
            return ThrottleDecision(
                allowed=True,
                retry_after=0.0,
                rate_limit=rate_limit,
            )

        now = time.monotonic()

        key = self._cache_key(
            user_id,
            event,
        )

        async with self._lock:
            self._cleanup_cache_locked(
                now
            )

            last_time = (
                self.activity_cache.get(
                    key
                )
            )

            if last_time is None:
                self.activity_cache[
                    key
                ] = now

                return ThrottleDecision(
                    allowed=True,
                    retry_after=0.0,
                    rate_limit=rate_limit,
                )

            elapsed = (
                now - last_time
            )

            if elapsed >= rate_limit:
                self.activity_cache[
                    key
                ] = now

                return ThrottleDecision(
                    allowed=True,
                    retry_after=0.0,
                    rate_limit=rate_limit,
                )

            retry_after = max(
                0.0,
                rate_limit - elapsed,
            )

            return ThrottleDecision(
                allowed=False,
                retry_after=retry_after,
                rate_limit=rate_limit,
            )

    # ========================================================
    # Reset
    # ========================================================

    async def reset_user(
        self,
        user_id: int,
        event_type: Optional[str] = None,
    ) -> None:
        """
        event_type=None:
            Message + Callback هر دو پاک می‌شوند.

        event_type="message" / "callback":
            فقط همان نوع پاک می‌شود.
        """
        user_id = self._normalize_user_id(
            user_id
        )

        if (
            event_type is not None
            and event_type
            not in {
                EVENT_MESSAGE,
                EVENT_CALLBACK,
                EVENT_OTHER,
            }
        ):
            raise ValueError(
                "Invalid event_type."
            )

        async with self._lock:
            for cache in (
                self.activity_cache,
                self.warning_cache,
            ):
                keys = [
                    key
                    for key in cache
                    if (
                        key[0] == user_id
                        and (
                            event_type is None
                            or key[1] == event_type
                        )
                    )
                ]

                for key in keys:
                    cache.pop(
                        key,
                        None,
                    )

        logger.info(
            "Throttling state reset | "
            "user=%s event_type=%s",
            user_id,
            event_type or "all",
        )

    async def reset_all(
        self,
    ) -> None:
        async with self._lock:
            self.activity_cache.clear()
            self.warning_cache.clear()

            self._last_cleanup = (
                time.monotonic()
            )

            self._accepted_count = 0
            self._throttled_count = 0
            self._whitelisted_count = 0

        logger.info(
            "All throttling state cleared."
        )

    # ========================================================
    # Statistics
    # ========================================================

    def get_stats(
        self,
    ) -> Dict[str, Any]:
        """
        برای Dashboard/Debug.
        """
        message_cache_size = sum(
            1
            for (
                _user_id,
                event_type,
            )
            in self.activity_cache
            if event_type
            == EVENT_MESSAGE
        )

        callback_cache_size = sum(
            1
            for (
                _user_id,
                event_type,
            )
            in self.activity_cache
            if event_type
            == EVENT_CALLBACK
        )

        return {
            # Legacy keys
            "rate_limit":
                self.rate_limit,
            "callback_rate_limit":
                self.callback_rate_limit,
            "cache_size":
                len(
                    self.activity_cache
                ),
            "warning_cache_size":
                len(
                    self.warning_cache
                ),
            "max_cache_size":
                self.max_cache_size,
            "cache_ttl":
                self.cache_ttl,
            "whitelist_size":
                len(
                    self.whitelist
                ),
            "warning_enabled":
                self.warning_enabled,

            # New metrics
            "message_cache_size":
                message_cache_size,
            "callback_cache_size":
                callback_cache_size,
            "accepted_count":
                self._accepted_count,
            "throttled_count":
                self._throttled_count,
            "whitelisted_count":
                self._whitelisted_count,
            "cleanup_interval":
                self.cleanup_interval,
            "warning_cooldown":
                self.warning_cooldown,
        }

    # ========================================================
    # Middleware
    # ========================================================

    async def __call__(
        self,
        handler: Callable[
            [
                TelegramObject,
                Dict[str, Any],
            ],
            Awaitable[Any],
        ],
        event: TelegramObject,
        data: Dict[str, Any],
    ) -> Any:
        user = self._get_user(
            event
        )

        # Event بدون User.
        if user is None:
            return await handler(
                event,
                data,
            )

        # Bot account.
        if getattr(
            user,
            "is_bot",
            False,
        ):
            return await handler(
                event,
                data,
            )

        try:
            user_id = int(
                user.id
            )
        except (
            TypeError,
            ValueError,
        ):
            return await handler(
                event,
                data,
            )

        event_type = (
            self._get_event_type(
                event
            )
        )

        # ----------------------------------------------------
        # Whitelist
        # ----------------------------------------------------

        if self.is_whitelisted(
            user_id
        ):
            self._whitelisted_count += 1

            data[
                "throttle_user_id"
            ] = user_id

            data[
                "throttle_event_type"
            ] = event_type

            data[
                "throttle_rate_limit"
            ] = 0.0

            data[
                "throttle_whitelisted"
            ] = True

            return await handler(
                event,
                data,
            )

        # ----------------------------------------------------
        # Rate Limit
        # ----------------------------------------------------

        decision = (
            await self.check_rate_limit(
                user_id=user_id,
                event=event,
            )
        )

        if not decision.allowed:
            self._throttled_count += 1

            if self.log_throttled_requests:
                logger.info(
                    "Rate limit exceeded | "
                    "user=%s event=%s retry_after=%.3fs",
                    user_id,
                    event_type,
                    decision.retry_after,
                )

            await self._send_throttle_warning(
                event=event,
                retry_after=(
                    decision.retry_after
                ),
            )

            # Handler اجرا نمی‌شود.
            return None

        self._accepted_count += 1

        # ----------------------------------------------------
        # Handler Context
        # ----------------------------------------------------

        data[
            "throttle_user_id"
        ] = user_id

        data[
            "throttle_event_type"
        ] = event_type

        data[
            "throttle_rate_limit"
        ] = decision.rate_limit

        data[
            "throttle_whitelisted"
        ] = False

        return await handler(
            event,
            data,
        )


# ============================================================
# Factory
# ============================================================

def create_throttling_middleware(
    rate_limit: float = 1.0,
    callback_rate_limit: Optional[float] = None,
    cleanup_interval: float = 60.0,
    cache_ttl: float = 300.0,
    max_cache_size: int = 10_000,
    warning_enabled: bool = True,
    warning_cooldown: float = 5.0,
    whitelist: Optional[set[int]] = None,
    *,
    whitelist_admins: bool = True,
    log_throttled_requests: bool = True,
) -> ThrottlingMiddleware:
    """
    Factory سازگار با API قبلی.
    """
    return ThrottlingMiddleware(
        rate_limit=rate_limit,
        callback_rate_limit=(
            callback_rate_limit
        ),
        cleanup_interval=(
            cleanup_interval
        ),
        cache_ttl=cache_ttl,
        max_cache_size=max_cache_size,
        warning_enabled=warning_enabled,
        warning_cooldown=(
            warning_cooldown
        ),
        whitelist=whitelist,
        whitelist_admins=(
            whitelist_admins
        ),
        log_throttled_requests=(
            log_throttled_requests
        ),
    )


__all__ = [
    "ThrottleDecision",
    "ThrottlingMiddleware",
    "create_throttling_middleware",
]