"""A Freelancer interface icon (.3db) from a PNG -- like the NeuroNet icons (NN_info.3db ...).

Retail icons are one textured quad: the mesh lives in DATA\\INTERFACE\\interface.generic.vms, the
material (DcDtOcOt, Dt_flags 64) and a 32-bit TGA with three mips are inside the .3db. The
texture is grey; the button tints it (white, yellow under the mouse). A new icon gets the same
quad in its own embedded VMeshLibrary, its own material name (material ids are global, keyed by
the name's CRC) and the PNG as grey + alpha, brightest pixel white.

    python tools/make_icon_3db.py <in.png> <out.3db> <name> [--color]

<name> becomes the material "UI_REM_<name>", the texture "REM_<name>.tga" and the mesh library
"rem_icon_<name>.lod0.vms". --color keeps the PNG's colours (not tinted like the originals).
"""
import os
import struct
import sys
import zlib

sys.path.insert(0, os.path.dirname(__file__))
import flcrc  # noqa: E402
import utf    # noqa: E402

# the retail NN_info quad (interface.generic-102.vms, start vertex 11026): x, y, z, u, v
QUAD_HALF = 0.020543185994029045
QUAD_Z = -2.1847961306775687e-06
QUAD = [(-QUAD_HALF, QUAD_HALF, QUAD_Z, 0.0, 1.0), (-QUAD_HALF, -QUAD_HALF, QUAD_Z, 0.0, 0.0),
        (QUAD_HALF, QUAD_HALF, QUAD_Z, 1.0, 1.0), (QUAD_HALF, -QUAD_HALF, QUAD_Z, 1.0, 0.0)]
QUAD_INDICES = [0, 1, 2, 1, 3, 2]
TEX_SIZE = 128


def read_png(path):
    """RGBA rows (top first) of an 8-bit PNG: grey, grey+alpha, RGB, RGBA or palette."""
    d = open(path, 'rb').read()
    assert d[:8] == b'\x89PNG\r\n\x1a\n', path
    pos = 8
    idat = b''
    palette = trns = None
    while pos < len(d):
        n, kind = struct.unpack('>I4s', d[pos:pos + 8])
        body = d[pos + 8:pos + 8 + n]
        if kind == b'IHDR':
            w, h, depth, ctype, _c, _f, interlace = struct.unpack('>IIBBBBB', body)
        elif kind == b'PLTE':
            palette = body
        elif kind == b'tRNS':
            trns = body
        elif kind == b'IDAT':
            idat += body
        pos += 12 + n
    assert depth in (8, 16) and interlace == 0, path + ': only 8/16-bit non-interlaced PNG'
    assert depth == 8 or ctype != 3, path + ': 16-bit palette PNG'
    # a 16-bit sample is filtered byte-wise as two bytes; its high byte is kept below
    channels = {0: 1, 2: 3, 3: 1, 4: 2, 6: 4}[ctype] * (depth // 8)
    raw = zlib.decompress(idat)
    stride = w * channels
    rows = []
    prev = bytearray(stride)
    p = 0
    for _ in range(h):
        f = raw[p]
        line = bytearray(raw[p + 1:p + 1 + stride])
        p += 1 + stride
        for i in range(stride):
            a = line[i - channels] if i >= channels else 0
            b = prev[i]
            c = prev[i - channels] if i >= channels else 0
            if f == 1:
                line[i] = (line[i] + a) & 0xFF
            elif f == 2:
                line[i] = (line[i] + b) & 0xFF
            elif f == 3:
                line[i] = (line[i] + ((a + b) >> 1)) & 0xFF
            elif f == 4:
                pa, pb, pc = abs(b - c), abs(a - c), abs(a + b - 2 * c)
                pr = a if pa <= pb and pa <= pc else (b if pb <= pc else c)
                line[i] = (line[i] + pr) & 0xFF
        prev = line
        rows.append(line if depth == 8 else line[0::2])
    channels //= depth // 8
    out = []
    for line in rows:
        px = []
        for x in range(w):
            s = line[x * channels:(x + 1) * channels]
            if ctype == 6:
                px.append(tuple(s))
            elif ctype == 2:
                px.append((s[0], s[1], s[2], 255))
            elif ctype == 4:
                px.append((s[0], s[0], s[0], s[1]))
            elif ctype == 0:
                px.append((s[0], s[0], s[0], 255))
            else:
                i = s[0]
                a = trns[i] if trns is not None and i < len(trns) else 255
                px.append((palette[3 * i], palette[3 * i + 1], palette[3 * i + 2], a))
        out.append(px)
    return w, h, out


def resample(w, h, rows, size):
    """Box filter to size x size, colour weighted by alpha (no dark fringes)."""
    out = []
    for y in range(size):
        y0, y1 = y * h // size, max((y + 1) * h // size, y * h // size + 1)
        line = []
        for x in range(size):
            x0, x1 = x * w // size, max((x + 1) * w // size, x * w // size + 1)
            r = g = b = a = n = 0
            for yy in range(y0, y1):
                for xx in range(x0, x1):
                    pr, pg, pb, pa = rows[yy][xx]
                    r += pr * pa
                    g += pg * pa
                    b += pb * pa
                    a += pa
                    n += 1
            line.append((r // a, g // a, b // a, a // n) if a else (0, 0, 0, 0))
        out.append(line)
    return out


def to_grey(rows, gamma=0.55):
    """Luminance, the brightest visible pixel white and the rest lifted (gamma < 1) -- the retail
    icons are near-white shapes the button tints."""
    lum = [[(0.2126 * r + 0.7152 * g + 0.0722 * b, a) for (r, g, b, a) in line] for line in rows]
    peak = max([l for line in lum for (l, a) in line if a > 32] or [255.0])
    out = []
    for line in lum:
        out.append([(min(int(255.0 * (min(l / max(peak, 1.0), 1.0) ** gamma) + 0.5), 255),) * 3 + (a,)
                    for (l, a) in line])
    return out


def crop_to_content(w, h, rows, margin=0.06):
    """The visible part (alpha > 8) as a centred square with a small margin, so the icon fills
    the quad like the retail ones."""
    xs = [x for line in rows for x, p in enumerate(line) if p[3] > 8]
    ys = [y for y, line in enumerate(rows) if any(p[3] > 8 for p in line)]
    if not xs:
        return w, h, rows
    x0, x1, y0, y1 = min(xs), max(xs) + 1, min(ys), max(ys) + 1
    side = int(max(x1 - x0, y1 - y0) * (1.0 + 2.0 * margin)) + 1
    cx, cy = (x0 + x1) // 2, (y0 + y1) // 2
    ox, oy = cx - side // 2, cy - side // 2
    empty = (0, 0, 0, 0)
    out = [[rows[y][x] if 0 <= y < h and 0 <= x < w else empty for x in range(ox, ox + side)]
           for y in range(oy, oy + side)]
    return side, side, out


def tga(rows):
    """32-bit uncompressed TGA, bottom row first, 8 alpha bits -- as the retail icons."""
    size = len(rows)
    head = struct.pack('<BBBHHBHHHHBB', 0, 0, 2, 0, 0, 0, 0, 0, size, size, 32, 8)
    body = bytearray()
    for line in reversed(rows):
        for r, g, b, a in line:
            body += bytes((b, g, r, a))
    return head + bytes(body)


def vmesh_data(material_crc):
    """One mesh, four XYZ+TEX1 vertices (FVF 0x102), six indices."""
    d = struct.pack('<IIHHHH', 1, 4, 1, len(QUAD_INDICES), 0x102, len(QUAD))
    d += struct.pack('<IHHHH', material_crc, 0, len(QUAD) - 1, len(QUAD_INDICES), 204)
    d += struct.pack('<%dH' % len(QUAD_INDICES), *QUAD_INDICES)
    for v in QUAD:
        d += struct.pack('<5f', *v)
    return d


def vmesh_ref(library_crc):
    """VMeshRef (60 bytes): the whole library, bounds and sphere of the quad."""
    return struct.pack('<IIHHHHHH10f', 60, library_crc, 0, len(QUAD), 0, len(QUAD_INDICES), 0, 1,
                       QUAD_HALF, -QUAD_HALF, QUAD_HALF, -QUAD_HALF, QUAD_Z, QUAD_Z, 0.0, 0.0, QUAD_Z,
                       QUAD_HALF * 1.5094)


def make(png, out, name, color=False):
    w, h, rows = read_png(png)
    w, h, rows = crop_to_content(w, h, rows)
    mips = []
    for size in (TEX_SIZE, TEX_SIZE // 2, TEX_SIZE // 4):
        img = resample(w, h, rows, size)
        mips.append(tga(img if color else to_grey(img)))
    material = 'UI_REM_' + name
    texture = 'REM_' + name + '.tga'
    library = 'rem_icon_' + name.lower() + '.lod0.vms'
    N = utf.Node
    root = N('\\', children=[
        N('Exporter Version', data=b'REM make_icon_3db\0'),
        N('VMeshLibrary', children=[N(library, children=[N('VMeshData', data=vmesh_data(flcrc.crc(material)))])]),
        N('VMeshPart', children=[N('VMeshRef', data=vmesh_ref(flcrc.crc(library)))]),
        N('Material library', children=[
            N('Material count', data=struct.pack('<I', 1)),
            N(material, children=[N('Type', data=b'DcDtOcOt\0'), N('Dt_name', data=texture.encode() + b'\0'),
                                  N('Dt_flags', data=struct.pack('<I', 64))])]),
        N('Texture library', children=[
            N(texture, children=[N('MIP%d' % i, data=m) for i, m in enumerate(mips)])]),
    ])
    utf.write(out, root)


if __name__ == '__main__':
    args = [a for a in sys.argv[1:] if not a.startswith('--')]
    if len(args) != 3:
        print(__doc__)
        sys.exit(1)
    make(args[0], args[1], args[2], color='--color' in sys.argv)
    print('wrote', args[1])