"""PSP theme (``.ptf``) introspection: name + preview image.

PTF themes embed a readable name in their header and a preview image (a GIM
bitmap) in object 0. This module decodes both without any external service:

* the name lives at offset 0x08 (128 UTF-8 bytes);
* the preview is object 0's largest image item, stored as a GIM that may be
  Sony-LZR compressed.

GIM decoding covers the common RGBA formats (5650/5551/4444/8888) and 8-bit
palettes; anything exotic simply yields no preview.
"""
from __future__ import annotations

import zlib
from pathlib import Path

PTF_MAGIC = b"\x00PTF"

# --------------------------------------------------------------------------- #
# Sony LZR/RLZ decompression (sceLzrDecompress), ported from PTF-Studio.
# --------------------------------------------------------------------------- #
_LZR_PROBABILITY_SLOTS = 2800
_LZR_INITIAL_PROBABILITY = 0x80
_LZR_RENORM_LIMIT = 0x00FFFFFF
_LZR_U32_MOD = 0x100000000


def _lzr_decompress(inp, max_output, offset=0):
    inp = bytes(inp)
    end = offset + len(inp) - offset
    out = bytearray(max_output)
    probs = [_LZR_INITIAL_PROBABILITY] * _LZR_PROBABILITY_SLOTS
    out_pos = 0
    in_pos = offset + 5
    code = int.from_bytes(inp[offset + 1:offset + 5], "big")
    range_ = 0xFFFFFFFF
    spec = 0
    corrupt = False
    buf_off = 0
    last_byte = 0
    last_number_flag = 0
    type_ = inp[offset]
    if type_ >= 0x80:
        type_ -= 0x100

    def next_byte():
        nonlocal in_pos, corrupt
        if in_pos >= end:
            corrupt = True
            return 0
        b = inp[in_pos]
        in_pos += 1
        return b

    if type_ > 7:
        raise ValueError("unsupported LZR literal-context mode")
    if type_ < 0:
        stored = code
        if stored > len(out) or in_pos + stored > end:
            raise ValueError("malformed stored LZR block")
        return inp[in_pos:in_pos + stored]

    def renormalize():
        nonlocal code, range_
        if range_ <= _LZR_RENORM_LIMIT:
            code = (code * 256 + next_byte()) % _LZR_U32_MOD
            range_ = (range_ * 256) % _LZR_U32_MOD

    def decode_bit(prob_index):
        nonlocal code, range_, corrupt
        if prob_index < 0 or prob_index >= _LZR_PROBABILITY_SLOTS:
            corrupt = True
            return 0
        renormalize()
        if corrupt:
            return 0
        p = probs[prob_index]
        bound = (range_ // 256) * p
        probs[prob_index] = p - (p // 8)
        if code < bound:
            range_ = bound
            probs[prob_index] += 31
            return 1
        code -= bound
        range_ -= bound
        return 0

    def decode_bit_speculative(prob_index):
        nonlocal code, range_, spec, corrupt
        if prob_index < 0 or prob_index >= _LZR_PROBABILITY_SLOTS:
            corrupt = True
            return 0
        if spec <= _LZR_RENORM_LIMIT:
            code = (code * 256 + next_byte()) % _LZR_U32_MOD
            range_ = (spec * 256) % _LZR_U32_MOD
        if corrupt:
            return 0
        p = probs[prob_index]
        bound = (range_ // 256) * p
        spec = bound
        probs[prob_index] = p - (p // 8)
        if code < bound:
            range_ = bound
            probs[prob_index] += 31
            return 1
        code -= bound
        range_ -= bound
        return 0

    def decode_number(bits, base, step):
        nonlocal code, range_, corrupt, last_number_flag
        number = 1
        if bits >= 3:
            number = number * 2 + decode_bit(base + 3 * step)
            if bits >= 4:
                number = number * 2 + decode_bit(base + 3 * step)
                if bits >= 5:
                    renormalize()
                    remaining = bits
                    while remaining >= 5:
                        if range_ == 0:
                            corrupt = True
                            return 0
                        number *= 2
                        range_ = range_ // 2
                        if code < range_:
                            number += 1
                        else:
                            code -= range_
                        remaining -= 1
        last_number_flag = decode_bit(base)
        number = number * 2 + last_number_flag
        if bits >= 1:
            number = number * 2 + decode_bit(base + step)
            if bits >= 2:
                number = number * 2 + decode_bit(base + 2 * step)
        return number

    def selector_context():
        return buf_off + 2488

    def decode_literal(shift):
        nonlocal out_pos, last_byte, corrupt, buf_off
        if buf_off > 0:
            buf_off -= 1
        if out_pos == len(out):
            return False
        context = ((((out_pos & 7) << 8) + last_byte) // (2 ** shift)) & 7
        tree_base = context * 0xFF - 1
        node = 1
        while node <= 0xFF:
            node = node * 2 + decode_bit(tree_base + node)
            if corrupt:
                return False
        out[out_pos] = node & 0xFF
        last_byte = node & 0xFF
        out_pos += 1
        return True

    def decode_match():
        nonlocal out_pos, last_byte, buf_off, spec, code, range_, corrupt, last_number_flag
        slot = selector_context()
        spec = range_
        length_class = -1
        bit = 0
        while True:
            slot += 8
            bit = decode_bit_speculative(slot)
            if corrupt:
                return -1
            length_class += bit
            if bit == 0 or length_class >= 6:
                break
        offset_group = length_class + 2033
        offset_bias = 64
        if bit != 0 or length_class >= 0:
            length_context = (length_class << 5) + (((out_pos * (2 ** length_class)) & 3) << 3) + buf_off + 2552
            match_length = decode_number(length_class, length_context, 8)
            if corrupt:
                return -1
            if match_length == 0xFF:
                return 1
            if last_number_flag != 0 or length_class > 0:
                offset_group += 56
                offset_bias = 352
        else:
            match_length = 1
        index = 1
        offset_class = 0
        while True:
            offset_class = (index << 4) - offset_bias
            prob_index = offset_group + (index << 3)
            if prob_index >= _LZR_PROBABILITY_SLOTS:
                return -1
            bit = decode_bit(prob_index)
            if corrupt:
                return -1
            index = (index << 1) | bit
            if offset_class >= 0:
                break
        if bit != 0 or offset_class > 0:
            if bit == 0:
                offset_class -= 8
            base = offset_class + 2344
            if base < 0 or base + 3 >= _LZR_PROBABILITY_SLOTS:
                return -1
            back_offset = decode_number(offset_class // 8, base, 1)
            if corrupt:
                return -1
        else:
            back_offset = 1
        if back_offset > out_pos:
            return -1
        copy_end = out_pos + match_length + 1
        if copy_end > len(out):
            return -1
        buf_off = ((copy_end + 1) & 1) + 6
        from_ = out_pos - back_offset
        while out_pos < copy_end:
            out[out_pos] = out[from_]
            out_pos += 1
            from_ += 1
        last_byte = out[out_pos - 1]
        return 0

    token_budget = max(1024, max_output * 2 + 1024)
    while True:
        token_budget -= 1
        if token_budget < 0:
            raise ValueError("malformed LZR stream did not terminate")
        token_is_match = decode_bit(selector_context())
        if corrupt:
            raise ValueError("truncated LZR stream")
        if token_is_match == 0:
            if not decode_literal(type_):
                raise ValueError("LZR stream exceeds declared output size")
        else:
            result = decode_match()
            if result == 1:
                return bytes(out[:out_pos])
            if result < 0:
                raise ValueError("malformed LZR back-reference")


# --------------------------------------------------------------------------- #
# PTF parsing
# --------------------------------------------------------------------------- #
def _decompress_item(blob, comp, ucsize):
    if comp == 2 and len(blob) < ucsize:
        return zlib.decompress(blob)
    if comp == 1 and len(blob) < ucsize:
        return _lzr_decompress(blob, ucsize)
    return blob


def parse_ptf(data):
    """Return ``(name, items)`` where items are decompressed item blobs."""
    if not isinstance(data, (bytes, bytearray)):
        data = data.read_bytes() if hasattr(data, "read_bytes") else bytes(data)
    data = bytes(data)
    if data[:4] != PTF_MAGIC:
        raise ValueError("not a PTF file")

    raw_name = data[8:8 + 128]
    name = raw_name.split(b"\x00", 1)[0].decode("utf-8", "replace").strip()

    def u16(o):
        return int.from_bytes(data[o:o + 2], "little")

    def u32(o):
        return int.from_bytes(data[o:o + 4], "little")

    items = []
    pos = 0x100
    while True:
        offset = u32(pos)
        pos += 4
        if offset == 0 or offset >= len(data):
            break
        count = u16(offset + 2)
        data_off = u32(offset + 8)
        p = data_off
        for _ in range(count):
            file_type = u16(p + 4)
            comp = u16(p + 6)
            size = u32(p + 8)
            ucsize = u32(p + 12)
            blob = data[p + 0x20:p + 0x20 + size]
            items.append((file_type, _decompress_item(blob, comp, ucsize)))
            p += 0x20 + size
    return name, items


def ptf_name(path):
    """Theme display name embedded in a ``.ptf`` file."""
    try:
        with open(path, "rb") as fh:
            head = fh.read(8 + 128)
    except OSError:
        return ""
    if head[:4] != PTF_MAGIC:
        return ""
    name = head[8:8 + 128].split(b"\x00", 1)[0].decode("utf-8", "replace").strip()
    return name or Path(path).stem


# --------------------------------------------------------------------------- #
# GIM decoding (RGBA formats only)
# --------------------------------------------------------------------------- #
def _overscan(value, tile):
    return value if value % tile == 0 else value + (tile - (value % tile))


def _untile(pixels, w, h, tile_w, tile_h):
    over_w = _overscan(w, tile_w)
    over_h = _overscan(h, tile_h)
    out = [0] * (w * h)
    tile_pos = 0
    origin_x = 0
    origin_y = 0
    for y in range(over_h):
        for x in range(over_w):
            dx = origin_x + (tile_pos % tile_w)
            dy = origin_y + (tile_pos // tile_w)
            if dx < w and dy < h:
                out[dy * w + dx] = pixels[y * over_w + x]
            tile_pos += 1
            if tile_pos == tile_w * tile_h:
                tile_pos = 0
                origin_x += tile_w
                if origin_x >= over_w:
                    origin_x = 0
                    origin_y += tile_h
    return out


def _gim_to_rgba(data):
    """Decode a GIM to ``(width, height, rgba_bytes)`` or None if unsupported."""
    data = bytes(data)
    if len(data) < 0x10:
        return None

    # Detect byte order from the magic. Sony GIMs store the magic as the byte
    # string "MIG.00.1PSP" with little-endian data (LibPSPThemes semantics).
    if data[:4] == b"MIG.":
        byteorder = "little"
    elif data[:4] == b".GIM":
        byteorder = "big"
    else:
        return None

    def rd(off, size):
        return int.from_bytes(data[off:off + size], byteorder)

    image = None
    palette = None
    pos = 0x10
    for _ in range(32):
        if pos + 16 > len(data):
            break
        btype = rd(pos, 2)
        next_block = rd(pos + 8, 4)
        if btype == 4 or btype == 5:
            h = pos + 0x10
            fmt = rd(h + 4, 2)
            pixel_order = rd(h + 6, 2)
            width = rd(h + 8, 2)
            height = rd(h + 10, 2)
            bpp = rd(h + 12, 2)
            next_index = rd(h + 24, 4)
            frame_data_end = rd(h + 32, 4)
            frame_offset = rd(pos + 0x10 + next_index, 4)
            block_content = pos + 0x10
            start = block_content + frame_offset
            end = block_content + frame_data_end
            raw = data[start:end]
            pixels = _decode_pixels(raw, fmt, bpp, width, height, pixel_order)
            if pixels is None:
                return None
            if btype == 4:
                image = (width, height, pixels)
            else:
                palette = (fmt, width, height, pixels)
        if next_block == 0:
            break
        pos += next_block

    if image is None:
        return None
    width, height, pixels = image
    if palette is not None:
        _, pw, ph, pvals = palette
        if pw * ph > 0:
            pixels = [pvals[p] if p < len(pvals) else 0 for p in pixels]

    rgba = bytearray(width * height * 4)
    for i, v in enumerate(pixels):
        # pixels are packed 0xAABBGGRR
        o = i * 4
        rgba[o] = v & 0xFF
        rgba[o + 1] = (v >> 8) & 0xFF
        rgba[o + 2] = (v >> 16) & 0xFF
        rgba[o + 3] = (v >> 24) & 0xFF
    return width, height, bytes(rgba)


def _decode_pixels(raw, fmt, bpp, width, height, pixel_order):
    """Decode a GIM frame into a list of packed 0xAABBGGRR ints."""
    if fmt in (0, 1, 2, 3, 5):
        pass
    else:
        return None  # indexed-16 / DXT etc. not supported

    if fmt == 5:  # p8 indexed
        bytes_per_pixel = 1
        tile_w = 0x80 // max(bpp, 8)
    else:
        bytes_per_pixel = max(1, (bpp + 7) // 8)
        tile_w = 0x80 // max(bpp, 1)
    tile_h = 8

    if pixel_order == 1:
        w = _overscan(width, tile_w)
        h = _overscan(height, tile_h)
    else:
        w, h = width, height

    count = w * h
    values = []
    pos = 0
    for _ in range(count):
        v = int.from_bytes(raw[pos:pos + bytes_per_pixel], "little")
        pos += bytes_per_pixel
        values.append(_expand(fmt, bpp, v))
        if pos >= len(raw):
            # pad missing trailing pixels with transparent
            values.extend([0] * (count - len(values)))
            break

    if pixel_order == 1:
        values = _untile(values, width, height, tile_w, tile_h)
    return values[:width * height]


def _expand(fmt, bpp, v):
    """Expand a raw GIM pixel value into packed 0xAABBGGRR."""
    if fmt == 3:  # rgba8888
        r = v & 0xFF
        g = (v >> 8) & 0xFF
        b = (v >> 16) & 0xFF
        a = (v >> 24) & 0xFF
        return (a << 24) | (b << 16) | (g << 8) | r
    if fmt == 0:  # rgba5650
        r = (v >> 11) & 0x1F
        g = (v >> 5) & 0x3F
        b = v & 0x1F
        r = (r << 3) | (r >> 2)
        g = (g << 2) | (g >> 4)
        b = (b << 3) | (b >> 2)
        return (0xFF << 24) | (b << 16) | (g << 8) | r
    if fmt == 1:  # rgba5551
        r = (v >> 11) & 0x1F
        g = (v >> 6) & 0x1F
        b = (v >> 1) & 0x1F
        a = 0xFF if v & 1 else 0
        r = (r << 3) | (r >> 2)
        g = (g << 3) | (g >> 2)
        b = (b << 3) | (b >> 2)
        return (a << 24) | (b << 16) | (g << 8) | r
    if fmt == 2:  # rgba4444
        r = (v >> 12) & 0xF
        g = (v >> 8) & 0xF
        b = (v >> 4) & 0xF
        a = v & 0xF
        r = (r << 4) | r
        g = (g << 4) | g
        b = (b << 4) | b
        a = (a << 4) | a
        return (a << 24) | (b << 16) | (g << 8) | r
    if fmt == 5:  # p8 — value is a palette index (expanded later)
        return v & 0xFF
    return 0


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #
def theme_preview_rgba(ptf_path):
    """Decode the preview image of a theme PTF.

    Returns ``(width, height, rgba_bytes)`` or None if it can't be decoded.
    """
    try:
        data = Path(ptf_path).read_bytes()
        _, items = parse_ptf(data)
    except (OSError, ValueError):
        return None

    # The preview image is the largest decodable GIM item.
    best = None
    for file_type, blob in items:
        if file_type != 5:
            continue
        decoded = _gim_to_rgba(blob)
        if decoded is None:
            continue
        w, h, rgba = decoded
        if best is None or w * h > best[0] * best[1]:
            best = (w, h, rgba)
    return best