"""The packet format every sender and receiver in this project uses. Given; do not change it.

    0      1              5        7          9
    +------+--------------+--------+----------+-----------------+
    | type |  seq (32)    | len    | checksum | payload (len)   |
    +------+--------------+--------+----------+-----------------+

All header fields are big-endian ("network order"). The checksum is the
Internet checksum (RFC 1071) over the whole packet with the checksum field set
to zero, the same one IPv4, UDP and TCP use.

References:
    https://docs.python.org/3/library/struct.html
    https://www.rfc-editor.org/rfc/rfc1071
"""

import struct

DATA = 0      # carries payload; seq numbers the packet, starting at 0
ACK = 1       # acknowledges; what seq means depends on the protocol (see README)
FIN = 2       # the sender is done; seq is the number of data packets it sent
FINACK = 3    # the receiver has everything and is closing

HEADER = struct.Struct("!BIHH")
MSS = 1000    # most payload bytes in one packet


def _checksum(data):
    if len(data) % 2:
        data += b"\0"
    total = sum(struct.unpack(f"!{len(data) // 2}H", data))
    while total >> 16:
        total = (total & 0xFFFF) + (total >> 16)
    return ~total & 0xFFFF


def encode(kind, seq, payload=b""):
    """Build one packet as bytes."""
    if len(payload) > MSS:
        raise ValueError(f"payload of {len(payload)} bytes is over the MSS of {MSS}")
    header = HEADER.pack(kind, seq, len(payload), 0)
    check = _checksum(header + payload)
    return HEADER.pack(kind, seq, len(payload), check) + payload


def decode(raw):
    """Return (kind, seq, payload), or None if the packet is short or corrupted."""
    if len(raw) < HEADER.size:
        return None
    kind, seq, length, check = HEADER.unpack_from(raw)
    payload = raw[HEADER.size:]
    if length != len(payload) or kind not in (DATA, ACK, FIN, FINACK):
        return None
    if _checksum(HEADER.pack(kind, seq, length, 0) + payload) != check:
        return None
    return kind, seq, payload


def split(data):
    """Cut a byte string into MSS-sized payloads, in order."""
    return [data[i:i + MSS] for i in range(0, len(data), MSS)]