Import('env')
env = env.Clone()

VK_FUNCTION_BLACKLIST = {'vkGetFramebufferTilePropertiesQCOM', 'vkGetDynamicRenderingTilePropertiesQCOM'}

def parse_cmd_node(node, vulkan_hpp):
    proto_tag = node.find('proto')
    cmd_name = proto_tag and proto_tag.find('name').text or node.attrib.get('name', '')
    if not cmd_name or cmd_name in VK_FUNCTION_BLACKLIST:
        return
    if proto_tag is None:
        vulkan_hpp['commands'][node.attrib['name']] = {
            'name': node.attrib['name'],
            'alias': node.attrib['alias'],
            'extension': '',
            'platform': ''
        }
        return
    type_tag = proto_tag.find('type')
    name_tag = proto_tag.find('name')

    params = []
    for param_tag in node.findall('param'):
        param_type_tag = param_tag.find('type')
        param_name_tag = param_tag.find('name')
        params.append({
            'type': ((param_tag.text or "") + param_type_tag.text + (param_type_tag.tail or "")).strip(),
            'name': param_name_tag.text.strip(),
            'name_type_suffix': (param_name_tag.tail or "").strip()
        })

    vulkan_hpp['commands'][name_tag.text.strip()] = {
        'name': name_tag.text.strip(),
        'return_type': type_tag.text,
        'params': params,
        'extension': '',
        'platform': ''
    }

def parse_type_node(node, vulkan_hpp):
    alias_for = node.attrib.get('alias')
    type_name = node.attrib.get('name', '')
    if alias_for is not None:
        vulkan_hpp['types'][type_name] = {
            'name': type_name,
            'alias': alias_for,
            'extension': '',
            'platform': ''
        }
        return
    category = node.attrib.get('category')
    if category in ('struct', 'union'):
        parse_struct_type_node(node, vulkan_hpp)
    elif category == 'handle':
        parse_handle_type_node(node, vulkan_hpp)

def parse_struct_type_node(node, vulkan_hpp):
    type_name = node.attrib['name']
    category = node.attrib.get('category')

    members = []
    sType = ''
    for member_tag in node.findall('member'):
        member_type_tag = member_tag.find('type')
        member_name_tag = member_tag.find('name')
        members.append({
            'type': ((member_tag.text or "") + member_type_tag.text + (member_type_tag.tail or "")).strip(),
            'name': member_name_tag.text.strip(),
            'name_type_suffix': (member_name_tag.tail or "").strip(),
            'length': member_tag.attrib.get('len', '')
        })
        if member_name_tag.text == 'sType':
            sType = member_tag.attrib.get('values', '').split(',')[0]
    vulkan_hpp['types'][type_name] = {
        'name': type_name,
        'category': category,
        'members': members,
        'sType': sType,
        'extension': '',
        'platform': ''
    }

def parse_handle_type_node(node, vulkan_hpp):
    name_tag = node.find('name')
    if name_tag is None:
        return
    type_name = name_tag.text
    vulkan_hpp['types'][type_name] = {
        'name': type_name,
        'category': 'handle',
        'objtypeenum': node.attrib['objtypeenum'],
        'extension': '',
        'platform': ''
    }


def parse_platform_node(platform_node, vulkan_hpp):
    platform_name = platform_node.attrib['name']
    vulkan_hpp['platforms'][platform_name] = {
        'protect': platform_node.attrib['protect']
    }

def parse_extension_node(extension_node, vulkan_hpp):
    extension_name = extension_node.attrib['name']
    platform_name = extension_node.attrib.get('platform', '')
    require_node = extension_node.find('require')
    for cmd_node in require_node.findall('command'):
        cmd_name = cmd_node.attrib['name']
        if cmd_name in vulkan_hpp['commands']:
            vulkan_hpp['commands'][cmd_name]['extension'] = extension_name
            vulkan_hpp['commands'][cmd_name]['platform'] = platform_name
    for type_node in require_node.findall('type'):
        type_name = type_node.attrib['name']
        if type_name in vulkan_hpp['types']:
            vulkan_hpp['types'][type_name]['extension'] = extension_name
            vulkan_hpp['types'][type_name]['platform'] = platform_name
    vulkan_hpp['extensions'][extension_name] = {
        'platform': platform_name
    }

def parse_vulkan_hpp():
    import xml.etree.ElementTree as ET

    vulkan_hpp = {
        'commands': {},
        'types': {},
        'extensions': {},
        'platforms': {}
    }

    tree = ET.parse('thirdparty/vulkan-docs/xml/vk.xml')
    root = tree.getroot()
    
    cmds_node = root.find('commands')
    for cmd_node in cmds_node:
        parse_cmd_node(cmd_node, vulkan_hpp)
    
    types_node = root.find('types')
    for type_node in types_node:
        parse_type_node(type_node, vulkan_hpp)
    
    platforms_node = root.find('platforms')
    for platform_node in platforms_node:
        parse_platform_node(platform_node, vulkan_hpp)
    
    extensions_node = root.find('extensions')
    for extension_node in extensions_node:
        parse_extension_node(extension_node, vulkan_hpp)

    return vulkan_hpp

vulkan_hpp = parse_vulkan_hpp()

def generate_source(target, source, env):
    import os
    import sys
    from importlib.util import module_from_spec, spec_from_file_location

    # allow the module to import from its current folder
    sys.path.append(os.path.dirname(source[0].abspath))
    python_spec = spec_from_file_location('source', source[0].abspath)
    python_mod = module_from_spec(python_spec)
    python_spec.loader.exec_module(python_mod)
    python_mod.vulkan_hpp = vulkan_hpp
    python_mod.generate(targets = [t.abspath for t in target])

    # restore path
    sys.path = sys.path[:-1]

    return None

pygen_builder = Builder(action = generate_source)
# env.Append(CCFLAGS = ' -Wall -Wextra -Werror -pedantic -std=c++20')
env.Append(BUILDERS = {'PyGen': pygen_builder})
env.Append(CPPPATH = ['include'])

source_files = Split("""
    source/data_pool.cpp
    source/dispatch_table.cpp
    source/functions.cpp
    source/function_ids.cpp
    source/layer.cpp
    source/record_list.cpp
    source/variant_pool.cpp
""")

# env.PyGen('source/functions.hpp', 'generators/functions.hpp.py')
env.PyGen('include/vk_function_ids.h', 'generators/vk_function_ids.h.py')
env.PyGen('source/function_ids.cpp', 'generators/function_ids.cpp.py')
env.PyGen('source/functions.cpp', 'generators/functions.cpp.py')
env.PyGen('source/variant_wrap.hpp', 'generators/variant_wrap.hpp.py')
env.SharedLibrary(
    target = 'VkLayer_capture',
    source = source_files
)