"""Regenerate src/RR/Patching/LogoArt.h - the Campus logo the patch swaps in for Rec Room's.

    python tools/make_logo_art.py <logo.png> [--game D:\\Puck\\Campus2025] [--check]

What it does (see src/RR/Patching/LogoSwap.h for why):
  1. Opens the game's UI bundle (91aca73acb86d6607f0efa1f803348be.bundle) with UnityPy and reads the raw
     bytes of the two logo textures, RecRoomLogo_Sprite_Shadow and RecRoomLogo_Sprite (1024x1024, DXT5, one mip),
     hashing each (FNV-1a over 64-bit words - the same function as LogoSwap::Fnv). Those hashes are how the
     patch recognises the originals.
  2. Encodes <logo.png> (must be 1024x1024, RGBA - transparent background) as BC3 with etcpak, in the memory
     order Unity uses (bottom row first, so the image is flipped before encoding).
  3. Writes LogoArt.h. --check only verifies the existing header already matches this png + bundle.

Needs: pip install UnityPy etcpak texture2ddecoder pillow numpy
"""
import argparse, os, struct, sys
import numpy as np
from PIL import Image, ImageOps
import UnityPy, etcpak, texture2ddecoder

BUNDLE = r'StreamingAssets\aa\StandaloneWindows64\91aca73acb86d6607f0efa1f803348be.bundle'
NAMES = ('RecRoomLogo_Sprite_Shadow', 'RecRoomLogo_Sprite')
HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.normpath(os.path.join(HERE, '..', 'src', 'RR', 'Patching', 'LogoArt.h'))


def fnv64(b: bytes) -> int:
    h = 0xcbf29ce484222325; M = (1 << 64) - 1
    n = len(b) // 8
    for (w,) in struct.iter_unpack('<Q', b[:n * 8]):
        h = ((h ^ w) * 0x100000001b3) & M
    for x in b[n * 8:]:
        h = ((h ^ x) * 0x100000001b3) & M
    return h


def originals(game):
    env = UnityPy.load(os.path.join(game, 'Recroom_Release_Data', BUNDLE))
    found = {}
    for o in env.objects:
        if o.type.name != 'Texture2D':
            continue
        d = o.read()
        if d.m_Name in NAMES:
            raw = bytes(d.get_image_data())
            assert (d.m_Width, d.m_Height, d.m_MipCount, len(raw)) == (1024, 1024, 1, 1048576), (d.m_Name, d.m_Width, d.m_Height, d.m_MipCount, len(raw))
            found[d.m_Name] = raw
    assert set(found) == set(NAMES), found.keys()
    return found


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('png'); ap.add_argument('--game', default=r'D:\Puck\Campus2025'); ap.add_argument('--check', action='store_true')
    a = ap.parse_args()

    orig = originals(a.game)
    # Unity keeps these textures bottom-row-first: confirm against UnityPy's own (flipped) decode, so a future
    # bundle with a different layout fails loudly instead of producing an upside-down logo.
    d = Image.fromarray(np.frombuffer(texture2ddecoder.decode_bc3(orig[NAMES[0]], 1024, 1024), np.uint8).reshape(1024, 1024, 4)[:, :, [2, 1, 0, 3]], 'RGBA')
    env = UnityPy.load(os.path.join(a.game, 'Recroom_Release_Data', BUNDLE))
    ref = next(o.read().image for o in env.objects if o.type.name == 'Texture2D' and o.read().m_Name == NAMES[0])
    assert np.abs(np.array(ImageOps.flip(d)).astype(int) - np.array(ref.convert('RGBA')).astype(int)).mean() < 1, 'texture is not bottom-row-first any more'

    art = Image.open(a.png).convert('RGBA')
    assert art.size == (1024, 1024), f'logo must be 1024x1024, got {art.size}'
    blob = bytes(etcpak.compress_bc3(ImageOps.flip(art).tobytes(), 1024, 1024))
    assert len(blob) == 1048576

    back = np.frombuffer(texture2ddecoder.decode_bc3(blob, 1024, 1024), np.uint8).reshape(1024, 1024, 4)[:, :, [2, 1, 0, 3]][::-1].astype(float)
    mse = ((np.array(art).astype(float) - back) ** 2).mean()
    print(f'encoded: PSNR {10 * np.log10(255 ** 2 / mse):.1f} dB')

    lines = ['#pragma once',
             '// GENERATED by tools/make_logo_art.py - do not edit by hand.',
             "// The Campus logo as one 1024x1024 BC3 (DXT5) image, in the memory order Unity uses for Rec Room's own",
             '// logo textures (bottom row first), plus the FNV-1a hash of each ORIGINAL texture the patch swaps.',
             '#include <cstdint>', 'namespace LogoArt {',
             '    inline constexpr int kSize = 1024;', '    inline constexpr uint32_t kBytes = 1024u * 1024u;']
    for n in NAMES:
        lines.append(f'    inline constexpr uint64_t k{n.replace("_", "")}Hash = 0x{fnv64(orig[n]):016x}ULL;   // original {n}')
    lines.append('    alignas(16) inline const unsigned char kBc3[1048576] = {')
    for i in range(0, len(blob), 32):
        lines.append('        ' + ','.join(str(x) for x in blob[i:i + 32]) + ',')
    lines += ['    };', '}']
    text = '\r\n'.join(lines) + '\r\n'
    if a.check:
        same = os.path.exists(OUT) and open(OUT, 'rb').read() == text.encode()
        print('LogoArt.h is', 'up to date' if same else 'DIFFERENT from this png/bundle'); sys.exit(0 if same else 1)
    open(OUT, 'wb').write(text.encode())
    print('wrote', OUT)


main()