from __future__ import annotations

import time
from dataclasses import dataclass
from threading import Lock
from typing import Optional


@dataclass(frozen=True)
class CooldownDecision:
    """Result of a cooldown check before relay."""

    allowed: bool
    should_ack: bool


class CooldownService:
    """Per-user cooldown. In-memory; enough for a single bot process.

    Cooldown is committed only after a successful relay via :meth:`commit`.
    Media albums share one ack and one cooldown window.
    Expired entries are pruned opportunistically so the maps stay bounded.
    """

    def __init__(self, seconds: float) -> None:
        self._seconds = max(0.0, seconds)
        self._last_allowed: dict[int, float] = {}
        # user_id -> (media_group_id, monotonic_ts when first item was accepted)
        self._active_albums: dict[int, tuple[str, float]] = {}
        self._services_last: dict[int, float] = {}
        self._lock = Lock()
        # Album items can arrive a few seconds apart.
        self._album_window = max(self._seconds, 3.0)
        self._services_seconds = max(self._seconds, 5.0)

    def _retention(self) -> float:
        return max(self._album_window, self._services_seconds, self._seconds, 1.0) * 2

    def _prune_unlocked(self, now: float) -> None:
        retain = self._retention()
        stale_users = [uid for uid, ts in self._last_allowed.items() if now - ts > retain]
        for uid in stale_users:
            del self._last_allowed[uid]

        stale_albums = [
            uid for uid, (_gid, ts) in self._active_albums.items() if now - ts > retain
        ]
        for uid in stale_albums:
            del self._active_albums[uid]

        stale_services = [
            uid for uid, ts in self._services_last.items() if now - ts > retain
        ]
        for uid in stale_services:
            del self._services_last[uid]

    def check(self, user_id: int, media_group_id: Optional[str] = None) -> CooldownDecision:
        """Return whether the user may relay, without consuming the cooldown."""
        if self._seconds <= 0 and media_group_id is None:
            return CooldownDecision(allowed=True, should_ack=True)

        now = time.monotonic()
        with self._lock:
            self._prune_unlocked(now)

            if media_group_id is not None:
                active = self._active_albums.get(user_id)
                if (
                    active is not None
                    and active[0] == media_group_id
                    and now - active[1] < self._album_window
                ):
                    return CooldownDecision(allowed=True, should_ack=False)

            last = self._last_allowed.get(user_id)
            if last is not None and now - last < self._seconds:
                return CooldownDecision(allowed=False, should_ack=False)

            return CooldownDecision(allowed=True, should_ack=True)

    def commit(self, user_id: int, media_group_id: Optional[str] = None) -> None:
        """Mark cooldown after a successful relay."""
        if self._seconds <= 0 and media_group_id is None:
            return

        now = time.monotonic()
        with self._lock:
            self._prune_unlocked(now)

            if media_group_id is not None:
                active = self._active_albums.get(user_id)
                if (
                    active is not None
                    and active[0] == media_group_id
                    and now - active[1] < self._album_window
                ):
                    # Continuation of the same album — cooldown already stamped.
                    return
                self._active_albums[user_id] = (media_group_id, now)

            self._last_allowed[user_id] = now

    def try_take_services(self, user_id: int) -> bool:
        """Return True once; silently False if tapped again within the services window."""
        if self._services_seconds <= 0:
            return True

        now = time.monotonic()
        with self._lock:
            self._prune_unlocked(now)
            last = self._services_last.get(user_id)
            if last is not None and now - last < self._services_seconds:
                return False
            self._services_last[user_id] = now
            return True