import socket
import struct
import zlib
from dataclasses import dataclass

from hegel.protocol.utils import (
    ChannelId,
    ConnectionClosedError,
    MessageId,
    ProtocolError,
)

# 5 unsigned 32-bit integers, big-endian:
# magic cookie, checksum, channel, message ID, payload length
PACKET_HEADER_FORMAT = ">5I"
# ASCII for "HEGL"
PACKET_MAGIC = 0x4845474C
PACKET_TERMINATOR = 0x0A  # '\n'
# If this is set in the message id, this packet is a reply to a previous packet
REPLY_BIT = 1 << 31

# Special payload that is sent on a channel when it is shut down. The shutdown
# is not acked and is handled specifially.
# Chosen to be invalid CBOR as per https://www.rfc-editor.org/rfc/rfc8949.html
# It is currently also not the prefix of any valid CBOR (this is a reserved)
# tag byte) but even if it became valid in future this would not be a problem.
CLOSE_CHANNEL_PAYLOAD = bytes([0b11111110])
CLOSE_CHANNEL_MESSAGE_ID = MessageId((1 << 31) - 1)


@dataclass(frozen=True, slots=True)
class Packet:
    """
    A logical message in the protocol.

    Packets are the only valid way to send bytes over the wire in the protocol. No "raw"
    bytes are ever sent.

    Wire format:

        0                   1                   2                   3
        0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                     Magic (0x4845474C)                        |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                     Checksum (CRC32)                          |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                     Channel id                              |S|
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |R|                   Message id                                |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                     Payload length                            |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        |                     Payload (variable length)                 |
        +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
        | Terminator 0x0A |
        +-+-+-+-+-+-+-+-+-+

    The first five fields comprise the header. Each field is a unsigned 32-bit big-endian
    integer:
    - Magic: The constant 0x4845474C (ASCII for "HEGL").
    - Checksum: CRC32 of the header with the checksum field zeroed, concatenated with
       the payload.
    - Channel id: The logical channel this packet is being sent over. The S (source) bit
       is 1 for channels created by the client, and 0 for channels created by the server.
       The S bit is only part of the protocol to allow both the client and server to
       create channels without coordination.
    - Message id: The id of the message. The R (reply) bit is set if this packet is a reply
       to a previous packet. The message id of a reply packet will be the same as the
       message id of a non-reply packet, but with the R bit set. The message id is
       included in the protocol to support out-of-order replies over the same channel.
    - Payload length: The length of the payload, in bytes.

    The header is followed by the variable-length payload field, and then a single
    terminator byte (0x0A).
    """

    channel_id: ChannelId
    message_id: MessageId
    is_reply: bool
    payload: bytes


def read_exact(sock: socket.socket, *, n: int) -> bytes:
    """Read exactly n bytes from the socket."""
    assert n >= 0
    if n == 0:
        return b""

    data = bytearray()
    while len(data) < n:
        chunk = sock.recv(n - len(data))
        if chunk:
            data.extend(chunk)
            continue

        if not data:
            raise ConnectionClosedError("Connection closed")
        raise ProtocolError(
            f"Connection closed during socket read (bytes read so far: {data!r})"
        )
    return bytes(data)


def read_packet(sock: socket.socket, *, timeout: float | None = None) -> Packet:
    sock.settimeout(timeout)
    header = read_exact(sock, n=struct.calcsize(PACKET_HEADER_FORMAT))
    sock.settimeout(None)
    magic, checksum, channel, message_id, length = struct.unpack(
        PACKET_HEADER_FORMAT, header
    )
    assert magic == PACKET_MAGIC

    is_reply = (message_id & REPLY_BIT) != 0
    if is_reply:
        message_id ^= REPLY_BIT

    payload = read_exact(sock, n=length)
    terminator = read_exact(sock, n=1)[0]
    assert terminator == PACKET_TERMINATOR

    # checksum is defined as crc(header + payload), where the header's checksum has
    # been zeroed
    zeroed_header = header[:4] + b"\x00\x00\x00\x00" + header[8:]
    assert zlib.crc32(zeroed_header + payload) == checksum

    return Packet(
        channel_id=channel,
        message_id=message_id,
        payload=payload,
        is_reply=is_reply,
    )


def write_packet(sock: socket.socket, packet: Packet) -> None:
    message_id: int = packet.message_id
    if packet.is_reply:
        message_id |= REPLY_BIT

    # checksum is defined as crc(header + payload), where the header's checksum has
    # been zeroed
    zeroed_header = struct.pack(
        ">5I", PACKET_MAGIC, 0, packet.channel_id, message_id, len(packet.payload)
    )
    checksum = zlib.crc32(zeroed_header + packet.payload)

    zeroed_header = struct.pack(
        ">5I",
        PACKET_MAGIC,
        checksum,
        packet.channel_id,
        message_id,
        len(packet.payload),
    )
    sock.sendall(zeroed_header + packet.payload + bytes([PACKET_TERMINATOR]))
