"""Check first-person weapon registration, including every firing pose."""
from pathlib import Path
import struct
from PIL import Image, ImageChops, ImageDraw

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / 'build/scratch'


def check_assets():
    OUT.mkdir(parents=True, exist_ok=True)
    sheet = Image.new('RGB', (5*288, 2*260), (40,44,50))
    draw = ImageDraw.Draw(sheet)
    errors = []
    for row, name in enumerate(('pistol','shotgun')):
        raw = (ROOT/f'assets/hell/{name}.bin').read_bytes()
        magic,w,h,count,size = struct.unpack_from('<5I',raw)
        assert magic == 0x314d4e41 and (w,h) == (320,160)
        assert size == w*h*count*4 == len(raw)-20
        for frame in range(count):
            start = 20+frame*w*h*4
            image = Image.frombytes('RGBA',(w,h),raw[start:start+w*h*4],'raw','BGRA')
            alpha = image.getchannel('A')
            bounds = alpha.getbbox()
            if bounds[0] == 0 or bounds[2] == w:
                errors.append(f'{name} frame {frame}: weapon touches horizontal canvas edge')
            if bounds[3] != h:
                errors.append(f'{name} frame {frame}: hands end {h-bounds[3]} px above panel')
            if frame == 0:
                # Upper tip of resting barrel must register at screen centre.
                tip = alpha.crop((0,bounds[1],w,bounds[1]+4)).getbbox()
                centre = (tip[0]+tip[2])/2
                if abs(centre-w/2) > 1:
                    errors.append(f'{name}: resting barrel x={centre:g}, expected {w/2:g}')
            # Crop equal empty margins for the compact contact sheet.
            preview = image.resize((480,240),Image.Resampling.NEAREST).crop((96,0,384,240))
            sheet.paste(preview,(frame*288,row*260),preview)
            draw.line((frame*288+144,row*260,frame*288+144,row*260+239), fill=(170,60,60))
            draw.line((frame*288,row*260+240,frame*288+287,row*260+240), fill=(200,180,110))
            draw.text((frame*288+4,row*260+245),f'{name} {frame}',fill='white')
    sheet.save(OUT/'hell-weapon-alignment.png')
    assert not errors, '\n'.join(errors)


def check_native():
    """Compare full ASM scene captures with fixed-point, nearest-neighbour overlays."""
    sheet = Image.new('RGB', (5*320, 2*240))
    draw = ImageDraw.Draw(sheet)
    index = 0
    for row, name in enumerate(('pistol','shotgun')):
        raw = (ROOT/f'assets/hell/{name}.bin').read_bytes()
        _,w,h,count,_ = struct.unpack_from('<5I',raw)
        base = Image.frombytes('RGBA',(640,480),
                              (OUT/f'hell-weapon-{name}-base.bgra').read_bytes(),
                              'raw','BGRA').convert('RGB')
        for frame in range(count):
            start = 20+frame*w*h*4
            source = Image.frombytes('RGBA',(w,h),raw[start:start+w*h*4],'raw','BGRA')
            expected = base.copy()
            pixels = expected.load()
            sprite = source.load()
            # Match DrawAnimSprite3D's truncated 16.16 steps, not Pillow's sampling.
            xstep,ystep = (w<<16)//480,(h<<16)//240
            for y in range(240):
                for x in range(480):
                    color = sprite[(x*xstep)>>16,(y*ystep)>>16]
                    if color[3]:
                        assert color[3] == 255
                        pixels[80+x,176+y] = color[:3]
            # Crosshair is drawn over the weapon, as in the native renderer.
            for bounds in ((318,234,322,248),(312,240,328,242)):
                expected.paste(base.crop(bounds),bounds[:2])
            actual = Image.frombytes('RGBA',(640,480),
                                    (OUT/f'hell-weapon-{index}.bgra').read_bytes(),
                                    'raw','BGRA').convert('RGB')
            assert ImageChops.difference(actual,expected).getbbox() is None, \
                f'{name} frame {frame}: ASM placement/pose differs from asset contract'
            assert ImageChops.difference(actual,base).crop((80,415,560,416)).getbbox(), \
                f'{name} frame {frame}: native hands do not meet HUD at y=416'
            sheet.paste(actual.resize((320,240),Image.Resampling.NEAREST),(frame*320,row*240))
            draw.text((frame*320+4,row*240+4),f'{name} {frame}',fill='white')
            index += 1
    sheet.save(OUT/'hell-weapon-native.png')


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--native',action='store_true',help='also validate captures from check_hell_fps.py')
    options = parser.parse_args()
    check_assets()
    if options.native:
        check_native()
    print('PASS: centred resting barrels; all weapon poses meet the status bar')