"""Reader for Ren'Py .rpa archives (RPA-2.0, RPA-3.x and common variants).

The index is a zlib-compressed pickle, so it is loaded with a restricted
unpickler that only allows plain containers — archives are untrusted input.
"""

import io
import pickle
import zlib
from collections.abc import Callable
from pathlib import Path


class RpaError(Exception):
    pass


class _SafeUnpickler(pickle.Unpickler):
    _ALLOWED = {("_codecs", "encode"), ("builtins", "set"), ("builtins", "frozenset"),
                ("__builtin__", "set"), ("builtins", "bytearray"), ("__builtin__", "bytearray"),
                ("builtins", "bytes"), ("__builtin__", "bytes")}

    def find_class(self, module, name):
        if (module, name) in self._ALLOWED:
            import builtins
            import codecs
            return codecs.encode if module == "_codecs" else getattr(builtins, name)
        raise pickle.UnpicklingError(f"запрещённый тип в индексе RPA: {module}.{name}")


def _text(v) -> str:
    return v.decode("utf-8", "replace") if isinstance(v, bytes) else str(v)


def read_index(path: Path) -> dict[str, list[tuple[int, int, bytes]]]:
    with path.open("rb") as f:
        header = f.readline(200).decode("latin-1").strip()
        parts = header.split()
        if not parts or not parts[0].startswith("RPA-"):
            raise RpaError(f"{path.name}: неизвестный формат архива ({header[:20]!r})")
        version = parts[0]
        try:
            offset = int(parts[1], 16)
            key = 0
            if version != "RPA-2.0":
                for token in parts[2:]:
                    key ^= int(token, 16)
        except (IndexError, ValueError) as e:
            raise RpaError(f"{path.name}: повреждён заголовок") from e
        f.seek(offset)
        try:
            raw = _SafeUnpickler(io.BytesIO(zlib.decompress(f.read())), encoding="bytes").load()
        except (zlib.error, pickle.UnpicklingError, EOFError) as e:
            raise RpaError(f"{path.name}: не удалось прочитать индекс: {e}") from e

    index: dict[str, list[tuple[int, int, bytes]]] = {}
    for name, entries in raw.items():
        fixed = []
        for entry in entries:
            off, length = entry[0] ^ key, entry[1] ^ key
            prefix = entry[2] if len(entry) > 2 else b""
            if isinstance(prefix, str):
                prefix = prefix.encode("latin-1")
            fixed.append((off, length, prefix))
        index[_text(name)] = fixed
    return index


def extract_rpa(path: Path, dest: Path, only: Callable[[str], bool] | None = None) -> int:
    """Extract files of an archive into dest (all, or those `only` accepts)."""
    index = read_index(path)
    root = dest.resolve()
    count = 0
    with path.open("rb") as f:
        for name, entries in index.items():
            if only is not None and not only(name):
                continue
            target = (dest / name).resolve()
            if not target.is_relative_to(root):
                continue
            target.parent.mkdir(parents=True, exist_ok=True)
            with target.open("wb") as out:
                for off, length, prefix in entries:
                    out.write(prefix)
                    f.seek(off)
                    remaining = length - len(prefix)
                    while remaining > 0:
                        chunk = f.read(min(remaining, 1 << 20))
                        if not chunk:
                            break
                        out.write(chunk)
                        remaining -= len(chunk)
            count += 1
    return count


def write_rpa(path: Path, files: dict[str, bytes], key: int = 0x42424242) -> None:
    """Write an RPA-3.0 archive (used by tests)."""
    with path.open("wb") as f:
        f.write(b"RPA-3.0 " + b"0" * 16 + b" " + b"0" * 8 + b"\n")
        index = {}
        for name, data in files.items():
            off = f.tell()
            f.write(data)
            index[name] = [(off ^ key, len(data) ^ key, b"")]
        idx_off = f.tell()
        f.write(zlib.compress(pickle.dumps(index, protocol=2)))
        f.seek(0)
        f.write(f"RPA-3.0 {idx_off:016x} {key:08x}\n".encode())