"""Implement this class to create
a workable bucket for Limiter to use
"""

import asyncio
import logging
from abc import ABC, abstractmethod
from collections import defaultdict
from inspect import isawaitable, iscoroutine
from threading import Event, Thread
from typing import Any, Awaitable, Dict, List, Optional, Tuple, Type, Union

from ..clocks import AbstractClock, MonotonicClock
from ..utils import enforce_rate_list
from .algorithm import Decision, LogAlgorithm, SlidingWindowLog
from .rate import Rate, RateItem

logger = logging.getLogger("pyrate_limiter")


class AbstractBucket(ABC):
    """Base bucket interface
    Assumption: len(rates) always > 0
    TODO: allow empty rates
    """

    _rates: List[Rate]
    failing_rate: Optional[Rate] = None
    _clock: AbstractClock = MonotonicClock()
    # Default policy; every built-in bucket takes an ``algorithm=`` argument.
    _algorithm: LogAlgorithm = SlidingWindowLog()
    # (weight, absolute timestamp at which that weight fits), from the last
    # put(). Absolute rather than a delay so it survives a later waiting() call.
    _last_wait: Optional[Tuple[int, int]] = None
    # Whether this bucket's operations return awaitables. ``None`` means
    # "unknown" - the Leaker then probes once by calling ``leak(0)`` and
    # checking for a coroutine. Built-in sync/async buckets declare this so no
    # side-effecting probe is needed; ``RedisBucket`` leaves it ``None`` because
    # it may wrap either a sync or an async client (issue #305).
    is_async: Optional[bool] = None

    @property
    def rates(self) -> List[Rate]:
        return self._rates

    @rates.setter
    def rates(self, value: List[Rate]) -> None:
        """Sort, validate, and store the rate list.

        Both sorting and validation are enforced here, at the abstract level,
        so every bucket (built-in or custom) shares the same contract:

        - Rates are sorted by ascending interval. Several backends assume
          ``rates[-1]`` is the largest interval when leaking (e.g. Redis,
          SQLite, Postgres), but historically only some buckets sorted -
          centralizing it fixes that latent leak bug.
        - The sorted list is then validated, rejecting ill-formed rate lists
          at construction instead of silently misbehaving (issue #239).
        """
        ordered = sorted(value, key=lambda r: r.interval)
        enforce_rate_list(ordered)
        self._rates = ordered

    def now(self):
        """Retrieve current timestamp from the clock backend."""
        return self._clock.now()

    def _record(self, item: RateItem, decision: Decision) -> bool:
        """Store a put()'s verdict; return whether it was admitted.

        Every put() path must funnel through here - including trivial admits -
        so a previous denial can never stay visible.
        """
        self.failing_rate = decision.failing_rate

        if decision.retry_after_ms is None:
            self._last_wait = None
        else:
            self._last_wait = (item.weight, item.timestamp + decision.retry_after_ms)

        return decision.allowed

    def _recorded_wait(self, item: RateItem) -> Optional[int]:
        """Last put()'s retry-after, if it was for this same weight."""
        recorded = self._last_wait

        if recorded is None or recorded[0] != item.weight:
            return None

        return max(0, recorded[1] - item.timestamp)

    @abstractmethod
    def put(self, item: RateItem) -> Union[bool, Awaitable[bool]]:
        """Put an item (typically the current time) in the bucket
        return true if successful, otherwise false
        """

    def put_decision(self, item: RateItem) -> Union[Decision, Awaitable[Decision]]:
        """``put()``, returning the full ``Decision`` rather than a bare bool.

        Buckets need not override it; the default reads back what ``put()``
        recorded. ``retry_after_ms`` is ``None`` for buckets that record none.
        """
        result = self.put(item)

        if isawaitable(result):

            async def _await_decision() -> Decision:
                await result
                return self._decision(item)

            return _await_decision()

        return self._decision(item)

    def _decision(self, item: RateItem) -> Decision:
        if item.weight == 0 or self.failing_rate is None:
            return Decision()

        return Decision(failing_rate=self.failing_rate, retry_after_ms=self._recorded_wait(item))

    @abstractmethod
    def leak(
        self,
        current_timestamp: Optional[int] = None,
    ) -> Union[int, Awaitable[int]]:
        """leaking bucket - removing items that are outdated"""

    @abstractmethod
    def flush(self) -> Union[None, Awaitable[None]]:
        """Flush the whole bucket
        - Must remove `failing-rate` after flushing
        """

    @abstractmethod
    def count(self) -> Union[int, Awaitable[int]]:
        """Count number of items in the bucket"""

    @abstractmethod
    def peek(self, index: int) -> Union[Optional[RateItem], Awaitable[Optional[RateItem]]]:
        """Peek at the rate-item at a specific index in latest-to-earliest order
        NOTE: The reason we cannot peek from the start of the queue(earliest-to-latest) is
        we can't really tell how many outdated items are still in the queue
        """

    def waiting(self, item: RateItem) -> Union[int, Awaitable[int]]:
        """Calculate time until bucket become availabe to consume an item again"""
        if self.failing_rate is None:
            return 0

        assert item.weight > 0, "Item's weight must > 0"

        if item.weight > self._algorithm.max_weight(self.failing_rate):
            return -1

        recorded = self._recorded_wait(item)

        if recorded is not None:
            return recorded

        # Fallback for buckets written against the pre-4.5 contract, which
        # record nothing: derive the wait with an extra lookup.
        offset = self._algorithm.blocking_offset(self.failing_rate, item.weight)

        if offset is None:
            # Policy's wait does not depend on a stored entry (e.g. FixedWindow).
            return self._algorithm.retry_after(self.failing_rate, item.timestamp, None)

        bound_item = self.peek(offset)

        def _calc_waiting(inner_bound_item: Optional[RateItem]) -> int:
            assert self.failing_rate is not None  # NOTE: silence mypy
            blocking = None if inner_bound_item is None else inner_bound_item.timestamp
            return self._algorithm.retry_after(self.failing_rate, item.timestamp, blocking)

        async def _calc_waiting_async() -> int:
            nonlocal bound_item

            while isawaitable(bound_item):
                bound_item = await bound_item

            assert bound_item is None or isinstance(bound_item, RateItem)
            return _calc_waiting(bound_item)

        if isawaitable(bound_item):
            return _calc_waiting_async()

        assert bound_item is None or isinstance(bound_item, RateItem)
        return _calc_waiting(bound_item)

    def limiter_lock(self) -> Optional[object]:  # type: ignore
        """An additional lock to be used by Limiter in-front of the thread lock.
        Intended for multiprocessing environments where a thread lock is insufficient.
        """
        return None

    def close(self) -> None:  # noqa: B027
        """Release any resources held by the bucket.

        Subclasses may override this method to perform any necessary cleanup
        (e.g., closing files, network connections, or releasing locks) when the
        bucket is no longer needed.
        """
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()


class Leaker:
    """Responsible for scheduling buckets' leaking at the background either
    through a daemon thread (for sync buckets) or a task using asyncio.Task.

    The sync worker is an *encapsulated* daemon thread rather than ``self``
    (this class used to subclass ``Thread``). A ``Thread`` can only be started
    once, so when the worker exits after every sync bucket has been disposed,
    re-registering a bucket and calling ``start()`` again would raise
    ``RuntimeError: threads can only be started once`` (issue #301). Holding the
    thread as an attribute lets ``start()`` spin up a fresh one on demand.
    """

    name = "PyrateLimiter's Leaker"
    sync_buckets: Dict[int, AbstractBucket]
    async_buckets: Dict[int, AbstractBucket]
    leak_interval: int = 10_000
    aio_leak_task: Optional[asyncio.Task] = None
    _stop_event: Any
    _thread: Optional[Thread] = None

    def __init__(self, leak_interval: int):
        self.sync_buckets = defaultdict()
        self.async_buckets = defaultdict()
        self.leak_interval = leak_interval
        self._stop_event = Event()  # <--- add here
        self._thread = None

    def register(self, bucket: AbstractBucket):
        """Register a new bucket, routing it to the sync or async leak loop.

        Prefers the bucket's declared ``is_async``; only falls back to the
        side-effecting ``leak(0)`` probe when that is ``None`` (issue #305).
        """
        assert self.sync_buckets is not None
        assert self.async_buckets is not None

        bucket_id = id(bucket)
        is_async = bucket.is_async

        if is_async is None:
            try_leak = bucket.leak(0)
            if iscoroutine(try_leak):
                try_leak.close()
                is_async = True
            else:
                is_async = False

        if is_async:
            self.async_buckets[bucket_id] = bucket
        else:
            self.sync_buckets[bucket_id] = bucket

    def deregister(self, bucket_id: int) -> bool:
        """Deregister a bucket"""
        if self.sync_buckets and bucket_id in self.sync_buckets:
            del self.sync_buckets[bucket_id]
            return True

        if self.async_buckets and bucket_id in self.async_buckets:
            del self.async_buckets[bucket_id]

            if not self.async_buckets and self.aio_leak_task:
                self.aio_leak_task.cancel()
                self.aio_leak_task = None

            return True

        return False

    async def _leak(self, buckets: Dict[int, AbstractBucket]) -> None:
        while not self._stop_event.is_set() and buckets:
            try:
                for _, bucket in tuple(buckets.items()):
                    now = bucket.now()

                    while isawaitable(now):
                        now = await now

                    assert isinstance(now, int)
                    leak = bucket.leak(now)

                    while isawaitable(leak):
                        leak = await leak

                    assert isinstance(leak, int)

                await asyncio.sleep(self.leak_interval / 1000)
            except RuntimeError as e:
                logger.debug("Leak task stopped due to event loop shutdown. %s", e)
                return

    def leak_async(self):
        if self.async_buckets and not self.aio_leak_task:
            self.aio_leak_task = asyncio.create_task(self._leak(self.async_buckets))

    def is_alive(self) -> bool:
        """Whether the sync-leak worker thread is currently running."""
        return self._thread is not None and self._thread.is_alive()

    def _run(self) -> None:
        """Worker-thread target; not meant to be called directly."""
        assert self.sync_buckets
        asyncio.run(self._leak(self.sync_buckets))

    def start(self) -> None:
        """Start (or restart) the daemon thread that leaks sync buckets.

        The worker exits once every sync bucket is disposed; a finished
        ``Thread`` cannot be restarted, so we create a fresh one here instead of
        re-starting the dead one (which would raise ``RuntimeError``).
        """
        if self.sync_buckets and not self.is_alive():
            self._stop_event.clear()
            self._thread = Thread(target=self._run, name=self.name, daemon=True)
            self._thread.start()

    def close(self):
        self._stop_event.set()
        self.sync_buckets.clear()
        self.async_buckets.clear()


class BucketFactory(ABC):
    """Asbtract BucketFactory class.
    It is reserved for user to implement/override this class with
    his own bucket-routing/creating logic
    """

    _leaker: Optional[Leaker] = None
    _leak_interval: int = 10_000

    @property
    def leak_interval(self) -> int:
        """Retrieve leak-interval from inner Leaker task"""
        if not self._leaker:
            return self._leak_interval
        return self._leaker.leak_interval

    @leak_interval.setter
    def leak_interval(self, value: int):
        """Set leak-interval for inner Leaker task"""
        if self._leaker:
            self._leaker.leak_interval = value
        self._leak_interval = value

    @abstractmethod
    def wrap_item(
        self,
        name: str,
        weight: int = 1,
    ) -> Union[RateItem, Awaitable[RateItem]]:
        """Add the current timestamp to the receiving item using any clock backend
        - Turn it into a RateItem
        - Can return either a coroutine or a RateItem instance
        """

    @abstractmethod
    def get(self, item: RateItem) -> Union[AbstractBucket, Awaitable[AbstractBucket]]:
        """Get the corresponding bucket to this item"""

    def create(
        self,
        bucket_class: Type[AbstractBucket],
        *args,
        **kwargs,
    ) -> AbstractBucket:
        """Creating a bucket dynamically"""
        bucket = bucket_class(*args, **kwargs)
        self.schedule_leak(bucket)
        return bucket

    def schedule_leak(self, new_bucket: AbstractBucket) -> None:
        """Schedule all the buckets' leak, reset bucket's failing rate"""
        assert new_bucket.rates, "Bucket rates are not set"

        if not self._leaker:
            self._leaker = Leaker(self.leak_interval)

        self._leaker.register(new_bucket)
        self._leaker.start()
        self._leaker.leak_async()

    def get_buckets(self) -> List[AbstractBucket]:
        """Iterator over all buckets in the factory"""
        if not self._leaker:
            return []

        buckets = []

        if self._leaker.sync_buckets:
            for _, bucket in self._leaker.sync_buckets.items():
                buckets.append(bucket)

        if self._leaker.async_buckets:
            for _, bucket in self._leaker.async_buckets.items():
                buckets.append(bucket)

        return buckets

    def dispose(self, bucket: Union[int, AbstractBucket]) -> bool:
        """Delete a bucket from the factory"""
        if isinstance(bucket, AbstractBucket):
            bucket = id(bucket)

        assert isinstance(bucket, int), "not valid bucket id"

        if not self._leaker:
            return False

        return self._leaker.deregister(bucket)

    def __del__(self):
        # Make sure all leakers are deregistered
        for bucket in self.get_buckets():
            try:
                self.dispose(bucket)
            except Exception as e:
                logger.debug("Exception %s (%s) deleting bucket %r", type(e).__name__, e, bucket)

    def close(self) -> None:
        try:
            if self._leaker is not None:
                self._leaker.close()
                self._leaker = None
        except Exception as e:
            logger.info("Exception %s (%s) deleting bucket %r", type(e).__name__, e, self._leaker)

        for bucket in self.get_buckets():
            try:
                logger.debug("Closing bucket %s", bucket)
                bucket.close()
            except Exception as e:
                logger.info("Exception %s (%s) deleting bucket %r", type(e).__name__, e, bucket)

    def __getstate__(self):
        """Get state for pickling"""
        state = self.__dict__.copy()
        state.pop("_leaker", None)
        return state

    def __setstate__(self, state):
        """Restore state after unpickling"""
        self.__dict__.update(state)
        self._leaker = None
