#!/usr/bin/python3
# Copyright 2019 The ANGLE Project Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
#
# gen_mtl_internal_shaders.py:
#   Code generation for Metal backend's default shaders.
#   NOTE: don't run this script directly. Run scripts/run_code_generation.py.

import json
import os
import subprocess
import sys

sys.path.append('../..')
import angle_format
import gen_angle_format_table

configs = [{
    "sdk": "macosx",
    "compile_flags": ["--std=macos-metal2.1", "-mmacosx-version-min=10.14"],
    "variable_name": "gDefaultMetallib",
    "header": "mtl_internal_shaders_macos_autogen.h"
}, {
    "sdk": "iphoneos",
    "compile_flags": ["--std=ios-metal2.1", "-mios-version-min=12"],
    "variable_name": "gDefaultMetallib",
    "header": "mtl_internal_shaders_ios_autogen.h"
}]

metal_source_output_header = "mtl_internal_shaders_src_autogen.h"

template_header_boilerplate = """// GENERATED FILE - DO NOT EDIT.
// Generated by {script_name}
//
// Copyright 2020 The ANGLE Project Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
//
"""

def gen_shader_enums_code(angle_formats):

    code = """// This file is similar to src/libANGLE/renderer/FormatID_autogen.h but is used by Metal default
// shaders instead of C++ code.
//
"""

    code += "namespace rx\n"
    code += "{\n"
    code += "namespace mtl_shader\n"
    code += "{\n"
    code += "\n"
    code += "namespace FormatID\n"
    code += "{\n"
    code += "enum\n"
    code += "{\n"
    code += gen_angle_format_table.gen_enum_string(angle_formats) + '\n'
    code += "};\n\n"
    code += "}\n"
    code += "\n"
    code += "}\n"
    code += "}\n"

    return code


def find_clang():
    if os.name == 'nt':
        binary = 'clang-cl.exe'
    else:
        binary = 'clang++'

    clang = os.path.join('..', '..', '..', '..', '..', 'third_party', 'llvm-build',
                         'Release+Asserts', 'bin', binary)

    if not os.path.isfile(clang):
        xcrun_clang = subprocess.run(["xcrun", "-f", binary], stdout=subprocess.PIPE, text=True)
        if xcrun_clang.returncode == 0:
            clang = xcrun_clang.stdout.strip()
    if (not os.path.isfile(clang)):
        raise Exception('Cannot find clang')

    return clang


def metal_to_air(metal_src, sdk, compile_flags, dest_air_file):
    temp_fname = 'temp_metal_src.metal'
    with open(temp_fname, 'wt') as out_file:
        out_file.write(metal_src.decode("utf-8"))

    result = subprocess.run(
        ["xcrun", "-sdk", sdk, "metal"] + compile_flags + ["-c", temp_fname, "-o", dest_air_file],
        stdout=subprocess.PIPE,
        text=True)
    os.remove(temp_fname)

    if result.returncode != 0:
        raise Exception('Failed to compile metal to air: ' + result.stdout.strip())


def air_to_mtllib(src_air_file, sdk, dest_mtllib_file):
    result = subprocess.run(
        ["xcrun", "-sdk", sdk, "metallib", src_air_file, "-o", dest_mtllib_file],
        stdout=subprocess.PIPE,
        text=True)
    if result.returncode != 0:
        raise Exception('Failed to compile metal to air: ' + result.stdout.strip())


def metal_to_metallib(metal_src, sdk, compile_flags):
    intermediate_air_file = "temp_metal.air"
    metal_to_air(metal_src, sdk, compile_flags, intermediate_air_file)

    intermediate_metallib_file = "temp_metal.metallib"
    air_to_mtllib(intermediate_air_file, sdk, intermediate_metallib_file)
    os.remove(intermediate_air_file)

    with open(intermediate_metallib_file, 'rb') as f:
        mtllib = f.read()
    os.remove(intermediate_metallib_file)

    return mtllib


def generate_metallib_header(metal_src, sdk, compile_flags, variable_name, dest_header_file):
    boilerplate_code = template_header_boilerplate.format(
        script_name=os.path.basename(sys.argv[0]))

    metallib_data = metal_to_metallib(metal_src, sdk, compile_flags)

    with open(dest_header_file, 'wt') as out_file:
        out_file.write(boilerplate_code)
        out_file.write('\n')
        out_file.write('// C++ string version of default shaders mtllib.\n\n')
        out_file.write('\n\nstatic constexpr uint8_t ' + variable_name + '[] = {\n')
        for byte in metallib_data:
            out_file.write(f"{byte}, ")
        out_file.write('\n')
        out_file.write('};\n')
        out_file.close()


def generate_metal_autogen_header(dest_metal_header):
    angle_to_gl = angle_format.load_inverse_table('../../angle_format_map.json')
    shader_autogen_header = gen_shader_enums_code(angle_to_gl.keys())

    with open(dest_metal_header, 'wt') as out_file:
        out_file.write(shader_autogen_header)
        out_file.close()


def generate_combined_metal_src(metal_src_files):
    autogen_header_file = "format_autogen.h"
    generate_metal_autogen_header(autogen_header_file)

    clang = find_clang()

    # Use clang to preprocess the combination source. "@@" token is used to prevent clang from
    # expanding the preprocessor directive
    temp_fname = 'temp_master_source.metal'
    with open(temp_fname, 'wb') as temp_file:
        for src_file in metal_src_files:
            include_str = '#include "' + src_file + '" \n'
            temp_file.write(include_str.encode('utf-8'))

    args = [clang]
    if not os.name == 'nt':
        args += ['-xc++']
    args += ['-E', temp_fname]

    combined_source = subprocess.check_output(args)
    os.remove(temp_fname)
    os.remove(autogen_header_file)

    # Remove '@@' tokens
    final_combined_src_string = combined_source.replace('@@'.encode('utf-8'), ''.encode('utf-8'))

    return final_combined_src_string


def generate_combined_metal_src_header(combined_metal_src, dest_header):
    boilerplate_code = template_header_boilerplate.format(
        script_name=os.path.basename(sys.argv[0]))

    with open(dest_header, 'wt') as out_file:
        out_file.write(boilerplate_code)
        out_file.write('\n')
        out_file.write('// C++ string version of combined Metal default shaders.\n\n')
        out_file.write('\n\nstatic char gDefaultMetallibSrc[] = R"(\n')
        out_file.write(combined_metal_src.decode("utf-8"))
        out_file.write('\n')
        out_file.write(')";\n')
        out_file.close()


def main():
    angle_format_script_files = [
        '../../angle_format_map.json', '../../angle_format.py', '../../gen_angle_format_table.py'
    ]
    src_files = [
        'blit.metal', 'clear.metal', 'gen_indices.metal', 'gen_mipmap.metal', 'copy_buffer.metal',
        'visibility.metal', 'rewrite_indices.metal'
    ]

    # auto_script parameters.
    if len(sys.argv) > 1:
        inputs = angle_format_script_files + src_files + ['common.h', 'constants.h']
        outputs = [config["header"] for config in configs] + [metal_source_output_header]

        if sys.argv[1] == 'inputs':
            print(','.join(inputs))
        elif sys.argv[1] == 'outputs':
            print(','.join(outputs))
        else:
            print('Invalid script parameters')
            return 1
        return 0

    os.chdir(sys.path[0])

    combined_metal_src = generate_combined_metal_src(src_files)

    generate_combined_metal_src_header(combined_metal_src, metal_source_output_header)

    for config in configs:
        generate_metallib_header(combined_metal_src, config["sdk"], config["compile_flags"],
                                 config["variable_name"], config["header"])


if __name__ == '__main__':
    sys.exit(main())
