"""A bucket using PostgreSQL as backend"""

from __future__ import annotations

import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Awaitable, List, Optional, Union

from ..abstracts import AbstractBucket, Rate, RateItem
from ..clocks import PostgresClock

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    from psycopg_pool import ConnectionPool  # type: ignore[import-untyped]


class Queries:
    CREATE_BUCKET_TABLE = """
    CREATE TABLE IF NOT EXISTS {table} (
        name VARCHAR,
        weight SMALLINT,
        item_timestamp TIMESTAMP
    )
    """
    CREATE_INDEX_ON_TIMESTAMP = """
    CREATE INDEX IF NOT EXISTS {index} ON {table} (item_timestamp)
    """
    LOCK_TABLE = """
    LOCK TABLE {table} IN EXCLUSIVE MODE NOWAIT
    """
    COUNT = """
    SELECT COUNT(*) FROM {table}
    """
    PUT = """
    INSERT INTO {table} (name, weight, item_timestamp) VALUES (%s, %s, TO_TIMESTAMP(%s))
    """
    FLUSH = """
    DELETE FROM {table}
    """
    PEEK = """
    SELECT name, weight, (extract(epoch FROM item_timestamp) * 1000) as item_timestamp
    FROM {table}
    ORDER BY item_timestamp DESC
    LIMIT 1
    OFFSET {offset}
    """
    LEAK = """
    DELETE FROM {table} WHERE item_timestamp < TO_TIMESTAMP({timestamp})
    """
    LEAK_COUNT = """
    SELECT COUNT(*) FROM {table} WHERE item_timestamp < TO_TIMESTAMP({timestamp})
    """


class PostgresBucket(AbstractBucket):
    table: str
    pool: ConnectionPool

    def __init__(self, pool: ConnectionPool, table: str, rates: List[Rate]):
        self._clock = PostgresClock(pool)
        self.table = table.lower()
        self.pool = pool
        assert rates
        self.rates = rates
        self._full_tbl = f"ratelimit___{self.table}"
        self._create_table()

    @contextmanager
    def _get_conn(self):
        with self.pool.connection() as conn:
            yield conn

    def _create_table(self):
        with self._get_conn() as conn:
            lock_id = hash(self._full_tbl) & 0x7FFFFFFF
            conn.execute("SELECT pg_advisory_xact_lock(%s)", (lock_id,))
            conn.execute(Queries.CREATE_BUCKET_TABLE.format(table=self._full_tbl))
            index_name = f"timestampIndex_{self.table}"
            conn.execute(Queries.CREATE_INDEX_ON_TIMESTAMP.format(table=self._full_tbl, index=index_name))

    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
        """
        from psycopg.errors import LockNotAvailable

        if item.weight == 0:
            return True

        item_ts_seconds = item.timestamp / 1000

        with self._get_conn() as conn:
            # Acquire an EXCLUSIVE MODE lock on the bucket table using NOWAIT.
            # This ensures the "check current count" + "insert new items" sequence
            # is atomic with respect to other writers, so rate limits cannot be
            # exceeded due to concurrent requests interleaving.
            #
            # Because we use NOWAIT, if the table is already locked by another
            # transaction, PostgreSQL raises LockNotAvailable and we immediately
            # reject this request (return False) instead of blocking or retrying.
            # This provides predictable, fail-fast behavior but may limit
            # throughput under high contention since only one writer can perform
            # the check-and-put at a time.
            try:
                conn.execute(Queries.LOCK_TABLE.format(table=self._full_tbl))
            except LockNotAvailable:
                logger.debug("LockNotAvailable")
                self.failing_rate = self.rates[0]
                return False

            for rate in self.rates:
                cur = conn.execute(
                    f"SELECT COUNT(*) FROM {self._full_tbl} "  # noqa: S608
                    f"WHERE item_timestamp >= TO_TIMESTAMP(%s) - (%s * INTERVAL '1 milliseconds')",
                    (item_ts_seconds, rate.interval),
                )
                count = int(cur.fetchone()[0])
                cur.close()

                if rate.limit - count < item.weight:
                    self.failing_rate = rate
                    return False

            self.failing_rate = None
            query = Queries.PUT.format(table=self._full_tbl)
            for _ in range(item.weight):
                conn.execute(query, (item.name, item.weight, item_ts_seconds))

        return True

    def leak(
        self,
        current_timestamp: Optional[int] = None,
    ) -> Union[int, Awaitable[int]]:
        """leaking bucket - removing items that are outdated"""
        assert current_timestamp is not None, "current-time must be passed on for leak"
        lower_bound = current_timestamp - self.rates[-1].interval

        if lower_bound <= 0:
            return 0

        count = 0

        with self._get_conn() as conn:
            conn = conn.execute(Queries.LEAK_COUNT.format(table=self._full_tbl, timestamp=lower_bound / 1000))
            result = conn.fetchone()

            if result:
                conn.execute(Queries.LEAK.format(table=self._full_tbl, timestamp=lower_bound / 1000))
                count = int(result[0])

        return count

    def flush(self) -> Union[None, Awaitable[None]]:
        """Flush the whole bucket
        - Must remove `failing-rate` after flushing
        """
        with self._get_conn() as conn:
            conn.execute(Queries.FLUSH.format(table=self._full_tbl))
            self.failing_rate = None

        return None

    def count(self) -> Union[int, Awaitable[int]]:
        """Count number of items in the bucket"""
        count = 0
        with self._get_conn() as conn:
            conn = conn.execute(Queries.COUNT.format(table=self._full_tbl))
            result = conn.fetchone()
            assert result
            count = int(result[0])

        return count

    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
        """
        item = None

        with self._get_conn() as conn:
            conn = conn.execute(Queries.PEEK.format(table=self._full_tbl, offset=index))
            result = conn.fetchone()
            if result:
                name, weight, timestamp = result[0], int(result[1]), int(result[2])
                item = RateItem(name=name, weight=weight, timestamp=timestamp)

        return item

    def close(self):
        if self.pool is not None and not self.pool.closed:
            try:
                self.pool.close()
            except Exception as e:
                logger.debug("Exception closing pool, %s", e)