"""
rbxl_binreader.py — низкоуровневое чтение бинарного потока Roblox.

Особенности формата:
1. Little-endian для всех целочисленных и float.
2. "Interleaved transformed" массивы: для N значений длиной L байт каждое,
   байты идут не AAAA BBBB CCCC, а ABCD ABCD ABCD ABCD (interleave).
   То есть сначала все first-byte'ы, потом все second-byte'ы, итд. Это
   делается потому что среди одинакового свойства часто меняется только
   младший байт, и interleave даёт длинные последовательности нулей и
   повторяющихся байт → лучше сжимается LZ4.
3. Signed int → zigzag: (n << 1) ^ (n >> 31) — для int32, аналогично для int64.
4. Float → "Roblox float encoding": сдвиг битов чтобы знаковый бит был в конце.
   Сделано опять же чтобы маленькие близкие к нулю float'ы выглядели похоже
   и сжимались.
5. String: uint32 length + UTF-8 bytes.

Полная спецификация: https://dom.rojo.space/binary
"""
import struct
import io


class BinReader:
    """Тонкая обёртка над BytesIO с методами для Roblox-типов."""

    def __init__(self, data: bytes):
        self.buf = data
        self.pos = 0

    # ─── атомарные ───

    def read(self, n: int) -> bytes:
        out = self.buf[self.pos:self.pos + n]
        if len(out) < n:
            raise IOError(f"unexpected EOF: wanted {n}, got {len(out)} at pos {self.pos}")
        self.pos += n
        return out

    def remaining(self) -> int:
        return len(self.buf) - self.pos

    def at_end(self) -> bool:
        return self.pos >= len(self.buf)

    def skip(self, n: int) -> None:
        self.pos += n

    # ─── простые числовые ───

    def uint8(self) -> int:
        return self.read(1)[0]

    def uint16(self) -> int:
        return struct.unpack('<H', self.read(2))[0]

    def uint32(self) -> int:
        return struct.unpack('<I', self.read(4))[0]

    def uint64(self) -> int:
        return struct.unpack('<Q', self.read(8))[0]

    def int32(self) -> int:
        return struct.unpack('<i', self.read(4))[0]

    def int64(self) -> int:
        return struct.unpack('<q', self.read(8))[0]

    def float32(self) -> float:
        return struct.unpack('<f', self.read(4))[0]

    def float64(self) -> float:
        return struct.unpack('<d', self.read(8))[0]

    def bool(self) -> bool:
        return self.read(1)[0] != 0

    def string(self) -> str:
        length = self.uint32()
        if length == 0:
            return ''
        return self.read(length).decode('utf-8', errors='replace')

    def bytes_with_len(self) -> bytes:
        length = self.uint32()
        return self.read(length)


# ──────────────────────────────────────────────────────────────────────
# Decoders для "interleaved transformed" массивов.
# Используются в PROP chunks где для класса с N инстансами идёт одно
# свойство, и значение лежит в interleaved виде.
# ──────────────────────────────────────────────────────────────────────


def deinterleave(data: bytes, count: int, element_size: int) -> bytes:
    """Снимает interleave. Получает count*element_size байт и возвращает
    нормальный массив байт где значения идут последовательно.

    interleaved:  b0[0] b1[0] b2[0] ... b0[1] b1[1] b2[1] ... b0[3] b1[3] b2[3]
    нормальный:   b0[0] b0[1] b0[2] b0[3]  b1[0] b1[1] b1[2] b1[3]  ...
    """
    if len(data) != count * element_size:
        raise ValueError(
            f"deinterleave: expected {count}*{element_size}={count*element_size} bytes, got {len(data)}"
        )
    if count == 0:
        return b''
    out = bytearray(count * element_size)
    for byte_idx in range(element_size):
        for value_idx in range(count):
            out[value_idx * element_size + byte_idx] = data[byte_idx * count + value_idx]
    return bytes(out)


def zigzag_decode_int32(n: int) -> int:
    """ZigZag decode для int32: (n >> 1) ^ -(n & 1).

    Используется для PROP с типом Int32 (свойство 'Size' MeshPart и т.п.).
    """
    return (n >> 1) ^ -(n & 1)


def zigzag_decode_int64(n: int) -> int:
    return (n >> 1) ^ -(n & 1)


def roblox_float_decode(raw: int) -> float:
    """Roblox float encoding для PROP-массивов.

    Roblox перекодирует float32 чтобы знаковый бит был в самом конце
    (это даёт лучшее сжатие при близких к нулю значениях).

    raw — uint32 little-endian.
    """
    # Развернём биты так чтобы знаковый бит вернулся в позицию 31.
    # Roblox: { mantissa_high(23 bits), exponent(8 bits), sign(1 bit) } → стандартный float
    # Просто (raw >> 1) | ((raw & 1) << 31)
    standard = (raw >> 1) | ((raw & 1) << 31)
    return struct.unpack('<f', struct.pack('<I', standard))[0]


# ──────────────────────────────────────────────────────────────────────
# Высокоуровневые readers для interleaved-массивов
# ──────────────────────────────────────────────────────────────────────


def read_interleaved_int32_array(reader: BinReader, count: int) -> list:
    """Читает count int32 в interleaved-zigzag форме."""
    raw = reader.read(count * 4)
    deint = deinterleave(raw, count, 4)
    out = []
    for i in range(count):
        u = struct.unpack('>I', deint[i*4:i*4+4])[0]  # ВНИМАНИЕ: после deinterleave порядок big-endian!
        out.append(zigzag_decode_int32(u))
    return out


def read_interleaved_uint32_array(reader: BinReader, count: int) -> list:
    """uint32 array (для Referent) — interleaved, но без zigzag."""
    raw = reader.read(count * 4)
    deint = deinterleave(raw, count, 4)
    out = []
    for i in range(count):
        out.append(struct.unpack('>I', deint[i*4:i*4+4])[0])
    return out


def read_interleaved_int64_array(reader: BinReader, count: int) -> list:
    raw = reader.read(count * 8)
    deint = deinterleave(raw, count, 8)
    out = []
    for i in range(count):
        u = struct.unpack('>Q', deint[i*8:i*8+8])[0]
        out.append(zigzag_decode_int64(u))
    return out


def read_interleaved_float_array(reader: BinReader, count: int) -> list:
    """Float32 array — interleaved с Roblox float encoding."""
    raw = reader.read(count * 4)
    deint = deinterleave(raw, count, 4)
    out = []
    for i in range(count):
        u = struct.unpack('>I', deint[i*4:i*4+4])[0]
        out.append(roblox_float_decode(u))
    return out


def read_referent_array(reader: BinReader, count: int) -> list:
    """Referent array — interleaved int32 array с накопительной разностью.

    Каждое следующее значение = предыдущее + delta. Это даёт длинные
    последовательности (1, 2, 3, 4, ...) которые после deinterleave +
    zigzag отлично сжимаются в нули.
    """
    deltas = read_interleaved_int32_array(reader, count)
    out = []
    accumulator = 0
    for d in deltas:
        accumulator += d
        out.append(accumulator)
    return out