#!/usr/bin/env python3
"""Read/write the TEXT chunk (textures) of level WADs.

Layout (from groove.exe 0x4067b0):
  u32 nPal; nPal x { u32 id, u32 n, u16 col[n], u8 extra[n] }
  u32 nPages; nPages x { u32 flags, u32 w, u32 h, data }
      flags==0      : w*h bytes (8-bit, palette)
      flags&0x80    : u32 length, RLE16 (s16 n: n>0 -> n literal u16, n<0 -> next u16 repeated -n times)
      otherwise     : w*h u16 raw
  u32 nTex; rest = texture entries (8 bytes) + appendix (see notes/textures.md)
"""
import struct
import numpy as np


def rle_decode(b, npix):
    out = np.zeros(npix, np.uint16)
    v = np.frombuffer(b, '<u2')
    i = p = 0
    while p < npix and i < len(v):
        n = int(v[i]) - 0x10000 if v[i] & 0x8000 else int(v[i]); i += 1
        if n > 0:
            out[p:p + n] = v[i:i + n]; i += n; p += n
        elif n < 0:
            out[p:p - n] = v[i]; i += 1; p -= n
    return out


def rle_encode(px):
    px = np.asarray(px, np.uint16)
    out, i, N = [], 0, len(px)
    while i < N:
        j = i
        while j + 1 < N and px[j + 1] == px[i] and j - i < 32766:
            j += 1
        if j > i:                       # run >= 2
            out += [(-(j - i + 1)) & 0xffff, int(px[i])]; i = j + 1
            continue
        k = i                           # literals up to the next run
        while k < N and k - i < 32767 and not (k + 1 < N and px[k + 1] == px[k]):
            k += 1
        out += [k - i] + [int(x) for x in px[i:k]]; i = k
    return struct.pack('<%dH' % len(out), *out)


def parse(d):
    o = 0
    pals = []
    n = struct.unpack_from('<I', d, o)[0]; o += 4
    for _ in range(n):
        pid, c = struct.unpack_from('<II', d, o); o += 8
        pals.append((pid, d[o:o + 2 * c], d[o + 2 * c:o + 3 * c])); o += 3 * c
    pages = []
    n = struct.unpack_from('<I', d, o)[0]; o += 4
    for _ in range(n):
        f, w, h = struct.unpack_from('<3I', d, o); o += 12
        if f == 0:
            raw = d[o:o + w * h]; o += w * h
            px = np.frombuffer(raw, np.uint8)
        elif f & 0x80:
            s = struct.unpack_from('<I', d, o)[0]; o += 4
            px = rle_decode(d[o:o + s], w * h); o += s
        else:
            px = np.frombuffer(d[o:o + 2 * w * h], '<u2').copy(); o += 2 * w * h
        pages.append(dict(flags=f, w=w, h=h, px=px.reshape(h, w)))
    return dict(pals=pals, pages=pages, tail=d[o:])


def build(t, compress=True):
    out = [struct.pack('<I', len(t['pals']))]
    for pid, col, ex in t['pals']:
        out.append(struct.pack('<II', pid, len(ex)) + col + ex)
    out.append(struct.pack('<I', len(t['pages'])))
    for p in t['pages']:
        f, w, h = p['flags'], p['w'], p['h']
        if f == 0:
            out.append(struct.pack('<3I', f, w, h) + p['px'].astype(np.uint8).tobytes())
        elif f & 0x80 and compress:
            b = rle_encode(p['px'].reshape(-1))
            out.append(struct.pack('<4I', f, w, h, len(b)) + b)
        else:
            out.append(struct.pack('<3I', f & ~0x80, w, h) + p['px'].astype('<u2').tobytes())
    out.append(t['tail'])
    return b''.join(out)


def abgr1555_to_rgba(px):
    """PS1 format: bits 0-4 R, 5-9 G, 10-14 B, bit 15 = STP.
    0x0000 = transparent (alpha 0). STP=1 -> alpha 255, STP=0 (pixel != 0) -> alpha 128
    (semi-transparent on pages with a blend mode)."""
    px = px.astype(np.uint32)
    r = (px & 31) * 255 // 31
    g = ((px >> 5) & 31) * 255 // 31
    b = ((px >> 10) & 31) * 255 // 31
    a = np.where(px == 0, 0, np.where(px & 0x8000, 255, 128))
    return np.dstack([r, g, b, a]).astype(np.uint8)


def rgba_to_abgr1555(img):
    """Inverse: alpha < 64 -> 0x0000, 64..191 -> STP=0, >= 192 (or RGB image) -> STP=1."""
    img = np.asarray(img).astype(np.uint32)
    v = ((img[..., 2] >> 3) << 10) | ((img[..., 1] >> 3) << 5) | (img[..., 0] >> 3)
    if img.shape[-1] == 4:
        a = img[..., 3]
        v = np.where(a >= 192, v | 0x8000, v)
        v = np.where((a >= 64) & (a < 192) & (v == 0), 0x0421, v)   # near-black instead of transparent
        v = np.where(a < 64, 0, v)
    else:
        v = v | 0x8000
    return v.astype(np.uint16)


def tex_entries(tail):
    """Texture entries (faces reference these indices): u16 flags, u8 page, u8 clut(0xff),
    u8 x0, u8 x1, u8 y0, u8 y1 (edges inclusive)."""
    n = struct.unpack_from('<I', tail, 0)[0]
    return [struct.unpack_from('<H6B', tail, 4 + 8 * i) for i in range(n)]