src/free_claude_code/messaging/limiter.py

"""Runtime-owned queued delivery for one messaging platform."""

import asyncio
from collections import deque
from collections.abc import Awaitable, Callable
from typing import Any

from loguru import logger

from free_claude_code.core.rate_limit import (
    StrictSlidingWindowLimiter as SlidingWindowLimiter,
)

from .safe_diagnostics import format_exception_for_log


class MessagingRateLimiter:
    """
    Rate limiter and compacting work queue for one messaging runtime.

    Uses a custom queue with task compaction (deduplication) to ensure
    only the latest version of a message update is processed.
    """

    def __init__(
        self,
        *,
        rate_limit: int,
        rate_window: float,
        log_error_details: bool = False,
    ) -> None:
        self.limiter = SlidingWindowLimiter(rate_limit, rate_window)
        self._log_error_details = log_error_details
        # Custom queue state - using deque for O(1) popleft
        self._queue_list: deque[str] = deque()  # Deque of dedup_keys in order
        self._queue_map: dict[
            str, tuple[Callable[[], Awaitable[Any]], list[asyncio.Future]]
        ] = {}
        self._condition = asyncio.Condition()
        self._shutdown = asyncio.Event()
        self._worker_task: asyncio.Task[None] | None = None
        self._background_tasks: set[asyncio.Task[None]] = set()
        self._active_futures: list[asyncio.Future[Any]] = []
        self._closed = False
        self._paused_until = 0

        logger.info(
            f"MessagingRateLimiter initialized ({rate_limit} req / {rate_window}s with Task Compaction)"
        )

    def start(self) -> None:
        """Start the owned worker on the current event loop."""
        if self._closed:
            raise RuntimeError("Messaging rate limiter is closed.")
        if self._worker_task and not self._worker_task.done():
            return
        self._worker_task = asyncio.create_task(
            self._worker(), name="msg-limiter-worker"
        )

    async def _worker(self) -> None:
        """Background worker that processes queued messaging tasks."""
        logger.info("MessagingRateLimiter worker started")
        while not self._shutdown.is_set():
            try:
                # Get a task from the queue
                async with self._condition:
                    while not self._queue_list and not self._shutdown.is_set():
                        await self._condition.wait()

                    if self._shutdown.is_set():
                        break

                    dedup_key = self._queue_list.popleft()
                    func, futures = self._queue_map.pop(dedup_key)
                    self._active_futures = futures

                # Check for manual pause (FloodWait)
                now = asyncio.get_event_loop().time()
                if self._paused_until > now:
                    wait_time = self._paused_until - now
                    logger.warning(
                        f"Limiter worker paused, waiting {wait_time:.1f}s more..."
                    )
                    await asyncio.sleep(wait_time)

                # Wait for rate limit capacity
                async with self.limiter:
                    try:
                        result = await func()
                        for f in futures:
                            if not f.done():
                                f.set_result(result)
                    except asyncio.CancelledError:
                        for f in futures:
                            if not f.done():
                                f.cancel()
                        worker = asyncio.current_task()
                        if self._shutdown.is_set() or (
                            worker is not None and worker.cancelling()
                        ):
                            raise
                        logger.debug(
                            "Messaging operation cancelled for key {}; worker remains active",
                            dedup_key,
                        )
                    except Exception as e:
                        # Report error to all futures and log it
                        for f in futures:
                            if not f.done():
                                f.set_exception(e)

                        error_msg = str(e).lower()
                        if "flood" in error_msg or "wait" in error_msg:
                            seconds = 30
                            try:
                                if hasattr(e, "seconds"):
                                    seconds = e.seconds
                                elif "after " in error_msg:
                                    # Try to parse "retry after X"
                                    parts = error_msg.split("after ")
                                    if len(parts) > 1:
                                        seconds = int(parts[1].split()[0])
                            except Exception:
                                pass

                            logger.error(
                                f"FloodWait detected! Pausing worker for {seconds}s"
                            )
                            wait_secs = (
                                float(seconds)
                                if isinstance(seconds, (int, float, str))
                                else 30.0
                            )
                            self._paused_until = (
                                asyncio.get_event_loop().time() + wait_secs
                            )
                        else:
                            logger.error(
                                "Error in limiter worker for key {}: {}",
                                dedup_key,
                                format_exception_for_log(
                                    e,
                                    log_full_message=self._log_error_details,
                                ),
                            )
                    finally:
                        self._active_futures = []
            except asyncio.CancelledError:
                for future in self._active_futures:
                    if not future.done():
                        future.cancel()
                self._active_futures = []
                if self._shutdown.is_set():
                    break
                raise
            except Exception as e:
                if self._log_error_details:
                    logger.error(
                        "MessagingRateLimiter worker critical error: {}",
                        e,
                        exc_info=True,
                    )
                else:
                    logger.error(
                        "MessagingRateLimiter worker critical error: exc_type={}",
                        type(e).__name__,
                    )
                await asyncio.sleep(1)

    async def shutdown(self, timeout: float | None = None) -> None:
        """Cancel queued work and stop every task owned by this limiter."""
        self._closed = True
        self._shutdown.set()
        async with self._condition:
            queued_futures = [
                future
                for _func, futures in self._queue_map.values()
                for future in futures
            ]
            self._queue_list.clear()
            self._queue_map.clear()
            for future in queued_futures:
                if not future.done():
                    future.cancel()
            for future in self._active_futures:
                if not future.done():
                    future.cancel()
            self._condition.notify_all()

        cancellation: asyncio.CancelledError | None = None
        timeout_error: TimeoutError | None = None
        task = self._worker_task
        if task and not task.done():
            task.cancel()
            try:
                drain = asyncio.gather(task, return_exceptions=True)
                if timeout is None:
                    await drain
                else:
                    await asyncio.wait_for(drain, timeout=timeout)
            except TimeoutError as exc:
                timeout_error = exc
            except asyncio.CancelledError as exc:
                cancellation = exc
        if task is None or task.done():
            self._worker_task = None

        background_tasks = tuple(self._background_tasks)
        for background_task in background_tasks:
            background_task.cancel()
        if background_tasks:
            try:
                await asyncio.gather(*background_tasks, return_exceptions=True)
            except asyncio.CancelledError as exc:
                cancellation = exc
        self._background_tasks.difference_update(
            task for task in background_tasks if task.done()
        )

        if cancellation is not None:
            raise cancellation
        if timeout_error is not None:
            raise TimeoutError(
                "MessagingRateLimiter worker did not stop before timeout"
            ) from timeout_error

    async def _enqueue_internal(
        self,
        func: Callable[[], Awaitable[Any]],
        future: asyncio.Future[Any],
        dedup_key: str,
        front: bool = False,
    ) -> None:
        await self._enqueue_internal_multi(func, [future], dedup_key, front)

    async def _enqueue_internal_multi(
        self,
        func: Callable[[], Awaitable[Any]],
        futures: list[asyncio.Future[Any]],
        dedup_key: str,
        front: bool = False,
    ) -> None:
        async with self._condition:
            if self._closed:
                raise RuntimeError("Messaging rate limiter is closed.")
            if self._worker_task is None or self._worker_task.done():
                raise RuntimeError("Messaging rate limiter has not been started.")
            if dedup_key in self._queue_map:
                # Compaction: Update existing task with new func, append new futures
                _old_func, old_futures = self._queue_map[dedup_key]
                old_futures.extend(futures)
                self._queue_map[dedup_key] = (func, old_futures)
                logger.debug(
                    f"Compacted task for key: {dedup_key} (now {len(old_futures)} futures)"
                )
            else:
                self._queue_map[dedup_key] = (func, futures)
                if front:
                    self._queue_list.appendleft(dedup_key)
                else:
                    self._queue_list.append(dedup_key)
                self._condition.notify_all()

    async def enqueue(
        self, func: Callable[[], Awaitable[Any]], dedup_key: str | None = None
    ) -> Any:
        """
        Enqueue a messaging task and return its future result.
        If dedup_key is provided, subsequent tasks with the same key will replace this one.
        """
        self._require_running()
        if dedup_key is None:
            # Unique key to avoid deduplication
            dedup_key = f"task_{id(func)}_{asyncio.get_running_loop().time()}"

        future = asyncio.get_running_loop().create_future()
        try:
            await self._enqueue_internal(func, future, dedup_key)
        except BaseException:
            future.cancel()
            raise
        return await future

    def fire_and_forget(
        self, func: Callable[[], Awaitable[Any]], dedup_key: str | None = None
    ) -> None:
        """Enqueue a task without waiting for the result."""
        self._require_running()
        if dedup_key is None:
            dedup_key = f"task_{id(func)}_{asyncio.get_running_loop().time()}"

        async def _wrapped() -> None:
            max_retries = 2
            for attempt in range(max_retries + 1):
                try:
                    await self.enqueue(func, dedup_key)
                    return
                except Exception as e:
                    error_msg = str(e).lower()
                    # Only retry transient connectivity issues that might have slipped through
                    # or occurred between platform checks.
                    if attempt < max_retries and any(
                        x in error_msg for x in ["connect", "timeout", "broken"]
                    ):
                        wait = 2**attempt
                        if self._log_error_details:
                            logger.warning(
                                "Limiter fire_and_forget transient error (attempt {}): {}. Retrying in {}s...",
                                attempt + 1,
                                e,
                                wait,
                            )
                        else:
                            logger.warning(
                                "Limiter fire_and_forget transient error (attempt {}): exc_type={}. Retrying in {}s...",
                                attempt + 1,
                                type(e).__name__,
                                wait,
                            )
                        await asyncio.sleep(wait)
                        continue

                    logger.error(
                        "Final error in fire_and_forget for key {}: {}",
                        dedup_key,
                        format_exception_for_log(
                            e,
                            log_full_message=self._log_error_details,
                        ),
                    )
                    break

        task = asyncio.create_task(_wrapped(), name=f"msg-limiter:{dedup_key}")
        self._background_tasks.add(task)
        task.add_done_callback(self._background_tasks.discard)

    def _require_running(self) -> None:
        if self._closed:
            raise RuntimeError("Messaging rate limiter is closed.")
        if self._worker_task is None or self._worker_task.done():
            raise RuntimeError("Messaging rate limiter has not been started.")