#!/usr/bin/env python3
# vim: set syntax=python ts=4 sw=4 sts=4 smarttab expandtab autoindent

DESCRIPTION = \
"Compile shader source code to a baked SPIR-V binary file."

EPILOG = \
"Applying a number of quirks and fixes to deal with drivers bugging out on certain" \
"SPIR-V instructions. Generate GLSL from spirv-opt-ified binary and recompile with" \
"that in order to have consistency between Nvidia and everyone else."

args = None
import subprocess

def glslc(src, src_is_a_filepath = True, emit_assembly = False):
    incdirs = []
    for i in args.incdirs:
        incdirs.append('-I')
        incdirs.append(i)

    rr = subprocess.run([
            "glslc",
            f"--target-env={args.env}",
            f"-std={args.std}",
            f"-fshader-stage={args.stage}"] + incdirs +
            ["-g", "-O", "-o", "-",] + (["-S"] if emit_assembly else []) +
            [(src if src_is_a_filepath else '-')],
        stdout=subprocess.PIPE,
        input = None if src_is_a_filepath else src)
    return rr.returncode, rr.stdout

def spirv_opt(src, strip_debug = True):
    rr = subprocess.run([
            "spirv-opt",
            '--preserve-bindings', '--preserve-spec-constants',
            '--strip-nonsemantic', '--strip-reflect'
            ] + (['--strip-debug'] if strip_debug else []) + [
            '--amd-ext-to-khr', '--workaround-1209',
            f"--target-env={args.env}",
            "-O", "-o", "-", "-" ],
        stdout=subprocess.PIPE, input = src )
    return rr.returncode, rr.stdout

def jfmt_glsl(src):
    srclines = [b'']
    for sl in [l for l in src.split(b'\n') if 0 < len(l)]:
        if sl.startswith(b'#'):
            srclines[-1] = sl
            srclines.append(b'')
        else:
            srclines[-1] += b' ' + sl
    srclines = [b' '.join(sl.split()) for sl in srclines]
    return b'\n'.join(srclines)

def filter_spirv(spirv, filename='<stdin>'):
    return b'\n'.join(
            [l for l in spirv.split(b'\n') if not
                (b'OpLine' in l or b'OpNoLine' in l)
            ]).replace(b'<stdin>', filename.encode())

def spirv_cross(src):
    rr = subprocess.run([
            "spirv-cross", '--no-es', '-'],
        stdout=subprocess.PIPE, input = src )
    return rr.returncode, rr.stdout

def spirv_dis(src):
    rr = subprocess.run(["spirv-dis", "-o", "-", "-"],
        stdout=subprocess.PIPE, input = src )
    return rr.returncode, rr.stdout

def spirv_as(src, dstpath = None):
    rr = subprocess.run(["spirv-as",
            "--target-env", args.env,
            "-o", ("-" if not dstpath else dstpath), "-"],
            input = src )
    return rr.returncode

def exit_on_error(ec):
    import sys
    if 0 != ec:
        sys.exit(ec)

if __name__ == '__main__':
    from argparse import ArgumentParser
    ap = ArgumentParser(
            prog = "shaderbake",
            description = DESCRIPTION,
            epilog = EPILOG )
    ap.add_argument('input', default='-', nargs='?')
    ap.add_argument('-I', '--add-include-directory', dest='incdirs', action='append', default=[])
    ap.add_argument('-o', '--output', dest='output', default=None)
    ap.add_argument('-s', '--std', dest='std', default='450core')
    ap.add_argument('-e', '--environment', dest='env', default='opengl4.5')
    ap.add_argument('-t', '--stage', dest='stage', default=None)

    import sys
    args = ap.parse_args(sys.argv[1:])

    args.stripped_input = args.input[:-5] if args.input.lower().endswith('.glsl') else args.input
    if not args.stage:
        if args.stripped_input.lower().endswith('.cs'):
            args.stage = 'comp'
        if args.stripped_input.lower().endswith('.vs'):
            args.stage = 'vert'
        if args.stripped_input.lower().endswith('.tc'):
            args.stage = 'tesc'
        if args.stripped_input.lower().endswith('.gs'):
            args.stage = 'geom'
        if args.stripped_input.lower().endswith('.te'):
            args.stage = 'tese'
        if args.stripped_input.lower().endswith('.fs'):
            args.stage = 'frag'

    ec,sc = glslc(args.input)
    exit_on_error(ec)

    if True:
        ec,sc = spirv_opt(sc, strip_debug=False)
        exit_on_error(ec)

    if True:
        ec,sc = spirv_cross(sc)
        exit_on_error(ec)
        ec,sc = glslc(jfmt_glsl(sc), src_is_a_filepath=False, emit_assembly=False)
        exit_on_error(ec)
        ec,sc = spirv_opt(sc, strip_debug=False)
        exit_on_error(ec)

    if True:
        ec,sc = spirv_dis(sc)
        exit_on_error(ec)
        ec = spirv_as(
                filter_spirv(sc, args.stripped_input),
                args.output)
        exit_on_error(ec)
    else:
        with open(args.output, 'wb') if args.output else sys.stdout as output:
            output.write(sc)