from ipaddress import IPv4Address, IPv4Network
from typing import Any, Dict, List, Tuple

import pytest

from primaite.simulator.network.container import Network
from primaite.simulator.network.hardware.nodes.computer import Computer
from primaite.simulator.network.hardware.nodes.server import Server
from primaite.simulator.network.hardware.nodes.switch import Switch
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.applications.application import Application
from primaite.simulator.system.services.service import Service


class BroadcastService(Service):
    """A service for sending broadcast and unicast messages over a network."""

    def __init__(self, **kwargs):
        # Set default service properties for broadcasting
        kwargs["name"] = "BroadcastService"
        kwargs["port"] = Port.HTTP
        kwargs["protocol"] = IPProtocol.TCP
        super().__init__(**kwargs)

    def describe_state(self) -> Dict:
        # Implement state description for the service
        pass

    def unicast(self, ip_address: IPv4Address):
        # Send a unicast payload to a specific IP address
        super().send(
            payload="unicast",
            dest_ip_address=ip_address,
            dest_port=Port.HTTP,
        )

    def broadcast(self, ip_network: IPv4Network):
        # Send a broadcast payload to an entire IP network
        super().send(
            payload="broadcast",
            dest_ip_address=ip_network,
            dest_port=Port.HTTP,
        )


class BroadcastClient(Application):
    """A client application to receive broadcast and unicast messages."""

    payloads_received: List = []

    def __init__(self, **kwargs):
        # Set default client properties
        kwargs["name"] = "BroadcastClient"
        kwargs["port"] = Port.HTTP
        kwargs["protocol"] = IPProtocol.TCP
        super().__init__(**kwargs)

    def describe_state(self) -> Dict:
        # Implement state description for the application
        pass

    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
        # Append received payloads to the list and print a message
        self.payloads_received.append(payload)
        print(f"Payload: {payload} received on node {self.sys_log.hostname}")


@pytest.fixture(scope="function")
def broadcast_network() -> Network:
    network = Network()

    client_1 = Computer(
        hostname="client_1",
        ip_address="192.168.1.2",
        subnet_mask="255.255.255.0",
        default_gateway="192.168.1.1",
        start_up_duration=0,
    )
    client_1.power_on()
    client_1.software_manager.install(BroadcastClient)
    application_1 = client_1.software_manager.software["BroadcastClient"]
    application_1.run()

    client_2 = Computer(
        hostname="client_2",
        ip_address="192.168.1.3",
        subnet_mask="255.255.255.0",
        default_gateway="192.168.1.1",
        start_up_duration=0,
    )
    client_2.power_on()
    client_2.software_manager.install(BroadcastClient)
    application_2 = client_2.software_manager.software["BroadcastClient"]
    application_2.run()

    server_1 = Server(
        hostname="server_1",
        ip_address="192.168.1.1",
        subnet_mask="255.255.255.0",
        default_gateway="192.168.1.1",
        start_up_duration=0,
    )
    server_1.power_on()

    server_1.software_manager.install(BroadcastService)
    service: BroadcastService = server_1.software_manager.software["BroadcastService"]
    service.start()

    switch_1 = Switch(hostname="switch_1", num_ports=6, start_up_duration=0)
    switch_1.power_on()

    network.connect(endpoint_a=client_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[1])
    network.connect(endpoint_a=client_2.ethernet_port[1], endpoint_b=switch_1.switch_ports[2])
    network.connect(endpoint_a=server_1.ethernet_port[1], endpoint_b=switch_1.switch_ports[3])

    return network


@pytest.fixture(scope="function")
def broadcast_service_and_clients(broadcast_network) -> Tuple[BroadcastService, BroadcastClient, BroadcastClient]:
    client_1: BroadcastClient = broadcast_network.get_node_by_hostname("client_1").software_manager.software[
        "BroadcastClient"
    ]
    client_2: BroadcastClient = broadcast_network.get_node_by_hostname("client_2").software_manager.software[
        "BroadcastClient"
    ]
    service: BroadcastService = broadcast_network.get_node_by_hostname("server_1").software_manager.software[
        "BroadcastService"
    ]

    return service, client_1, client_2


def test_broadcast_correct_subnet(broadcast_service_and_clients):
    service, client_1, client_2 = broadcast_service_and_clients

    assert not client_1.payloads_received
    assert not client_2.payloads_received

    service.broadcast(IPv4Network("192.168.1.0/24"))

    assert client_1.payloads_received == ["broadcast"]
    assert client_2.payloads_received == ["broadcast"]


def test_broadcast_incorrect_subnet(broadcast_service_and_clients):
    service, client_1, client_2 = broadcast_service_and_clients

    assert not client_1.payloads_received
    assert not client_2.payloads_received

    service.broadcast(IPv4Network("192.168.2.0/24"))

    assert not client_1.payloads_received
    assert not client_2.payloads_received


def test_unicast_correct_address(broadcast_service_and_clients):
    service, client_1, client_2 = broadcast_service_and_clients

    assert not client_1.payloads_received
    assert not client_2.payloads_received

    service.unicast(IPv4Address("192.168.1.2"))

    assert client_1.payloads_received == ["unicast"]
    assert not client_2.payloads_received


def test_unicast_incorrect_address(broadcast_service_and_clients):
    service, client_1, client_2 = broadcast_service_and_clients

    assert not client_1.payloads_received
    assert not client_2.payloads_received

    service.unicast(IPv4Address("192.168.2.2"))

    assert not client_1.payloads_received
    assert not client_2.payloads_received