# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import socket
import threading

import pytest
import zmq

from vllm.utils import network_utils
from vllm.utils.network_utils import (
    get_file_store_init_method,
    get_open_port,
    get_open_ports_list,
    get_tcp_uri,
    join_host_port,
    make_zmq_path,
    make_zmq_socket,
    split_host_port,
    split_zmq_path,
)


def test_get_file_store_init_method_is_unique():
    init_methods = {get_file_store_init_method() for _ in range(2)}

    assert len(init_methods) == 2
    assert all(method.startswith("file://") for method in init_methods)


def _call_with_timeout(func, timeout: float = 10.0):
    """Run func in a daemon thread so a livelock regression fails the test
    quickly instead of hanging the suite."""
    result: dict = {}

    def target():
        try:
            result["value"] = func()
        except BaseException as e:
            result["error"] = e

    thread = threading.Thread(target=target, daemon=True)
    thread.start()
    thread.join(timeout)
    assert not thread.is_alive(), f"call did not finish within {timeout}s"
    if "error" in result:
        raise result["error"]
    return result["value"]


def test_get_open_port(monkeypatch: pytest.MonkeyPatch):
    with monkeypatch.context() as m:
        m.setenv("VLLM_PORT", "5678")
        # make sure we can get multiple ports, even if the env var is set
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s1:
            s1.bind(("localhost", get_open_port()))
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s2:
                s2.bind(("localhost", get_open_port()))
                with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s3:
                    s3.bind(("localhost", get_open_port()))


def test_get_open_port_vllm_port_in_dp_reserved_range(
    monkeypatch: pytest.MonkeyPatch,
):
    # VLLM_PORT falling inside the data-parallel reserved window used to make
    # get_open_port() loop forever (issue #50024). It must instead return a
    # port outside the reserved range.
    with monkeypatch.context() as m:
        m.setenv("VLLM_DP_MASTER_PORT", "5680")
        # 5682 is inside [5680, 5690).
        m.setenv("VLLM_PORT", "5682")
        port = _call_with_timeout(get_open_port)
        assert port not in range(5680, 5690)


def test_get_open_ports_list_with_vllm_port(monkeypatch: pytest.MonkeyPatch):
    with monkeypatch.context() as m:
        m.setenv("VLLM_PORT", "5678")
        ports = get_open_ports_list(5)
        assert len(ports) == 5
        assert len(set(ports)) == 5, "ports must be unique"

        # verify every port is actually bindable
        sockets = []
        try:
            for p in ports:
                s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                s.bind(("localhost", p))
                sockets.append(s)
        finally:
            for s in sockets:
                s.close()


def test_get_open_port_skips_reserved_dp_master_ports(
    monkeypatch: pytest.MonkeyPatch,
):
    """VLLM_PORT == VLLM_DP_MASTER_PORT must not livelock by returning the
    same free-but-reserved port forever (upstream issue #50024)."""
    base = 42123
    with monkeypatch.context() as m:
        m.setenv("VLLM_PORT", str(base))
        m.setenv("VLLM_DP_MASTER_PORT", str(base))
        port = _call_with_timeout(get_open_port)
        assert port >= base + 10


def test_get_open_ports_list_skips_reserved_dp_master_ports(
    monkeypatch: pytest.MonkeyPatch,
):
    """The VLLM_PORT band scan must not hand out ports inside the window
    reserved for the data parallel master process."""
    base = 42223
    reserved = range(base, base + 10)
    with monkeypatch.context() as m:
        m.setenv("VLLM_PORT", str(base))
        m.setenv("VLLM_DP_MASTER_PORT", str(base))
        ports = _call_with_timeout(lambda: get_open_ports_list(5))
        assert len(ports) == 5
        assert len(set(ports)) == 5, "ports must be unique"
        assert all(p not in reserved for p in ports)


def test_get_open_port_ephemeral_skips_reserved_range(
    monkeypatch: pytest.MonkeyPatch,
):
    """An ephemeral port that happens to fall in the reserved range is
    replaced by a rescan starting past the range."""
    base = 42323
    with monkeypatch.context() as m:
        m.delenv("VLLM_PORT", raising=False)
        m.setenv("VLLM_DP_MASTER_PORT", str(base))
        m.setattr(
            network_utils,
            "_get_open_port",
            lambda start_port=None, max_attempts=None: (
                base if start_port is None else start_port
            ),
        )
        assert _call_with_timeout(get_open_port) == base + 10


def test_get_open_port_ephemeral_without_dp_master_port(
    monkeypatch: pytest.MonkeyPatch,
):
    """Without VLLM_PORT and VLLM_DP_MASTER_PORT, an ephemeral port is
    returned as before."""
    with monkeypatch.context() as m:
        m.delenv("VLLM_PORT", raising=False)
        m.delenv("VLLM_DP_MASTER_PORT", raising=False)
        port = _call_with_timeout(get_open_port)
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            s.bind(("localhost", port))


@pytest.mark.parametrize(
    "path,expected",
    [
        ("ipc://some_path", ("ipc", "some_path", "")),
        ("tcp://127.0.0.1:5555", ("tcp", "127.0.0.1", "5555")),
        ("tcp://[::1]:5555", ("tcp", "::1", "5555")),  # IPv6 address
        ("inproc://some_identifier", ("inproc", "some_identifier", "")),
    ],
)
def test_split_zmq_path(path, expected):
    assert split_zmq_path(path) == expected


@pytest.mark.parametrize(
    "invalid_path",
    [
        "invalid_path",  # Missing scheme
        "tcp://127.0.0.1",  # Missing port
        "tcp://[::1]",  # Missing port for IPv6
        "tcp://:5555",  # Missing host
    ],
)
def test_split_zmq_path_invalid(invalid_path):
    with pytest.raises(ValueError):
        split_zmq_path(invalid_path)


def test_make_zmq_socket_ipv6():
    # Check if IPv6 is supported by trying to create an IPv6 socket
    try:
        sock = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
        sock.close()
    except OSError:
        pytest.skip("IPv6 is not supported on this system")

    ctx = zmq.Context()
    ipv6_path = "tcp://[::]:5555"  # IPv6 loopback address
    socket_type = zmq.REP  # Example socket type

    # Create the socket
    zsock: zmq.Socket = make_zmq_socket(ctx, ipv6_path, socket_type)

    # Verify that the IPV6 option is set
    assert zsock.getsockopt(zmq.IPV6) == 1, (
        "IPV6 option should be enabled for IPv6 addresses"
    )

    # Clean up
    zsock.close()
    ctx.term()


def test_make_zmq_path():
    assert make_zmq_path("tcp", "127.0.0.1", "5555") == "tcp://127.0.0.1:5555"
    assert make_zmq_path("tcp", "::1", "5555") == "tcp://[::1]:5555"


def test_get_tcp_uri():
    assert get_tcp_uri("127.0.0.1", 5555) == "tcp://127.0.0.1:5555"
    assert get_tcp_uri("::1", 5555) == "tcp://[::1]:5555"


def test_split_host_port():
    # valid ipv4
    assert split_host_port("127.0.0.1:5555") == ("127.0.0.1", 5555)
    # invalid ipv4
    with pytest.raises(ValueError):
        # multi colon
        assert split_host_port("127.0.0.1::5555")
    with pytest.raises(ValueError):
        # tailing colon
        assert split_host_port("127.0.0.1:5555:")
    with pytest.raises(ValueError):
        # no colon
        assert split_host_port("127.0.0.15555")
    with pytest.raises(ValueError):
        # none int port
        assert split_host_port("127.0.0.1:5555a")

    # valid ipv6
    assert split_host_port("[::1]:5555") == ("::1", 5555)
    # invalid ipv6
    with pytest.raises(ValueError):
        # multi colon
        assert split_host_port("[::1]::5555")
    with pytest.raises(IndexError):
        # no colon
        assert split_host_port("[::1]5555")
    with pytest.raises(ValueError):
        # none int port
        assert split_host_port("[::1]:5555a")


def test_join_host_port():
    assert join_host_port("127.0.0.1", 5555) == "127.0.0.1:5555"
    assert join_host_port("::1", 5555) == "[::1]:5555"