#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import io
import json
import logging
import socket
from contextlib import contextmanager
from unittest.mock import MagicMock, patch

import pytest
import torch

from vllm_ascend.model_loader.netloader.interaction import elastic
from vllm_ascend.model_loader.netloader.interaction.elastic import ElasticClient, ElasticServer


# Simulate server's normal response
def mock_server_response(data):
    return json.dumps({"label": "JOIN_ACK", "content": {"name": "mocked_name"}}).encode("utf-8")


# Simulate server's error response
def mock_server_error_response(data):
    return json.dumps({"label": "JOIN_ACK", "content": None}).encode("utf-8")


# Simulated server's abnormal response
def mock_server_exception_response(data):
    raise Exception("Mocked server exception")


@contextmanager
def capture_elastic_logs(level=logging.DEBUG):
    log_capture_string = io.StringIO()
    handler = logging.StreamHandler(log_capture_string)
    handler.setLevel(level)
    original_level = elastic.logger.level
    elastic.logger.setLevel(level)
    elastic.logger.addHandler(handler)
    try:
        yield log_capture_string
    finally:
        elastic.logger.removeHandler(handler)
        elastic.logger.setLevel(original_level)
        log_capture_string.close()


# Test the initialization of ElasticClient
def test_elastic_client_init():
    sources = ["127.0.0.1:12345"]
    device_id = 0
    model_path = "mocked_model_path"
    tp = 1
    pp = 1

    with patch("socket.socket") as mock_socket:
        mock_socket_instance = MagicMock()
        mock_socket.return_value = mock_socket_instance
        mock_socket_instance.recv.return_value = mock_server_response(None)

        mock_socket_instance.getsockname.return_value = ("127.0.0.1", 12346)
        mock_socket_instance.__enter__.return_value = mock_socket_instance

        with ElasticClient(sources, device_id, model_path, tp, pp) as client:
            assert client.server_addr == "127.0.0.1"
            assert client.server_port == 12345
            assert client.ack == ("mocked_name", 12346)
        mock_socket_instance.close.assert_called_once()


# Test the register method of ElasticClient
def test_elastic_client_register():
    sources = ["127.0.0.1:12345"]
    device_id = 0
    model_path = "mocked_model_path"
    tp = 1
    pp = 1

    with patch("socket.socket") as mock_socket:
        mock_socket_instance = MagicMock()
        mock_socket.return_value = mock_socket_instance
        mock_socket_instance.connect.return_value = None
        mock_socket_instance.recv.return_value = mock_server_response(None)

        mock_socket_instance.getsockname.return_value = ("127.0.0.1", 12346)
        mock_socket_instance.__enter__.return_value = mock_socket_instance

        client = ElasticClient(sources, device_id, model_path, tp, pp)
        assert client.register(device_id, model_path, tp, pp) == ("mocked_name", 12346)


# Test the behavior of the `register` method of ElasticClient when the server returns an error response.
def test_elastic_client_register_error_response():
    sources = ["127.0.0.1:12345"]
    device_id = 0
    model_path = "mocked_model_path"
    tp = 1
    pp = 1

    with patch("socket.socket") as mock_socket:
        mock_socket_instance = MagicMock()
        mock_socket.return_value = mock_socket_instance
        mock_socket_instance.connect.return_value = None
        mock_socket_instance.recv.return_value = mock_server_error_response(None)

        with ElasticClient(sources, device_id, model_path, tp, pp) as client, pytest.raises(RuntimeError):
            client.register(device_id, model_path, tp, pp)
        mock_socket_instance.close.assert_called_once()


# Test the behavior of the `register` method of ElasticClient when an exception is thrown on the server.
def test_elastic_client_register_exception():
    sources = ["127.0.0.1:12345"]
    device_id = 0
    model_path = "mocked_model_path"
    tp = 1
    pp = 1

    with patch("socket.socket") as mock_socket:
        mock_socket_instance = MagicMock()
        mock_socket.return_value = mock_socket_instance
        mock_socket_instance.connect.return_value = None
        mock_socket_instance.recv.side_effect = mock_server_exception_response
        mock_socket_instance.__enter__.return_value = mock_socket_instance
        mock_socket_instance.__exit__.return_value = None

        with ElasticClient(sources, device_id, model_path, tp, pp) as client, pytest.raises(RuntimeError):
            client.register(device_id, model_path, tp, pp)
        mock_socket_instance.close.assert_called_once()


class FakeInt8Param:
    def __init__(self, name="param", device="npu", dtype=torch.int8):
        self.dtype = dtype
        self.device = torch.device(device)

    @property
    def data(self):
        return self  # Simulate .data returning self so .cpu() etc. can be chained

    def clone(self):
        return self

    def detach(self):
        return self

    def cpu(self):
        self.device = torch.device("cpu")
        return self


class FakeModel:
    def __init__(self):
        self.params = {
            "param1": MagicMock(dtype=torch.float32),  # This will be ignored
            "param2": FakeInt8Param(),  # This simulates a real int8 param
        }

    def named_parameters(self):
        return self.params.items()


@pytest.fixture
def mock_model():
    return FakeModel()


@pytest.fixture
def server_config():
    return {
        "addr": "127.0.0.1",
        "port": 8080,
        "model": MagicMock(),
        "device_id": 0,
        "model_path": "/test/model",
        "tp": 1,
        "pp": 1,
        "int8_cache": "dram",
        "int8_cache_name": None,
    }


# Test server initialization
def test_server_initialization(server_config, mock_model):
    server_config["model"] = mock_model
    with patch("socket.socket") as mock_socket, capture_elastic_logs() as log_capture_string:
        server = ElasticServer(**server_config)

        # Check the socket configuration
        mock_socket.assert_called_with(socket.AF_INET, socket.SOCK_STREAM)
        mock_socket.return_value.setsockopt.assert_called_with(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        mock_socket.return_value.bind.assert_called_with(("127.0.0.1", 8080))
        mock_socket.return_value.listen.assert_called_with(256)

        # Check int8 cache
        assert "param2" in server.original_int8
        assert server.original_int8["param2"].device.type == "cpu"  # Verifying DRAM Cache

        assert server.addr == server_config["addr"]
        assert server.port == server_config["port"]
        assert server.device_id == server_config["device_id"]
        assert server.model_path == server_config["model_path"]
        assert server.tp == server_config["tp"]
        assert server.pp == server_config["pp"]

        # Get captured logs
        log_output = log_capture_string.getvalue()

        # Check output
        assert "Server 127.0.0.1:8080 starts" in log_output


# Test the int8 cache option
@pytest.mark.parametrize("cache_option,expected_device", [("dram", "cpu"), ("no", None), ("invalid", None)])
def test_int8_cache_handling(server_config, mock_model, cache_option, expected_device):
    server_config["int8_cache"] = cache_option
    server_config["model"] = mock_model

    with patch("socket.socket"), capture_elastic_logs() as log_capture_string:
        server = ElasticServer(**server_config)

        log_output = log_capture_string.getvalue()

        if cache_option == "invalid":
            assert "int8_cache should be selected in [HBM, DRAM]" in log_output

        if expected_device is None:
            assert len(server.original_int8) == 0
        else:
            assert server.original_int8["param2"].device.type == expected_device


# Test client processing
def test_client_handler_valid_join(server_config, mock_model):
    server_config["model"] = mock_model
    with patch("vllm_ascend.model_loader.netloader.interaction.elastic.P2PSend") as mock_p2p_send:
        # Create a simulated connection
        mock_conn = MagicMock()
        mock_addr = ("192.168.1.1", 12345)

        # Configuring Client Data
        valid_data = {
            "label": "JOIN",
            "content": {"device_id": 0, "model_path": "/test/model", "tp": 1, "pp": 1, "port": 9090},
        }
        mock_conn.recv.return_value = json.dumps(valid_data).encode("utf-8")

        # Start the server
        server = ElasticServer(**server_config)
        server.register_handler(mock_conn, mock_addr)

        # Verify response
        expected_ack = {"label": "JOIN_ACK", "content": {"name": "192.168.1.1:12345"}}
        mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
        mock_p2p_send.assert_called_once_with("127.0.0.1", 9090, "192.168.1.1:12345", "netloader")
        mock_conn.close.assert_called_once()


# Test mismatched JOIN requests
def test_client_handler_mismatch(server_config):
    with patch("socket.socket"):
        server = ElasticServer(**server_config)
        mock_conn = MagicMock()
        mock_addr = ("192.168.1.1", 12345)

        # Send mismatched data
        mismatch_data = {
            "label": "JOIN",
            "content": {
                "device_id": 1,  # 不匹配的ID
                "model_path": "/wrong/model",
                "tp": 2,
                "pp": 2,
                "port": 9090,
            },
        }
        mock_conn.recv.return_value = json.dumps(mismatch_data).encode("utf-8")

        server.register_handler(mock_conn, mock_addr)

        assert isinstance(mismatch_data["content"], dict)

        # Verify response
        mismatch_tuple = (
            mismatch_data["content"]["device_id"],
            mismatch_data["content"]["model_path"],
            mismatch_data["content"]["tp"],
            mismatch_data["content"]["pp"],
        )

        server_tuple = (
            server_config["device_id"],
            server_config["model_path"],
            server_config["tp"],
            server_config["pp"],
        )

        expected_ack = {
            "label": "JOIN_NACK",
            "content": (f"Received data {mismatch_tuple} does not consist with this server {server_tuple}"),
        }
        mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
        mock_conn.close.assert_called_once()


# Test Invalid Request
@pytest.mark.parametrize(
    "invalid_data,should_send",
    [
        ({"label": "WRONG_LABEL"}, True),  # Incorrect label, can be decoded as JSON, but the content is invalid.
        (
            {"content": {"missing_fields": True}},
            True,
        ),  # Missing field, can be decoded as JSON, but the content is invalid.
        ("plain text", False),  # Non-JSON data, json.loads failed
        (b"invalid_bytes", False),  # Invalid byte, decode or json.loads failed
    ],
)
def test_client_handler_invalid_requests(server_config, invalid_data, should_send):
    with patch("socket.socket"), capture_elastic_logs() as log_capture_string:
        server = ElasticServer(**server_config)
        mock_conn = MagicMock()
        mock_addr = ("192.168.1.1", 12345)

        if isinstance(invalid_data, (str, bytes)):
            mock_conn.recv.return_value = invalid_data if isinstance(invalid_data, bytes) else invalid_data.encode()
        else:
            mock_conn.recv.return_value = json.dumps(invalid_data).encode("utf-8")

        server.register_handler(mock_conn, mock_addr)

        if should_send:
            expected_ack = {
                "label": "JOIN_NACK",
                "content": f"Received data does not contain required fields: {invalid_data}",
            }
            mock_conn.send.assert_called_once_with(json.dumps(expected_ack).encode("utf-8"))
        else:
            mock_conn.send.assert_not_called()

        log_output = log_capture_string.getvalue()

        # Any warning in the log is acceptable
        assert "Failed to load" in log_output or "does not contain" in log_output
        mock_conn.close.assert_called_once()


# Test the thread startup.
def test_server_start(server_config):
    with patch("socket.socket"), patch("threading.Thread") as mock_thread:
        handler_thread_instance = mock_thread.return_value

        server = ElasticServer(**server_config)
        server.start()

        # Assert that the correct target parameter was passed when instantiating the Thread instance.
        mock_thread.assert_called_once()
        args, kwargs = mock_thread.call_args
        assert kwargs["target"] == server.elastic_client_handler

        # Verify the daemon attribute is set to True (the attribute value will be recorded after MagicMock assignment).
        assert handler_thread_instance.daemon is True

        # Check if the start() method is called.
        handler_thread_instance.start.assert_called_once()


# Test resource clearing
def test_server_cleanup(server_config):
    with patch("socket.socket") as mock_socket:
        server = ElasticServer(**server_config)
        del server
        mock_socket.return_value.close.assert_called_once()


def test_draft_group_name_in_client_register():
    sent_payloads = []

    with (
        patch("socket.socket") as mock_socket,
        patch("vllm_ascend.model_loader.netloader.interaction.elastic.find_free_port", return_value=12346),
    ):
        mock_socket_instance = MagicMock()
        mock_socket.return_value = mock_socket_instance
        mock_socket_instance.recv.return_value = mock_server_response(None)
        mock_socket_instance.send.side_effect = lambda data: sent_payloads.append(json.loads(data.decode()))

        client = ElasticClient(["127.0.0.1:12345"], 0, "draft-model", 1, 1, "netloader_draft")
        client.s = mock_socket_instance
        client.register(0, "draft-model", 1, 1)

    assert client.group_name == "netloader_draft"
    assert sent_payloads[0]["content"]["group_name"] == "netloader_draft"


def test_draft_group_name_in_server_p2p_send(server_config, mock_model):
    server_config["model"] = mock_model
    join_data = {
        "label": "JOIN",
        "content": {
            "device_id": 0,
            "model_path": "/test/model",
            "tp": 1,
            "pp": 1,
            "port": 9090,
            "group_name": "netloader_draft",
        },
    }

    with (
        patch("socket.socket"),
        patch("vllm_ascend.model_loader.netloader.interaction.elastic.P2PSend") as mock_p2p_send,
    ):
        mock_conn = MagicMock()
        mock_conn.recv.return_value = json.dumps(join_data).encode("utf-8")

        ElasticServer(**server_config, group_name="netloader_draft").register_handler(mock_conn, ("192.168.1.1", 12345))

        mock_p2p_send.assert_called_once_with("127.0.0.1", 9090, "192.168.1.1:12345", "netloader_draft")


if __name__ == "__main__":
    pytest.main()