from __future__ import annotations

import socket
import struct


DDP_DEFAULT_PORT = 4048
DDP_MAX_DATA_LENGTH = 1440
DDP_VERSION_1 = 0x40
DDP_PUSH_FLAG = 0x01
DDP_RGB888 = 0x0B
DDP_DESTINATION_ID = 1


class DDPClient:
    def __init__(self, port: int = DDP_DEFAULT_PORT) -> None:
        self.port = int(port)
        self._sequence = 1

    def count_packets(self, host_payloads: dict[str, bytes]) -> int:
        packet_count = 0
        for payload in host_payloads.values():
            chunks, remainder = divmod(len(payload), DDP_MAX_DATA_LENGTH)
            packet_count += chunks + (1 if remainder else 0)
        return packet_count

    def build_packets(self, host_payloads: dict[str, bytes]) -> list[tuple[str, bytes]]:
        packets: list[tuple[str, bytes]] = []
        sequence = self._next_sequence()
        for host, payload in host_payloads.items():
            for offset in range(0, len(payload), DDP_MAX_DATA_LENGTH):
                chunk = payload[offset : offset + DDP_MAX_DATA_LENGTH]
                last = offset + DDP_MAX_DATA_LENGTH >= len(payload)
                header = struct.pack(
                    "!BBBBLH",
                    DDP_VERSION_1 | (DDP_PUSH_FLAG if last else 0),
                    sequence,
                    DDP_RGB888,
                    DDP_DESTINATION_ID,
                    offset,
                    len(chunk),
                )
                packets.append((host, header + chunk))
        return packets

    def send(self, host_payloads: dict[str, bytes]) -> int:
        packets = self.build_packets(host_payloads)
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as sock:
            for host, packet in packets:
                sock.sendto(packet, (host, self.port))
        return len(packets)

    def _next_sequence(self) -> int:
        sequence = self._sequence
        self._sequence = 1 if self._sequence >= 15 else self._sequence + 1
        return sequence