#!/usr/bin/env python3
# The host's own Metal shaders (dm/gfx/msl/<name>.metal: clears, copies, the present, index conversions, geometry packing, ...), compiled when the host is built, so the host
# never compiles MSL source at run time. MTL4Compiler's newLibraryWithDescriptor: with a source compiles it through the MTLDevice (it calls newLibraryWithSource:), which
# the Metal performance HUD reports ("Compiled 2 shaders ... Consider adopting MTL4Compiler") and which held the submitting thread about 50 ms per library on an iPhone.
# Writes one C header (the host Makefile includes its directory; dm/gfx/nvn/host_msl.m reads it) with, per library, its name, its source (the fallback when the library
# is missing) and the metallib compiled for the SDK. Without the Metal toolchain the header has the sources only: the host then compiles them and says so in a [RED] line.
# Each run rechecks the compiler, SDK, deployment target and source contents. Matching inputs reuse the header without compiling; a toolchain installed after a source-only
# build regenerates the libraries. The header is written only when it changes, so an unchanged run does not rebuild the host.
# Usage: host_msl.py --sdk macosx|iphoneos [--min <OS version>] <header> <name.metal>...
import argparse, hashlib, os, subprocess, sys, tempfile

ap = argparse.ArgumentParser()
ap.add_argument('--sdk', required=True, choices=['macosx', 'iphoneos'])
ap.add_argument('--min', help='the lowest OS version the libraries must load on (iphoneos: the app\'s deployment target)')
ap.add_argument('out')
ap.add_argument('sources', nargs='+')
a = ap.parse_args()

def c_str(s):
    lines = s.split('\n')
    esc = lambda l: l.replace('\\', '\\\\').replace('"', '\\"')
    return '\n'.join(f'    "{esc(l)}\\n"' for l in lines[:-1]) + (f'\n    "{esc(lines[-1])}"' if lines[-1] else '') or '    ""'

ver = subprocess.run(['xcrun', '-sdk', a.sdk, 'metal', '--version'], capture_output=True, text=True)
ver = ver.stdout.splitlines()[0].strip() if ver.returncode == 0 and ver.stdout else ''
sdk = subprocess.run(['xcrun', '-sdk', a.sdk, '--show-sdk-path'], capture_output=True, text=True).stdout.strip()
sources = [(src, open(src).read()) for src in sorted(a.sources)]
# The generator itself is an input too: a changed compile flag or header format must invalidate an existing library.
inputs = hashlib.sha256(repr((a.sdk, a.min, ver, sdk, sources)).encode() + open(__file__, 'rb').read()).hexdigest()
stamp = f'// Inputs: {inputs}'
if os.path.exists(a.out):
    with open(a.out) as f:
        f.readline()
        if f.readline().rstrip('\n') == stamp: sys.exit(0)
if not ver: print(f'host_msl.py: no Metal compiler for {a.sdk} (xcodebuild -downloadComponent MetalToolchain): the host compiles its shaders at run time', file=sys.stderr)
out = ['// Generated by dm/gfx/tools/host_msl.py from dm/gfx/msl/*.metal: do not edit', stamp, f'#define HOST_MSL_SDK "{a.sdk}"', f'#define HOST_MSL_COMPILER "{ver or "none"}"']
ents = []
with tempfile.TemporaryDirectory() as td:
    for src, text in sources:
        name = os.path.splitext(os.path.basename(src))[0]
        lib = b''
        if ver:
            ml = os.path.join(td, name + '.metallib')
            cmd = ['xcrun', '-sdk', a.sdk, 'metal'] + ([f'-m{"ios" if a.sdk == "iphoneos" else "macosx"}-version-min={a.min}'] if a.min else []) + ['-o', ml, src]
            r = subprocess.run(cmd)
            if r.returncode: sys.exit(f'host_msl.py: {src} does not compile ({" ".join(cmd)})')
            lib = open(ml, 'rb').read()
        out.append(f'static const char k_host_msl_src_{name}[] =\n{c_str(text)};')
        out.append(f'static const unsigned char k_host_msl_lib_{name}[] = {{' + ','.join(str(b) for b in lib or b'\0') + '};')
        ents.append(f'    {{"{name}", k_host_msl_src_{name}, k_host_msl_lib_{name}, {len(lib)}}},')
out.append('static const struct { const char *name, *src; const unsigned char *lib; unsigned long lib_len; } k_host_msl[] = {\n' + '\n'.join(ents) + '\n};\n')
text = '\n'.join(out)
if not os.path.exists(a.out) or open(a.out).read() != text:
    os.makedirs(os.path.dirname(a.out) or '.', exist_ok=True)
    with open(a.out + '.tmp', 'w') as f: f.write(text)
    os.replace(a.out + '.tmp', a.out)
print(f'host_msl.py: {len(ents)} libraries for {a.sdk}' + (f' ({ver})' if ver else ', sources only') + f' in {a.out}', file=sys.stderr)