"""
Фрагментация и сборка изображений SDR Rover Link.

Изображение рассматривается как обычная последовательность байтов.
Модулю неважно, используется JPEG, PNG или другой формат.
"""

from dataclasses import dataclass
import struct
from typing import Iterable

from protocol.packet import MAX_PAYLOAD_SIZE


# ============================================================
# Формат заголовка фрагмента изображения
# ============================================================

# I — идентификатор изображения, 4 байта
# H — номер фрагмента, 2 байта
# H — общее количество фрагментов, 2 байта
# I — размер исходного изображения, 4 байта
FRAGMENT_HEADER_FORMAT = ">IHHI"

FRAGMENT_HEADER_SIZE = struct.calcsize(
    FRAGMENT_HEADER_FORMAT
)

MAX_FRAGMENT_DATA_SIZE = (
    MAX_PAYLOAD_SIZE - FRAGMENT_HEADER_SIZE
)


# ============================================================
# Ошибки
# ============================================================

class ImageFragmentError(ValueError):
    """Ошибка структуры фрагмента изображения."""


class MissingFragmentsError(ImageFragmentError):
    """Отсутствуют необходимые фрагменты изображения."""


# ============================================================
# Представление одного фрагмента
# ============================================================

@dataclass(frozen=True)
class ImageFragment:
    """
    Один фрагмент изображения.

    Атрибуты
    --------
    image_id:
        Идентификатор изображения.

    fragment_index:
        Номер фрагмента. Нумерация начинается с нуля.

    fragment_count:
        Общее количество фрагментов.

    original_size:
        Размер исходного изображения в байтах.

    data:
        Часть исходного изображения.
    """

    image_id: int
    fragment_index: int
    fragment_count: int
    original_size: int
    data: bytes


# ============================================================
# Разбиение изображения
# ============================================================

def split_image_bytes(
    image_bytes: bytes,
    image_id: int,
    fragment_data_size: int = 512,
) -> list[ImageFragment]:
    """
    Разбить изображение на фрагменты.

    Параметры
    --------
    image_bytes:
        Полное содержимое файла изображения.

    image_id:
        Идентификатор изображения от 0 до 4294967295.

    fragment_data_size:
        Максимальный размер данных одного фрагмента.

    Возвращает
    ----------
    Список объектов ImageFragment.
    """

    if not isinstance(image_bytes, (bytes, bytearray)):
        raise TypeError(
            "image_bytes должен иметь тип bytes или bytearray"
        )

    image_bytes = bytes(image_bytes)

    if not image_bytes:
        raise ValueError(
            "Изображение не должно быть пустым"
        )

    if not 0 <= image_id <= 0xFFFFFFFF:
        raise ValueError(
            "image_id должен находиться "
            "в диапазоне 0...4294967295"
        )

    if not 1 <= fragment_data_size <= MAX_FRAGMENT_DATA_SIZE:
        raise ValueError(
            "Некорректный размер фрагмента"
        )

    fragment_count = (
        len(image_bytes)
        + fragment_data_size
        - 1
    ) // fragment_data_size

    if fragment_count > 65535:
        raise ValueError(
            "Изображение требует более 65535 фрагментов"
        )

    fragments = []

    for fragment_index in range(fragment_count):

        start = fragment_index * fragment_data_size
        end = start + fragment_data_size

        fragment_data = image_bytes[start:end]

        fragments.append(
            ImageFragment(
                image_id=image_id,
                fragment_index=fragment_index,
                fragment_count=fragment_count,
                original_size=len(image_bytes),
                data=fragment_data,
            )
        )

    return fragments


# ============================================================
# Преобразование фрагмента в байты
# ============================================================

def encode_image_fragment(
    fragment: ImageFragment,
) -> bytes:
    """
    Преобразовать ImageFragment в полезную нагрузку пакета.
    """

    if not isinstance(fragment, ImageFragment):
        raise TypeError(
            "fragment должен иметь тип ImageFragment"
        )

    if not 0 <= fragment.fragment_index < fragment.fragment_count:
        raise ImageFragmentError(
            "Некорректный номер фрагмента"
        )

    if not fragment.data:
        raise ImageFragmentError(
            "Фрагмент не должен быть пустым"
        )

    if len(fragment.data) > MAX_FRAGMENT_DATA_SIZE:
        raise ImageFragmentError(
            "Данные фрагмента слишком велики"
        )

    header = struct.pack(
        FRAGMENT_HEADER_FORMAT,
        fragment.image_id,
        fragment.fragment_index,
        fragment.fragment_count,
        fragment.original_size,
    )

    return header + fragment.data


# ============================================================
# Преобразование байтов обратно во фрагмент
# ============================================================

def decode_image_fragment(
    payload: bytes,
) -> ImageFragment:
    """
    Разобрать полезную нагрузку пакета изображения.
    """

    if not isinstance(payload, (bytes, bytearray)):
        raise TypeError(
            "payload должен иметь тип bytes или bytearray"
        )

    payload = bytes(payload)

    if len(payload) <= FRAGMENT_HEADER_SIZE:
        raise ImageFragmentError(
            "Фрагмент слишком короткий"
        )

    header = payload[:FRAGMENT_HEADER_SIZE]
    fragment_data = payload[FRAGMENT_HEADER_SIZE:]

    (
        image_id,
        fragment_index,
        fragment_count,
        original_size,
    ) = struct.unpack(
        FRAGMENT_HEADER_FORMAT,
        header,
    )

    if fragment_count == 0:
        raise ImageFragmentError(
            "Количество фрагментов не может быть равно нулю"
        )

    if fragment_index >= fragment_count:
        raise ImageFragmentError(
            "Номер фрагмента выходит за допустимый диапазон"
        )

    if original_size == 0:
        raise ImageFragmentError(
            "Размер исходного изображения равен нулю"
        )

    return ImageFragment(
        image_id=image_id,
        fragment_index=fragment_index,
        fragment_count=fragment_count,
        original_size=original_size,
        data=fragment_data,
    )


# ============================================================
# Сборка изображения
# ============================================================

def reassemble_image(
    fragments: Iterable[ImageFragment],
) -> bytes:
    """
    Собрать исходное изображение из фрагментов.

    Фрагменты могут поступать в произвольном порядке.
    """

    fragments = list(fragments)

    if not fragments:
        raise MissingFragmentsError(
            "Список фрагментов пуст"
        )

    first_fragment = fragments[0]

    image_id = first_fragment.image_id
    fragment_count = first_fragment.fragment_count
    original_size = first_fragment.original_size

    fragments_by_index = {}

    for fragment in fragments:

        if fragment.image_id != image_id:
            raise ImageFragmentError(
                "Фрагменты относятся к разным изображениям"
            )

        if fragment.fragment_count != fragment_count:
            raise ImageFragmentError(
                "Различается общее количество фрагментов"
            )

        if fragment.original_size != original_size:
            raise ImageFragmentError(
                "Различается размер исходного изображения"
            )

        existing_fragment = fragments_by_index.get(
            fragment.fragment_index
        )

        if existing_fragment is not None:

            if existing_fragment.data != fragment.data:
                raise ImageFragmentError(
                    "Получены разные данные "
                    "для одного номера фрагмента"
                )

            # Полностью одинаковый повторный фрагмент
            # можно безопасно проигнорировать.
            continue

        fragments_by_index[
            fragment.fragment_index
        ] = fragment

    missing_indexes = [
        index
        for index in range(fragment_count)
        if index not in fragments_by_index
    ]

    if missing_indexes:
        raise MissingFragmentsError(
            "Отсутствуют фрагменты: "
            + ", ".join(map(str, missing_indexes))
        )

    assembled_bytes = b"".join(
        fragments_by_index[index].data
        for index in range(fragment_count)
    )

    if len(assembled_bytes) != original_size:
        raise ImageFragmentError(
            "Размер собранного изображения "
            "не совпадает с исходным размером"
        )

    return assembled_bytes