"""Executor group management for HA / failover.

An executor group is a set of qRemoteX instances that share the same
``QRAPTOR_EXECUTOR_GROUP`` name.  The **gateway** performs group‐aware
routing (round‐robin / least‐loaded among group members, failover when
a member drops), but the **client** has a few responsibilities:

1. Declare its group in the ``hello`` message (handled by config + ConnectionManager).
2. Handle the ``drain`` command from the gateway — stop accepting new
   tasks, let in‐flight jobs finish, then disconnect (or stay connected
   as a warm standby).
3. Report drain state in heartbeats so the gateway does not route new
   work to a draining executor.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Callable, Coroutine

logger = logging.getLogger("qremotex.executor_group")


class DrainController:
    """Manages the drain lifecycle for an executor.

    When the gateway sends a ``drain`` message the controller:
    1. Sets the draining flag (exposed to connection manager).
    2. Waits until all in‐flight jobs complete (polled via ``active_jobs_fn``).
    3. Calls ``on_drained`` callback so the agent can disconnect or idle.

    The gateway will stop routing new tasks once it sees
    ``"draining": true`` in subsequent heartbeats.
    """

    def __init__(
        self,
        active_jobs_fn: Callable[[], int],
        on_drained: Callable[[], Coroutine[Any, Any, None]],
        *,
        poll_interval_sec: float = 2.0,
    ) -> None:
        self._active_jobs_fn = active_jobs_fn
        self._on_drained = on_drained
        self._poll_interval = poll_interval_sec
        self._draining = False
        self._drain_task: asyncio.Task | None = None  # type: ignore[type-arg]

    @property
    def is_draining(self) -> bool:
        return self._draining

    def handle_drain_message(self, msg: dict[str, Any]) -> None:
        """Called when the gateway sends a ``drain`` message."""
        if self._draining:
            logger.info("Already draining — ignoring duplicate drain command")
            return

        self._draining = True
        reason = msg.get("reason", "gateway request")
        logger.info("Drain command received (reason: %s). Waiting for in-flight jobs to finish.", reason)
        self._drain_task = asyncio.create_task(self._wait_and_notify())

    async def _wait_and_notify(self) -> None:
        """Poll until all active jobs complete, then fire callback."""
        while self._active_jobs_fn() > 0:
            active = self._active_jobs_fn()
            logger.info("Draining — %d job(s) still active", active)
            await asyncio.sleep(self._poll_interval)

        logger.info("All jobs complete — executor fully drained")
        try:
            await self._on_drained()
        except Exception:
            logger.exception("Error in on_drained callback")
