"""Garlic encryption types — DeliveryInstructions, GarlicClove, GarlicMessage."""

import enum
import struct


class DeliveryType(enum.IntEnum):
    LOCAL = 0
    DESTINATION = 1
    ROUTER = 2
    TUNNEL = 3


class DeliveryInstructions:
    """Delivery instructions for a garlic clove."""

    def __init__(self, delivery_type: DeliveryType, *,
                 destination: bytes | None = None,
                 router: bytes | None = None,
                 tunnel_id: int | None = None):
        self.delivery_type = delivery_type
        self.destination = destination
        self.router = router
        self.tunnel_id = tunnel_id

    def to_bytes(self) -> bytes:
        # Flag byte: delivery type in bits 5-6
        flag = (self.delivery_type & 0x03) << 5
        parts = [struct.pack("!B", flag)]

        if self.delivery_type == DeliveryType.DESTINATION:
            assert self.destination is not None
            parts.append(self.destination)
        elif self.delivery_type == DeliveryType.ROUTER:
            assert self.router is not None
            parts.append(self.router)
        elif self.delivery_type == DeliveryType.TUNNEL:
            assert self.router is not None
            parts.append(self.router)
            parts.append(struct.pack("!I", self.tunnel_id))

        return b"".join(parts)

    @classmethod
    def from_bytes(cls, data: bytes) -> tuple["DeliveryInstructions", int]:
        flag = data[0]
        delivery_type = DeliveryType((flag >> 5) & 0x03)
        offset = 1
        destination = None
        router = None
        tunnel_id = None

        if delivery_type == DeliveryType.DESTINATION:
            destination = data[offset:offset + 32]
            offset += 32
        elif delivery_type == DeliveryType.ROUTER:
            router = data[offset:offset + 32]
            offset += 32
        elif delivery_type == DeliveryType.TUNNEL:
            router = data[offset:offset + 32]
            offset += 32
            tunnel_id = struct.unpack("!I", data[offset:offset + 4])[0]
            offset += 4

        return cls(delivery_type, destination=destination, router=router,
                   tunnel_id=tunnel_id), offset


class GarlicClove:
    """A single garlic clove: delivery instructions + message + id + expiration."""

    def __init__(self, delivery_instructions: DeliveryInstructions,
                 message_data: bytes, clove_id: int, expiration: int):
        self.delivery_instructions = delivery_instructions
        self.message_data = message_data
        self.clove_id = clove_id
        self.expiration = expiration
        self._size: int = 0

    def to_bytes(self) -> bytes:
        di_bytes = self.delivery_instructions.to_bytes()
        return (di_bytes +
                struct.pack("!I", len(self.message_data)) +
                self.message_data +
                struct.pack("!IQ", self.clove_id, self.expiration))

    @classmethod
    def from_bytes(cls, data: bytes, offset: int = 0) -> "GarlicClove":
        di, di_len = DeliveryInstructions.from_bytes(data[offset:])
        pos = offset + di_len
        msg_len = struct.unpack("!I", data[pos:pos + 4])[0]
        pos += 4
        message_data = data[pos:pos + msg_len]
        pos += msg_len
        clove_id, expiration = struct.unpack("!IQ", data[pos:pos + 12])
        pos += 12
        clove = cls(di, message_data, clove_id, expiration)
        clove._size = pos - offset
        return clove


class GarlicMessage:
    """Bundle of garlic cloves."""

    def __init__(self, cloves: list[GarlicClove]):
        self.cloves = cloves

    def clove_count(self) -> int:
        return len(self.cloves)

    def to_bytes(self) -> bytes:
        parts = [struct.pack("!B", len(self.cloves))]
        for clove in self.cloves:
            parts.append(clove.to_bytes())
        return b"".join(parts)

    @classmethod
    def from_bytes(cls, data: bytes) -> "GarlicMessage":
        count = data[0]
        offset = 1
        cloves = []
        for _ in range(count):
            clove = GarlicClove.from_bytes(data, offset)
            cloves.append(clove)
            offset += clove._size
        return cls(cloves)