"""Import pinned Freedoom assets into native ANM1 packs (offline build tool)."""
import argparse
import hashlib
import io
import json
import struct
import urllib.request
import wave
import zipfile
from pathlib import Path

from PIL import Image, ImageDraw

ROOT = Path(__file__).resolve().parents[2]
SHA256 = "3f9b264f3e3ce503b4fb7f6bdcb1f419d93c7b546f4df3e874dd878db9688f59"
URL = "https://github.com/freedoom/freedoom/releases/download/v0.13.0/freedoom-0.13.0.zip"


class Wad:
    def __init__(self, data):
        magic, count, offset = struct.unpack_from('<4sII', data)
        assert magic == b'IWAD'
        self.lumps = {}
        for i in range(count):
            pos, size, name = struct.unpack_from('<II8s', data, offset + i * 16)
            assert pos + size <= len(data)
            self.lumps[name.rstrip(b'\0').decode()] = data[pos:pos + size]
        self.palette = self.lumps['PLAYPAL'][:768]

    def patch(self, name):
        data = self.lumps[name]
        w, h, left, top = struct.unpack_from('<HHhh', data)
        image = Image.new('RGBA', (w, h))
        pixels = image.load()
        for x in range(w):
            pos = struct.unpack_from('<I', data, 8 + x * 4)[0]
            previous = -1
            while data[pos] != 255:
                y, n = data[pos], data[pos + 1]
                if y <= previous:
                    y += previous
                previous = y
                for dy, index in enumerate(data[pos + 3:pos + 3 + n]):
                    if y + dy < h:
                        pixels[x, y + dy] = (*self.palette[index * 3:index * 3 + 3], 255)
                pos += n + 4
        return image, left, top

    def flat(self, name):
        image = Image.frombytes('P', (64, 64), self.lumps[name])
        image.putpalette(self.palette)
        return image.convert('RGBA')

    def texture(self, name):
        pnames = self.lumps['PNAMES']
        patches = [pnames[4+i*8:12+i*8].rstrip(b'\0').decode()
                   for i in range(struct.unpack_from('<I', pnames)[0])]
        for table in ('TEXTURE1', 'TEXTURE2'):
            data = self.lumps.get(table)
            if not data:
                continue
            for i in range(struct.unpack_from('<I', data)[0]):
                pos = struct.unpack_from('<I', data, 4+i*4)[0]
                texname = data[pos:pos+8].rstrip(b'\0').decode()
                if texname != name:
                    continue
                w, h = struct.unpack_from('<HH', data, pos+12)
                image = Image.new('RGBA', (w, h), (0, 0, 0, 255))
                count = struct.unpack_from('<H', data, pos+20)[0]
                for j in range(count):
                    x, y, patch = struct.unpack_from('<hhH', data, pos+22+j*10)
                    sprite, _, _ = self.patch(patches[patch])
                    image.paste(sprite, (x, y), sprite)
                return image
        raise KeyError(name)


def pack(path, images):
    w, h = images[0].size
    assert all(im.size == (w, h) for im in images)
    payload = b''.join(im.tobytes('raw', 'BGRA') for im in images)
    path.write_bytes(struct.pack('<5I', 0x314d4e41, w, h, len(images), len(payload)) + payload)


def aligned(wad, names, size=(128, 128)):
    frames = [wad.patch(n) for n in names]
    # Shared Doom origin, shared scale: no per-frame trim/resize jitter.
    x0 = min(-left for _, left, _ in frames)
    x1 = max(im.width-left for im, left, _ in frames)
    y0 = min(-top for _, _, top in frames)
    y1 = max(im.height-top for im, _, top in frames)
    scale = min(size[0] / (x1-x0), size[1] / (y1-y0))
    result = []
    for im, left, top in frames:
        canvas = Image.new('RGBA', size)
        im = im.resize((max(1, round(im.width*scale)), max(1, round(im.height*scale))), Image.Resampling.NEAREST)
        x = round((-left-x0)*scale + (size[0]-(x1-x0)*scale)/2)
        y = round((-top-y0)*scale + size[1]-(y1-y0)*scale)
        canvas.paste(im, (x, y))
        result.append(canvas)
    return result


def weapon_frames(wad, names):
    """Register by resting barrel, not animation bounds; keep wrists at HUD edge."""
    frames = aligned(wad, names, (192,160))
    alpha = frames[0].getchannel('A')
    top = alpha.getbbox()[1]
    tip = alpha.crop((0,top,192,top+4)).getbbox()
    shift_x = round(160-(tip[0]+tip[2])/2)
    result = []
    for image in frames:
        bottom = image.getchannel('A').getbbox()[3]
        canvas = Image.new('RGBA',(320,160))
        # Extra horizontal canvas retains the shotgun's leftward pump motion.
        # Same scale and horizontal translation for every frame of each weapon.
        canvas.paste(image,(shift_x,160-bottom))
        result.append(canvas)
    return result


def interface_assets(wad, out, manifest):
    """Doom-layout status bar, native-size glyphs and common-origin face frames."""
    pack(out/'hud.bin', [wad.patch('STBAR')[0]])
    names = [f'STTNUM{i}' for i in range(10)] + ['STTPRCNT']
    digits = []
    for name in names:
        image = wad.patch(name)[0]
        canvas = Image.new('RGBA', (16,16))
        canvas.paste(image, (16-image.width,0))
        digits.append(canvas)
    pack(out/'digits.bin', digits)
    names_face = [name for health in range(5)
                  for name in [f'STFST{health}0', f'STFST{health}1',
                               f'STFST{health}2', f'STFOUCH{health}']] + ['STFDEAD0']
    faces = []
    for name in names_face:
        image = wad.patch(name)[0]
        canvas = Image.new('RGBA', (24,32))
        canvas.paste(image, ((24-image.width)//2,32-image.height))
        faces.append(canvas)
    pack(out/'faces.bin', faces)
    impacts = ['PUFFA0','PUFFB0','PUFFC0','PUFFD0']
    pack(out/'impacts.bin', aligned(wad, impacts, (32,32)))
    manifest['packs'].update(hud=['STBAR'], digits=names, faces=names_face, impacts=impacts)

    # Authored landscape, entirely composed from Freedoom textures and sprites.
    # Opaque power-of-two 512x256 BGRA surface; wall renderer wraps with bit masks.
    image = Image.new('RGBA',(512,192),(20,9,8,255))
    sky = wad.texture('SKY3').resize((488,166),Image.Resampling.NEAREST)
    image.paste(sky,(12,12))
    draw = ImageDraw.Draw(image)
    draw.polygon([(12,136),(65,98),(96,126),(139,86),(174,125),(222,104),
                  (273,142),(321,90),(376,123),(423,94),(500,130),(500,180),(12,180)],
                 fill=(34,20,22,255))
    lava = wad.flat('LAVA1').resize((488,48),Image.Resampling.NEAREST)
    image.paste(lava,(12,132))
    for x,h in [(42,70),(92,52),(362,62),(428,84)]:
        stone = wad.texture('GSTONE1').resize((28,h),Image.Resampling.NEAREST)
        image.paste(stone,(x,164-h))
        draw.rectangle((x-3,160-h,x+30,165-h),fill=(107,72,45,255))
    for name,x,y,w,h in [('BOSSA1',224,48,70,116),('TROOA1',145,104,34,58),
                          ('SARGA1',330,112,42,52)]:
        actor = wad.patch(name)[0]
        actor.thumbnail((w,h),Image.Resampling.NEAREST)
        image.alpha_composite(actor,(x,y))
    # Layered gilt moulding and dark inner bevel: a painting, not an open portal.
    for inset,color in [(0,(42,25,12)),(2,(199,147,64)),(4,(104,61,22)),
                        (6,(236,191,99)),(8,(116,74,28)),(10,(26,13,8))]:
        draw.rectangle((inset,inset,511-inset,191-inset),outline=(*color,255),width=2)
    for x in range(18,495,16):
        draw.line((x,4,x+3,7),fill=(255,216,126,255),width=1)
        draw.line((x,184,x+3,187),fill=(255,216,126,255),width=1)
    pack(out/'painting.bin',[image.resize((512,256),Image.Resampling.NEAREST)])
    image.save(out/'painting-preview.png')
    manifest['packs']['painting'] = ['SKY3','LAVA1','GSTONE1','BOSSA1','TROOA1','SARGA1']


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('archive', type=Path)
    args = parser.parse_args()
    data = args.archive.read_bytes()
    assert hashlib.sha256(data).hexdigest() == SHA256, 'Freedoom archive hash mismatch'
    archive = zipfile.ZipFile(io.BytesIO(data))
    wad = Wad(archive.read(next(n for n in archive.namelist() if n.endswith('/freedoom2.wad'))))
    out = ROOT / 'assets/hell'
    out.mkdir(parents=True, exist_ok=True)
    manifest = {'version': '0.13.0', 'url': URL, 'archive_sha256': SHA256, 'packs': {}}
    monsters = {
        'soldier': ['POSSA1','POSSB1','POSSC1','POSSD1','POSSE1','POSSF1','POSSG1','POSSH0','POSSI0','POSSJ0','POSSK0','POSSL0'],
        'demon': ['SARGA1','SARGB1','SARGC1','SARGD1','SARGE1','SARGF1','SARGH1','SARGI0','SARGJ0','SARGK0','SARGL0','SARGN0'],
        'imp': ['TROOA1','TROOB1','TROOC1','TROOD1','TROOE1','TROOF1','TROOH1','TROOI0','TROOJ0','TROOK0','TROOL0','TROOM0'],
        'boss': ['BOSSA1','BOSSB1','BOSSC1','BOSSD1','BOSSE1','BOSSF1','BOSSH1','BOSSI0','BOSSJ0','BOSSK0','BOSSL0','BOSSO0'],
        'pistol': ['PISGA0','PISGB0','PISGC0','PISGD0','PISGE0'],
        'shotgun': ['SHTGA0','SHTGB0','SHTGC0','SHTGD0'],
        'fire': ['BAL1A0','BAL1B0','BAL1C0','BAL1D0','BAL1E0'],
        'portal': ['TFOGA0','TFOGB0','TFOGC0','TFOGD0','TFOGE0','TFOGF0','TFOGG0','TFOGH0'],
    }
    for name, names in monsters.items():
        # Some Doom rotations are packed in mirrored pairs (A2A8 etc.).
        for i, lump in enumerate(names):
            if lump not in wad.lumps:
                choices = [n for n in wad.lumps if n.startswith(lump[:5]) and len(n) in (6,8)]
                if not choices:
                    raise KeyError(lump)
                names[i] = sorted(choices)[0]
        images = weapon_frames(wad,names) if name in ('pistol','shotgun') else aligned(wad,names)
        if name == 'portal':
            # Recolour original Freedoom teleport animation to infernal orange/red.
            for image in images:
                image.putdata([(max(r,g), min(180,g//2), b//4, a)
                               for r,g,b,a in image.getdata()])
        pack(out / (name+'.bin'), images)
        manifest['packs'][name] = names
    items = ['MEDIA0','ARM1A0','CLIPA0','SHELA0','SHOTA0','BKEYA0','BON2A0']
    pack(out/'items.bin', [aligned(wad, [n])[0] for n in items])
    manifest['packs']['items'] = items
    textures = ['MARBLE1','GSTONE1','SKIN2','BIGDOOR2','SW1GARG','SP_HOT1']
    flats = ['FLOOR7_2','F_SKY1','LAVA1','LAVA2','LAVA3','LAVA4']
    images = [wad.texture(n).resize((128,128),Image.Resampling.NEAREST) for n in textures]
    # F_SKY1 is a marker, use a real stone flat for ceiling.
    flats[1] = 'CEIL5_2'
    images += [wad.flat(n).resize((128,128),Image.Resampling.NEAREST) for n in flats]
    pack(out/'materials.bin', images)
    manifest['packs']['materials'] = textures+flats
    interface_assets(wad,out,manifest)
    for name in ['DSPISTOL','DSSHOTGN','DSSGCOCK','DSPLPAIN','DSPOPAIN','DSFIRSHT','DSPODTH1','DSITEMUP','DSDOROPN','DSTELEPT']:
        raw = wad.lumps[name]
        fmt, rate, count = struct.unpack_from('<HHI', raw)
        assert fmt == 3 and count <= len(raw)-8
        with wave.open(str(out/(name.lower()+'.wav')), 'wb') as wav:
            wav.setparams((1,1,rate,0,'NONE','not compressed'))
            wav.writeframes(raw[24:8+count-16])
    for basename in ('COPYING','CREDITS'):
        names = [n for n in archive.namelist() if Path(n).name.startswith(basename)]
        if names:
            (out/(basename+'.txt')).write_bytes(archive.read(names[0]))
        else:
            url = f'https://raw.githubusercontent.com/freedoom/freedoom/v0.13.0/{basename}.adoc'
            (out/(basename+'.txt')).write_bytes(urllib.request.urlopen(url).read())
    (out/'manifest.json').write_text(json.dumps(manifest,indent=2)+'\n',encoding='utf-8')
    sheet = Image.new('RGBA',(6*128,3*128),(30,25,25,255))
    for i, name in enumerate(monsters):
        raw = (out/(name+'.bin')).read_bytes()
        _,w,h,_,_ = struct.unpack_from('<5I',raw)
        im = Image.frombytes('RGBA',(w,h),raw[20:20+w*h*4],'raw','BGRA')
        im.thumbnail((128,128),Image.Resampling.NEAREST)
        sheet.paste(im,((i%6)*128,(i//6)*128),im)
    sheet.save(out/'preview.png')
    print(f'Imported {len(manifest["packs"])} packs to {out}')


if __name__ == '__main__':
    main()