"""
rbxl_parser.py — главный парсер Roblox Binary Level (.rbxl) v1.

Использование:
    from rbxl_parser import parse, RobloxModel
    with open('map.rbxl', 'rb') as f:
        model = parse(f.read())
    for inst in model.instances:
        print(inst.class_name, inst.properties.get('Name'))

Архитектура:
  1. parse() читает все chunks (META, SSTR, INST, PROP, PRNT, END).
  2. Из INST chunks извлекаем какие классы (Part, Script, ...) есть в файле
     и сколько инстансов у каждого.
  3. Из PROP chunks извлекаем значения свойств (по одному chunk на свойство
     для каждого класса).
  4. Из PRNT chunk — иерархия parent-ссылок (referent->referent).
  5. Собираем плоский список Instance'ов с привязкой properties + parent.

Полная документация формата: https://dom.rojo.space/binary
"""
import struct
import io
import lz4.block
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any

from rbxl_binreader import (
    BinReader,
    read_interleaved_uint32_array,
    read_referent_array,
)
from rbxl_types import (
    decode_prop_chunk,
    PropChunk,
)


RBXL_SIGNATURE = b'<roblox!\x89\xff\r\n\x1a\n'


class RbxlParseError(Exception):
    pass


# ──────────────────────────────────────────────────────────────────────
# Промежуточные структуры (raw chunks)
# ──────────────────────────────────────────────────────────────────────

@dataclass
class RawChunk:
    name: str
    payload: bytes


@dataclass
class RawHeader:
    version: int
    class_count: int
    instance_count: int


@dataclass
class InstChunk:
    """Описание одного класса (берётся из INST chunk)."""
    class_id: int              # порядковый индекс (=index в массиве)
    class_name: str            # 'Part', 'Script', 'MeshPart'
    is_service: bool           # сервис (Workspace, Lighting) или обычный объект
    referent_ids: List[int]    # для каждого инстанса его referent (referent — уникальный int32 id)


# ──────────────────────────────────────────────────────────────────────
# Высокоуровневая модель: Instance + RobloxModel
# ──────────────────────────────────────────────────────────────────────


@dataclass
class Instance:
    """Один объект Roblox-сцены, готовый к преобразованию в Rublox."""
    referent: int                            # уникальный id из rbxl
    class_name: str                          # 'Part', 'Script', ...
    properties: Dict[str, Any] = field(default_factory=dict)
    parent_referent: Optional[int] = None    # None == root (DataModel)
    children: List['Instance'] = field(default_factory=list)


@dataclass
class RobloxModel:
    """Полная разобранная модель."""
    version: int
    class_count: int
    instance_count: int
    instances: List[Instance]                # плоский список
    by_referent: Dict[int, Instance]         # быстрый лукап по referent
    roots: List[Instance]                    # инстансы без parent (children of DataModel)
    shared_strings: List[bytes]              # из SSTR chunk
    meta: Dict[str, str]                     # из META chunk
    warnings: List[str]                      # некритичные проблемы парсинга
    # сырые chunks на случай если что-то надо допарсить отдельно
    raw_chunks: List[RawChunk] = field(default_factory=list)


# ──────────────────────────────────────────────────────────────────────
# Read chunks (без decode значений)
# ──────────────────────────────────────────────────────────────────────


def _read_chunks(blob: bytes) -> tuple:
    """Возвращает (header, [RawChunk]).

    Каждый chunk:
        name (4 bytes ASCII), compressed_len (uint32), uncompressed_len (uint32),
        reserved (uint32), payload (compressed_len bytes if compressed_len>0 else uncompressed_len bytes).

    Если compressed_len == 0 — payload не сжат.
    """
    if not blob.startswith(RBXL_SIGNATURE):
        raise RbxlParseError(
            f"signature mismatch: got {blob[:14].hex()}, expected {RBXL_SIGNATURE.hex()}"
        )

    stream = io.BytesIO(blob[len(RBXL_SIGNATURE):])

    # Header
    version, class_count, instance_count = struct.unpack('<HII', stream.read(10))
    stream.read(8)  # reserved
    header = RawHeader(version, class_count, instance_count)

    chunks = []
    while True:
        hdr = stream.read(16)
        if len(hdr) < 16:
            break
        name_raw, compressed_len, uncompressed_len, _ = struct.unpack('<4sIII', hdr)
        name = name_raw.decode('ascii', errors='replace').rstrip('\x00').rstrip()

        if compressed_len == 0:
            payload = stream.read(uncompressed_len)
        else:
            data = stream.read(compressed_len)
            try:
                payload = lz4.block.decompress(data, uncompressed_size=uncompressed_len)
            except Exception as e:
                raise RbxlParseError(f"LZ4 decompress failed for chunk {name!r}: {e}")

        chunks.append(RawChunk(name, payload))

        if name.startswith('END'):
            break

    return header, chunks


# ──────────────────────────────────────────────────────────────────────
# Decode INST chunk
# ──────────────────────────────────────────────────────────────────────


def _decode_inst_chunk(payload: bytes) -> InstChunk:
    """Декодирует INST chunk.

    Структура:
        class_id    (uint32)
        class_name  (string)
        is_service  (uint8)
        count       (uint32)  — число инстансов класса
        referents   (interleaved-zigzag-累加 int32, count штук)
        if is_service:
            services_flags (count байт)
    """
    r = BinReader(payload)
    class_id = r.uint32()
    class_name = r.string()
    is_service = bool(r.uint8())
    count = r.uint32()
    referents = read_referent_array(r, count)

    if is_service:
        # Пропускаем сервис-флаги (один байт на инстанс)
        r.read(count)

    return InstChunk(
        class_id=class_id,
        class_name=class_name,
        is_service=is_service,
        referent_ids=referents,
    )


# ──────────────────────────────────────────────────────────────────────
# Decode SSTR chunk
# ──────────────────────────────────────────────────────────────────────


def _decode_sstr_chunk(payload: bytes) -> List[bytes]:
    """SSTR (Shared Strings) chunk.

    Структура:
        version (uint32)
        count   (uint32)
        для каждого:
            md5_hash (16 bytes)  — игнорируется
            length   (uint32)
            data     (length bytes)
    """
    r = BinReader(payload)
    _version = r.uint32()
    count = r.uint32()
    strings = []
    for _ in range(count):
        r.read(16)  # md5
        length = r.uint32()
        strings.append(r.read(length))
    return strings


# ──────────────────────────────────────────────────────────────────────
# Decode PRNT chunk (parent hierarchy)
# ──────────────────────────────────────────────────────────────────────


def _decode_prnt_chunk(payload: bytes, instance_count: int) -> List[tuple]:
    """PRNT (Parents) chunk: для каждого instance его parent.

    Структура:
        version (uint8) — обычно 0
        count   (uint32) — instance_count
        child_referents:   interleaved int32 array (length=count, cumulative)
        parent_referents:  interleaved int32 array (length=count, cumulative)

    Возвращает [(child_referent, parent_referent)]. parent_referent == -1
    означает что parent — root (Workspace/DataModel).
    """
    r = BinReader(payload)
    _version = r.uint8()
    count = r.uint32()
    children = read_referent_array(r, count)
    parents = read_referent_array(r, count)
    return list(zip(children, parents))


# ──────────────────────────────────────────────────────────────────────
# Decode META chunk
# ──────────────────────────────────────────────────────────────────────


def _decode_meta_chunk(payload: bytes) -> Dict[str, str]:
    """META chunk: пары ключ-значение."""
    r = BinReader(payload)
    count = r.uint32()
    meta = {}
    for _ in range(count):
        k = r.string()
        v = r.string()
        meta[k] = v
    return meta


# ──────────────────────────────────────────────────────────────────────
# Главная функция parse()
# ──────────────────────────────────────────────────────────────────────


def parse(blob: bytes) -> RobloxModel:
    """Парсит .rbxl-байты в RobloxModel со списком Instance'ов."""
    header, raw_chunks = _read_chunks(blob)
    warnings: List[str] = []

    # 1. Извлекаем META
    meta = {}
    for c in raw_chunks:
        if c.name == 'META':
            try:
                meta = _decode_meta_chunk(c.payload)
            except Exception as e:
                warnings.append(f"META decode failed: {e}")

    # 2. SSTR
    shared_strings: List[bytes] = []
    for c in raw_chunks:
        if c.name == 'SSTR':
            try:
                shared_strings = _decode_sstr_chunk(c.payload)
            except Exception as e:
                warnings.append(f"SSTR decode failed: {e}")
            break

    # 3. INST — описания классов
    inst_chunks: Dict[int, InstChunk] = {}
    for c in raw_chunks:
        if c.name == 'INST':
            try:
                ic = _decode_inst_chunk(c.payload)
                inst_chunks[ic.class_id] = ic
            except Exception as e:
                warnings.append(f"INST decode failed: {e}")

    # 4. Создаём пустые Instance'ы по референтам
    by_referent: Dict[int, Instance] = {}
    instances: List[Instance] = []
    for class_id, ic in inst_chunks.items():
        for ref in ic.referent_ids:
            inst = Instance(referent=ref, class_name=ic.class_name)
            by_referent[ref] = inst
            instances.append(inst)

    # 5. PROP — заполняем свойства
    for c in raw_chunks:
        if c.name == 'PROP':
            try:
                # Узнаём class_id (первые 4 байта чанка) чтобы взять count
                class_id = struct.unpack('<I', c.payload[:4])[0]
                ic = inst_chunks.get(class_id)
                if not ic:
                    warnings.append(f"PROP for unknown class_id={class_id}")
                    continue
                pc = decode_prop_chunk(c.payload, len(ic.referent_ids), shared_strings)
                if len(pc.values) != len(ic.referent_ids):
                    warnings.append(
                        f"PROP {ic.class_name}.{pc.prop_name}: value count "
                        f"{len(pc.values)} != instance count {len(ic.referent_ids)}"
                    )
                for ref, val in zip(ic.referent_ids, pc.values):
                    inst = by_referent.get(ref)
                    if inst:
                        inst.properties[pc.prop_name] = val
            except Exception as e:
                warnings.append(f"PROP decode failed: {e!r}")

    # 6. PRNT — собираем дерево
    for c in raw_chunks:
        if c.name == 'PRNT':
            try:
                links = _decode_prnt_chunk(c.payload, header.instance_count)
                for child_ref, parent_ref in links:
                    child = by_referent.get(child_ref)
                    if not child:
                        continue
                    if parent_ref == -1:
                        child.parent_referent = None
                    else:
                        child.parent_referent = parent_ref
                        parent = by_referent.get(parent_ref)
                        if parent:
                            parent.children.append(child)
            except Exception as e:
                warnings.append(f"PRNT decode failed: {e!r}")

    # 7. roots
    roots = [i for i in instances if i.parent_referent is None]

    return RobloxModel(
        version=header.version,
        class_count=header.class_count,
        instance_count=header.instance_count,
        instances=instances,
        by_referent=by_referent,
        roots=roots,
        shared_strings=shared_strings,
        meta=meta,
        warnings=warnings,
        raw_chunks=raw_chunks,
    )


# ──────────────────────────────────────────────────────────────────────
# Утилиты для отчёта
# ──────────────────────────────────────────────────────────────────────


def class_histogram(model: RobloxModel) -> Dict[str, int]:
    """Возвращает {class_name: count} — сколько объектов каждого типа."""
    out: Dict[str, int] = {}
    for inst in model.instances:
        out[inst.class_name] = out.get(inst.class_name, 0) + 1
    return out


def summarize(model: RobloxModel, top_classes: int = 30) -> str:
    """Печатает человекочитаемый отчёт о модели."""
    lines = []
    lines.append(f"=== Roblox model ===")
    lines.append(f"  version:        {model.version}")
    lines.append(f"  class count:    {model.class_count}")
    lines.append(f"  instance count: {model.instance_count}")
    lines.append(f"  shared strings: {len(model.shared_strings)}")
    lines.append(f"  warnings:       {len(model.warnings)}")
    if model.meta:
        lines.append("  meta:")
        for k, v in list(model.meta.items())[:5]:
            lines.append(f"    {k}: {v}")
    lines.append("")
    lines.append(f"=== Top {top_classes} classes ===")
    hist = class_histogram(model)
    for cls, n in sorted(hist.items(), key=lambda x: -x[1])[:top_classes]:
        lines.append(f"  {n:>5d}  {cls}")
    lines.append("")
    if model.warnings:
        lines.append(f"=== Warnings (first 20) ===")
        for w in model.warnings[:20]:
            lines.append(f"  ! {w}")
    return "\n".join(lines)


if __name__ == '__main__':
    import sys
    path = sys.argv[1] if len(sys.argv) > 1 else '/opt/rbxl-importer/reference_files/easy_obby.rbxl'
    with open(path, 'rb') as f:
        blob = f.read()
    print(f"file size: {len(blob)} bytes\n")
    model = parse(blob)
    print(summarize(model, top_classes=40))

    # Дополнительная статистика: что внутри у Part
    print("\n=== Sample Part properties ===")
    parts = [i for i in model.instances if i.class_name == 'Part']
    if parts:
        p = parts[0]
        print(f"  Sample Part (referent={p.referent}):")
        for k, v in p.properties.items():
            sv = str(v)
            if len(sv) > 80:
                sv = sv[:77] + '...'
            print(f"    {k}: {sv}")
        print(f"\n  Total Parts: {len(parts)}")