#!/usr/bin/python3

import argparse
import pathlib
import os
import subprocess
from datetime import datetime
import sys


def main():
    parser = argparse.ArgumentParser(
        description='Compile glsl source to spv bytes and generate a C file that export the spv bytes as an array.')
    parser.add_argument('source_file', type=pathlib.Path)
    parser.add_argument('target_file', type=str)
    parser.add_argument('export_symbol', type=str)
    args = parser.parse_args()
    vulkan_sdk_path = os.getenv('VULKAN_SDK')
    if vulkan_sdk_path == None:
        print("Environment variable VULKAN_SDK not set. Please install the LunarG Vulkan SDK and set VULKAN_SDK accordingly")
        return
    vulkan_sdk_path = pathlib.Path(vulkan_sdk_path)
    glslc_path = vulkan_sdk_path / 'bin' / 'glslc'
    if os.name == 'nt':
        glslc_path = glslc_path.with_suffix('.exe')
    if not glslc_path.exists():
        print("Can't find " + str(glslc_path))
        return
    if not args.source_file.exists():
        print("Can't find source file: " + str(args.source_file))
        return

    with subprocess.Popen([str(glslc_path), str(args.source_file), '-o', '-'], stdout=subprocess.PIPE, stderr=subprocess.PIPE) as proc, open(args.target_file, mode='wt') as target_file:
        if proc.wait() != 0:
            print(proc.stderr.read().decode('utf-8'))
            return
        spv_bytes = proc.stdout.read()
        spv_bytes = [int.from_bytes(spv_bytes[i:i + 4], byteorder="little")
                     for i in range(0, len(spv_bytes), 4)]
        chunk_size = 4
        spv_bytes_chunks = [spv_bytes[i:i + chunk_size]
                            for i in range(0, len(spv_bytes), chunk_size)]
        spv_bytes_in_c_array = ',\n'.join([', '.join(
            [f'{spv_byte:#010x}' for spv_byte in spv_bytes_chunk]) for spv_bytes_chunk in spv_bytes_chunks])
        spv_bytes_in_c_array = f'const uint32_t {args.export_symbol}[] = ' + \
            '{\n' + spv_bytes_in_c_array + '\n};'
        comments = f"""// Copyright (C) {datetime.today().year} The Android Open Source Project
// Copyright (C) {datetime.today().year} Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Autogenerated module {args.target_file}
// generated by {"python3 " + " ".join(sys.argv)}
// Please do not modify directly.\n\n"""
        prelude = "#include <stdint.h>\n\n"
        target_file.write(comments)
        target_file.write(prelude)
        target_file.write(spv_bytes_in_c_array)


if __name__ == '__main__':
    main()