# Copyright 2025 Acute Technology Inc. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
"""Session to connect to the server."""

import logging
import threading
import time
from collections.abc import Callable
from typing import Any, ClassVar

import zmq
import zmq.asyncio

import aqpxlib.message as pxmsg
from aqpxlib.channels import SynchronousSocketWrapper
from aqpxlib.exceptions import (
    BackendError,
)

_LOGGER = logging.getLogger(__name__)


class PxSession:
    """A connection interface to the PX server via ZeroMQ channels.

    There are two channels associated with the server listed below:

    * client: for request/reply calls to the server.
    * subscriber: for the server to publish events to the client.
    """

    _socket_channel_type: ClassVar[dict[str, int]] = {
        "client": zmq.DEALER,
        "subscriber": zmq.SUB,
    }

    def __init__(self) -> None:
        """Initialize session variables."""
        self._ctx: zmq.Context = zmq.Context()
        self._client: zmq.Socket | None = None
        self._subscriber: zmq.Socket | None = None

        self._event_monitor_start_event: threading.Event | None = None
        self._event_monitor_stop_event: threading.Event | None = None
        self._event_monitor_thread: threading.Thread | None = None
        self._event_handlers: dict[str, Callable] = {}

    def _create_connected_socket(self, identity: str, endpoint: str) -> zmq.Socket:
        socket = self._ctx.socket(self._socket_channel_type[identity])

        # Set linger to 10 seconds to ensure the socket is closed properly
        socket.linger = 10000

        # Set the identity of the socket to the identity of the channel
        socket.setsockopt(zmq.IDENTITY, identity.encode("utf-8"))

        # Connect the socket to the endpoint
        socket.connect(endpoint)

        return socket

    @property
    def client(self) -> SynchronousSocketWrapper[pxmsg.PxProtobufMsg, pxmsg.PxProtobufMsg]:
        """A wrapper of the ZeroMQ DEALER socket."""
        if self._client is None:
            err_msg = "Not connected to server"
            raise BackendError(err_msg)

        return SynchronousSocketWrapper[pxmsg.PxProtobufMsg, pxmsg.PxProtobufMsg](self._client)

    @property
    def subscriber(self) -> SynchronousSocketWrapper[pxmsg.PxProtobufMsg, pxmsg.PxEventMsg]:
        """A wrapper of the ZeroMQ SUB socket."""
        if self._subscriber is None:
            err_msg = "Not connected to server"
            raise BackendError(err_msg)

        return SynchronousSocketWrapper[pxmsg.PxProtobufMsg, pxmsg.PxEventMsg](self._subscriber)

    def connect(self, host: str = "localhost", port: int = 60600):
        """Connect to the server on the given host and port.

        Args:
            host (str, optional): The hostname or IP address to connect to. Defaults to "localhost".
            port (int, optional): The port to connect to. Defaults to 60600.

        """
        _LOGGER.info(
            "Connecting to server on %s:%d, subscribing to %s:%d",
            host,
            port,
            host,
            port + 1,
        )

        self._client = self._create_connected_socket(
            "client",
            f"tcp://{host}:{port}",
        )

        self._subscriber = self._create_connected_socket(
            "subscriber",
            f"tcp://{host}:{port + 1}",
        )
        self._subscriber.subscribe(b"")

        _LOGGER.info("Starting event monitor async")
        self._start_event_monitor_async()

    def disconnect(self):
        """Disconnect from the server."""
        _LOGGER.info("Disconnecting from server")

        self._client.close()
        self._subscriber.close()

        if self._event_monitor_thread is not None:
            self._event_monitor_stop_event.set()
            self._event_monitor_thread.join()
            self._event_monitor_thread = None
            self._event_monitor_stop_event = None
            self._event_monitor_start_event = None

        self._client = None
        self._subscriber = None

    def request(
        self,
        data: pxmsg.PxProtobufMsg,
        timeout: int = 5000,
    ) -> pxmsg.PxProtobufMsg | None:
        """Send a request to the server and wait for a response.

        Args:
            data (bytes): The data to send to the server
            timeout (int, optional): The timeout in milliseconds. Defaults to 1000.

        """
        return self.client.request(data, timeout=timeout).result

    def subscribe(self, event_topic: str):
        """Subscribe to events from the server."""
        _LOGGER.info("Subscribing to events topic: %s", event_topic)
        self._subscriber.subscribe(event_topic.encode("utf-8"))

    def _start_event_monitor_async(self) -> None:
        def event_monitor_thread_async(stop_event: threading.Event) -> None:
            _LOGGER.debug("Starting event loop...")
            self._event_monitor_start_event.set()
            try:
                while not stop_event.is_set():
                    try:
                        response = self.subscriber.recv(timeout=1000)
                    except BackendError:
                        continue

                    if response is not None:
                        _LOGGER.debug("Received event: %s", response.topic)
                        _LOGGER.debug("Event data: %s", response.result)

                        handler = self._event_handlers.get(response.topic)
                        if handler is not None:
                            handler(response.result)
                        else:
                            _LOGGER.debug(
                                "No handler found for event: %s",
                                response.topic,
                            )

            except KeyboardInterrupt:
                time.sleep(1)

        self._event_monitor_start_event = threading.Event()
        self._event_monitor_stop_event = threading.Event()
        self._event_monitor_thread = threading.Thread(
            target=event_monitor_thread_async,
            args=(self._event_monitor_stop_event,),
        )
        self._event_monitor_thread.daemon = True
        self._event_monitor_thread.start()
        if not self._event_monitor_start_event.wait(timeout=10):
            err_msg = "Event monitor thread failed to start"
            raise BackendError(err_msg)

    def add_event_handler(self, event: str, handler: Callable[[Any], Any]):
        """Add an event handler."""
        self._event_handlers[event] = handler
        self.subscribe(event)

    def remove_event_handler(self, event: str):
        """Remove an event handler."""
        self._event_handlers.pop(event, None)

    def get_event_handlers(
        self,
        filter_by: Callable[[str], bool] | None = None,
    ) -> dict[str, Callable]:
        """Get all available event handlers."""
        if filter_by is None:
            return self._event_handlers

        return {k: v for k, v in self._event_handlers.items() if filter_by(k)}
