from __future__ import annotations

from collections.abc import Mapping

from liverylab_core.proto import ProtoField, build_fields, parse_fields


VISUAL_SELECTOR_FIELD = 15
CARFINALSTATE_SELECTOR_FIELD = 51


def selector_payload(data: bytes, field_number: int) -> bytes | None:
    """Return the first selector message payload, preserving an empty message."""
    for field in parse_fields(data):
        if field.number == int(field_number) and field.wire == 2:
            return bytes(field.value)
    return None


def selector_values(data: bytes, field_number: int) -> dict[int, int]:
    payload = selector_payload(data, field_number)
    if payload is None:
        return {}
    return {
        int(field.number): int(field.value)
        for field in parse_fields(payload)
        if field.wire == 0 and int(field.value) > 0
    }


def build_selector_payload(selector: Mapping[int, int] | None) -> bytes:
    normalized = {
        int(surface): int(index)
        for surface, index in (selector or {}).items()
        if int(index) > 0
    }
    return build_fields(
        [ProtoField(surface, 0, index) for surface, index in sorted(normalized.items())]
    )


def replace_selector_payload(
    data: bytes,
    field_number: int,
    payload: bytes,
) -> tuple[bytes, bool]:
    """Set one selector message and keep it present even when its payload is empty."""
    fields = parse_fields(data)
    replacement = bytes(payload)
    new_fields: list[ProtoField] = []
    written = False
    changed = False
    for field in fields:
        if field.number == int(field_number) and field.wire == 2:
            if written:
                changed = True
                continue
            new_fields.append(ProtoField(int(field_number), 2, replacement))
            written = True
            if bytes(field.value) != replacement:
                changed = True
            continue
        new_fields.append(field)
    if not written:
        new_fields.append(ProtoField(int(field_number), 2, replacement))
        changed = True
    return (build_fields(new_fields), True) if changed else (data, False)


def set_selector_values(
    data: bytes,
    field_number: int,
    selector: Mapping[int, int] | None,
) -> tuple[bytes, bool]:
    return replace_selector_payload(data, field_number, build_selector_payload(selector))


def set_visual_selector(
    visual_data: bytes,
    selector: Mapping[int, int] | None,
) -> tuple[bytes, bool]:
    return set_selector_values(visual_data, VISUAL_SELECTOR_FIELD, selector)


def sync_carfinalstate_from_visual(
    carfinalstate_data: bytes,
    visual_data: bytes,
) -> tuple[bytes, bool]:
    payload = selector_payload(visual_data, VISUAL_SELECTOR_FIELD)
    if payload is None:
        return carfinalstate_data, False
    return replace_selector_payload(
        carfinalstate_data,
        CARFINALSTATE_SELECTOR_FIELD,
        payload,
    )
