from unittest.mock import patch, MagicMock
import csle_common.constants.constants as constants
from csle_common.controllers.five_g_cu_controller import FiveGCUController
from csle_common.dao.emulation_config.emulation_env_config import EmulationEnvConfig
from csle_common.dao.emulation_config.node_container_config import NodeContainerConfig
from csle_common.dao.emulation_config.five_g_config import FiveGConfig
import pytest


class TestFiveGCUControllerSuite:
    """
    Test five_g_cu_controller
    """

    @pytest.fixture(autouse=True)
    def logger_setup(self) -> None:
        """
        Set up logger

        :return: None
        """
        self.logger = MagicMock()

    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.start_five_g_cu")
    def test_start_five_g_cus(self, mock_start_five_g_cu, example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for starting the 5G CUs

        :param mock_start_five_g_cu: mock_start_five_g_cu
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        constants.CONTAINER_IMAGES.FIVE_G_CU_IMAGES = \
            [example_emulation_env_config.containers_config.containers[0].name]
        FiveGCUController.start_five_g_cus(
            emulation_env_config=example_emulation_env_config,
            physical_server_ip=example_emulation_env_config.containers_config.containers[0].physical_host_ip,
            logger=self.logger)
        mock_start_five_g_cu.assert_called_once_with(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip,
            logger=self.logger)

    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.stop_five_g_cu")
    def test_stop_five_g_cus(self, mock_stop_five_g_cu, example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for stopping the 5G CUs

        :param mock_stop_five_g_cu: mock_stop_five_g_cu
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        constants.CONTAINER_IMAGES.FIVE_G_CU_IMAGES = \
            [example_emulation_env_config.containers_config.containers[0].name]
        FiveGCUController.stop_five_g_cus(
            emulation_env_config=example_emulation_env_config,
            physical_server_ip=example_emulation_env_config.containers_config.containers[0].physical_host_ip,
            logger=self.logger)
        mock_stop_five_g_cu.assert_called_once_with(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip,
            logger=self.logger)

    @patch("csle_collector.five_g_cu_manager.query_five_g_cu_manager.start_five_g_cu")
    @patch("grpc.insecure_channel")
    def test_start_five_g_cu(self, mock_insecure_channel, mock_start_five_g_cu,
                             example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for starting the 5G CU on a specific IP

        :param mock_insecure_channel: mock_insecure_channel
        :param mock_start_five_g_cu: mock_start_five_g_cu
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        constants.GRPC_SERVERS.GRPC_OPTIONS = []
        mock_channel = MagicMock()
        mock_insecure_channel.return_value = mock_channel
        mock_stub = MagicMock()
        mock_channel.__enter__.return_value = mock_stub
        FiveGCUController.start_five_g_cu(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip,
            logger=self.logger)
        self.logger.info.assert_called_once_with(
            f"Starting the 5G CU on container with ip "
            f"{example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip} in "
            f"execution {example_emulation_env_config.execution_id} "
            f"of emulation: {example_emulation_env_config.name}")
        mock_insecure_channel.assert_called()
        mock_start_five_g_cu.assert_called()

    @patch("csle_collector.five_g_cu_manager.query_five_g_cu_manager.stop_five_g_cu")
    @patch("grpc.insecure_channel")
    def test_stop_five_g_cu(self, mock_insecure_channel, mock_stop_five_g_cu,
                            example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for stopping the 5G CU on a specific IP

        :param mock_insecure_channel: mock_insecure_channel
        :param mock_stop_five_g_cu: mock_stop_five_g_cu
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        constants.GRPC_SERVERS.GRPC_OPTIONS = []
        mock_channel = MagicMock()
        mock_insecure_channel.return_value = mock_channel
        mock_stub = MagicMock()
        mock_channel.__enter__.return_value = mock_stub
        FiveGCUController.stop_five_g_cu(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip,
            logger=self.logger)
        self.logger.info.assert_called_once_with(
            f"Stopping the 5G CU on container with ip "
            f"{example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip} in "
            f"execution {example_emulation_env_config.execution_id} "
            f"of emulation: {example_emulation_env_config.name}")
        mock_insecure_channel.assert_called()
        mock_stop_five_g_cu.assert_called()

    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.start_five_g_cu_manager")
    def test_start_five_g_cu_managers(self, start_five_g_cu_manager,
                                      example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for starting 5G CU managers

        :param start_five_g_cu_manager: start_five_g_cu_manager
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        FiveGCUController.start_five_g_cu_managers(
            emulation_env_config=example_emulation_env_config,
            physical_server_ip=example_emulation_env_config.containers_config.containers[0].physical_host_ip,
            logger=self.logger)
        start_five_g_cu_manager.assert_called()

    @patch("csle_common.util.emulation_util.EmulationUtil.connect_admin")
    @patch("csle_common.util.emulation_util.EmulationUtil.execute_ssh_cmd")
    @patch("time.sleep", return_value=None)
    def test_start_five_g_cu_manager(self, mock_sleep, mock_execute_ssh_cmd, mock_connect_admin,
                                     example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for starting the 5G CU manager on a specific IP

        :param mock_execute_ssh_cmd: mock_execute_ssh_cmd
        :param mock_connect_admin: mock_connect_admin
        :param example_emulation_env_config: example_emulation_env_config
        :param mock_sleep: mock_sleep
        :return: None
        """
        mock_connection = MagicMock()
        example_emulation_env_config.get_connection = MagicMock()  # type: ignore[method-assign]
        example_emulation_env_config.get_connection.return_value = mock_connection
        mock_execute_ssh_cmd.side_effect = [
            (b"", b"", 0),  # Output for stopping old background job
            (b"", b"", 0),  # Output for starting the five_g_cu_manager
        ]
        FiveGCUController.start_five_g_cu_manager(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip, logger=self.logger)
        mock_connect_admin.assert_called_once_with(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip)
        mock_execute_ssh_cmd.assert_called()
        self.logger.info.assert_any_call(
            f"Starting 5G CU manager on node "
            f"{example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip}")

    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.stop_five_g_cu_manager")
    def test_stop_five_g_cu_managers(self, stop_five_g_cu_manager,
                                     example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for stopping 5G CU managers

        :param stop_five_g_cu_manager: stop_five_g_cu_manager
        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        FiveGCUController.stop_five_g_cu_managers(
            emulation_env_config=example_emulation_env_config,
            physical_server_ip=example_emulation_env_config.containers_config.containers[0].physical_host_ip,
            logger=self.logger)
        stop_five_g_cu_manager.assert_called()

    @patch("csle_common.util.emulation_util.EmulationUtil.connect_admin")
    @patch("csle_common.util.emulation_util.EmulationUtil.execute_ssh_cmd")
    @patch("time.sleep", return_value=None)
    def test_stop_five_g_cu_manager(self, mock_sleep, mock_execute_ssh_cmd, mock_connect_admin,
                                    example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test utility function for stopping the 5G CU manager on a specific IP

        :param mock_execute_ssh_cmd: mock_execute_ssh_cmd
        :param mock_connect_admin: mock_connect_admin
        :param example_emulation_env_config: example_emulation_env_config
        :param mock_sleep: mock_sleep
        :return: None
        """
        mock_connection = MagicMock()
        example_emulation_env_config.get_connection = MagicMock()  # type: ignore[method-assign]
        example_emulation_env_config.get_connection.return_value = mock_connection
        mock_execute_ssh_cmd.side_effect = [(b"", b"", 0)]
        FiveGCUController.stop_five_g_cu_manager(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip, logger=self.logger)
        mock_connect_admin.assert_called_once_with(
            emulation_env_config=example_emulation_env_config,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip)
        mock_execute_ssh_cmd.assert_called()
        self.logger.info.assert_any_call(
            f"Stopping 5G CU manager on node "
            f"{example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip}")

    def test_get_five_g_cu_managers_ips(self, example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test a method that extracts the IPs of the 5G Cu managers in a given emulation

        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        ips = FiveGCUController.get_five_g_cu_managers_ips(emulation_env_config=example_emulation_env_config)
        expected_ips = [example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip]
        assert ips == expected_ips

    def test_get_five_g_cu_managers_ports(self, example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test a method that extracts the ports of the 5G CU managers in a given emulation

        :param example_emulation_env_config: example_emulation_env_config
        :return: None
        """
        ports = FiveGCUController.get_five_g_cu_managers_ports(emulation_env_config=example_emulation_env_config)
        expected_ports = [50053]
        assert ports == expected_ports

    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.get_five_g_cu_managers_ips")
    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController.get_five_g_cu_managers_ports")
    @patch("csle_common.controllers.five_g_cu_controller.FiveGCUController."
           "get_five_g_cu_status_by_ip_and_port")
    @patch("csle_common.util.emulation_util.EmulationUtil.physical_ip_match")
    def test_get_five_g_cu_managers_info(self, mock_physical_ip_match, mock_get_statuses, mock_get_ports,
                                         mock_get_ips, example_emulation_env_config: EmulationEnvConfig) -> None:
        """
        Test the method that extracts the information of the 5G CU managers for a given emulation

        :param mock_physical_ip_match: mock_physical_ip_match
        :param mock_get_statuses: mock_get_statuses
        :param mock_get_ports:mock_get_ports
        :param mock_get_ips: mock_get_ips
        :return: None
        """
        mock_get_ips.return_value = [example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip]
        mock_get_ports.return_value = [example_emulation_env_config.five_g_config.five_g_cu_manager_port]
        mock_status = MagicMock()
        mock_get_statuses.side_effect = [mock_status, Exception("Test exception")]
        mock_physical_ip_match.side_effect = [True, False]
        active_ips = [example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip]
        physical_server_ip = example_emulation_env_config.containers_config.containers[0].physical_host_ip
        FiveGCUController.get_five_g_cu_managers_info(
            emulation_env_config=example_emulation_env_config, active_ips=active_ips, logger=self.logger,
            physical_server_ip=physical_server_ip)
        mock_get_ips.assert_called_once_with(emulation_env_config=example_emulation_env_config)
        mock_get_ports.assert_called_once_with(emulation_env_config=example_emulation_env_config)
        mock_get_statuses.assert_any_call(
            port=example_emulation_env_config.five_g_config.five_g_cu_manager_port,
            ip=example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip)

    @patch("csle_collector.five_g_cu_manager.query_five_g_cu_manager.init_five_g_cu")
    @patch("grpc.insecure_channel")
    def test_init_five_g_cu(self, mock_insecure_channel, mock_init_five_g_cu,
                            example_emulation_env_config: EmulationEnvConfig,
                            example_containers_config_five_g: NodeContainerConfig,
                            example_five_g_config_two: FiveGConfig) -> None:
        """
        Test utility function for initializing the 5G UE on a specific IP

        :param mock_insecure_channel: mock_insecure_channel
        :param mock_init_five_g_cu: mock_init_five_g_cu
        :param example_emulation_env_config: example_emulation_env_config
        :param example_containers_config_five_g: example_containers_config_five_g
        :param example_five_g_config_two: example_five_g_config_two
        :return: None
        """
        constants.GRPC_SERVERS.GRPC_OPTIONS = []
        mock_channel = MagicMock()
        mock_insecure_channel.return_value = mock_channel
        mock_stub = MagicMock()
        mock_channel.__enter__.return_value = mock_stub
        result = FiveGCUController.init_five_g_cu(
            emulation_env_config=example_emulation_env_config,
            container=example_emulation_env_config.containers_config.containers[0],
            logger=self.logger)
        self.logger.info.assert_called_once_with(
            f"Initializing the 5G CU on container with ip "
            f"{example_emulation_env_config.containers_config.containers[0].docker_gw_bridge_ip} "
            f"in execution {example_emulation_env_config.execution_id} "
            f"of emulation: {example_emulation_env_config.name}")
        assert result is None
        example_emulation_env_config.containers_config = example_containers_config_five_g
        example_emulation_env_config.five_g_config = example_five_g_config_two
        result = FiveGCUController.init_five_g_cu(
            emulation_env_config=example_emulation_env_config,
            container=example_emulation_env_config.containers_config.containers[4],
            logger=self.logger)
        assert result is not None
        mock_insecure_channel.assert_called()
        mock_init_five_g_cu.assert_called()
