import asyncio
import functools
import logging
from typing import TYPE_CHECKING, Any, Awaitable, Callable

from redis.asyncio.observability.recorder import (
    record_connection_handoff,
    record_connection_relaxed_timeout,
    record_maint_notification_count,
)
from redis.maint_notifications import (
    MaintenanceNotification,
    MaintenanceState,
    MaintNotificationsConfig,
    NodeMovingNotification,
    OSSNodeMigratedNotification,
    OSSNodeMigratingNotification,
    _get_maintenance_notification_name,
    _get_maintenance_notification_type,
    _should_skip_connection_timeout_update,
)
from redis.observability.attributes import get_pool_name

if TYPE_CHECKING:
    from redis.asyncio.cluster import RedisCluster
    from redis.asyncio.connection import AsyncMaintNotificationsAbstractConnection

logger = logging.getLogger(__name__)

_ScheduledCallback = Callable[..., Awaitable[None]]


def _log_task_exception(message: str, task: "asyncio.Task[Any]") -> None:
    """Task done-callback that surfaces a failed task's exception in the logs.

    Without it, an unhandled exception in a fire-and-forget task is only
    reported by asyncio as a noisy "Task exception was never retrieved"
    warning. Retrieving the exception here suppresses that warning and logs a
    meaningful error instead. Cancellation is expected and left unlogged.

    Bind ``message`` with ``functools.partial`` before passing to
    ``add_done_callback`` (which supplies ``task``).
    """
    try:
        exc = task.exception()
    except asyncio.CancelledError:
        return
    if exc:
        logger.error(message, exc_info=exc)


def add_debug_log_for_notification(
    connection: object,
    notification: str | MaintenanceNotification,
) -> None:
    if not logger.isEnabledFor(logging.DEBUG):
        return

    socket_address = None
    try:
        writer = getattr(connection, "_writer", None)
        socket_name = writer.get_extra_info("sockname") if writer else None
        socket_address = (
            socket_name[1] if socket_name and len(socket_name) > 1 else None
        )
    except (AttributeError, OSError, TypeError):
        pass

    resolved_ip = None
    try:
        get_resolved_ip = getattr(connection, "get_resolved_ip", None)
        if callable(get_resolved_ip):
            resolved_ip = get_resolved_ip()
    except (AttributeError, OSError):
        pass

    logger.debug(
        f"Handling maintenance notification: {notification}, "
        f"with connection: {connection}, connected to ip {resolved_ip}, "
        f"local socket port: {socket_address}",
    )


class AsyncMaintNotificationsPoolHandler:
    def __init__(
        self,
        pool: Any,
        config: MaintNotificationsConfig,
    ) -> None:
        self.pool = pool
        self.config = config
        self._processed_notifications: set[MaintenanceNotification] = set()
        self._scheduled_tasks: set[asyncio.Task[None]] = set()
        self._lock = asyncio.Lock()
        self.connection: Any | None = None

    def set_connection(
        self, connection: "AsyncMaintNotificationsAbstractConnection"
    ) -> None:
        self.connection = connection

    def get_handler_for_connection(self) -> "AsyncMaintNotificationsPoolHandler":
        # Copy all data that should be shared between connections, while each
        # connection gets its own handler instance and current connection state.
        copy = AsyncMaintNotificationsPoolHandler(self.pool, self.config)
        copy._processed_notifications = self._processed_notifications
        copy._scheduled_tasks = self._scheduled_tasks
        copy._lock = self._lock
        copy.connection = None
        return copy

    async def remove_expired_notifications(self) -> None:
        async with self._lock:
            for notification in tuple(self._processed_notifications):
                if notification.is_expired():
                    self._processed_notifications.remove(notification)

    async def handle_notification(self, notification: MaintenanceNotification) -> None:
        await self.remove_expired_notifications()

        if isinstance(notification, NodeMovingNotification):
            await self.handle_node_moving_notification(notification)
        else:
            logger.error(f"Unhandled notification type: {notification}")

    async def handle_node_moving_notification(
        self, notification: NodeMovingNotification
    ) -> None:
        if (
            not self.config.proactive_reconnect
            and not self.config.is_relaxed_timeouts_enabled()
        ):
            return

        async with self._lock:
            if notification in self._processed_notifications:
                # nothing to do in the connection pool handling
                # the notification has already been handled or is expired
                # just return
                return
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(
                    f"Handling node MOVING notification: {notification}, "
                    f"with connection: {self.connection}, connected to ip "
                    f"{self.connection.get_resolved_ip() if self.connection else None}"
                )
            # Get the current connected address - if any
            # This is the address that is being moved
            # and we need to handle only connections
            # connected to the same address
            moving_address_src = (
                self.connection.getpeername() if self.connection else None
            )

            # The async pool owns the active/free connection collections and
            # asyncio.Lock is not reentrant, so the whole MOVING pool mutation
            # has to be one pool-owned atomic operation. The handler still owns
            # the notification policy and passes the already-decided inputs.
            await self.pool.apply_moving_notification(
                notification=notification,
                config=self.config,
                moving_address_src=moving_address_src,
                run_proactive_reconnect=(
                    self.config.proactive_reconnect
                    and notification.new_node_host is not None
                ),
            )

            if self.config.proactive_reconnect and notification.new_node_host is None:
                self._schedule(
                    notification.ttl / 2,
                    self.run_proactive_reconnect,
                    moving_address_src,
                )

            self._schedule(
                notification.ttl,
                self.handle_node_moved_notification,
                notification,
            )

            await record_connection_handoff(
                pool_name=get_pool_name(self.pool),
            )

            self._processed_notifications.add(notification)

    async def run_proactive_reconnect(
        self, moving_address_src: str | None = None
    ) -> None:
        """
        Run proactive reconnect for the pool.
        Active connections are marked for reconnect after they complete the current command.
        Inactive connections are disconnected and will be connected on next use.
        """
        async with self._lock:
            # This delayed reconnect must be atomic for the same reason as the
            # initial MOVING mutation: a connection can move between active and
            # free lists while another task is acquiring or releasing it.
            await self.pool.run_proactive_reconnect(
                moving_address_src=moving_address_src,
            )

    async def handle_node_moved_notification(
        self, notification: NodeMovingNotification
    ) -> None:
        """
        Handle the cleanup after a node moving notification expires.
        """
        notification_hash = hash(notification)

        async with self._lock:
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(
                    f"Reverting temporary changes related to notification: {notification}, "
                    f"with connection: {self.connection}, connected to ip "
                    f"{self.connection.get_resolved_ip() if self.connection else None}"
                )
            reset_relaxed_timeout = self.config.is_relaxed_timeouts_enabled()
            reset_host_address = self.config.proactive_reconnect

            # Cleanup has to reset future connection kwargs and existing
            # matching connections together under the pool lock. Splitting it
            # lets an acquire/release interleave and leaves stale MOVING state.
            await self.pool.cleanup_moving_notification(
                notification_hash=notification_hash,
                reset_relaxed_timeout=reset_relaxed_timeout,
                reset_host_address=reset_host_address,
            )

    def _schedule(
        self,
        delay: float,
        callback: _ScheduledCallback,
        *args: Any,
    ) -> None:
        # Record the absolute deadline now so that any lag between create_task
        # and the task's first execution slice does not push the fire time out.
        deadline = asyncio.get_running_loop().time() + delay
        task = asyncio.create_task(self._run_after(deadline, callback, *args))
        self._scheduled_tasks.add(task)
        task.add_done_callback(self._scheduled_tasks.discard)
        task.add_done_callback(
            functools.partial(
                _log_task_exception,
                "Error handling scheduled maintenance notification",
            )
        )

    async def _run_after(
        self,
        deadline: float,
        callback: _ScheduledCallback,
        *args: Any,
    ) -> None:
        remaining = deadline - asyncio.get_running_loop().time()
        if remaining > 0:
            await asyncio.sleep(remaining)
        await callback(*args)

    async def cancel_scheduled_tasks(self) -> None:
        if not self._scheduled_tasks:
            return
        tasks = tuple(self._scheduled_tasks)
        for task in tasks:
            task.cancel()
        await asyncio.gather(*tasks, return_exceptions=True)


class AsyncMaintNotificationsConnectionHandler:
    def __init__(
        self,
        connection: "AsyncMaintNotificationsAbstractConnection",
        config: MaintNotificationsConfig,
    ) -> None:
        self.connection = connection
        self.config = config

    def _get_pool_name(self) -> str:
        """
        Get the pool name from the connection's pool handler.
        Falls back to connection representation if pool is not available.
        """
        pool_handler = getattr(
            self.connection, "_maint_notifications_pool_handler", None
        )
        if pool_handler and getattr(pool_handler, "pool", None):
            return get_pool_name(pool_handler.pool)
        # Fallback for standalone connections without a pool
        return repr(self.connection)

    async def handle_notification(self, notification: MaintenanceNotification) -> None:
        # 1 for start, 0 for end notification type, None for unknown.
        notification_type = _get_maintenance_notification_type(notification)
        maint_notification = _get_maintenance_notification_name(notification)

        await record_maint_notification_count(
            server_address=self.connection.host,
            server_port=self.connection.port,
            network_peer_address=self.connection.host,
            network_peer_port=self.connection.port,
            maint_notification=maint_notification,
        )

        if notification_type is None:
            logger.error(f"Unhandled notification type: {notification}")
            return

        if notification_type:
            await self.handle_maintenance_start_notification(
                MaintenanceState.MAINTENANCE, notification
            )
        else:
            await self.handle_maintenance_completed_notification(
                notification=notification
            )

    async def handle_maintenance_start_notification(
        self,
        maintenance_state: MaintenanceState,
        notification: MaintenanceNotification,
    ) -> None:
        add_debug_log_for_notification(self.connection, notification)

        if _should_skip_connection_timeout_update(
            self.connection.maintenance_state, self.config
        ):
            return

        self.connection.maintenance_state = maintenance_state
        self.connection.set_tmp_settings(
            tmp_relaxed_timeout=self.config.relaxed_timeout
        )
        self.connection.update_current_socket_timeout(self.config.relaxed_timeout)
        if isinstance(notification, OSSNodeMigratingNotification):
            # add the notification id to the set of processed start maint notifications
            # this is used to skip the unrelaxing of the timeouts if we have received more than
            # one start notification before the the final end notification
            self.connection.add_maint_start_notification(notification.id)

        maint_notification = _get_maintenance_notification_name(notification)
        await record_connection_relaxed_timeout(
            connection_name=self._get_pool_name(),
            maint_notification=maint_notification,
            relaxed=True,
        )

    async def handle_maintenance_completed_notification(self, **kwargs: Any) -> None:
        # Only reset timeouts if state is not MOVING and relaxed timeouts are enabled
        if _should_skip_connection_timeout_update(
            self.connection.maintenance_state, self.config
        ):
            return

        notification = None
        if kwargs.get("notification"):
            notification = kwargs["notification"]
        add_debug_log_for_notification(
            self.connection, notification if notification else "MAINTENANCE_COMPLETED"
        )
        self.connection.reset_tmp_settings(reset_relaxed_timeout=True)
        # Maintenance completed - reset the connection
        # timeouts by providing -1 as the relaxed timeout
        self.connection.update_current_socket_timeout(-1)
        self.connection.maintenance_state = MaintenanceState.NONE
        # reset the sets that keep track of received start maint
        # notifications and skipped end maint notifications
        self.connection.reset_received_notifications()

        if notification:
            maint_notification = _get_maintenance_notification_name(notification)
            await record_connection_relaxed_timeout(
                connection_name=self._get_pool_name(),
                maint_notification=maint_notification,
                relaxed=False,
            )


class AsyncOSSMaintNotificationsHandler:
    """
    Cluster-wide handler for OSS (open-source) cluster maintenance notifications.

    Reacts to SMIGRATED (slot migration completed) push notifications and
    triggers topology re-initialization via nodes_manager.initialize().

    Lock discipline: _lock is held across the whole pool mutation, including
    await initialize(), so the topology refresh and the subsequent connection
    marking/disconnect happen as one atomic operation — releasing the lock
    mid-mutation would let other tasks observe partial state.

    Re-entrancy is not a deadlock risk here even though asyncio.Lock is
    non-reentrant: initialize() may dispatch commands whose responses carry new
    push notifications, but handle_notification schedules the actual handling as
    a separate background task rather than calling into it inline, so the
    re-entrant arrival never tries to re-acquire the lock on this call stack.
    The cheap _in_progress/_processed dedup that gates that scheduling runs
    without the lock — those sets are only mutated from the single event loop.
    """

    def __init__(
        self,
        cluster_client: "RedisCluster",
        config: MaintNotificationsConfig,
    ) -> None:
        self.cluster_client = cluster_client
        self.config = config
        self._processed_notifications: set[MaintenanceNotification] = set()
        self._in_progress: set[MaintenanceNotification] = set()
        self._lock = asyncio.Lock()
        self._background_tasks: set[asyncio.Task] = set()

    async def remove_expired_notifications(self) -> None:
        async with self._lock:
            for n in tuple(self._processed_notifications):
                if n.is_expired():
                    self._processed_notifications.remove(n)

    async def handle_notification(self, notification: MaintenanceNotification) -> None:
        # Synchronous pre-dedup BEFORE scheduling a task.
        #
        # The same SMIGRATED notification is delivered by the server on every
        # connection, so under load this callback fires many times for the same
        # notification. Without this guard we would schedule one background task
        # per arrival - a flood of tasks that each acquire the lock and dedup,
        # saturating the single event loop. Reserving the notification in
        # _in_progress here (and skipping if already in-progress/processed) caps
        # it at exactly ONE handling task per unique notification.
        #
        # These sets are only mutated from the single event loop, so this
        # check-and-add is race-free without holding the lock.
        if (
            notification in self._in_progress
            or notification in self._processed_notifications
        ):
            return
        self._in_progress.add(notification)

        # Schedule as a background task so the parser's read path is not blocked.
        # This also breaks the inline call chain that would otherwise deadlock:
        # initialize() dispatches commands whose responses may carry more push
        # notifications; as a separate task it can run while this one awaits.
        #
        # If scheduling the task (or wiring its callbacks) fails, release the
        # in-progress reservation made above. Otherwise the notification would be
        # stuck in _in_progress forever - the dedup guard would skip it on every
        # future arrival, and _do_handle_notification's finally (which normally
        # clears it) never runs because the task was never started.
        try:
            task = asyncio.get_running_loop().create_task(
                self._do_handle_notification(notification)
            )
            self._background_tasks.add(task)
            task.add_done_callback(self._background_tasks.discard)
            task.add_done_callback(
                functools.partial(
                    _log_task_exception,
                    "Error handling maintenance notification background task",
                )
            )
        except Exception:
            self._in_progress.discard(notification)
            raise

    async def _do_handle_notification(
        self, notification: MaintenanceNotification
    ) -> None:
        try:
            if isinstance(notification, OSSNodeMigratedNotification):
                await self.handle_oss_maintenance_completed_notification(notification)
            else:
                logger.error(f"Unhandled notification type: {notification}")
        finally:
            # Release the in-progress reservation. On success the notification is
            # also in _processed_notifications (so it won't be re-handled); on
            # failure it is not, allowing a later retry.
            self._in_progress.discard(notification)

    async def handle_oss_maintenance_completed_notification(
        self, notification: OSSNodeMigratedNotification
    ) -> None:
        await self.remove_expired_notifications()

        async with self._lock:
            # handle_notification already reserved this notification in
            # _in_progress and guaranteed uniqueness; the processed check here is
            # defensive (e.g. across handler copies sharing the same sets).
            if notification in self._processed_notifications:
                return
            if logger.isEnabledFor(logging.DEBUG):
                logger.debug(f"Handling SMIGRATED notification: {notification}")

            # Extract the information about the src and destination nodes that are
            # affected by the maintenance. nodes_to_slots_mapping structure:
            # {
            #     "src_host:port": [
            #         {"dest_host:port": "slot_range"},
            #         ...
            #     ],
            #     ...
            # }
            additional_startup_nodes_info = []
            affected_nodes = set()
            for (
                src_address,
                dest_mappings,
            ) in notification.nodes_to_slots_mapping.items():
                src_host, src_port = src_address.rsplit(":", 1)
                src_node = self.cluster_client.nodes_manager.get_node(
                    host=src_host, port=int(src_port)
                )
                if src_node is not None:
                    affected_nodes.add(src_node)
                for dest_mapping in dest_mappings:
                    for dest_address in dest_mapping.keys():
                        dest_host, dest_port = dest_address.rsplit(":", 1)
                        additional_startup_nodes_info.append(
                            (dest_host, int(dest_port))
                        )
            # Updates the cluster slots cache with the new slots mapping
            # This will also update the nodes cache with the new nodes mapping
            await self.cluster_client.nodes_manager.initialize(
                additional_startup_nodes_info=additional_startup_nodes_info,
            )

            all_nodes = set(affected_nodes)
            all_nodes = all_nodes.union(
                self.cluster_client.nodes_manager.nodes_cache.values()
            )
            for current_node in all_nodes:
                handoff_recorded = False
                if current_node in affected_nodes:
                    # mark for reconnect all in-use connections to the node — this
                    # forces them to disconnect after completing their current command
                    free_set = set(current_node._free)
                    for conn in current_node._connections:
                        if conn not in free_set:
                            add_debug_log_for_notification(
                                conn, "SMIGRATED - mark for reconnect"
                            )
                            conn.mark_for_reconnect()
                    await record_connection_handoff(
                        pool_name=f"{current_node.host}:{current_node.port}"
                    )
                    handoff_recorded = True
                else:
                    if logger.isEnabledFor(logging.DEBUG):
                        logger.debug(
                            f"SMIGRATED: Node {current_node.name} not affected "
                            f"by maintenance, skipping mark for reconnect"
                        )
                if (
                    current_node
                    not in self.cluster_client.nodes_manager.nodes_cache.values()
                ):
                    task = asyncio.get_running_loop().create_task(
                        current_node.disconnect_free_connections()
                    )
                    self._background_tasks.add(task)
                    task.add_done_callback(self._background_tasks.discard)
                    task.add_done_callback(
                        functools.partial(
                            _log_task_exception,
                            "Error disconnecting free connections after "
                            "maintenance notification",
                        )
                    )
                    if not handoff_recorded:
                        await record_connection_handoff(
                            pool_name=f"{current_node.host}:{current_node.port}"
                        )

            self._processed_notifications.add(notification)
