"""Passthrough geometry as one Metal mesh shader (<vertex key>_bm<geometry key>).

A lean passthrough kernel (dmsh_geometry.lean) only decides, per input primitive, whether
the rasterizer gets it. A mesh threadgroup takes PRIMS triangles of a list draw: each thread
runs the vertex shader's pulling entry (dmsh_pull) for one index position and writes the
vertex out, keeps the words the kernel reads in threadgroup memory, and the first PRIMS
threads run the kernel body on them and cull their triangle when viewport 0's mask bit is
clear. No capture pass, no compute pass, no index compaction, no render pass break.

The kernel's own storage writes (and the subgroup operations that order them) run as in the
compute kernel: one triangle per thread, PRIMS consecutive triangles per SIMD group. When the
vertices of PRIMS triangles exceed the mesh output limit (MESH_BYTES), a threadgroup takes 16,
8, ... (prims_for) and its SIMD group runs with the rest of its lanes inactive, which subgroup
operations of the guest's own (ballots, scans over active lanes) already allow.

Each entry has a run twin (dm_mesh_run, dm_mesh_run_z, ..._p16) that draws several such draws in
one dispatch: buffer 12 holds a table of dm_mesh_rec, one a draw, each with the draw's arguments,
the end of its threadgroups (a running total) and the addresses bound at the vertex table's slots
for it. A threadgroup finds its draw by halving, takes its resources through those addresses and
runs the same body with the draw's own threadgroup and primitive numbers. A mesh draw on an M1
costs about 2.7 us however few triangles it has; Donkey Kong Bananza's depth prepass draws 420-550
of them a frame (nvn_geometry.inc, geo_mesh_run).
"""
import re
import dmsh_geometry, dmsh_pull

PRIMS = 32   # triangles a threadgroup takes, fewer when their vertices exceed the mesh output limit (prims_for); the entry names any other count (dm_mesh_p16)
MESH_BYTES = 32768   # Metal's limit on a mesh threadgroup's output: vertices and primitives
DRAW_SLOT, GS_SLOT = 12, 13   # argument-table slots the vertex stage leaves free in a pulling pipeline (a run entry takes its table at DRAW_SLOT)
RUN_SLOTS = 31   # vertex table slots a run record holds (= nvn_geometry.inc geo_mesh_rec)
_SCALAR = {'float': 4, 'int': 4, 'uint': 4, 'half': 2, 'short': 2, 'ushort': 2}


def out_bytes(src):
    """The bytes of one vertex the mesh writes: the members of main0_out, packed as Metal counts
    the mesh output (Donkey Kong Bananza: 24 float4 and the position, 388 bytes, with the point size)."""
    m = re.search(r'struct main0_out\s*\{(.*?)\};', src, re.S)
    if not m: raise ValueError('mesh: no main0_out')
    n = 0
    for ln in m[1].split(';'):
        a = re.fullmatch(r'\s*(\w+?)([234]?) \w+(?: \[(\d+)\])?(?: \[\[.*\]\])?\s*', ln)
        if not ln.strip(): continue
        if not a or a[1] not in _SCALAR: raise ValueError('mesh: vertex output member ' + ln.strip())
        n += _SCALAR[a[1]] * int(a[2] or 1) * int(a[3] or 1)
    return n


def prims_for(vbytes):
    """The most triangles (32, 16, ... 1) whose vertices and primitives fit one mesh threadgroup."""
    for p in (32, 16, 8, 4, 2, 1):
        if 3 * p * vbytes + 4 * p <= MESH_BYTES: return p
    raise ValueError(f'mesh: one triangle of {vbytes}-byte vertices exceeds the mesh output limit')
_BUILTIN_V = {'vertex_id': 'dm_vid', 'instance_id': 'dm_iid', 'base_instance': 'dm_bi', 'base_vertex': 'dm_bv'}
_BUILTIN_G = {'thread_position_in_grid': 'uint3(dm_p0 + dm_t, 0u, 0u)', 'thread_index_in_simdgroup': 'dm_sg'}


def _params(s):
    """Split a parameter list at top-level commas."""
    out, depth, cur = [], 0, ''
    for ch in s:
        depth += ch in '<([' ; depth -= ch in '>)]'
        if ch == ',' and not depth: out.append(cur.strip()); cur = ''
        else: cur += ch
    return out + ([cur.strip()] if cur.strip() else [])


def _strip_head(src):
    keep = [l for l in src.split('\n') if not re.match(r'\s*#include\b|\s*using namespace metal;', l)]
    return '\n'.join(keep)


def _entry(src, kind, name):
    m = re.search(r'\n%s (\w+) %s\(([^\n]*)\)\n\{\n' % (kind, name), src)
    if not m: raise ValueError(f'mesh: no {kind} {name}')
    return m


def compose(vs_src, gs_src, reads):
    """vs_src: the vertex shader's library source (<key>_b.metal); gs_src: the lean kernel
    (<geometry key>_bl3.metal) and the capture words it reads, in its record order."""
    if 2 in reads or not reads: raise ValueError('mesh: the kernel reads clip z or nothing')
    pulled, info = dmsh_pull.variant(vs_src)   # keeps main0 / main0_z only
    if pulled is None or not isinstance(info, dict): raise ValueError(f'mesh: no pulling entry: {info}')
    words = dmsh_geometry.output_words(pulled)
    pulled = pulled.replace(' [[point_size]];', ';')   # a triangle mesh has no point size
    P = prims_for(out_bytes(pulled))
    # The vertex entries as functions; their resource parameters become the mesh entry's
    entries, res = {}, None
    for name in [info['entry']] + ([info['entry_z']] if info.get('entry_z') else []):
        m = _entry(pulled, 'vertex', name)
        args, fparams, rparams = [], [], []
        for p in _params(m[2]):
            a = re.fullmatch(r'(.+?) (\w+) \[\[(\w+)(?:\((\d+)\))?\]\]', p)
            if not a: raise ValueError('mesh: vertex parameter ' + p)
            fparams.append(f'{a[1]} {a[2]}')
            if a[3] in _BUILTIN_V: args.append(_BUILTIN_V[a[3]])
            elif a[3] == 'buffer' and int(a[4]) not in (DRAW_SLOT, GS_SLOT): args.append(a[2]); rparams.append(p)
            else: raise ValueError('mesh: vertex parameter ' + p)
        if res is not None and res != rparams: raise ValueError('mesh: entries differ in resources')
        res = rparams
        pulled = pulled[:m.start()] + f'\nstatic inline main0_out dm_vs_{name}(' + ', '.join(fparams) + ')\n{\n' + pulled[m.end():]
        entries[name] = args
    # The kernel as a function: capture reads from threadgroup memory, the topology map as the identity
    # (triangle t's vertices are 3t..3t+2), the mask store into dm_mask, index stores dropped
    k = _entry(gs_src, 'kernel', 'main0')
    gargs = []
    for p in _params(k[2]):
        a = re.fullmatch(r'(.+?) (\w+) \[\[(\w+)(?:\((\d+)\))?\]\]', p)
        if a and a[3] in _BUILTIN_G: gargs.append(_BUILTIN_G[a[3]])
        elif a and a[0] == 'constant spvDescriptorSetBuffer0& spvDescriptorSet0 [[buffer(0)]]': gargs.append('dm_gset')
        else: raise ValueError('mesh: kernel parameter ' + p)
    fparams = [re.sub(r' \[\[.*\]\]$', '', p) for p in _params(k[2])]
    body = gs_src[k.end():]
    if re.search(r'\bthreadgroup\b|\bthreadgroup_barrier\b', body): raise ValueError('mesh: kernel uses threadgroup memory')
    # GeometryToCompute uses GlobalId.x both for the guest PrimitiveId and for
    # addressing its capture. Keep that ID draw-wide; only the capture address
    # is relative to this mesh group's records. Replacing GlobalId with the
    # local lane silently repeats guest IDs every PRIMS primitives.
    body = re.sub(r'\(\*spvDescriptorSet0\.vertex_output\)\._m0\[(\w+)\]',
                  r'dm_rec[uint(\1) - dm_capture_base]', body)
    prefix = 'spvDescriptorSet0.trb_data['
    while prefix in body:
        s = body.index(prefix); at = e = s + len(prefix); depth = 1
        while depth:
            depth += (body[e] == '[') - (body[e] == ']'); e += 1
        body = body[:s] + 'uint(' + body[at:e - 1] + ')' + body[e:]
    body, n = re.subn(r'\(\*spvDescriptorSet0\.geometry_vb_output\)\._m0\[\w+\] = (\w+);', r'dm_mask = as_type<uint>(\1);', body)
    if not n: raise ValueError('mesh: no mask store')
    body = re.sub(r'\n\s*\(\*spvDescriptorSet0\.geometry_ib_output\)\._m0\[\w+\] = [^;]+;', '', body)
    if re.search(r'vertex_output|geometry_vb_output|geometry_ib_output|trb_data\[', body): raise ValueError('mesh: kernel output not understood')
    gfn = gs_src[:k.start()] + '\nstatic inline void dm_gs(' + ', '.join(fparams + ['threadgroup float* dm_rec', 'uint dm_capture_base', 'thread uint& dm_mask']) + ')\n{\n' + body
    # The words the kernel reads, from the vertex's outputs (missing ones as the capture leaves them: 0, w and point size 1)
    def word(w):
        if w in words: return 'dm_o.' + words[w][4:]
        return '0x3f800000u' if w in (3, 4) or (w >= 19 and w % 4 == 3) else '0u'
    n = len(reads)
    stores = ' '.join(f'dm_rec[dm_t * {n}u + {i}u] = as_type<float>({word(w)});' for i, w in enumerate(reads))
    out = ['#include <metal_stdlib>', '#include <simd/simd.h>', '#include <metal_atomic>', 'using namespace metal;',
           'namespace dm_v {', _strip_head(pulled), '}', 'namespace dm_g {', _strip_head(gfn), '}', 'using namespace dm_v;',
           'struct dm_mesh_draw { device const uchar* ib; uint itype; uint first; uint count; int base_vertex; uint base_instance; };',
           'struct dm_mesh_prim { bool culled [[primitive_culled]]; };',
           f'struct dm_mesh_rec {{ constant uchar* slot[{RUN_SLOTS}]; dm_mesh_draw d; uint group_end; uint n; }};']
    if res is None: raise ValueError('mesh: no vertex entries')
    # The run twin takes each resource of the vertex entries from the draw's record (constant references only: a run entry is left out otherwise)
    run_res = [m for m in (re.fullmatch(r'constant (.+?)& (\w+) \[\[buffer\((\d+)\)\]\]', p) for p in res) if m and int(m[3]) < RUN_SLOTS]
    run_ok = len(run_res) == len(res)
    head = f'metal::mesh<main0_out, dm_mesh_prim, {3 * P}, {P}, metal::topology::triangle> dm_m'
    for name, args in entries.items():
        sfx = f'{name[5:]}{"" if P == PRIMS else f"_p{P}"}'
        body = f'''
    threadgroup float dm_rec[{3 * P * n}];
    uint dm_p0 = dm_tg * {P}u, dm_np = min({P}u, dm_d.count / 3u - dm_p0);
    if (dm_t < dm_np * 3u) {{
        uint dm_i = dm_d.first + dm_p0 * 3u + dm_t;
        uint dm_ix = dm_d.itype == 0u ? dm_i : dm_d.itype == 1u ? uint(dm_d.ib[dm_i]) : dm_d.itype == 2u ? uint(((device const ushort*)dm_d.ib)[dm_i]) : ((device const uint*)dm_d.ib)[dm_i];
        uint dm_vid = dm_ix + uint(dm_d.base_vertex), dm_bv = uint(dm_d.base_vertex), dm_bi = dm_d.base_instance, dm_iid = dm_bi;
        main0_out dm_o = dm_vs_{name}({', '.join(args)});
        dm_m.set_vertex(dm_t, dm_o);
        {stores}
    }}
    if (dm_t == 0u) dm_m.set_primitive_count(dm_np);
    threadgroup_barrier(mem_flags::mem_threadgroup);
    if (dm_t < dm_np) {{
        uint dm_mask = 0u;
        dm_g::dm_gs({', '.join(gargs)}, dm_rec, dm_p0 * {3 * n}u, dm_mask);
        dm_m.set_index(dm_t * 3u, dm_t * 3u); dm_m.set_index(dm_t * 3u + 1u, dm_t * 3u + 1u); dm_m.set_index(dm_t * 3u + 2u, dm_t * 3u + 2u);
        dm_m.set_primitive(dm_t, dm_mesh_prim{{(dm_mask & 1u) == 0u}});
    }}
}}'''
        out.append(f'''[[mesh]] void dm_mesh{sfx}({head}, {', '.join(res)},
    constant dm_mesh_draw& dm_d [[buffer({DRAW_SLOT})]], constant dm_g::spvDescriptorSetBuffer0& dm_gset [[buffer({GS_SLOT})]],
    uint dm_t [[thread_index_in_threadgroup]], uint dm_tg [[threadgroup_position_in_grid]], uint dm_sg [[thread_index_in_simdgroup]])
{{{body}''')
        if not run_ok: continue
        takes = ''.join(f'\n    constant {m[1]}& {m[2]} = *(constant {m[1]}*)dm_r.slot[{m[3]}];' for m in run_res)
        out.append(f'''[[mesh]] void dm_mesh_run{sfx}({head}, constant dm_mesh_rec* dm_run [[buffer({DRAW_SLOT})]],
    uint dm_t [[thread_index_in_threadgroup]], uint dm_tgr [[threadgroup_position_in_grid]], uint dm_sg [[thread_index_in_simdgroup]])
{{
    uint dm_lo = 0u, dm_hi = dm_run[0].n - 1u;   // the draw this threadgroup belongs to: the first whose group_end is past it
    while (dm_lo < dm_hi) {{ uint dm_mid = (dm_lo + dm_hi) / 2u; if (dm_run[dm_mid].group_end <= dm_tgr) dm_lo = dm_mid + 1u; else dm_hi = dm_mid; }}
    constant dm_mesh_rec& dm_r = dm_run[dm_lo];
    uint dm_tg = dm_tgr - (dm_lo ? dm_run[dm_lo - 1u].group_end : 0u);
    constant dm_mesh_draw& dm_d = dm_r.d;
    constant dm_g::spvDescriptorSetBuffer0& dm_gset = *(constant dm_g::spvDescriptorSetBuffer0*)dm_r.slot[{GS_SLOT}];{takes}{body}''')
    return '\n'.join(out) + '\n'


if __name__ == '__main__':
    import json, sys
    vs, gs, out = sys.argv[1:4]
    reads = json.load(open(gs[:-len('.metal')] + '.json'))['geometry']['capture_reads']
    open(out, 'w').write(compose(open(vs).read(), open(gs).read(), reads))