"""Native MASM kiosk/lantern regression and framebuffer checks; no game simulation."""
import argparse
from pathlib import Path
import re
import struct
import subprocess
from PIL import Image, ImageChops
from check_card_exit_build import Coff

ROOT = Path(__file__).resolve().parents[2]
OUT = ROOT / "build/terminals"


def read(path):
    return (ROOT / path).read_text(encoding="cp1251")


def proc(text, name):
    return re.search(rf"^{name} PROC\b.*?^{name} ENDP", text, re.M | re.S).group()


def signature(coff, name, tail_branches=False):
    result = []
    for code, refs, target in coff.signature(name, tail_branches=tail_branches):
        if code[:1] == b"\xe8" and len(refs) == 1 and refs[0][1][1] == 20:
            target = refs[0][1][0]
            refs = ()
        result.append((code, refs, target))
    return result


def prepare():
    OUT.mkdir(exist_ok=True)
    terminal = read("src_graphics/terminal.inc")
    city = read("src_graphics/city_bins.inc")
    city_names = re.findall(r"^(\w+) PROC", city, re.M)[:7]
    code = terminal + "\n" + "\n".join(proc(city, name) for name in city_names)
    code += "\n" + proc(read("src_graphics/world_dressing.inc"), "RenderHallDressing3D")
    # The shared surface dispatcher also knows the separately packed upper lounge.
    code += "\n" + proc(read("src_graphics/hall_expansion.inc"), "ResolveUpperMaterial3D")
    code += "\n" + proc(read("src_graphics/fortune_room.inc"), "ResolveFortuneMaterial3D")
    trigger = proc(read("casino3d.inc"), "Trigger3DInteraction")
    code += "\n" + trigger
    if "_InitHellTower3D@0" not in Coff(ROOT / "build/actions/actions-test.obj").functions:
        # This unrelated initializer is newer than the reusable helper object.
        # The kiosk fixture must never enter it; fail immediately if it does.
        code += "\nInitHellTower3D PROC\ninvoke ExitProcess,98\nret\nInitHellTower3D ENDP\n"
    own_functions = set(re.findall(r"^(\w+) PROC", code, re.M))
    own_data = set(re.findall(r"^(\w+) (?:dd|db|AnimPack3D)", read("src_graphics/terminal_data.inc"), re.M))
    own_data.update(("cityRayX3D", "cityRayY3D"))
    own_data.update(("fortuneMaterialArt3D", "fortuneArt3D"))
    own_data.update(("upperMaterialArt3D","upperFeltPixels3D","upperKeyPixels3D",
                     "upperGlowPixels3D","upperPortalArt3D","upperPortal3D",
                     "portalRendering3D","portalClipDepth3D","basementDoorArt3D",
                     "basementWheelTime3D"))
    coff = Coff(ROOT / "build/actions/actions-test.obj")
    data = bytearray(coff.data)
    at = struct.unpack_from("<I", data, 8)[0]
    count = struct.unpack_from("<H", data, 2)[0]
    data_sections = set()
    for section in range(1, count + 1):
        offset = 20 + (section - 1) * 40
        name = bytes(data[offset:offset+8]).rstrip(b"\0")
        if name in (b".data", b".bss", b".const"):
            data_sections.add(section)
    externs = []
    protos = []
    for index, name in coff.symbols.items():
        pos = at + index * 18
        _, section, kind, storage, _ = struct.unpack_from("<IhHBB", data, pos+8)
        if section in data_sections and re.fullmatch(r"[A-Za-z_]\w*", name) and storage in (2, 3):
            if name in own_data:
                data[pos+16] = 3
            else:
                data[pos+16] = 2
                typ = "AnimPack3D" if name in ("cityDecorPack3D","hallProps3D") else "DWORD"
                externs.append(f"EXTERN {name}:{typ}")
        match = re.fullmatch(r"_(\w+)@(\d+)", name)
        if match and kind == 0x20 and section > 0:
            function, size = match[1], int(match[2])
            if function in own_functions:
                data[pos+16] = 3
            else:
                protos.append(function + " PROTO" + (" " + ",".join([":DWORD"] * (size//4)) if size else ""))
    # Visibility-only helper copy. Never patch instructions or the shipping object.
    allowed = {at + index * 18 + 16 for index in coff.symbols}
    assert all(i in allowed for i, (a, b) in enumerate(zip(data, coff.data)) if a != b)
    (OUT / "helpers.obj").write_bytes(data)
    header = """.386
.model flat
option casemap:none
FP_1 equ 65536
FP_HALF equ 32768
SCREEN_W equ 640
SCREEN_H equ 480
MODE_WALK equ 0
MODE_ATM equ 5
WORLD_WALL_ART_COUNT equ 6
AnimPack3D STRUCT
 magic dd ?
 imageWidth dd ?
 height dd ?
 frames dd ?
 bytes dd ?
 pixels dd ?
AnimPack3D ENDS
CityProp3D STRUCT
 posX dd ?
 posY dd ?
 worldW dd ?
 worldH dd ?
 frame dd ?
 asset dd ?
 groundZ dd 0
CityProp3D ENDS
SpriteObject STRUCT
 posX dd ?
 posY dd ?
 baseTexId dd ?
 frameCount dd ?
 curFrame dd ?
 animTimer dd ?
 animSpeed dd ?
 interactId dd ?
 promptPtr dd ?
 distSq dd ?
SpriteObject ENDS
NpcWaypoint STRUCT
 posX dd ?
 posY dd ?
 neighbor0 dd ?
 neighbor1 dd ?
 neighbor2 dd ?
 neighbor3 dd ?
 zoneType dd ?
NpcWaypoint ENDS
"""
    for name in city_names:
        line = proc(city, name).splitlines()[0]
        size = line.count(":DWORD")
        protos.append(name + " PROTO" + (" " + ",".join([":DWORD"] * size) if size else ""))
    text = header + "\n".join(externs) + "\noption language:stdcall\nExitProcess PROTO :DWORD\n" + "\n".join(protos)
    constants = {}
    for path in [ROOT / "casino3d.inc", *sorted((ROOT / "src_graphics").glob("*_data.inc"))]:
        constants.update(re.findall(r"^(\w+)\s+equ\s+([^;\n]+)", path.read_text(encoding="cp1251"), re.M))
    for name in sorted(set(re.findall(r"\b(?:MODE_|INTERACT_)\w+", trigger)) - {"MODE_WALK", "MODE_ATM"}):
        text += f"\n{name} equ {constants[name]}"
    text += "\nResolveUpperMaterial3D PROTO :DWORD\ninclude src_graphics/terminal_data.inc\n.data\nupperMaterialArt3D AnimPack3D <>\nupperPortalArt3D AnimPack3D <>\nupperPortal3D CityProp3D <>\nbasementDoorArt3D AnimPack3D <>\nbasementWheelTime3D dd 0\nportalRendering3D dd 0\n.data?\nupperFeltPixels3D dd 1024 dup(?)\nupperKeyPixels3D dd 1024 dup(?)\nupperGlowPixels3D dd 1024 dup(?)\nportalClipDepth3D dd 640 dup(?)\ncityRayX3D dd 640 dup(?)\ncityRayY3D dd 640 dup(?)\n.code\n"
    text += "\nResolveFortuneMaterial3D PROTO :DWORD\n.data\nfortuneMaterialArt3D AnimPack3D <>\nfortuneArt3D AnimPack3D <>\n.code\n"
    text += code + "\ninclude tools/tests/terminal_test.inc\nend TerminalTestStart\n"
    (OUT / "module.asm").write_text(text, encoding="ascii")


def verify():
    for path, expected in (("terminals-v2.bin", (256, 512, 4)), ("terminal-screen-v1.bin", (640, 480, 1))):
        data = (ROOT / "assets/animations/hall" / path).read_bytes()
        magic, w, h, frames, size = struct.unpack_from("<4s4I", data)
        assert magic == b"ANM1" and (w, h, frames) == expected
        assert size == w*h*frames*4 == len(data)-20
        assert set(data[23::4]) == {255}
    data = (ROOT / "assets/animations/hall/terminals-v2.bin").read_bytes()
    old = (ROOT / "assets/animations/hall/terminals-v1.bin").read_bytes()
    assert data[20:20+3*256*512*4] == old[20:]
    body = Image.open(ROOT / "assets/sprites/terminals/body-side-rear-v2.png").convert("RGBA").resize((256,512),Image.Resampling.LANCZOS)
    assert data[20+3*256*512*4:] == body.tobytes("raw","BGRA")
    cabinet = proc(read("src_graphics/terminal.inc"),"DrawTerminalCabinet3D")
    assert "FP_1*8/5,121" in cabinet and "FP_1*8/5,86" not in cabinet
    images = {}
    for name in ("cash", "ach-0", "ach-1", "ach-2", "ach-locked", "cabinet-front", "cabinet-side", "cabinet-rear", "trophy-front", "trophy-side", "trophy-rear", "wall-placement", "lantern-front", "lantern-side"):
        image = Image.open(OUT / (name + ".bmp")).convert("RGB")
        assert image.size == (640, 480)
        assert len(image.getcolors(640*480)) > 1000, name
        image.save(OUT / (name + ".png"))
        images[name] = image
    for a, b in (("cash", "ach-0"), ("ach-0", "ach-1"), ("ach-1", "ach-2"),
                 ("cabinet-front", "cabinet-side"), ("lantern-front", "lantern-side")):
        assert ImageChops.difference(images[a], images[b]).getbbox(), (a, b)
    print("PASS: wall placement, real front interaction, old footprints clear, dedicated enclosure/bar independence, native save/load/input, fourteen captures")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--parity", action="store_true")
    parser.add_argument("--prove-regression", action="store_true")
    args = parser.parse_args()
    if args.parity:
        a, b = Coff(OUT / "module.obj"), Coff(ROOT / "build/graphics/poker.obj")
        names = re.findall(r"^(\w+) PROC[^\n]*", read("src_graphics/terminal.inc"), re.M)
        for name in names + ["Trigger3DInteraction", "RenderHallDressing3D", "ResolveCitySurface3D", "DrawCityVertical3D", "DrawCityHorizontal3D", "DrawCityBoxSides3D"]:
            symbol = next(key for key in a.functions if key.startswith("_" + name + "@"))
            assert signature(a,symbol) == signature(b,symbol), name
        verify()
        print("PASS: shipping terminal/lantern renderer matches fresh native tests")
        return
    prepare()
    stem = "module"
    if args.prove_regression:
        stem = "old-body"
        source = (OUT / "module.asm").read_text(encoding="ascii")
        before = "invoke DrawCityBoxSides3D,x0,x1,y0,y1,0,FP_1*8/5,121"
        assert source.count(before) == 1
        source = source.replace(before, before[:-3] + "86")
        (OUT / "old-body.asm").write_text(source, encoding="ascii")
    masm = Path("F:/MASM32/masm32")
    subprocess.run([str(masm / "bin/ml.exe"), "/c", "/coff", "/Cp", "/nologo", "/I.",
                    f"/Fobuild/terminals/{stem}.obj", f"build/terminals/{stem}.asm"], cwd=ROOT, check=True)
    sdk = sorted(Path("C:/Program Files (x86)/Windows Kits/10/Lib").glob("*/um/x86"))[-1]
    subprocess.run([str(masm / "bin/link.exe"), "/nologo", "/SUBSYSTEM:CONSOLE", "/ENTRY:TerminalTestStart@0", f"/LIBPATH:{sdk}",
                    f"/LIBPATH:{masm / 'lib'}", f"/OUT:build/terminals/{stem}.exe", f"build/terminals/{stem}.obj",
                    "build/terminals/helpers.obj", "kernel32.lib", "user32.lib", "gdi32.lib", "winmm.lib", "ws2_32.lib"], cwd=ROOT, check=True)
    result = subprocess.run([str(OUT / (stem + ".exe"))], cwd=ROOT, check=False)
    if args.prove_regression:
        assert result.returncode == 6, result.returncode
        print("PASS: restoring the old bar material is caught by the native texture-dependency regression")
        return
    result.check_returncode()
    verify()


if __name__ == "__main__":
    main()