# Copyright 2026 Marimo. All rights reserved.
"""ZeroMQ connection management for inter-process communication."""

from __future__ import annotations

import dataclasses
import queue
import sys
import typing

import msgspec

from marimo import _loggers
from marimo._ipc.queue_proxy import PushQueue, start_receiver_thread
from marimo._ipc.types import ConnectionInfo
from marimo._messaging.types import KernelMessage
from marimo._runtime.commands import (
    BatchableCommand,
    CommandMessage,
    OutOfBandCommand,
)
from marimo._session.queue import QueueType

if typing.TYPE_CHECKING:
    import zmq

LOGGER = _loggers.marimo_logger()
ADDR = "tcp://127.0.0.1"

T = typing.TypeVar("T")


@dataclasses.dataclass
class Channel(typing.Generic[T]):
    """A typed communication channel wrapping a ZeroMQ socket and queue.

    Channels can be either:
    - Push: Send-only, uses PushQueue to immediately send via socket
    - Pull: Receive-only, uses Queue to buffer received messages
    """

    kind: typing.Literal["push", "pull"]
    socket: zmq.Socket[bytes]
    queue: QueueType[T]
    decoder: msgspec.msgpack.Decoder[T] | None = None

    @classmethod
    def Push(
        cls, context: zmq.Context[zmq.Socket[bytes]], *, maxsize: int = 0
    ) -> Channel[T]:
        """Create a push (send-only) channel.

        Note: maxsize is ignored for push channels as ZeroMQ handles buffering.
        """
        import zmq

        socket = context.socket(zmq.PUSH)
        return cls(
            kind="push",
            socket=socket,
            queue=PushQueue(socket, maxsize=maxsize),
        )

    @typing.overload
    @classmethod
    def Pull(
        cls,
        context: zmq.Context[zmq.Socket[bytes]],
        *,
        msg_type: type[T],
        maxsize: int = 0,
    ) -> Channel[T]: ...

    @typing.overload
    @classmethod
    def Pull(
        cls,
        context: zmq.Context[zmq.Socket[bytes]],
        *,
        msg_type: typing.Any,
        maxsize: int = 0,
    ) -> Channel[typing.Any]: ...

    @classmethod
    def Pull(
        cls,
        context: zmq.Context[zmq.Socket[bytes]],
        *,
        msg_type: typing.Any,
        maxsize: int = 0,
    ) -> Channel[typing.Any]:
        """Create a pull (receive-only) channel.

        Args:
            context: ZeroMQ context for creating sockets
            msg_type: The type to decode incoming messages as. Accepts
                concrete types, union types (e.g. `A | B`), and NewTypes;
                msgspec's `Decoder` handles all three.
            maxsize: Maximum queue size (0 = unlimited)
        """
        import zmq

        socket = context.socket(zmq.PULL)
        return cls(
            kind="pull",
            socket=socket,
            queue=queue.Queue(maxsize=maxsize),
            decoder=msgspec.msgpack.Decoder(msg_type),
        )


@dataclasses.dataclass
class Connection:
    """Manages all ZeroMQ sockets for marimo IPC communication."""

    context: zmq.Context[zmq.Socket[bytes]]

    control: Channel[CommandMessage]
    ui_element: Channel[BatchableCommand]
    completion: Channel[OutOfBandCommand]
    win32_interrupt: Channel[bool] | None

    input: Channel[str]
    stream: Channel[KernelMessage]

    def __post_init__(self) -> None:
        """Start receiver threads for all pull channels."""
        pull_channels: list[Channel[typing.Any]] = []
        if self.control.kind == "pull":
            pull_channels.append(self.control)
        if self.ui_element.kind == "pull":
            pull_channels.append(self.ui_element)
        if self.completion.kind == "pull":
            pull_channels.append(self.completion)
        if self.win32_interrupt and self.win32_interrupt.kind == "pull":
            pull_channels.append(self.win32_interrupt)
        if self.input.kind == "pull":
            pull_channels.append(self.input)
        if self.stream.kind == "pull":
            pull_channels.append(self.stream)

        self._stop_event, self._receiver_thread = start_receiver_thread(
            pull_channels
        )

    @classmethod
    def create(cls) -> tuple[Connection, ConnectionInfo]:
        """Create host-side connection with all sockets bound to random ports.

        Returns:
            Tuple of (Connection instance, ConnectionInfo with port numbers)
        """
        import zmq

        context = zmq.Context()
        conn = cls(
            context=context,
            control=Channel.Push(context),
            ui_element=Channel.Push(context),
            completion=Channel.Push(context),
            win32_interrupt=(
                Channel.Push(context) if sys.platform == "win32" else None
            ),
            input=Channel.Push(context, maxsize=1),
            stream=Channel.Pull(context, msg_type=KernelMessage),
        )
        info = ConnectionInfo(
            control=conn.control.socket.bind_to_random_port(ADDR),
            ui_element=conn.ui_element.socket.bind_to_random_port(ADDR),
            completion=conn.completion.socket.bind_to_random_port(ADDR),
            input=conn.input.socket.bind_to_random_port(ADDR),
            stream=conn.stream.socket.bind_to_random_port(ADDR),
            win32_interrupt=conn.win32_interrupt.socket.bind_to_random_port(
                ADDR
            )
            if conn.win32_interrupt
            else None,
        )
        return conn, info

    @classmethod
    def connect(cls, connection_info: ConnectionInfo) -> Connection:
        """Connect to host with all sockets and start receivers.

        Args:
            connection_info: Port information from host

        Returns:
            Connected Connection instance
        """
        import zmq

        context = zmq.Context()

        conn = cls(
            context=context,
            control=Channel.Pull(context, msg_type=CommandMessage),
            ui_element=Channel.Pull(context, msg_type=BatchableCommand),
            completion=Channel.Pull(context, msg_type=OutOfBandCommand),
            win32_interrupt=Channel.Pull(context, msg_type=bool)
            if connection_info.win32_interrupt
            else None,
            input=Channel.Pull(context, msg_type=str, maxsize=1),
            stream=Channel.Push(context),
        )

        # Attach to existing ports
        conn.control.socket.connect(f"{ADDR}:{connection_info.control}")
        conn.ui_element.socket.connect(f"{ADDR}:{connection_info.ui_element}")
        conn.completion.socket.connect(f"{ADDR}:{connection_info.completion}")
        if conn.win32_interrupt:
            conn.win32_interrupt.socket.connect(
                f"{ADDR}:{connection_info.win32_interrupt}"
            )
        conn.input.socket.connect(f"{ADDR}:{connection_info.input}")
        conn.stream.socket.connect(f"{ADDR}:{connection_info.stream}")

        return conn

    def close(self) -> None:
        """Close all sockets and cleanup resources."""
        # Stop receiver thread
        self._stop_event.set()
        if self._receiver_thread.is_alive():
            self._receiver_thread.join(timeout=1)

        # Close all sockets and terminate the context.
        # linger=0 drops any unsent messages immediately, preventing
        # hangs during shutdown when the remote end is already gone.
        self.context.destroy(linger=0)
