"""
Lab008. Передача JPEG через виртуальный канал с ARQ.

Этапы:
1. Чтение тестового JPEG из Lab007.
2. Разбиение файла на фрагменты.
3. Передача каждого фрагмента через виртуальный канал.
4. Проверка CRC на приёмнике.
5. Передача ACK по обратному каналу.
6. Повторная передача при потере или повреждении.
7. Сборка принятого JPEG.
8. Сравнение SHA-256 исходного и принятого файлов.
"""

from hashlib import sha256
from pathlib import Path
import random

from PIL import Image

from protocol.image_fragments import (
    ImageFragmentError,
    decode_image_fragment,
    encode_image_fragment,
    reassemble_image,
    split_image_bytes,
)

from protocol.packet import (
    CRCError,
    MESSAGE_TYPE_ACK,
    MESSAGE_TYPE_IMAGE_FRAGMENT,
    PacketError,
    build_packet,
    parse_packet,
)

from tools.virtual_channel import transmit


# ============================================================
# Настройки эксперимента
# ============================================================

IMAGE_ID = 20260713
FRAGMENT_DATA_SIZE = 512

# Максимальное количество передач одного фрагмента.
MAX_ATTEMPTS = 12

# Прямой канал: передатчик → приёмник.
DATA_DROP_PROBABILITY = 0.10
DATA_BIT_ERROR_PROBABILITY = 0.0002

# Обратный канал: приёмник → передатчик.
ACK_DROP_PROBABILITY = 0.10
ACK_BIT_ERROR_PROBABILITY = 0.0002

# Фиксированное значение делает результат воспроизводимым.
random_generator = random.Random(2026)


# ============================================================
# Пути к файлам
# ============================================================

source_path = Path(
    "data/raw/lab007_source.jpg"
)

received_path = Path(
    "data/received/lab008_received.jpg"
)

if not source_path.exists():
    raise FileNotFoundError(
        "Не найден файл data/raw/lab007_source.jpg. "
        "Сначала необходимо выполнить Lab007."
    )

received_path.parent.mkdir(
    parents=True,
    exist_ok=True,
)


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

source_bytes = source_path.read_bytes()

source_hash = sha256(
    source_bytes
).hexdigest()

source_fragments = split_image_bytes(
    image_bytes=source_bytes,
    image_id=IMAGE_ID,
    fragment_data_size=FRAGMENT_DATA_SIZE,
)


# ============================================================
# Счётчики
# ============================================================

data_transmissions = 0
ack_transmissions = 0

dropped_data_packets = 0
invalid_data_packets = 0

dropped_ack_packets = 0
invalid_ack_packets = 0

duplicate_data_packets = 0

flipped_data_bits = 0
flipped_ack_bits = 0

confirmed_fragments = 0
failed_fragment_indexes = []

# Принятые фрагменты хранятся по их номеру.
received_fragments = {}

# Число попыток передачи каждого фрагмента.
attempts_per_fragment = {}


# ============================================================
# Передача всех фрагментов
# ============================================================

for source_fragment in source_fragments:

    fragment_index = source_fragment.fragment_index

    fragment_payload = encode_image_fragment(
        source_fragment
    )

    data_packet = build_packet(
        payload=fragment_payload,
        message_type=MESSAGE_TYPE_IMAGE_FRAGMENT,
        sequence_number=fragment_index,
    )

    fragment_confirmed = False

    for attempt_number in range(
        1,
        MAX_ATTEMPTS + 1,
    ):
        data_transmissions += 1

        received_data_packet, changed_bits = transmit(
            packet=data_packet,
            drop_probability=DATA_DROP_PROBABILITY,
            bit_error_probability=DATA_BIT_ERROR_PROBABILITY,
            random_generator=random_generator,
        )

        flipped_data_bits += changed_bits

        # DATA полностью потерян.
        if received_data_packet is None:
            dropped_data_packets += 1
            continue

        # Проверка структуры и CRC.
        try:
            parsed_data_packet = parse_packet(
                received_data_packet
            )

        except (CRCError, PacketError):
            invalid_data_packets += 1
            continue

        if (
            parsed_data_packet.message_type
            != MESSAGE_TYPE_IMAGE_FRAGMENT
        ):
            invalid_data_packets += 1
            continue

        # Разбор заголовка фрагмента изображения.
        try:
            received_fragment = decode_image_fragment(
                parsed_data_packet.payload
            )

        except ImageFragmentError:
            invalid_data_packets += 1
            continue

        # Проверяем согласованность номеров.
        if (
            parsed_data_packet.sequence_number
            != received_fragment.fragment_index
        ):
            invalid_data_packets += 1
            continue

        if received_fragment.image_id != IMAGE_ID:
            invalid_data_packets += 1
            continue

        # Сохраняем новый фрагмент либо распознаём дубликат.
        existing_fragment = received_fragments.get(
            received_fragment.fragment_index
        )

        if existing_fragment is None:
            received_fragments[
                received_fragment.fragment_index
            ] = received_fragment

        else:
            # DATA мог прийти повторно, если предыдущий ACK потерялся.
            assert existing_fragment.data == received_fragment.data
            duplicate_data_packets += 1

        # ====================================================
        # Формирование ACK
        # ====================================================

        ack_packet = build_packet(
            payload=b"",
            message_type=MESSAGE_TYPE_ACK,
            sequence_number=received_fragment.fragment_index,
        )

        ack_transmissions += 1

        received_ack_packet, changed_ack_bits = transmit(
            packet=ack_packet,
            drop_probability=ACK_DROP_PROBABILITY,
            bit_error_probability=ACK_BIT_ERROR_PROBABILITY,
            random_generator=random_generator,
        )

        flipped_ack_bits += changed_ack_bits

        # ACK полностью потерян.
        if received_ack_packet is None:
            dropped_ack_packets += 1
            continue

        # Передатчик проверяет ACK.
        try:
            parsed_ack = parse_packet(
                received_ack_packet
            )

        except (CRCError, PacketError):
            invalid_ack_packets += 1
            continue

        if (
            parsed_ack.message_type != MESSAGE_TYPE_ACK
            or parsed_ack.sequence_number != fragment_index
            or parsed_ack.payload != b""
        ):
            invalid_ack_packets += 1
            continue

        # Фрагмент успешно подтверждён.
        fragment_confirmed = True
        confirmed_fragments += 1

        attempts_per_fragment[
            fragment_index
        ] = attempt_number

        break

    if not fragment_confirmed:
        failed_fragment_indexes.append(
            fragment_index
        )

        attempts_per_fragment[
            fragment_index
        ] = MAX_ATTEMPTS


# ============================================================
# Результаты передачи
# ============================================================

print("=== Lab008. Передача JPEG с ARQ ===")

print("\nИсходный файл:")
print(source_path)

print("\nРазмер JPEG:")
print(len(source_bytes), "байт")

print("\nКоличество фрагментов:")
print(len(source_fragments))

print("\nПодтверждено фрагментов:")
print(confirmed_fragments)

print("\nПолучено уникальных фрагментов:")
print(len(received_fragments))

print("\nНе подтверждены фрагменты:")
print(failed_fragment_indexes)


print("\nПередачи:")

print(
    "Всего передач DATA:",
    data_transmissions,
)

print(
    "Всего передач ACK:",
    ack_transmissions,
)


print("\nОшибки прямого канала:")

print(
    "Полностью потеряно DATA:",
    dropped_data_packets,
)

print(
    "Повреждено DATA:",
    invalid_data_packets,
)

print(
    "Изменено битов DATA:",
    flipped_data_bits,
)


print("\nОшибки обратного канала:")

print(
    "Полностью потеряно ACK:",
    dropped_ack_packets,
)

print(
    "Повреждено ACK:",
    invalid_ack_packets,
)

print(
    "Изменено битов ACK:",
    flipped_ack_bits,
)

print(
    "Получено повторных DATA:",
    duplicate_data_packets,
)


print("\nКоличество попыток по фрагментам:")

for fragment_index in sorted(
    attempts_per_fragment
):
    print(
        f"Фрагмент {fragment_index}: "
        f"{attempts_per_fragment[fragment_index]}"
    )


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

assert not failed_fragment_indexes

assert (
    len(received_fragments)
    == len(source_fragments)
)

received_bytes = reassemble_image(
    received_fragments.values()
)

received_path.write_bytes(
    received_bytes
)

received_hash = sha256(
    received_bytes
).hexdigest()


# ============================================================
# Финальные проверки
# ============================================================

assert received_bytes == source_bytes
assert received_hash == source_hash

with Image.open(received_path) as image:
    image.verify()


print("\nSHA-256 исходного JPEG:")
print(source_hash)

print("\nSHA-256 принятого JPEG:")
print(received_hash)

print("\nПринятый файл:")
print(received_path)

print(
    "\nПроверка пройдена: "
    "все фрагменты подтверждены, "
    "JPEG восстановлен байт в байт."
)