#!/usr/bin/env python3
"""Texture contact sheet: all textures of a level WAD (or only those used by the world
geometry) as a grid labelled with their index -> for picking textures for levelbuild.py.

  texsheet.py level.wad out.png [--all] [--size 64] [--cols 16] [--only 12,40,77]

Without --all: only textures that occur in TRAK (world tiles), sorted by frequency.
Label: index, below it usage  O=top (floor) / S=side (wall) / U=bottom, size in px.
Flat colours (flag bit 0) are shown as a colour swatch ("F").
"""
import argparse
import collections
import os
import sys

from PIL import Image, ImageDraw

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import geomlib as G  # noqa: E402


def tex_image(pages_rgba, t, size):
    if t['flags'] & 1 or t['page'] < 0:
        im = Image.new('RGBA', (size, size), (t['u0'], t['u1'], t['v0'], 255))
        return im
    pg = pages_rgba[t['page']]
    crop = pg.crop((t['u0'], t['v0'], t['u1'] + 1, t['v1'] + 1))
    return crop.resize((size, size), Image.NEAREST)


def usage(trak):
    up, side, down = collections.Counter(), collections.Counter(), collections.Counter()
    for m in trak:
        for f in m.faces:
            ny = f['plane'][1]
            (up if ny > 0.7 else down if ny < -0.7 else side)[f['tex']] += 1
    return up, side, down


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument('wad'); ap.add_argument('out')
    ap.add_argument('--all', action='store_true')
    ap.add_argument('--size', type=int, default=64)
    ap.add_argument('--cols', type=int, default=16)
    ap.add_argument('--only', help='comma-separated indices')
    a = ap.parse_args()
    ch = G.wad_read(a.wad)
    pages, texs = G.parse_text(G.wad_get(ch, 'TEXT'))
    prgba = [Image.fromarray(G.page_rgba(p), 'RGBA') for p in pages]
    up, side, down = usage(G.parse_trak(G.wad_get(ch, 'TRAK')))
    tot = up + side + down
    if a.only:
        idx = [int(x) for x in a.only.split(',')]
    elif a.all:
        idx = list(range(len(texs)))
    else:
        idx = [i for i, _ in tot.most_common()]
    S, C = a.size, a.cols
    cw, chh = S + 8, S + 30
    rows = (len(idx) + C - 1) // C
    sheet = Image.new('RGB', (C * cw, rows * chh), (40, 40, 48))
    d = ImageDraw.Draw(sheet)
    for k, i in enumerate(idx):
        t = texs[i]
        x, y = (k % C) * cw + 4, (k // C) * chh + 2
        # checkerboard for transparency
        bg = Image.new('RGB', (S, S), (90, 90, 90))
        bd = ImageDraw.Draw(bg)
        for yy in range(0, S, 8):
            for xx in range(0, S, 8):
                if (xx + yy) // 8 % 2:
                    bd.rectangle((xx, yy, xx + 7, yy + 7), fill=(130, 130, 130))
        im = tex_image(prgba, t, S)
        bg.paste(im, (0, 0), im)
        sheet.paste(bg, (x, y))
        w, h = t['u1'] - t['u0'] + 1, t['v1'] - t['v0'] + 1
        tag = 'F' if t['flags'] & 1 else ''
        if up[i] or side[i] or down[i]:
            tag += 'O' if up[i] >= max(side[i], down[i]) else ('S' if side[i] >= down[i] else 'U')
        d.text((x, y + S + 1), f'{i}', fill=(255, 255, 120))
        d.text((x + S // 2 + 4, y + S + 1), tag, fill=(140, 220, 255))
        d.text((x, y + S + 14), f'{w}x{h}', fill=(170, 170, 170))
    sheet.save(a.out)
    print(f'{a.out}: {len(idx)} textures, {sheet.size[0]}x{sheet.size[1]}')


if __name__ == '__main__':
    main()