# ai_hydra/server/HydraServer.py
#
#    AI Hydra
#    Author: Nadim-Daniel Ghaznavi
#    Copyright: (c) 2025-2026 Nadim-Daniel Ghaznavi
#    GitHub: https://github.com/NadimGhaznavi/ai_hydra
#    Website: https://ai-hydra.readthedocs.io/en/latest
#    License: GPL 3.0

import sys
import zmq
import asyncio
import signal
from typing import Callable, Optional

from ai_hydra.constants.DHydra import (
    DHydraLog,
    DHydraLogDef,
    DHydraRouterDef,
    DHydraServerDef,
    DMethod,
    DModule,
)
from ai_hydra.utils.HydraLog import HydraLog
from ai_hydra.zmq.HydraMsg import HydraMsg
from ai_hydra.zmq.HydraServerMQ import HydraServerMQ


class HydraServer:
    """
    Abstract base class for HydraServer implementations.

    Provides common ZeroMQ-based server functionality that binds to a port
    and handles client requests using the REQ/REP pattern. Subclasses must
    implement application-specific message handling logic.
    """

    def __init__(
        self,
        address: str = "*",
        port: int = DHydraServerDef.PORT,
        pub_port: int = DHydraServerDef.PUB_PORT,
        router_address: str = DHydraRouterDef.HOSTNAME,
        router_port: int = DHydraRouterDef.PORT,
        router_hb_port: int = DHydraRouterDef.HEARTBEAT_PORT,
        identity: str = DModule.HYDRA_SERVER,
        log_level: DHydraLog = DHydraLogDef.DEFAULT_LOG_LEVEL,
    ):
        """
        Initialize the HydraServer with binding parameters.

        Args:
            address (str): The address to bind to (default: "*" for all
                interfaces)
            port (int): The port to bind to
            server_id (str): Identifier for logging purposes
        """

        self.address = address
        self.port = port
        self.pub_port = pub_port
        self.router_address = router_address
        self.router_port = router_port
        self.router_hb_port = router_hb_port
        self.identity = identity
        self.log_level = log_level

        self._methods: dict[str, Callable[[HydraMsg], object]] = {}

        # Messaging stub
        self.mq: Optional[HydraServerMQ] = None

        # Structured console logs
        self.log = HydraLog(
            client_id=self.identity, log_level=log_level, to_console=True
        )

        self.log.info(f"Binding to network interface: {self.address}")
        self.log.info(f"Set control port: {self.port}")
        self.log.info(f"Set ZeroMQ PUB port: {self.pub_port}")

        # Shutdown coordination
        self._stop_event: asyncio.Event | None = None
        self._main_task: asyncio.Task[None] | None = None

    def _install_signal_handlers(self, stop_event: asyncio.Event) -> None:
        """
        Prefer signal handlers when available (Unix). On Windows,
        add_signal_handler may not be implemented; Ctrl-C will still raise
        KeyboardInterrupt.
        """
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            return

        def _request_stop() -> None:
            if not stop_event.is_set():
                self.log.info("Shutdown requested (signal)")
                stop_event.set()

        for sig in (signal.SIGINT, signal.SIGTERM):
            try:
                loop.add_signal_handler(sig, _request_stop)
            except (NotImplementedError, RuntimeError):
                # Not supported on some platforms/event loops
                pass

    def loglevel(self, log_level: DHydraLog) -> None:
        """
        Initialize console logging for the server instance.
        """
        self.log = HydraLog(
            client_id=self.identity, log_level=log_level, to_console=True
        )

    async def quit(self):
        if self.mq is not None:
            await self.mq.quit()

    async def _main_loop(self, stop_event: asyncio.Event) -> None:
        try:
            self.mq = HydraServerMQ(
                router_address=self.router_address,
                router_port=self.router_port,
                router_hb_port=self.router_hb_port,
                identity=self.identity,
                srv_methods=self._methods,
                pub_port=self.pub_port,
                log_level=self.log_level,
            )
            self.mq.start()
        except Exception:
            if self.mq is not None:
                await self.mq.quit()
                self.mq = None
            raise

        try:
            while not stop_event.is_set():
                await asyncio.sleep(1)

        except asyncio.CancelledError:
            # Normal during shutdown
            raise

    async def ping(self, msg: HydraMsg) -> None:
        self.log.info(f"Received ping from {msg.sender}")
        reply_msg = HydraMsg(
            sender=DModule.HYDRA_SERVER,
            target=msg.sender,
            method=DMethod.PONG,
        )
        if self.mq is not None:
            await self.mq.send(reply_msg)
            self.log.info(f"Sent pong to {msg.sender}")

    async def run(self) -> None:
        self._stop_event = asyncio.Event()
        stop_event = self._stop_event
        self._install_signal_handlers(stop_event)

        self.log.info("Initialized")
        self._main_task = asyncio.create_task(
            self._main_loop(stop_event), name="hydra-server-main"
        )
        wait_stop_task = asyncio.create_task(
            stop_event.wait(), name="wait-stop"
        )

        try:
            done, pending = await asyncio.wait(
                [wait_stop_task, self._main_task],
                return_when=asyncio.FIRST_COMPLETED,
            )

            if self._main_task in done:
                exc = self._main_task.exception()
                if exc is not None:
                    self.log.critical(
                        f"Main task failed: {exc!r}",
                    )
                    raise exc

        finally:
            wait_stop_task.cancel()
            await asyncio.gather(wait_stop_task, return_exceptions=True)
            await self._shutdown()

    async def _shutdown(self) -> None:
        """Graceful shutdown: stop main task, close MQ, cancel stragglers."""
        self.log.info("Shutting down...")

        main_task = self._main_task
        self._main_task = None

        # Drain the main task, but do not re-raise its exception here.
        if main_task is not None:
            if not main_task.done():
                main_task.cancel()
            try:
                await main_task
            except asyncio.CancelledError:
                pass
            except Exception as e:
                self.log.debug(
                    f"Ignoring main task exception during shutdown: {e!r}"
                )

        # Close MQ if still present
        if self.mq is not None:
            try:
                await self.mq.quit()
            except Exception as e:
                self.log.debug(f"Ignoring MQ shutdown exception: {e!r}")
            finally:
                self.mq = None

        # Cancel any other pending tasks, excluding this one
        current = asyncio.current_task()
        pending = [
            t for t in asyncio.all_tasks() if t is not current and not t.done()
        ]

        for t in pending:
            t.cancel()

        if pending:
            await asyncio.gather(*pending, return_exceptions=True)

        self.log.info("Shutdown complete.")
