#!/usr/bin/env python3
"""TEXT-Chunk (Texturen) von Level-WADs lesen/schreiben.

Aufbau (aus 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, Daten }
      flags==0      : w*h Byte (8-Bit, Palette)
      flags&0x80    : u32 Länge, RLE16 (s16 n: n>0 -> n Literal-u16, n<0 -> nächstes u16 (-n)-mal)
      sonst         : w*h u16 roh
  u32 nTex; Rest = Textur-Einträge (8 Byte) + Anhang (siehe 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:                       # Lauf >= 2
            out += [(-(j - i + 1)) & 0xffff, int(px[i])]; i = j + 1
            continue
        k = i                           # Literale bis zum nächsten Lauf
        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 = durchsichtig (Alpha 0). STP=1 -> Alpha 255, STP=0 (Pixel != 0) -> Alpha 128
    (auf Seiten mit Blendmodus = halbtransparent)."""
    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):
    """Umkehrung: Alpha < 64 -> 0x0000, 64..191 -> STP=0, >= 192 (oder RGB-Bild) -> 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)   # fast schwarz statt durchsichtig
        v = np.where(a < 64, 0, v)
    else:
        v = v | 0x8000
    return v.astype(np.uint16)


def tex_entries(tail):
    """Textur-Einträge (Faces referenzieren diese Indizes): u16 flags, u8 page, u8 clut(0xff),
    u8 x0, u8 x1, u8 y0, u8 y1 (inklusive Kanten)."""
    n = struct.unpack_from('<I', tail, 0)[0]
    return [struct.unpack_from('<H6B', tail, 4 + 8 * i) for i in range(n)]