from __future__ import annotations

import asyncio
import json
import threading
import uuid
from typing import Any
from typing import Callable
from typing import Coroutine
from typing import Union

import zmq
import zmq.asyncio

from egse.env import bool_env
from egse.log import logging
from egse.registry import DEFAULT_RS_HB_PORT
from egse.registry import DEFAULT_RS_PUB_PORT
from egse.registry import DEFAULT_RS_REQ_PORT
from egse.registry import MessageType
from egse.system import do_every

VERBOSE_DEBUG = bool_env("VERBOSE_DEBUG")

REQUEST_TIMEOUT = 0.5
"""time to wait for a request-reply [seconds]. For responsive message handling."""
EVENT_POLL_TIMEOUT = 0.5
"""time to wait for while listening for events [seconds]."""
TASK_COMPLETION_TIMEOUT = 5.0
"""time to wait for Tasks to complete [seconds]."""
HEALTH_CHECK_TIMEOUT = 2.0
"""Health checks shouldn't take too long. Recommended: 1-5s."""
HEART_BEAT_TIMEOUT = 1.0
"""Heart beat should fail fast to detect network issues. Recommended: 0.5-2s."""
HEART_BEAT_RECONNECT = 1.0
"""Time to wait between a disconnect an a connect to the heartbeat socket."""
THREAD_JOIN_TIMEOUT = 10.0
"""Joining threads should be longer, but the key is balancing graceful cleanup with reasonable shutdown times."""
PROCESS_SHUTDOWN_TIMEOUT = 10.0
"""Waiting for a process shutdown confirmation message."""


class RegistryClient:
    """
    Synchronous client for the service registry.
    """

    def __init__(
        self,
        registry_req_endpoint: str | None = None,
        registry_sub_endpoint: str | None = None,
        registry_hb_endpoint: str | None = None,
        timeout: float = REQUEST_TIMEOUT,
        client_id: str = "registry-client",
    ):
        """
        Initialize the async registry client.

        Args:
            registry_req_endpoint: ZeroMQ endpoint for REQ-REP socket, defaults to DEFAULT_RS_REQ_PORT on localhost.
            registry_sub_endpoint: ZeroMQ endpoint for SUB socket, defaults to DEFAULT_RS_PUB_PORT on localhost.
            registry_hb_endpoint: ZeroMQ endpoint for Heartbeat socket, defaults to DEFAULT_RS_HB_PORT on localhost.
            timeout: Timeout for requests in seconds, defaults to 0.5s.
            client_id: client identification, default='registry-client'
        """
        self.registry_req_endpoint = registry_req_endpoint or f"tcp://localhost:{DEFAULT_RS_REQ_PORT}"
        self.registry_sub_endpoint = registry_sub_endpoint or f"tcp://localhost:{DEFAULT_RS_PUB_PORT}"
        self.registry_hb_endpoint = registry_hb_endpoint or f"tcp://localhost:{DEFAULT_RS_HB_PORT}"

        self.timeout = timeout
        self.timeout_ms = int(timeout * 1000)
        self.logger = logging.getLogger("egse.registry.client")

        # Service state
        self._service_id = None
        self._service_info = None
        self._ttl = None

        self._hb_stop_event = None
        self._hb_thread: threading.Thread | None = None

        self._client_id = f"{client_id}-{uuid.uuid4()}".encode()

        self.context: zmq.Context = zmq.Context.instance()

        self.req_socket: zmq.Socket | None = None
        self.sub_socket: zmq.Socket | None = None
        self.hb_socket: zmq.Socket | None = None

        self.poller = zmq.Poller()

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

    def connect(self):
        if VERBOSE_DEBUG:
            self.logger.debug("Connecting to service registry...")

        # REQ socket for request-reply pattern
        self.req_socket = self.context.socket(zmq.DEALER)
        self.req_socket.setsockopt(zmq.LINGER, 0)  # Don't wait for unsent messages on close()
        self.req_socket.setsockopt(zmq.IDENTITY, self._client_id)
        self.req_socket.connect(self.registry_req_endpoint)

        # SUB socket for receiving events
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.connect(self.registry_sub_endpoint)
        # Default to receiving all events
        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")

        self.poller.register(self.req_socket, zmq.POLLIN)
        self.poller.register(self.sub_socket, zmq.POLLIN)

    def _connect_hb_socket(self):
        self.hb_socket = self.context.socket(zmq.DEALER)
        self.hb_socket.setsockopt(zmq.LINGER, 0)  # Don't wait for unsent messages on close()
        self.hb_socket.setsockopt(zmq.IDENTITY, self._client_id)
        self.hb_socket.connect(self.registry_hb_endpoint)

    def disconnect(self):
        if VERBOSE_DEBUG:
            self.logger.debug("Disconnecting from service registry...")

        if self.req_socket:
            self.poller.unregister(self.req_socket)
            self.req_socket.close(linger=0)
        self.req_socket = None

        if self.sub_socket:
            self.poller.unregister(self.sub_socket)
            self.sub_socket.close(linger=0)
        self.sub_socket = None

    def _disconnect_hb_socket(self):
        if self.hb_socket:
            self.hb_socket.setsockopt(zmq.LINGER, 0)
            self.hb_socket.close()
        self.hb_socket = None

    def _send_request(self, msg_type: MessageType, request: dict[str, Any], timeout: float = None) -> dict[str, Any]:
        """
        Send a request to the registry and get the response.

        Args:
            msg_type: the type of message and reply
            request: The request to send
            timeout: The number of seconds to wait, if None, instance variable is used.

        Returns:
            The response from the registry.
        """
        timeout_ms = int(timeout * 1000) if timeout else self.timeout_ms

        try:
            if VERBOSE_DEBUG:
                self.logger.debug(f"Sending request: {request}")

            self.req_socket.send_multipart([msg_type.value, json.dumps(request).encode()])

            if msg_type == MessageType.REQUEST_NO_REPLY:
                return {"success": True}

            if self.poller.poll(timeout=timeout_ms):
                message_parts = self.req_socket.recv_multipart()

                if len(message_parts) >= 2:
                    message_type = MessageType(message_parts[0])
                    message_data = message_parts[1]

                    if message_type == MessageType.RESPONSE:
                        response = json.loads(message_data)
                        if VERBOSE_DEBUG:
                            self.logger.debug(f"Received response: {response}")
                        return response
                    else:
                        return {
                            "success": False,
                            "error": f"unexpected MessageType received: {message_type.name}, {message_data = }",
                        }
                else:
                    return {
                        "success": False,
                        "error": f"not enough parts received: {len(message_parts)}",
                        "data": message_parts,
                    }
            else:
                self.logger.error(f"Request timed out after {self.timeout:.2f}s")
                return {"success": False, "error": "Request timed out"}
        except zmq.ZMQError as exc:
            self.logger.error(f"ZMQ error: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}
        except Exception as exc:
            self.logger.error(f"Error sending request: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}

    def _send_heartbeat(self, request: dict[str, Any]) -> dict[str, Any]:
        """
        Send a heartbeat to the registry and get the response.

        Args:
            request: The request to send

        Returns:
            The response from the registry.
        """

        try:
            if VERBOSE_DEBUG:
                self.logger.debug(f"Sending heartbeat request: {request}")

            self.hb_socket.send_string(json.dumps(request))

            ready_sockets, _, _ = zmq.select([self.hb_socket], [], [], timeout=HEART_BEAT_TIMEOUT)

            if ready_sockets:
                response_json = self.hb_socket.recv_string()
                response = json.loads(response_json)
                if VERBOSE_DEBUG:
                    self.logger.debug(f"Received response: {response}")
                return response
            else:
                self.logger.error(f"Request timed out after {HEART_BEAT_TIMEOUT:.2f}s")
                return {"success": False, "error": "Request timed out"}

        except zmq.ZMQError as exc:
            self.logger.error(f"ZMQ error: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}
        except Exception as exc:
            self.logger.error(f"Error sending request: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}

    def register(
        self,
        name: str,
        host: str,
        port: int,
        service_type: str | None = None,
        metadata: dict[str, Any] | None = None,
        ttl: int = 30,
    ) -> str | None:
        """
        Register this service with the registry.

        Args:
            name: Service name
            host: Service host/IP
            port: Service port
            service_type: Service type (for discovery)
            metadata: Additional service metadata
            ttl: Time-to-live in seconds

        Returns:
            The service ID if successful, None otherwise
        """
        # Prepare service info
        service_info: dict[str, str | int | dict | list] = {"name": name, "host": host, "port": port}

        # Add optional fields
        if service_type:
            service_info["type"] = service_type

        if metadata:
            service_info["metadata"] = metadata

        # Prepare tags for easier discovery
        tags = []
        if service_type:
            tags.append(service_type)
        service_info["tags"] = tags

        # Send registration request
        request = {"action": "register", "service_info": service_info, "ttl": ttl}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            # Store service information for later use
            self._service_id = response.get("service_id")
            self._service_info = service_info
            self._ttl = ttl

            self.logger.info(f"Service registered with ID: {self._service_id}")
            return self._service_id
        else:
            self.logger.error(f"Failed to register service: {response.get('error')}")
            return None

    def deregister(self, service_id: str = None) -> bool:
        """
        Deregister this service from the registry.

        Returns:
            True if successful, False otherwise.
        """

        service_id = service_id or self._service_id

        if not service_id:
            self.logger.warning("Cannot deregister: no service is registered")
            return False

        request = {"action": "deregister", "service_id": service_id}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            self.logger.info(f"Service deregistered: {service_id}")
            self._service_id = None
            self._service_info = None
            self._ttl = None
            return True
        else:
            self.logger.error(f"Failed to deregister service: {response.get('error')}")
            return False

    def reregister(self) -> str | None:
        if not self._service_id:
            self.logger.warning("Cannot reregister: no service is registered")
            return None

        if not self._service_info:
            self.logger.warning(
                "Cannot reregister: no service info was saved by this registry client or service already deregistered."
            )
            return None

        return self.register(
            name=self._service_info["name"],
            host=self._service_info["host"],
            port=self._service_info["port"],
            service_type=self._service_info["type"],
            metadata=self._service_info["metadata"],
        )

    def discover_service(self, service_type: str) -> dict[str, Any] | None:
        """
        Discover a service of the specified type. The service is guaranteed to be healthy at the time of discovery.

        The returned information contains:

        - name: the name of the service
        - host: the ip address or hostname of the service
        - port: the port number for requests to the microservice

        Args:
            service_type: Type of service to discover

        Returns:
            Service information if found, None otherwise
        """
        request = {"action": "discover", "service_type": service_type}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if VERBOSE_DEBUG:
            self.logger.debug(f"{response = }")

        if response.get("success"):
            service = response.get("service")
            return service
        else:
            if VERBOSE_DEBUG:
                self.logger.debug(f"Service discovery failed: {response.get('error')}")
            return None

    def get_service(self, service_id: str | None = None) -> dict[str, Any] | None:
        """
        Get information about a specific service. When no service_id is given,
        the service_id known to this client will be used.

        Args:
            service_id: ID of the service to get [default=None]

        Returns:
            Service information if found, None otherwise.
        """
        service_id = service_id or self._service_id

        if not service_id:
            self.logger.warning("Cannot get service: no service id is given or known.")
            return None

        request = {"action": "get", "service_id": service_id}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            service = response.get("service")
            return service
        else:
            self.logger.warning(f"Get service failed: {response.get('error')}")
            return None

    def list_services(self, service_type: str | None = None) -> list[dict[str, Any]]:
        """
        List all registered services, optionally filtered by type.

        Args:
            service_type: Type of services to list

        Returns:
            List of service information.
        """
        request = {"action": "list", "service_type": service_type}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            services = response.get("services", [])
            return services
        else:
            self.logger.warning(f"List services failed: {response.get('error')}")
            return []

    def get_endpoint(self, service_type) -> str | None:
        """Returns the endpoint for the given service type or None if the service_type is not registered."""
        service = self.discover_service(service_type)

        if service:
            protocol = service.get("protocol", "tcp")
            hostname = service["host"]
            port = service["port"]

            return f"{protocol}://{hostname}:{port}"
        else:
            return None

    def health_check(self) -> bool:
        """
        Check if the registry server is healthy.

        Returns:
            True if healthy, False otherwise
        """
        request = {"action": "health"}

        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request, timeout=HEALTH_CHECK_TIMEOUT)
        return response.get("success", False)

    def terminate_registry_server(self) -> bool:
        """
        Send a terminate request to the service registry server. Returns True when successful.
        """
        request = {"action": "terminate"}
        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request, timeout=PROCESS_SHUTDOWN_TIMEOUT)
        return response.get("success", False)

    def server_status(self) -> dict[str, Any]:
        """
        Requests the status information from the service registry.
        """
        request = {
            "action": "info",
        }
        response = self._send_request(MessageType.REQUEST_WITH_REPLY, request)
        return response

    def start_heartbeat(self, interval: int | None = None) -> None:
        if not self._service_id:
            self.logger.warning("Cannot start heartbeat: no service is registered")
            return None

        # Cancel an existing heartbeat thread if present
        self.stop_heartbeat()

        # If interval not specified, use 1/3 of TTL
        if interval is None:
            interval = max(1, self._ttl // 3)

        def send_heartbeat():
            try:
                request = {"action": "renew", "service_id": self._service_id}

                response = self._send_heartbeat(request)

                if not response.get("success", False):
                    self.logger.warning(f"Heartbeat failed: {response.get('error')}")

                    # Do a health check
                    if not self.health_check():
                        self.logger.warning("Heartbeat failed: ServiceRegistry not responding.")
                        return
                    else:
                        self.logger.info("Heartbeat failed, but health check succeeded, reregistering...")
                        self.reregister()

                else:
                    if VERBOSE_DEBUG:
                        self.logger.debug(f"Heartbeat succeeded: {response.get('message')}")

            except Exception as exc:
                self.logger.error(f"Error sending heartbeat: {exc}")

        self._hb_stop_event = threading.Event()
        self._hb_thread = threading.Thread(
            target=do_every,
            args=(interval, send_heartbeat),
            kwargs={
                "stop_event": self._hb_stop_event,
                "setup_func": self._connect_hb_socket,
                "teardown_func": self._disconnect_hb_socket,
            },
        )
        self._hb_thread.daemon = True
        self._hb_thread.start()
        return None

    def stop_heartbeat(self) -> None:
        if self._hb_stop_event:
            self._hb_stop_event.set()
        if self._hb_thread:
            if self._hb_thread.is_alive():
                self.logger.debug(f"Heartbeat thread is alive, joining with timeout={THREAD_JOIN_TIMEOUT}s")
                self._hb_thread.join(timeout=THREAD_JOIN_TIMEOUT)
            else:
                self.logger.debug("Heartbeat thread is not alive.")
                self._hb_thread = None
        else:
            self.logger.debug("No heartbeat thread defined.")

    def close(self) -> None:
        """Clean up resources."""
        try:
            if hasattr(self, "req_socket") and self.req_socket:
                self.req_socket.close()

            if hasattr(self, "sub_socket") and self.sub_socket:
                self.sub_socket.close()

            if hasattr(self, "context") and self.context:
                self.context.term()
        except Exception as exc:
            self.logger.error(f"Error during cleanup: {exc}")


class AsyncRegistryClient:
    """
    Asynchronous client for interacting with the ZeroMQ-based service registry.

    This class uses asyncio and the ZeroMQ asyncio API for non-blocking operations.
    """

    def __init__(
        self,
        registry_req_endpoint: str = None,
        registry_sub_endpoint: str = None,
        registry_hb_endpoint: str = None,
        timeout: float = REQUEST_TIMEOUT,
        client_id: str = "registry-client",
    ):
        """
        Initialize the async registry client.

        Args:
            registry_req_endpoint: ZeroMQ endpoint for REQ-REP socket, defaults to DEFAULT_RS_REQ_PORT on localhost.
            registry_sub_endpoint: ZeroMQ endpoint for SUB socket, defaults to DEFAULT_RS_PUB_PORT on localhost.
            registry_hb_endpoint: ZeroMQ endpoint for heartbeat socket, defaults to DEFAULT_RS_HB_PORT on localhost.
            timeout: Timeout for requests in seconds, defaults to 0.5.
            client_id: client identification, default='registry-client'
        """
        self.registry_req_endpoint = registry_req_endpoint or f"tcp://localhost:{DEFAULT_RS_REQ_PORT}"
        self.registry_sub_endpoint = registry_sub_endpoint or f"tcp://localhost:{DEFAULT_RS_PUB_PORT}"
        self.registry_hb_endpoint = registry_hb_endpoint or f"tcp://localhost:{DEFAULT_RS_HB_PORT}"

        self.timeout = timeout
        self.timeout_ms = int(timeout * 1000)
        self.logger = logging.getLogger("egse.registry.client")

        # Service state
        self._service_id = None
        self._service_info = None
        self._ttl = None

        self._client_id = f"{client_id}-{uuid.uuid4()}".encode()

        self.context = zmq.asyncio.Context.instance()

        self._running = False
        self._tasks = set()
        self._heartbeat_task = None
        self._event_listener_task = None

        self._event_handlers = {}

        self.req_socket: zmq.asyncio.Socket | None = None
        self.sub_socket: zmq.asyncio.Socket | None = None
        self.hb_socket: zmq.asyncio.Socket | None = None

    def __enter__(self):
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.disconnect()

    def connect(self):
        if VERBOSE_DEBUG:
            self.logger.debug("Connecting to service registry...")

        # REQ socket for request-reply pattern
        self.req_socket = self.context.socket(zmq.DEALER)
        self.req_socket.setsockopt(zmq.LINGER, 0)  # Don't wait for unsent messages on close()
        self.req_socket.setsockopt(zmq.IDENTITY, self._client_id)
        self.req_socket.connect(self.registry_req_endpoint)

        # SUB socket for receiving events
        self.sub_socket = self.context.socket(zmq.SUB)
        self.sub_socket.connect(self.registry_sub_endpoint)
        # Default to receiving all events
        self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, "")

    def _connect_hb_socket(self):
        self.hb_socket = self.context.socket(zmq.DEALER)
        self.hb_socket.setsockopt(zmq.LINGER, 0)  # Don't wait for unsent messages on close()
        self.hb_socket.setsockopt(zmq.IDENTITY, self._client_id)
        self.hb_socket.connect(self.registry_hb_endpoint)

    def disconnect(self):
        if VERBOSE_DEBUG:
            self.logger.debug("Disconnecting from service registry...")

        if self.req_socket:
            self.req_socket.close(linger=0)
        self.req_socket = None

        if self.sub_socket:
            self.sub_socket.close(linger=0)
        self.sub_socket = None

    def _disconnect_hb_socket(self):
        if self.hb_socket:
            self.hb_socket.setsockopt(zmq.LINGER, 0)
            self.hb_socket.close()
        self.hb_socket = None

    async def _send_request(
        self, msg_type: MessageType, request: dict[str, Any], timeout: float = None
    ) -> dict[str, Any]:
        """
        Send a request to the registry and get the response.

        Args:
            msg_type: the type of message and reply
            request: The request to send to the service registry server.
            timeout: The number of seconds to wait, if None, instance variable is used.

        Returns:
            The response from the registry as a dictionary.
        """

        assert self.req_socket is not None, "REQ socket is not connected, cannot send request."

        timeout = timeout or self.timeout
        try:
            if VERBOSE_DEBUG:
                self.logger.debug(f"Sending request: {request}")

            await self.req_socket.send_multipart([msg_type.value, json.dumps(request).encode()])

            try:
                message_parts = await asyncio.wait_for(self.req_socket.recv_multipart(), timeout=timeout)

                if len(message_parts) >= 2:
                    message_type = MessageType(message_parts[0])
                    message_data = message_parts[1]

                    if message_type == MessageType.RESPONSE:
                        response = json.loads(message_data)
                        if VERBOSE_DEBUG:
                            self.logger.debug(f"Received response: {response}")
                        return response
                    else:
                        return {
                            "success": False,
                            "error": f"unexpected MessageType received: {message_type.name}, {message_data = }",
                        }
                else:
                    return {
                        "success": False,
                        "error": f"not enough parts received: {len(message_parts)}",
                        "data": message_parts,
                    }
            except asyncio.TimeoutError:
                self.logger.error(f"Request timed out after {timeout:.2f}s")
                return {"success": False, "error": "Request timed out"}
        except zmq.ZMQError as exc:
            self.logger.error(f"ZMQ error: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}
        except Exception as exc:
            self.logger.error(f"Error sending request: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}

    async def _send_heartbeat(self, request: dict[str, Any]) -> dict[str, Any]:
        """
        Send a heartbreak request to the registry and get the response.

        Args:
            request: The request to send to the service registry server.

        Returns:
            The response from the registry as a dictionary.
        """

        assert self.hb_socket is not None, "HB socket is not connected, cannot send heartbeat request."

        try:
            if VERBOSE_DEBUG:
                self.logger.debug(f"Sending heartbeat request: {request}")

            await self.hb_socket.send_string(json.dumps(request))

            try:
                response_json = await asyncio.wait_for(self.hb_socket.recv_string(), timeout=HEART_BEAT_TIMEOUT)
                response = json.loads(response_json)
                if VERBOSE_DEBUG:
                    self.logger.debug(f"Received heartbeat response: {response = }")
                return response
            except asyncio.TimeoutError:
                self.logger.error(f"Heartbeat request timed out after {HEART_BEAT_TIMEOUT:.2f}s")
                return {"success": False, "error": "Heartbeat request timed out"}

        except zmq.ZMQError as exc:
            self.logger.error(f"ZMQ error: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}
        except Exception as exc:
            self.logger.error(f"Error sending heartbeat request: {exc}", exc_info=True)
            return {"success": False, "error": str(exc)}

    async def register(
        self,
        name: str,
        host: str,
        port: int,
        service_type: str | None = None,
        metadata: dict[str, Any] | None = None,
        ttl: int = 30,
    ) -> str | None:
        """
        Register this service with the registry.

        Args:
            name: Service name
            host: Service host/IP
            port: Service port
            service_type: Service type (for discovery)
            metadata: Additional service metadata
            ttl: Time-to-live in seconds

        Returns:
            The service ID if successful, None otherwise
        """
        # Prepare service info
        service_info: dict[str, Any] = {"name": name, "host": host, "port": port}

        # Add optional fields
        if service_type:
            service_info["type"] = service_type

        if metadata:
            service_info["metadata"] = metadata

        # Prepare tags for easier discovery
        tags = []
        if service_type:
            tags.append(service_type)
        service_info["tags"] = tags

        # Send registration request
        request = {"action": "register", "service_info": service_info, "ttl": ttl}

        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            # Store service information for later use
            self._service_id = response.get("service_id")
            self._service_info = service_info
            self._ttl = ttl

            self.logger.info(f"Service registered with ID: {self._service_id}")
            return self._service_id
        else:
            self.logger.error(f"Failed to register service: {response.get('error')}")
            return None

    async def deregister(self, service_id: str | None = None) -> bool:
        """
        Deregister this service or the service with service_id from the registry.

        When you register and deregister with the same client instance, you don't have to provide
        the service_id.

        NOTE:
            When a registry client is running and sending out heartbeats, it will reregister if it finds the
            registry server to be healthy. If you want to permanently deregister the service, you will need
            to stop the registry client for the service or terminate the service.

        Args:
            service_id: the service identifier that was previously handed out by the ServiceRegistry after
                registration. If not provided, the service_id that was saved when registering is used.

        Returns:
            True if successful, False otherwise
        """

        service_id = service_id or self._service_id

        if not service_id:
            self.logger.warning("Cannot deregister: no service is registered")
            return False

        request = {"action": "deregister", "service_id": service_id}

        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            self.logger.info(f"Service deregistered: {service_id}")
            self._service_id = None
            self._service_info = None
            self._ttl = None
            return True
        else:
            self.logger.error(f"Failed to deregister service: {response.get('error')}")
            return False

    async def reregister(self) -> str | None:
        if not self._service_id:
            self.logger.warning("Cannot reregister: no service is registered")
            return None

        if not self._service_info:
            self.logger.warning(
                "Cannot reregister: no service info was saved by this registry client or service already deregistered."
            )
            return None

        return await self.register(
            name=self._service_info["name"],
            host=self._service_info["host"],
            port=self._service_info["port"],
            service_type=self._service_info["type"],
            metadata=self._service_info["metadata"],
        )

    async def start_heartbeat(self, interval: int | None = None) -> asyncio.Task | None:
        """
        Start sending heartbeats to the registry.

        Args:
            interval: Heartbeat interval in seconds (default: 1/3 of TTL)

        Returns:
            The heartbeat task
        """
        if not self._service_id:
            self.logger.warning("Cannot start heartbeat: no service is registered")
            return None

        # Cancel existing heartbeat task if present
        await self.stop_heartbeat()

        # If interval not specified, use 1/3 of TTL
        if interval is None:
            interval = max(1, self._ttl // 3)

        self._running = True

        async def heartbeat_loop():
            self._connect_hb_socket()

            try:
                while self._running and self._service_id:
                    try:
                        request = {"action": "renew", "service_id": self._service_id}

                        response = await self._send_heartbeat(request)

                        if not response.get("success"):
                            self.logger.warning(f"Heartbeat failed: {response.get('error')}")

                            # Do a health check
                            if not await self.health_check():
                                self.logger.warning("Heartbeat failed: ServiceRegistry not responding.")
                                continue
                            else:
                                self.logger.info("Heartbeat failed, but health check succeeded, reregistering...")
                                await self.reregister()

                        else:
                            if VERBOSE_DEBUG:
                                self.logger.debug(f"Heartbeat succeeded: {response.get('message')}")

                    except Exception as exc:
                        self.logger.error(f"Error in heartbeat loop: {exc}", exc_info=True)

                    # Sleep until next heartbeat
                    await asyncio.sleep(interval)
            except asyncio.CancelledError:
                self.logger.info("Heartbeat task cancelled")
            finally:
                self._disconnect_hb_socket()

        # Start the heartbeat task

        self.logger.info(f"Starting heartbeat task with interval {interval}s")

        self._heartbeat_task = task = asyncio.create_task(heartbeat_loop())
        self._tasks.add(task)
        task.add_done_callback(lambda t: self._tasks.discard(t))

        return task

    async def stop_heartbeat(self) -> None:
        """Stop the running heartbeat task."""

        if self._heartbeat_task is None:
            if VERBOSE_DEBUG:
                self.logger.debug("Couldn't stop heartbeat, heartbeat_task is None")
            return

        self._heartbeat_task.cancel()
        try:
            await self._heartbeat_task
        except asyncio.CancelledError:
            self.logger.info("Heartbeat task cancelled")
            pass
        self._tasks.discard(self._heartbeat_task)
        self._heartbeat_task = None
        self.logger.info("Stopped heartbeat task")

    async def stop_event_listener(self) -> None:
        """Stop the running event listener task."""

        if self._event_listener_task is None:
            if VERBOSE_DEBUG:
                self.logger.debug("Couldn't stop event_listener, event_listener_task is None")
            return

        self._event_listener_task.cancel()
        try:
            await self._event_listener_task
        except asyncio.CancelledError:
            pass
        self._tasks.discard(self._event_listener_task)
        self._event_listener_task = None
        self.logger.info("Stopped event listener task")

    async def stop_all_tasks(self) -> None:
        self._running = False

        # Cancel all tasks
        for task in self._tasks:
            task.cancel()

        # Wait for tasks to complete (with timeout)
        if self._tasks:
            try:
                await asyncio.wait(self._tasks, timeout=TASK_COMPLETION_TIMEOUT)
            except asyncio.CancelledError:
                pass

        self._tasks.clear()
        self.logger.info("Stopped all background tasks")

    def on_event(self, event_type: str, handler: Callable[[dict[str, Any]], Union[None, Coroutine]]) -> None:
        """
        Register a handler for a specific event type.

        Args:
            event_type: Type of event to handle (register, deregister, expire)
            handler: Function or coroutine to call with event data
        """
        if event_type not in self._event_handlers:
            self._event_handlers[event_type] = []

            # Subscribe to this specific event type
            self.sub_socket.setsockopt_string(zmq.SUBSCRIBE, event_type)

        self._event_handlers[event_type].append(handler)
        self.logger.debug(f"Registered handler for {event_type} events")

    async def start_event_listener(self) -> asyncio.Task:
        """
        Start listening for registry events.

        Returns:
            The event listener task
        """
        # Cancel existing event listener task if present
        await self.stop_event_listener()

        self._running = True

        async def subscription_loop():
            try:
                while self._running:
                    try:
                        # Use a timeout to allow for clean shutdown
                        try:
                            message = await asyncio.wait_for(
                                self.sub_socket.recv_multipart(), timeout=EVENT_POLL_TIMEOUT
                            )
                        except asyncio.TimeoutError:
                            continue

                        # Parse the message
                        event_type_bytes, event_json_bytes = message
                        event_type = event_type_bytes.decode("utf-8")
                        event = json.loads(event_json_bytes.decode("utf-8"))

                        self.logger.debug(f"Received event: {event_type}")

                        # Call registered handlers
                        handlers = self._event_handlers.get(event_type, [])
                        for handler in handlers:
                            try:
                                # Check if handler is a coroutine function
                                if asyncio.iscoroutinefunction(handler):
                                    await handler(event["data"])
                                else:
                                    handler(event["data"])
                            except Exception as exc:
                                self.logger.error(f"Error in event handler: {exc}")
                    except zmq.ZMQError as exc:
                        self.logger.error(f"ZMQ error in event listener: {exc}")
                    except Exception as exc:
                        self.logger.error(f"Error in event listener: {exc}")
                        await asyncio.sleep(1)  # Prevent tight loop on error
            except asyncio.CancelledError:
                self.logger.info("Event listener task cancelled")

        # Start the subscription task
        self._event_listener_task = task = asyncio.create_task(subscription_loop())
        self._tasks.add(task)
        task.add_done_callback(lambda t: self._tasks.discard(t))

        self.logger.info("Started event listener task")
        return task

    async def discover_service(self, service_type: str) -> dict[str, Any] | None:
        """
        Discover a service of the specified type. The service is guaranteed to be healthy at the time of discovery.

        The returned information contains:

        - name: the name of the service
        - host: the ip address or hostname of the service
        - port: the port number for requests to the microservice

        Args:
            service_type: Type of service to discover

        Returns:
            Service information if found, None otherwise
        """

        request = {"action": "discover", "service_type": service_type}

        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            service = response.get("service")

            return service
        else:
            self.logger.warning(f"Service discovery failed: {response.get('error')}")
            return None

    async def get_service(self, service_id: str | None = None) -> dict[str, Any] | None:
        """
        Get information about a specific service.  When no service_id is given,
        the service_id known to this client will be used.

        Args:
            service_id: ID of the service to get

        Returns:
            Service information if found, None otherwise
        """

        service_id = service_id or self._service_id

        request = {"action": "get", "service_id": service_id}

        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            service = response.get("service")

            return service
        else:
            self.logger.warning(f"Get service failed: {response.get('error')}")
            return None

    async def list_services(self, service_type: str | None = None) -> list[dict[str, Any]]:
        """
        List all registered services, optionally filtered by type.

        Args:
            service_type: Type of services to list

        Returns:
            List of service information
        """
        request = {"action": "list", "service_type": service_type}

        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)

        if response.get("success"):
            services = response.get("services", [])

            return services
        else:
            self.logger.warning(f"List services failed: {response.get('error')}")
            return []

    async def get_endpoint(self, service_type) -> str | None:
        """Returns the endpoint for the given service type."""
        service = await self.discover_service(service_type)

        if service:
            protocol = service.get("protocol", "tcp")
            hostname = service["host"]
            port = service["port"]

            return f"{protocol}://{hostname}:{port}"
        else:
            return None

    async def health_check(self) -> bool:
        """
        Check if the registry server is healthy.

        Returns:
            True if healthy, False otherwise
        """
        request = {"action": "health"}
        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request, timeout=HEALTH_CHECK_TIMEOUT)
        return response.get("success", False)

    async def terminate_registry_server(self) -> bool:
        """
        Send a terminate request to the service registry server. Returns True when successful.
        """
        request = {"action": "terminate"}
        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request, timeout=PROCESS_SHUTDOWN_TIMEOUT)
        return response.get("success", False)

    async def server_status(self) -> dict[str, Any]:
        """
        Requests the status information from the service registry.
        """
        request = {"action": "info"}
        response = await self._send_request(MessageType.REQUEST_WITH_REPLY, request)
        return response

    async def close(self) -> None:
        """Clean up resources."""
        await self.stop_heartbeat()  # This stops all tasks

        try:
            if hasattr(self, "req_socket") and self.req_socket:
                self.req_socket.close()

            if hasattr(self, "sub_socket") and self.sub_socket:
                self.sub_socket.close()

            # We can not terminate the context, because we use a global instance, i.e. a singleton context.
            # When we try to terminate it, even after checking if it was closed, it raises an exception.
            if hasattr(self, "context") and self.context:
                if VERBOSE_DEBUG:
                    self.logger.debug(f"{self.context = !r}")
                    self.logger.debug(f"{self.context._sockets = !r}")
                # The zmq context instance is the global singleton instance.
                # Terminating it here would affect other parts of the application using zmq.
                # if not self.context.closed:
                #     self.context.term()
        except Exception as exc:
            self.logger.error(f"Error during cleanup: {exc}")


def is_service_registered(service_type: str):
    """Convenience function to check if a service is registered."""
    with RegistryClient() as reg:
        response = reg.discover_service(service_type)

    return False if response is None else True
