from __future__ import annotations

from typing import Any


LITE_TAGS = {
    "k5": 5,
    "k14": 14,
    "k22": 22,
    "k23": 23,
    "k27": 27,
    "k29": 29,
    "k31": 31,
    "k34": 34,
    "k35": 35,
    "k36": 36,
    "k39": 39,
    "k40": 40,
    "k46": 46,
    "k57": 57,
    "k61": 61,
    "k64": 64,
    "k66": 66,
    "k68": 68,
    "k83": 83,
    "k86": 86,
    "k93": 93,
    "k97": 97,
    "k101": 101,
    "k102": 102,
    "k105": 105,
    "k106": 106,
    "k107": 107,
    "k108": 108,
    "k109": 109,
    "k110": 110,
    "k111": 111,
    "k112": 112,
    "k113": 113,
}

FULL_TAGS = {f"k{index}": index for index in range(1, 120)}


def _varint(value: int) -> bytes:
    if value < 0:
        raise ValueError("varint value must be >= 0")
    out = bytearray()
    while value >= 0x80:
        out.append((value & 0x7F) | 0x80)
        value >>= 7
    out.append(value)
    return bytes(out)


def _read_varint(raw: bytes, offset: int) -> tuple[int, int]:
    shift = 0
    value = 0
    while True:
        if offset >= len(raw):
            raise ValueError("truncated varint")
        byte = raw[offset]
        offset += 1
        value |= (byte & 0x7F) << shift
        if byte < 0x80:
            return value, offset
        shift += 7
        if shift > 63:
            raise ValueError("varint too long")


def _tags_for_mode(mode: str) -> dict[str, int]:
    if mode == "lite":
        return LITE_TAGS
    if mode == "full":
        return FULL_TAGS
    raise ValueError(f"unsupported sq0 mode: {mode}")


def encode_sq0_device_info(values: dict[str, str], mode: str) -> bytes:
    tags = _tags_for_mode(mode)
    encoded = bytearray()
    for key, value in values.items():
        if key not in tags:
            raise KeyError(key)
        text = "" if value is None else str(value)
        if text == "":
            continue
        payload = text.encode("utf-8")
        encoded += _varint((tags[key] << 3) | 2)
        encoded += _varint(len(payload))
        encoded += payload
    return bytes(encoded)


def decode_sq0_string_fields(raw: bytes) -> list[dict[str, Any]]:
    fields: list[dict[str, Any]] = []
    offset = 0
    while offset < len(raw):
        key, offset = _read_varint(raw, offset)
        proto_tag = key >> 3
        wire_type = key & 7
        if wire_type != 2:
            raise ValueError(f"unsupported wire type: {wire_type}")
        size, offset = _read_varint(raw, offset)
        end = offset + size
        if end > len(raw):
            raise ValueError("truncated length-delimited field")
        payload = raw[offset:end]
        offset = end
        fields.append(
            {
                "proto_tag": proto_tag,
                "wire_type": wire_type,
                "value": payload.decode("utf-8", errors="replace"),
                "raw_hex": payload.hex(),
            }
        )
    return fields


__all__ = [
    "FULL_TAGS",
    "LITE_TAGS",
    "decode_sq0_string_fields",
    "encode_sq0_device_info",
]