# © Crown-owned copyright 2024, Defence Science and Technology Laboratory UK
from __future__ import annotations

from abc import abstractmethod
from typing import Any, Dict, Optional, Union

from prettytable import MARKDOWN, PrettyTable

from primaite.simulator.network.hardware.base import NetworkInterface
from primaite.simulator.network.protocols.arp import ARPEntry, ARPPacket
from primaite.simulator.network.transmission.network_layer import IPProtocol
from primaite.simulator.network.transmission.transport_layer import Port
from primaite.simulator.system.services.service import Service
from primaite.utils.validators import IPV4Address


class ARP(Service):
    """
    The ARP (Address Resolution Protocol) Service.

    Manages ARP for resolving network layer addresses into link layer addresses. It maintains an ARP cache,
    sends ARP requests and replies, and processes incoming ARP packets.
    """

    arp: Dict[IPV4Address, ARPEntry] = {}

    def __init__(self, **kwargs):
        kwargs["name"] = "ARP"
        kwargs["port"] = Port.ARP
        kwargs["protocol"] = IPProtocol.UDP
        super().__init__(**kwargs)

    def describe_state(self) -> Dict:
        """
        Produce a dictionary describing the current state of this object.

        :return: Current state of this object and child objects.
        """
        state = super().describe_state()
        state.update({str(ip): arp_entry.mac_address for ip, arp_entry in self.arp.items()})

        return super().describe_state()

    def show(self, markdown: bool = False):
        """
        Prints the current state of the ARP cache in a table format.

        :param markdown: If True, format the output as Markdown. Otherwise, use plain text.
        """
        table = PrettyTable(["IP Address", "MAC Address", "Via"])
        if markdown:
            table.set_style(MARKDOWN)
        table.align = "l"
        table.title = f"{self.sys_log.hostname} ARP Cache"
        for ip, arp in self.arp.items():
            table.add_row(
                [
                    str(ip),
                    arp.mac_address,
                    self.software_manager.node.network_interfaces[arp.network_interface_uuid].mac_address,
                ]
            )
        print(table)

    def clear(self):
        """Clears the arp cache."""
        self.arp.clear()

    def get_default_gateway_network_interface(self) -> Optional[NetworkInterface]:
        """Not used at the parent ARP level. Should return None when there is no override by child class."""
        return None

    def add_arp_cache_entry(
        self, ip_address: IPV4Address, mac_address: str, network_interface: NetworkInterface, override: bool = False
    ):
        """
        Add an ARP entry to the cache.

        If an entry for the given IP address already exists, the entry is only updated if the `override` parameter is
        set to True.

        :param ip_address: The IP address to be added to the cache.
        :param mac_address: The MAC address associated with the IP address.
        :param network_interface: The NIC through which the NIC with the IP address is reachable.
        :param override: If True, an existing entry for the IP address will be overridden. Default is False.
        """
        for _network_interface in self.software_manager.node.network_interfaces.values():
            if _network_interface.ip_address == ip_address:
                return
        if override or not self.arp.get(ip_address):
            self.sys_log.info(f"Adding ARP cache entry for {mac_address}/{ip_address} via NIC {network_interface}")
            arp_entry = ARPEntry(mac_address=mac_address, network_interface_uuid=network_interface.uuid)

            self.arp[ip_address] = arp_entry

    @abstractmethod
    def get_arp_cache_mac_address(self, ip_address: IPV4Address) -> Optional[str]:
        """
        Retrieves the MAC address associated with a given IP address from the ARP cache.

        :param ip_address: The IP address to look up.
        :return: The associated MAC address, if found. Otherwise, returns None.
        """
        pass

    @abstractmethod
    def get_arp_cache_network_interface(self, ip_address: IPV4Address) -> Optional[NetworkInterface]:
        """
        Retrieves the NIC associated with a given IP address from the ARP cache.

        :param ip_address: The IP address to look up.
        :return: The associated NIC, if found. Otherwise, returns None.
        """
        pass

    def send_arp_request(self, target_ip_address: Union[IPV4Address, str]):
        """
        Sends an ARP request to resolve the MAC address of a target IP address.

        :param target_ip_address: The target IP address for which the MAC address is being requested.
        """
        if target_ip_address in self.arp:
            return

        use_default_gateway = True
        for network_interface in self.software_manager.node.network_interfaces.values():
            if target_ip_address in network_interface.ip_network:
                use_default_gateway = False
                break

        if use_default_gateway:
            if self.software_manager.node.default_gateway:
                target_ip_address = self.software_manager.node.default_gateway
            else:
                return

        outbound_network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(
            target_ip_address
        )
        if outbound_network_interface:
            self.sys_log.info(f"Sending ARP request from NIC {outbound_network_interface} for ip {target_ip_address}")
            arp_packet = ARPPacket(
                sender_ip_address=outbound_network_interface.ip_address,
                sender_mac_addr=outbound_network_interface.mac_address,
                target_ip_address=target_ip_address,
            )
            self.software_manager.session_manager.receive_payload_from_software_manager(
                payload=arp_packet, dst_ip_address=target_ip_address, dst_port=self.port, ip_protocol=self.protocol
            )
        else:
            self.sys_log.warning(
                "Cannot send ARP request as there is no outbound Network Interface to use. Try configuring the default "
                "gateway."
            )

    def send_arp_reply(self, arp_reply: ARPPacket):
        """
        Sends an ARP reply in response to an ARP request.

        :param arp_reply: The ARP packet containing the reply.
        """
        outbound_network_interface = self.software_manager.session_manager.resolve_outbound_network_interface(
            arp_reply.target_ip_address
        )
        if outbound_network_interface:
            self.sys_log.info(
                f"Sending ARP reply from {arp_reply.sender_mac_addr}/{arp_reply.sender_ip_address} "
                f"to {arp_reply.target_ip_address}/{arp_reply.target_mac_addr} "
            )
            self.software_manager.session_manager.receive_payload_from_software_manager(
                payload=arp_reply,
                dst_ip_address=arp_reply.target_ip_address,
                dst_port=self.port,
                ip_protocol=self.protocol,
            )
        else:
            self.sys_log.warning(
                "Cannot send ARP reply as there is no outbound Network Interface to use. Try configuring the default "
                "gateway."
            )

    @abstractmethod
    def _process_arp_request(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface):
        """
        Processes an incoming ARP request.

        :param arp_packet: The ARP packet containing the request.
        :param from_network_interface: The NIC that received the ARP request.
        """
        self.sys_log.info(
            f"Received ARP request for {arp_packet.target_ip_address} from "
            f"{arp_packet.sender_mac_addr}/{arp_packet.sender_ip_address} "
        )

    def _process_arp_reply(self, arp_packet: ARPPacket, from_network_interface: NetworkInterface):
        """
        Processes an incoming ARP reply.

        :param arp_packet: The ARP packet containing the reply.
        :param from_network_interface: The NIC that received the ARP reply.
        """
        self.sys_log.info(
            f"Received ARP response for {arp_packet.sender_ip_address} "
            f"from {arp_packet.sender_mac_addr} via Network Interface {from_network_interface}"
        )
        self.add_arp_cache_entry(
            ip_address=arp_packet.sender_ip_address,
            mac_address=arp_packet.sender_mac_addr,
            network_interface=from_network_interface,
        )

    def receive(self, payload: Any, session_id: str, **kwargs) -> bool:
        """
        Processes received data, handling ARP packets.

        :param payload: The payload received.
        :param session_id: The session ID associated with the received data.
        :param kwargs: Additional keyword arguments.
        :return: True if the payload was processed successfully, otherwise False.
        """
        if not super().receive(payload, session_id, **kwargs):
            return False

        from_network_interface = kwargs["from_network_interface"]
        if payload.request:
            self._process_arp_request(arp_packet=payload, from_network_interface=from_network_interface)
        else:
            self._process_arp_reply(arp_packet=payload, from_network_interface=from_network_interface)
        return True

    def __contains__(self, item: Any) -> bool:
        """
        Checks if an item is in the ARP cache.

        :param item: The item to check.
        :return: True if the item is in the cache, otherwise False.
        """
        return item in self.arp