"""
rbxl_types.py — определения типов значений Roblox и парсинг PROP-chunks.

Roblox-формат хранит свойства per-class: для каждого класса свой набор
PROP chunks (по одному на каждое свойство). Внутри PROP лежат значения
этого свойства для ВСЕХ инстансов данного класса разом.

Schema PROP chunk:
   class_id     (uint32)   — индекс INST chunk'а (1:1 с порядком)
   prop_name    (string)   — например "Color", "Size", "CFrame"
   type_id      (uint8)    — тип значения (см. PROP_TYPE_* константы)
   values...    (массив длиной N где N = instance_count для класса class_id)

Полная спецификация PROP типов: https://dom.rojo.space/binary#type-id

Эта реализация — голый struct без логики приведения к Rublox-формату.
Конвертация в твою сцену — в rbxl_converter.py.
"""
from dataclasses import dataclass, field
from typing import List, Optional, Union, Any
import struct

from rbxl_binreader import (
    BinReader,
    deinterleave,
    zigzag_decode_int32,
    zigzag_decode_int64,
    roblox_float_decode,
    read_interleaved_int32_array,
    read_interleaved_uint32_array,
    read_interleaved_int64_array,
    read_interleaved_float_array,
    read_referent_array,
)


# ──────────────────────────────────────────────────────────────────────
# Тип-ID коды (используется в первом байте PROP chunk после prop_name).
# ──────────────────────────────────────────────────────────────────────

PROP_STRING            = 0x01
PROP_BOOL              = 0x02
PROP_INT32             = 0x03
PROP_FLOAT             = 0x04
PROP_DOUBLE            = 0x05
PROP_UDIM              = 0x06
PROP_UDIM2             = 0x07
PROP_RAY               = 0x08
PROP_FACES             = 0x09
PROP_AXES              = 0x0A
PROP_BRICKCOLOR        = 0x0B
PROP_COLOR3            = 0x0C
PROP_VECTOR2           = 0x0D
PROP_VECTOR3           = 0x0E
# 0x0F не используется
PROP_CFRAME            = 0x10
PROP_QUATERNION        = 0x11   # есть в спеке, но обычно cf хранится в CFrame
PROP_ENUM              = 0x12
PROP_REFERENT          = 0x13
PROP_VECTOR3INT16      = 0x14
PROP_NUMBERSEQUENCE    = 0x15
PROP_COLORSEQUENCE     = 0x16
PROP_NUMBERRANGE       = 0x17
PROP_RECT              = 0x18
PROP_PHYSICALPROPS     = 0x19
PROP_COLOR3UINT8       = 0x1A
PROP_INT64             = 0x1B
PROP_SHAREDSTRING      = 0x1C
PROP_BYTECODE          = 0x1D
PROP_OPTIONALCFRAME    = 0x1E
PROP_UNIQUEID          = 0x1F
PROP_FONT              = 0x20

# ──────────────────────────────────────────────────────────────────────
# Дата-классы для сложных типов (Vector3, CFrame, и т.д.)
# ──────────────────────────────────────────────────────────────────────


@dataclass
class Vector3:
    x: float
    y: float
    z: float


@dataclass
class Vector2:
    x: float
    y: float


@dataclass
class Color3:
    r: float
    g: float
    b: float

    def to_hex(self) -> str:
        return '#{:02x}{:02x}{:02x}'.format(
            max(0, min(255, int(self.r * 255))),
            max(0, min(255, int(self.g * 255))),
            max(0, min(255, int(self.b * 255))),
        )


@dataclass
class CFrame:
    """4x3 матрица Roblox: 3 позиции + 9 значений матрицы ротации (3x3, row-major).

    Хранится как (px, py, pz, r00, r01, r02, r10, r11, r12, r20, r21, r22).
    """
    position: Vector3
    matrix: tuple  # (r00, r01, r02, r10, r11, r12, r20, r21, r22)

    def to_euler_xyz(self) -> tuple:
        """Конверт 3x3 rotation matrix в Euler YXZ (Babylon convention).

        Babylon mesh.rotation = Vector3(rx, ry, rz) применяется в порядке YXZ
        (rotate Y first, then X, then Z). Чтобы извлечь Euler из матрицы под
        этот convention, используем формулу YXZ-extraction:
            Rx = asin(-r12)
            Ry = atan2(r02, r22)
            Rz = atan2(r10, r11)
        (имя метода to_euler_xyz сохраняем для совместимости вызовов.)
        """
        import math
        r00, r01, r02, r10, r11, r12, r20, r21, r22 = self.matrix
        # Edge case: r12 близко к ±1 (gimbal lock на X = ±90°)
        clamped = max(-1.0, min(1.0, -r12))
        rx = math.asin(clamped)
        if abs(clamped) > 0.99999:
            # Gimbal lock — z = 0, y = atan2(-r20, r00)
            ry = math.atan2(-r20, r00)
            rz = 0.0
        else:
            ry = math.atan2(r02, r22)
            rz = math.atan2(r10, r11)
        return (rx, ry, rz)


@dataclass
class UDim:
    scale: float
    offset: int


@dataclass
class UDim2:
    x: UDim
    y: UDim


@dataclass
class Ray:
    origin: Vector3
    direction: Vector3


@dataclass
class Rect:
    min: Vector2
    max: Vector2


@dataclass
class PhysicalProperties:
    custom: bool
    density: float = 0.0
    friction: float = 0.0
    elasticity: float = 0.0
    friction_weight: float = 0.0
    elasticity_weight: float = 0.0


@dataclass
class BrickColor:
    """Цвет из палитры BrickColor (старый формат Roblox).

    Хранится как int32 — код цвета из BrickColor.palette().
    Для нас можно мапить в Color3 по таблице.
    """
    code: int


@dataclass
class Vector3int16:
    x: int
    y: int
    z: int


@dataclass
class NumberSequenceKeypoint:
    time: float
    value: float
    envelope: float


@dataclass
class NumberSequence:
    keypoints: List[NumberSequenceKeypoint]


@dataclass
class ColorSequenceKeypoint:
    time: float
    value: Color3
    envelope: float = 0.0   # игнорируется Roblox для ColorSequence


@dataclass
class ColorSequence:
    keypoints: List[ColorSequenceKeypoint]


@dataclass
class NumberRange:
    min: float
    max: float


@dataclass
class Faces:
    """6 bool флагов: Top, Bottom, Left, Right, Front, Back."""
    flags: int  # bitmask


@dataclass
class Axes:
    """3 bool флага."""
    flags: int


@dataclass
class SharedStringRef:
    """Ссылка на запись в SSTR chunk (по индексу)."""
    index: int


@dataclass
class EnumValue:
    value: int


@dataclass
class UniqueId:
    """128-bit ID (uuid-like)."""
    raw: bytes


@dataclass
class Font:
    family: str       # font family asset id or name
    weight: int
    style: int
    cached_face_id: str


@dataclass
class OptionalCFrame:
    """CFrame который может быть null."""
    value: Optional[CFrame]


# ──────────────────────────────────────────────────────────────────────
# Главный декодер PROP chunk
# ──────────────────────────────────────────────────────────────────────


@dataclass
class PropChunk:
    """Распарсенный PROP chunk: значения свойства для всех инстансов одного класса."""
    class_id: int
    prop_name: str
    type_id: int
    values: List[Any]  # длина = count instances класса


def decode_prop_chunk(payload: bytes, instance_count_for_class: int, shared_strings: List[bytes]) -> PropChunk:
    """Декодирует PROP chunk.

    Args:
        payload: уже декомпрессированные байты chunk'а
        instance_count_for_class: сколько инстансов у класса (определяется по INST)
        shared_strings: список из SSTR chunk для PROP_SHAREDSTRING

    Returns:
        PropChunk со списком значений (длина = instance_count_for_class).
    """
    r = BinReader(payload)
    class_id = r.uint32()
    prop_name = r.string()
    type_id = r.uint8()

    count = instance_count_for_class

    if type_id == PROP_STRING:
        # N строк, каждая uint32+bytes
        values = [r.string() for _ in range(count)]

    elif type_id == PROP_BOOL:
        # N bool — НЕ interleaved, просто N байт
        values = [r.bool() for _ in range(count)]

    elif type_id == PROP_INT32:
        values = read_interleaved_int32_array(r, count)

    elif type_id == PROP_FLOAT:
        values = read_interleaved_float_array(r, count)

    elif type_id == PROP_DOUBLE:
        # 8 байт little-endian, БЕЗ interleave
        values = [r.float64() for _ in range(count)]

    elif type_id == PROP_UDIM:
        scales = read_interleaved_float_array(r, count)
        offsets = read_interleaved_int32_array(r, count)
        values = [UDim(s, o) for s, o in zip(scales, offsets)]

    elif type_id == PROP_UDIM2:
        xs_scale = read_interleaved_float_array(r, count)
        ys_scale = read_interleaved_float_array(r, count)
        xs_offset = read_interleaved_int32_array(r, count)
        ys_offset = read_interleaved_int32_array(r, count)
        values = [
            UDim2(UDim(xs_scale[i], xs_offset[i]), UDim(ys_scale[i], ys_offset[i]))
            for i in range(count)
        ]

    elif type_id == PROP_RAY:
        # 6 floats per value, БЕЗ interleave
        values = []
        for _ in range(count):
            ox, oy, oz = r.float32(), r.float32(), r.float32()
            dx, dy, dz = r.float32(), r.float32(), r.float32()
            values.append(Ray(Vector3(ox, oy, oz), Vector3(dx, dy, dz)))

    elif type_id == PROP_FACES:
        # 1 байт на каждый
        values = [Faces(r.uint8()) for _ in range(count)]

    elif type_id == PROP_AXES:
        values = [Axes(r.uint8()) for _ in range(count)]

    elif type_id == PROP_BRICKCOLOR:
        codes = read_interleaved_int32_array(r, count)
        values = [BrickColor(c) for c in codes]

    elif type_id == PROP_COLOR3:
        # 3 float channels, каждый interleaved отдельно
        rs = read_interleaved_float_array(r, count)
        gs = read_interleaved_float_array(r, count)
        bs = read_interleaved_float_array(r, count)
        values = [Color3(rs[i], gs[i], bs[i]) for i in range(count)]

    elif type_id == PROP_VECTOR2:
        xs = read_interleaved_float_array(r, count)
        ys = read_interleaved_float_array(r, count)
        values = [Vector2(xs[i], ys[i]) for i in range(count)]

    elif type_id == PROP_VECTOR3:
        xs = read_interleaved_float_array(r, count)
        ys = read_interleaved_float_array(r, count)
        zs = read_interleaved_float_array(r, count)
        values = [Vector3(xs[i], ys[i], zs[i]) for i in range(count)]

    elif type_id == PROP_CFRAME:
        # Сначала идут rotation matrices, для каждого инстанса:
        #   1 byte: orientation_id (0 = custom 9 float, иначе — шаблон из enum)
        #   если 0: следующие 36 байт = 9 float32 (row-major)
        # ПОТОМ 3 interleaved float-массива (positions x, y, z)
        rotations = []
        for _ in range(count):
            orientation_id = r.uint8()
            if orientation_id == 0:
                m = struct.unpack('<9f', r.read(36))
                rotations.append(m)
            else:
                # Стандартные orientations (см. ROBLOX_CFRAME_ORIENTATIONS)
                rotations.append(_cframe_orientation_to_matrix(orientation_id))
        xs = read_interleaved_float_array(r, count)
        ys = read_interleaved_float_array(r, count)
        zs = read_interleaved_float_array(r, count)
        values = [
            CFrame(Vector3(xs[i], ys[i], zs[i]), rotations[i])
            for i in range(count)
        ]

    elif type_id == PROP_QUATERNION:
        # Редко используется отдельно, обычно в составе CFrame
        # Структура: 4 floats * count, БЕЗ interleave (на всякий)
        values = []
        for _ in range(count):
            qx, qy, qz, qw = r.float32(), r.float32(), r.float32(), r.float32()
            # Конверт quaternion → 3x3 matrix
            values.append(CFrame(Vector3(0, 0, 0), _quat_to_matrix(qx, qy, qz, qw)))

    elif type_id == PROP_ENUM:
        # interleaved int32 (zigzag не применяется)
        # на самом деле в спеке — это uint32 interleaved
        vals = read_interleaved_uint32_array(r, count)
        values = [EnumValue(v) for v in vals]

    elif type_id == PROP_REFERENT:
        values = read_referent_array(r, count)

    elif type_id == PROP_VECTOR3INT16:
        # 3 * int16 per value, БЕЗ interleave
        values = []
        for _ in range(count):
            x = struct.unpack('<h', r.read(2))[0]
            y = struct.unpack('<h', r.read(2))[0]
            z = struct.unpack('<h', r.read(2))[0]
            values.append(Vector3int16(x, y, z))

    elif type_id == PROP_NUMBERSEQUENCE:
        # Для каждого instance: uint32 keypoint_count, потом keypoint_count * 12 bytes
        # (time, value, envelope) each float32
        values = []
        for _ in range(count):
            kp_count = r.uint32()
            keypoints = []
            for _ in range(kp_count):
                t = r.float32()
                v = r.float32()
                env = r.float32()
                keypoints.append(NumberSequenceKeypoint(t, v, env))
            values.append(NumberSequence(keypoints))

    elif type_id == PROP_COLORSEQUENCE:
        values = []
        for _ in range(count):
            kp_count = r.uint32()
            keypoints = []
            for _ in range(kp_count):
                t = r.float32()
                cr = r.float32()
                cg = r.float32()
                cb = r.float32()
                # envelope (всегда 0 для ColorSequence, но 4 байта в файле)
                _env = r.float32()
                keypoints.append(ColorSequenceKeypoint(t, Color3(cr, cg, cb), 0.0))
            values.append(ColorSequence(keypoints))

    elif type_id == PROP_NUMBERRANGE:
        # 2 float per value, БЕЗ interleave
        values = []
        for _ in range(count):
            mn = r.float32()
            mx = r.float32()
            values.append(NumberRange(mn, mx))

    elif type_id == PROP_RECT:
        xmin = read_interleaved_float_array(r, count)
        ymin = read_interleaved_float_array(r, count)
        xmax = read_interleaved_float_array(r, count)
        ymax = read_interleaved_float_array(r, count)
        values = [
            Rect(Vector2(xmin[i], ymin[i]), Vector2(xmax[i], ymax[i]))
            for i in range(count)
        ]

    elif type_id == PROP_PHYSICALPROPS:
        values = []
        for _ in range(count):
            is_custom = r.bool()
            if is_custom:
                density = r.float32()
                friction = r.float32()
                elasticity = r.float32()
                fw = r.float32()
                ew = r.float32()
                values.append(PhysicalProperties(True, density, friction, elasticity, fw, ew))
            else:
                values.append(PhysicalProperties(False))

    elif type_id == PROP_COLOR3UINT8:
        # 3 байта на цвет (uint8 каждый), интерливированные
        rs = r.read(count)
        gs = r.read(count)
        bs = r.read(count)
        values = [Color3(rs[i]/255.0, gs[i]/255.0, bs[i]/255.0) for i in range(count)]

    elif type_id == PROP_INT64:
        values = read_interleaved_int64_array(r, count)

    elif type_id == PROP_SHAREDSTRING:
        # interleaved uint32 индексы в SSTR
        idxs = read_interleaved_uint32_array(r, count)
        values = []
        for i in idxs:
            if i < len(shared_strings):
                values.append(shared_strings[i])
            else:
                values.append(b'')

    elif type_id == PROP_BYTECODE:
        # Для каждого instance: uint32 length + bytes (compiled Lua bytecode)
        values = []
        for _ in range(count):
            length = r.uint32()
            values.append(r.read(length))

    elif type_id == PROP_OPTIONALCFRAME:
        # OptionalCFrame в rbx-dom spec:
        #   1 byte:  type_id (всегда 0x10 = CFrame)
        #   далее структура как у обычного PROP_CFRAME:
        #       N orientation_ids (по 1 байту каждый)
        #       N rotations (если orientation_id=0 → 36 байт row-major)
        #       interleaved x[N], y[N], z[N] float32
        #   ПОТОМ:
        #   1 byte:  type_id (0x02 = Bool — флаги has_value)
        #   N bytes: has_value flags
        #
        # Источник: https://dom.rojo.space/binary#optionalcframe
        _inner_type_cframe = r.uint8()  # 0x10
        rotations = []
        for _ in range(count):
            orientation_id = r.uint8()
            if orientation_id == 0:
                m = struct.unpack('<9f', r.read(36))
                rotations.append(m)
            else:
                rotations.append(_cframe_orientation_to_matrix(orientation_id))
        xs = read_interleaved_float_array(r, count)
        ys = read_interleaved_float_array(r, count)
        zs = read_interleaved_float_array(r, count)
        _inner_type_bool = r.uint8()    # 0x02
        has_values = [r.bool() for _ in range(count)]
        values = []
        for i in range(count):
            if has_values[i]:
                cf = CFrame(Vector3(xs[i], ys[i], zs[i]), rotations[i])
                values.append(OptionalCFrame(cf))
            else:
                values.append(OptionalCFrame(None))

    elif type_id == PROP_UNIQUEID:
        # 16 байт на каждый, БЕЗ interleave (предположительно)
        values = [UniqueId(r.read(16)) for _ in range(count)]

    elif type_id == PROP_FONT:
        # Для каждого: family (string), weight (uint16), style (uint8), cached_face_id (string)
        values = []
        for _ in range(count):
            family = r.string()
            weight = r.uint16()
            style = r.uint8()
            cached = r.string()
            values.append(Font(family, weight, style, cached))

    else:
        # Неизвестный тип — пропускаем, ставим заглушки
        values = [None] * count

    return PropChunk(class_id=class_id, prop_name=prop_name, type_id=type_id, values=values)


def _cframe_orientation_to_matrix(orientation_id: int) -> tuple:
    """Roblox использует 24 стандартных orientations (вращения по 90°).

    Источник: https://dom.rojo.space/binary#cframe-orientation-ids

    Формула из rbx-dom:
      orientation_id = (rx_axis * 6) + ry_axis + 1
    где rx_axis, ry_axis ∈ {0..5} = (R0, R1, R2, R3, R4, R5):
      R0 = +X, R1 = +Y, R2 = +Z, R3 = -X, R4 = -Y, R5 = -Z

    rx — это направление куда смотрит локальная +X ось куба (правая грань),
    ry — направление куда смотрит локальная +Y ось (верхняя грань).

    Возвращает (r00, r01, r02, r10, r11, r12, r20, r21, r22) row-major.
    Матрица собирается так: rx, ry, rz это столбцы.
    """
    # Правильный порядок axes (rbx-dom):
    # 0=+X, 1=+Y, 2=+Z, 3=-X, 4=-Y, 5=-Z
    AXES = [
        (1, 0, 0),   # +X
        (0, 1, 0),   # +Y
        (0, 0, 1),   # +Z
        (-1, 0, 0),  # -X
        (0, -1, 0),  # -Y
        (0, 0, -1),  # -Z
    ]
    # orientation_id = 1..36 (некоторые комбинации rx==ry невалидны, в файлах
    # не встречаются — но id может доходить до 6*6 = 36, не 24).
    if not (1 <= orientation_id <= 36):
        return (1, 0, 0, 0, 1, 0, 0, 0, 1)

    idx = orientation_id - 1
    rx_idx = idx // 6
    ry_idx = idx % 6
    rx = AXES[rx_idx]
    ry = AXES[ry_idx]
    # rz = rx × ry (cross product) — третий столбец
    rz = (
        rx[1] * ry[2] - rx[2] * ry[1],
        rx[2] * ry[0] - rx[0] * ry[2],
        rx[0] * ry[1] - rx[1] * ry[0],
    )
    # rx, ry, rz — это СТОЛБЦЫ матрицы.
    # row-major: [r00=rx[0], r01=ry[0], r02=rz[0], r10=rx[1], r11=ry[1], r12=rz[1], ...]
    r00, r10, r20 = rx
    r01, r11, r21 = ry
    r02, r12, r22 = rz
    return (r00, r01, r02, r10, r11, r12, r20, r21, r22)


def _quat_to_matrix(qx, qy, qz, qw) -> tuple:
    """Кватернион в 3x3 row-major matrix."""
    xx = qx * qx
    yy = qy * qy
    zz = qz * qz
    xy = qx * qy
    xz = qx * qz
    yz = qy * qz
    wx = qw * qx
    wy = qw * qy
    wz = qw * qz
    return (
        1 - 2*(yy + zz),  2*(xy - wz),       2*(xz + wy),
        2*(xy + wz),      1 - 2*(xx + zz),   2*(yz - wx),
        2*(xz - wy),      2*(yz + wx),       1 - 2*(xx + yy),
    )
