#!/usr/bin/python3

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

BYTES_PER_UINT32 = 4

UINT32_PER_LINE = 8

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

    spv_bytes = None
    with subprocess.Popen([str(glslc_path), str(args.source_file), '-o', '-'], stdout=subprocess.PIPE, stderr=subprocess.PIPE) as proc:
        if proc.wait() != 0:
            print(proc.stderr.read().decode('utf-8'))
        spv_bytes = proc.stdout.read()
    spv_bytes = [int.from_bytes(spv_bytes[i:i + BYTES_PER_UINT32], byteorder="little")
                 for i in range(0, len(spv_bytes), BYTES_PER_UINT32)]
    spv_bytes_chunks = [spv_bytes[i:i + UINT32_PER_LINE]
                        for i in range(0, len(spv_bytes), UINT32_PER_LINE)]
    spv_bytes_in_vector = ',\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_vector = f'const std::vector<uint32_t> {args.export_symbol} = ' + \
            '{\n' + spv_bytes_in_vector + '\n};'

    source_file_lines = None
    with open(args.source_file, mode='r') as source_file:
        source_file_lines = source_file.readlines()

    with open(args.target_file, mode='wt') as target_file:
        header = 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.

// Please do not modify directly.

// Autogenerated module {args.target_file} generated by:
//   {"python3 " + " ".join(sys.argv)}

#include <stdint.h>
#include <vector>

// From {args.source_file}:

"""
        target_file.write(header)

        for source_file_line in source_file_lines:
            target_file.write("// " + source_file_line)

        target_file.write("\n\n")

        target_file.write(spv_bytes_in_vector)


if __name__ == '__main__':
    main()