"""
Listen for UDP broadcasts on a port and transmit packet information to other relays

This uses the standard asyncio.DatagramProtocol base class so most of the initial
management of the UDP packet is already done for us.
"""

import array
import asyncio
import ipaddress
import logging
import socket
import struct
from copy import deepcopy
from typing import Any

from .configure import ConfigArgs
from .netutils import get_broadcast_from_iface, get_macaddress_from_iface
from .packet import BadPacketException, Packet
from .pva_packet import log_pvaccess

logger = logging.getLogger(__name__)


class UDPRelayReceive(asyncio.DatagramProtocol):
    """Listen to UDP messages from remote relays and forward them as broadcasts on the local net"""

    def __init__(
        self,
        local_addr: tuple[ipaddress.IPv4Address | ipaddress.IPv6Address | str, int],
        broadcast_port: int,
        config: ConfigArgs | None = None,
    ) -> None:
        super().__init__()

        self.local_addr = (str(local_addr[0]), local_addr[1])  # Get typing right
        self.broadcast_port = broadcast_port
        self.transport: asyncio.DatagramTransport
        self._rebroad_sock: socket.socket

        if config:
            self._iface = config.target_interface
            self._decode_pvaccess = config.decode_pvaccess
        else:
            self._iface = "eth0"
            self._decode_pvaccess = False

        # Assume the MAC address is immutable
        self._mac = get_macaddress_from_iface(self._iface)

        # Also assume the broadcast address associated with the
        # interface is immutable
        self._broadcast_addr = get_broadcast_from_iface(self._iface)

        # A way to manage the forever loop for unit testing and other purposes
        self._loop_forever = True

    def connection_made(self, transport: asyncio.DatagramTransport) -> None:
        """Handle a connection being established"""
        self.transport = transport

    def connection_lost(self, exc: Exception | None) -> None:
        """Handle a connection being lost"""
        # What does connection lost even mean for UDP?
        # Seems only necessary to stop some spurious errors on server shutdown

    def recalculate_udp_checksum(self, ip_packet) -> bytes:
        """Calculate UDP checksum, using the IP and UDP parts of the packet,
        and change the existing packet UDP checksum with the newly calculcated
        checksum"""
        logger.debug("Recalculating UDP checksum")

        # The UDP checksum algorithm is defined in RFC768
        # https://www.rfc-editor.org/rfc/rfc768.txt
        # "Checksum is the 16-bit one's complement of the one's complement sum of a
        #  pseudo header of information from the IP header, the UDP header, and the
        #  data, padded with zero octets at the end (if necessary) to make a
        #  multiple of two octets"

        # Extract the data needed to form the pseudo packet and header
        # Got this from https://dev.to/cwprogram/python-networking-tcp-and-udp-4i3l

        # Make a deepcopy of the UDP portion of the whole ip_packet so that we don't
        # accidentally modify it. Then zero the part that contains the UDP checksum.
        # A zero checksum is valid but we'll calculate the correct value
        pseudo_packet = bytearray(deepcopy(ip_packet[20:]))
        pseudo_packet[6] = 0x0
        pseudo_packet[7] = 0x0

        # We need information from the IP header to construct the pseudo-header
        # needed in turn to calculate the UDP checksum. Specifically we need the
        # source and destination IP addresses
        ip_header = struct.unpack("!BBHHHBBH4s4s", ip_packet[0:20])
        pseudo_header = struct.pack("!4s4sHH", ip_header[8], ip_header[9], socket.IPPROTO_UDP, len(pseudo_packet))

        # Combine the pseudo header and pseudo packet to form a complete pseudo packet
        # that we'll perform the checksum calculations on
        checksum_packet = pseudo_header + pseudo_packet

        # If there is an odd number of bytes in the checksum packet we need to
        # pad it to an even number of bytes
        if len(checksum_packet) % 2 == 1:
            checksum_packet += b"\0"

        # The checksum calculation proceeds by summing the one’s complement where
        # all binary 0s become 1s, of all 16-bit words in these components.
        onecompsum = sum(array.array("H", checksum_packet))
        onecompsum = (onecompsum >> 16) + (onecompsum & 0xFFFF)
        onecompsum += onecompsum >> 16
        onecompsum = ~onecompsum  # Finally invert the bits

        # Test endianness and do some magic if we're on a little endian system
        if struct.pack("H", 1) != b"\x00\x01":
            onecompsum = ((onecompsum >> 8) & 0xFF) | onecompsum << 8

        # If checksum is 0 change it to 0xFFFF to signal it has been calculated
        udp_checksum = onecompsum & 0xFFFF

        # Insert the calculated checksum into the IP + UDP packet
        ip_packet = ip_packet[:26] + udp_checksum.to_bytes(2, "big") + ip_packet[28:]

        return ip_packet

    def datagram_received(self, data: bytes, addr: tuple[str | Any, int]) -> None:
        """Receive a UDP message and forward it any listeners on our local broadcast network segment"""
        logger.debug(
            "Received from %s for rebroadcast on port %i message: %r",
            addr,
            self.broadcast_port,
            data,
        )

        # Simple verification of the received payload, and remove the bytes
        # confirming that this is for us. We also remove the ethernet frame as the
        # sendto() below will take care of that part
        if data[0:2] == b"SS":
            data = data[2:]
        else:
            logger.debug("Malformed packet received")
            return

        # Check if this is a UDP packet and if it is recalculate the UDP checksum
        try:
            packet = Packet(data)
            packet.decode_ip()
            packet.decode_udp()
            data = self.recalculate_udp_checksum(data[14:])
        except BadPacketException:
            pass

        # TODO: Apply any filters

        # TODO: The code above does not change the IP source address
        # If we're on a different network segment then we should switch the
        # broadcast IP address to use get_broadcast_from_iface(). This will
        # then require recomputing checksums. Note: is this required? We
        # are sending to the broadcast address in the sendto() below

        # TODO: Logic to validate what we're receiving a PVAccess message
        # Note that although doing the validation on receipt means we're doing
        # it for every relay (instead of once if we did it on send), it's much
        # safer to do it on receipt since it means we don't have to trust the
        # sender as much

        # TODO: We should be using self.transport.sendto() in order to make
        # this asynchronous but unfortunately that wouldn't allow us to
        # control the IP headers. Is there another way to resolve that?

        # Use this unusual conditional in order to avoid expensive
        # decoding operations when we're not debugging
        if self._decode_pvaccess and logger.isEnabledFor(logging.INFO):
            # Construct a fake ethernet header so we can decode the IP address using existing
            # functionality. This will only work if the original source used IPv4.
            # TODO: This indicates a flaw in the current logic. The received broadcast and the
            # rebroadcast must use the same IP version (IPv4 or IPv6) or the message will
            # be mangled
            packet = Packet(b"\xff\xff\xff\xff\xff\xff\x02B\xac\x16\x00\x02\x08\x00" + data)
            packet.decode_ip()
            packet.decode_udp()
            logger.debug(Packet)

            packet_src_ip = packet.ip_src_addr

            log_pvaccess(data[28:], packet_src_ip)

        # Finally broadcast the new packet
        # It doesn't feel much simpler but we're not using fully raw sockets here
        # but instead letting Python do the work of handling the Ethernet frames
        sendbytes = self._rebroad_sock.sendto(data, (self._broadcast_addr, self.broadcast_port))
        logger.debug(
            "Broadcast UDP packet of length %d to (%s,%s) on iface %s: %s",
            sendbytes,
            self._broadcast_addr,
            self.broadcast_port,
            self._iface,
            data,
        )

    async def start(self) -> None:
        """Start the UDP server that listens for messages from other relays and broadcasts them"""

        logger.info(
            "Starting UDP server listening on %s; will rebroadcast on port %i",
            self.local_addr,
            self.broadcast_port,
        )

        # Open the socket we'll rebroadcast messages on
        self._rebroad_sock = socket.socket(socket.AF_INET, socket.SOCK_RAW, socket.IPPROTO_UDP)
        self._rebroad_sock.setsockopt(socket.IPPROTO_IP, socket.IP_HDRINCL, True)
        self._rebroad_sock.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, True)
        self._rebroad_sock.setblocking(False)

        # Get a reference to the event loop as we plan to use
        # low-level APIs.
        loop = asyncio.get_running_loop()

        # One protocol instance will be created to serve all
        # client requests.
        transport, _ = await loop.create_datagram_endpoint(
            lambda: self, local_addr=self.local_addr, allow_broadcast=True
        )

        try:
            while self._loop_forever:
                # Basically sleep forever!
                await asyncio.sleep(1)
        finally:
            transport.close()
            self._rebroad_sock.close()
