import time
from unittest.mock import MagicMock, call, patch

import zmq
import pytest

from zeromqthreadpool import WorkerThread, ThreadPoolParameters, ZeroMqThreadPool


class ConcreteWorker(WorkerThread):
    def do_work(self, data):
        return b"ok"

    def dispose(self):
        pass


def _make_mock_socket(poll_returns=0):
    """Create a mock ZMQ socket that returns no data on poll."""
    sock = MagicMock()
    sock.poll.return_value = poll_returns
    return sock


class TestWorkerSocketCleanup:
    """Tests that worker threads properly disconnect and close sockets on shutdown."""

    def test_disconnect_called_before_close_on_stop(self):
        address = "inproc://test-cleanup"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket

        worker = ConcreteWorker(ctx, address, idle_timeout=60000, identity="test-1")
        worker.start()
        time.sleep(0.3)
        worker.stop()
        worker.join()

        mock_socket.disconnect.assert_called_once_with(address)
        mock_socket.close.assert_called_once_with(linger=0)

    def test_disconnect_called_before_close_ordering(self):
        """Verify disconnect happens before close (order matters for TCP teardown)."""
        address = "inproc://test-order"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket

        call_order = []
        mock_socket.disconnect.side_effect = lambda *a: call_order.append("disconnect")
        original_close = mock_socket.close
        mock_socket.close.side_effect = lambda **kw: call_order.append("close")

        worker = ConcreteWorker(ctx, address, idle_timeout=60000, identity="test-2")
        worker.start()
        time.sleep(0.3)
        worker.stop()
        worker.join()

        assert call_order == ["disconnect", "close"]

    def test_close_linger_zero(self):
        """Socket must close with linger=0 so resources are freed immediately."""
        address = "inproc://test-linger"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket

        worker = ConcreteWorker(ctx, address, idle_timeout=60000, identity="test-3")
        worker.start()
        time.sleep(0.3)
        worker.stop()
        worker.join()

        mock_socket.close.assert_called_once_with(linger=0)

    def test_close_still_called_if_disconnect_raises(self):
        """If disconnect fails, close must still be called to release the socket."""
        address = "inproc://test-disconnect-fail"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket
        mock_socket.disconnect.side_effect = zmq.ZMQError(msg="not connected")

        worker = ConcreteWorker(ctx, address, idle_timeout=60000, identity="test-4")
        worker.start()
        time.sleep(0.3)
        worker.stop()
        worker.join()

        mock_socket.disconnect.assert_called_once_with(address)
        mock_socket.close.assert_called_once_with(linger=0)

    def test_idle_worker_cleanup_no_leaked_connections(self):
        """Simulate the monitor stopping an idle worker — socket must be fully cleaned up."""
        address = "inproc://test-idle"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket

        worker = ConcreteWorker(ctx, address, idle_timeout=500, identity="test-5")
        worker.start()
        # Wait for idle timeout to trigger
        time.sleep(1.5)
        assert worker.is_idle
        # Monitor would call stop()
        worker.stop()
        worker.join()

        mock_socket.disconnect.assert_called_once_with(address)
        mock_socket.close.assert_called_once_with(linger=0)

    def test_died_sent_before_disconnect(self):
        """DIED signal should be sent before disconnect so the ROUTER knows the worker is gone."""
        address = "inproc://test-died-order"
        ctx = MagicMock()
        mock_socket = _make_mock_socket()
        ctx.socket.return_value = mock_socket

        call_order = []
        original_send = mock_socket.send

        def track_send(data, **kwargs):
            if data == b"DIED":
                call_order.append("DIED")

        mock_socket.send.side_effect = track_send
        mock_socket.disconnect.side_effect = lambda *a: call_order.append("disconnect")
        mock_socket.close.side_effect = lambda **kw: call_order.append("close")

        worker = ConcreteWorker(ctx, address, idle_timeout=60000, identity="test-6")
        worker.start()
        time.sleep(0.3)
        worker.stop()
        worker.join()

        assert "DIED" in call_order
        assert "disconnect" in call_order
        assert "close" in call_order
        assert call_order.index("DIED") < call_order.index("disconnect") < call_order.index("close")


class TestThreadPoolWorkerChurn:
    """Tests that the pool's scale-down path triggers proper socket cleanup."""

    def test_pool_scaledown_disconnects_idle_workers(self):
        """When the monitor removes idle workers, their sockets must be disconnected."""
        address = "inproc://test-pool"
        ctx = MagicMock()
        sockets = []

        def make_socket(socket_type):
            sock = _make_mock_socket()
            sockets.append(sock)
            return sock

        ctx.socket.side_effect = make_socket

        params = ThreadPoolParameters(
            name="test",
            initial_count=1,
            max_count=3,
            increment=1,
            timeout=500,       # idle_timeout — workers idle after 500ms
            monitor_timeout=200  # monitor checks every 200ms
        )

        pool = ZeroMqThreadPool(ctx, params, address, is_daemon=True,
                                thread_class=ConcreteWorker)

        # Manually add extra workers beyond initial_count to trigger scale-down
        pool.create_worker_thread()
        assert len(pool.threads) == 2

        # Wait for idle timeout + monitor cycle to trigger scale-down
        time.sleep(2.0)

        # Pool should have scaled back to initial_count
        assert len(pool.threads) <= params.initial_count

        pool.dispose()

        # All created sockets should have been disconnected and closed with linger=0
        for sock in sockets:
            sock.disconnect.assert_called_with(address)
            sock.close.assert_called_with(linger=0)