import sys
import math
import random
from omg import WAD, MapEditor # omgifol
from PIL import Image, ImageDraw

SIDE_MARGIN = 40
ICON_SIZE = 16
TILE_SIZE = ICON_SIZE + 2
TILE_ROW = 50

g_index = 0
g_type2def = dict()
g_name2def = dict()


class S:
    pass


class ThingDef:
    def __init__(self, name, type_id, icon):
        self.name = name
        self.type_id = type_id
        self.icon = icon

        global g_index
        g_index += 1
        self.index = g_index

        global g_type2def, g_name2def
        g_type2def[type_id] = self
        g_name2def[name] = self


g_things = [
    ThingDef("zombie",              3004,   (1,0)),
    ThingDef("sergeant",            9,      (2,0)),
    ThingDef("commando",            65,     (3,0)),
    ThingDef("imp",                 3001,   (4,0)),
    ThingDef("demon",               3002,   (5,0)),
    ThingDef("spectre",             58,     (6,0)),
    ThingDef("lost soul",           3006,   (7,0)),
    ThingDef("cacodemon",           3005,   (8,0)),
    ThingDef("hell knight",         69,     (9,0)),
    ThingDef("baron of hell",       3003,   (10,0)),
    ThingDef("revenant",            66,     (11,0)),
    ThingDef("mancubus",            67,     (12,0)),
    ThingDef("arachnotron",         68,     (13,0)),
    ThingDef("pain elemental",      71,     (14,0)),
    ThingDef("archvile",            64,     (15,0)),
    ThingDef("cyberdemon",          16,     (16,0)),
    ThingDef("spider mastermind",   7,      (17,0)),

    ThingDef("ss guy",              84,     (18,0)),
    ThingDef("commander keen",      72,     (19,0)),

    # ThingDef("spawn target", 87, (0,0)),
    # ThingDef("spawn shooter", 89, (0,0)),
    # ThingDef("romero head", 88, (0,0)),

    # ThingDef("sunder hive mother", 72, (0,0)),
    # ThingDef("ancient aliens stealthy", 23, (0,0)),
    # ThingDef("ancient aliens guardian", 72, (0,0)),
    # ThingDef("ancient aliens guardian (2)", 12, (0,0)),
]


# Construct icons dict
g_icon_image = Image.open("icons.png")
g_icons = dict()
for y in range(g_icon_image.height // ICON_SIZE):
    for x in range(g_icon_image.width // ICON_SIZE):
        x0, y0 = x * ICON_SIZE, y * ICON_SIZE
        x1, y1 = x0 + ICON_SIZE, y0 + ICON_SIZE
        crop = g_icon_image.crop([x0,y0,x1,y1])
        g_icons[(x,y)] = crop


# Get list of all things in a specific map in g_type2def
def mapGetThings(wad, map_id, skill='hard'):
    edit = MapEditor(wad.maps[map_id])

    result = []

    for thing in edit.things:
        if thing.multiplayer: continue
        # I'm sure there's a better way to do this
        if skill == 'easy' and not thing.easy: continue
        if skill == 'medium' and not thing.medium: continue
        if skill == 'hard' and not thing.hard: continue
        type_id = thing.type
        if type_id not in g_type2def: continue
        thing = g_type2def[type_id]
        result.append(thing)

    result = sorted(result, key=lambda t: t.index)
    return result



# Draw the thing icons from a particular map_info
def processMap(wad, im, map_info, y_off):

    draw = ImageDraw.Draw(im)
    y_text = y_off
    for line in map_info.text:
        draw.text((4, y_text), line, fill='black')
        y_text += 12

    for i, thing in enumerate(map_info.things):
        x = SIDE_MARGIN + (i // map_info.rows) * TILE_SIZE
        y = y_off + (i % map_info.rows) * TILE_SIZE
        x0, y0 = x+1, y+1
        x1, y1 = x0 + ICON_SIZE, y0 + ICON_SIZE
        im.paste(g_icons[thing.icon], [x0,y0,x1,y1])



def main():
    # Args
    wad = WAD(sys.argv[1])
    skill_arg = 'hard' # 'easy', 'medium' or 'hard'
    if len(sys.argv) >= 3: skill_arg = sys.argv[2]
    if skill_arg == 'all':
        skills = ['easy', 'medium', 'hard']
    else:
        skills = [skill_arg]

    # Get list of sub-plots to draw (typically one per map)
    map_ids = []
    for map_id in wad.maps:
        map_ids.append(map_id)
    map_ids = sorted(map_ids)
    subs = []
    for map_id in map_ids:
        for skill in skills:
            subs.append((map_id, skill))

    # Get the necessarily map_infos
    map_infos = []
    im_height = 0
    for map_id, skill in subs:
        s = S()
        s.text = [map_id]
        if skill_arg == 'all':
            s.text.append(skill)
        print(f"scanning {s.text}")
        s.things = mapGetThings(wad, map_id, skill)
        s.rows = math.ceil(len(s.things) / TILE_ROW)
        if len(s.things) == 0: s.rows += 1
        map_infos.append(s)
        im_height += TILE_SIZE * (s.rows + 1)

    # Create image
    y_off = TILE_SIZE // 2 # Start half a tile lower
    im_width = SIDE_MARGIN + TILE_SIZE * TILE_ROW
    im = Image.new('RGB', (im_width, im_height), 'grey')
    for map_info in map_infos:
        print(f"drawing {map_info.text} ({len(map_info.things)} things)")
        processMap(wad, im, map_info, y_off)
        y_off += (map_info.rows + 1) * TILE_SIZE
    im.save("out.png")


if __name__ == "__main__":
    main()