import binascii
from common import struct_name_to_enum, write_epilogue, write_preamble, FIRST_PARAM_BLACKLIST, TYPE_BLACKLIST

HEADER_GUARD = 'VK_CAPTURE_VARIANT_WRAP_HPP_INCLUDED'

LENGTH_EXCEPTIONS = {
    ('VkShaderModuleCreateInfo', 'pCode'): 'value.codeSize / 4',
    ('VkPipelineMultisampleStateCreateInfo', 'pSampleMask'): '(value.rasterizationSamples + 31) / 32',
    ('VkAccelerationStructureVersionInfoKHR', 'pVersionData'): '2 * VK_UUID_SIZE'
}

def write_copyfunc_definition(f, struct):
    if 'alias' in struct or struct['name'] in TYPE_BLACKLIST or struct['category'] != 'struct':
        return
    f.write(f"""inline void copyIndirectData({struct['name']}& value);
""")
    

def write_copyfunc(f, struct):
    if 'alias' in struct or struct['name'] in TYPE_BLACKLIST or struct['category'] != 'struct':
        return
    f.write(f"""
inline void copyIndirectData([[maybe_unused]] {struct['name']}& value)
{{""")
    for member in struct['members']:
        if 'void*' in member['type']:
            if member['name'] == 'pNext':
                f.write("""
    value.pNext = copyNextPtr(value.pNext);""")
            continue
        if member['type'].endswith('**'):
            pass
        elif member['type'].endswith('*'):
            f.write(f"""
    if (value.{member['name']} != nullptr)
    {{
        using base_type_t = std::decay_t<{member['type'].strip('*')}>;""")
            length_exception = LENGTH_EXCEPTIONS.get((struct['name'], member['name']))
            if length_exception:
                f.write(f"""
        const std::size_t numElements = {length_exception};
                """)
            elif member['type'].endswith('char*'):
                f.write(f"""
        const std::size_t numElements = std::strlen(value.{member['name']}) + 1;""")
            elif member['length']:
                f.write(f"""
        const std::size_t numElements = value.{member['length'].split(',')[0]};""")
            else:
                f.write(f"""
        static constexpr std::size_t numElements = 1; // TODO""")
            f.write(f"""
        auto copy = allocType<base_type_t>(numElements);
        std::memcpy(copy, value.{member['name']}, numElements * sizeof(base_type_t));

        for (std::size_t ele = 0; ele < numElements; ++ele) {{
            copyIndirectData(copy[ele]);
        }}
        value.{member['name']} = copy;
    }}
""")
    f.write(f"""
}}
""")

def write_structwrapper(f, struct):
    if 'alias' in struct or struct['name'] in TYPE_BLACKLIST:
        return
    if(len(struct['members']) < 2 or struct['members'][1]['name'] != 'pNext'):
        return # TODO
    is_output = "const" not in struct['members'][1]['type']
    f.write(f"""
inline void variantWrap(const {struct['name']}& value, VkVariantMWN& outVariant)
{{
    outVariant.type = VK_VARIANT_TYPE_{is_output and "OUT" or "IN"}_STRUCTURE_MWN;

    auto valueCopy = allocType<{struct['name']}>();
    outVariant.{is_output and "out" or "in"}StructureValue = reinterpret_cast<VkBase{is_output and "Out" or "In"}Structure*>(valueCopy);

    std::memcpy(valueCopy, &value, sizeof({struct['name']}));
    copyIndirectData(*valueCopy);
}}
""")

def write_handlewrapper(f, handle):
    if 'alias' in handle or handle['name'] in TYPE_BLACKLIST:
        return
    f.write(f"""
inline void variantWrap({handle['name']} value, VkVariantMWN& outVariant)
{{
    outVariant.type = VK_VARIANT_TYPE_OBJECT_MWN;
    outVariant.objectValue.type = {handle['objtypeenum']};
    outVariant.objectValue.value = value;
}}
""")

def write_nextptrcase(f, struct):
    if 'alias' in struct or struct['name'] in TYPE_BLACKLIST or struct['name'] in ('VkBaseOutStructure', 'VkBaseInStructure'):
        return
    if(len(struct['members']) < 2 or struct['members'][0]['name'] != 'sType'):
        return
    f.write(f"""
        case {(struct['sType'])}:
            pCopy = allocType<{struct['name']}>();
            std::memcpy(pCopy, pNext, sizeof({struct['name']}));
            copyIndirectData(*static_cast<{struct['name']}*>(pCopy));
            break;""")

def generate(targets):
    assert(len(targets) == 1)
    with open(targets[0], 'w') as f:
        write_preamble(f, includes = ['<cstring>', '<type_traits>', '<vulkan/vulkan.h>', '"common.hpp"', '"vk_capture.h"'], header_guard=HEADER_GUARD)
        
        f.write("""
inline void* copyNextPtr(const void* pNext);

// do nothing by default
template<typename T>
inline void copyIndirectData(const T&) {}

inline void copyIndirectData(const char*& value)
{
    const std::size_t len = std::strlen(value) + 1;
    char* copy = allocType<char>(len);
    std::memcpy(copy, value, len);
    value = copy;
}

template<typename TStruct>
inline void copyIndirectData(const TStruct*& value)
{
    TStruct* copy = allocType<TStruct>();
    std::memcpy(copy, value, sizeof(TStruct));
    copyIndirectData(*copy);
    value = copy;
}
""")

        types = list(vulkan_hpp['types'].values())
        types.sort(key=lambda tp: tp['platform'])
        plat = ''
        for tp in types:
            if plat != tp['platform']:
                if plat != '':
                    f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
                plat = tp['platform']
                f.write(f"""
#if defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
            if tp.get('category') in ('struct', 'union'):
                write_copyfunc_definition(f, tp)
        # finish the final platform
        f.write(f"""#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")

        plat = ''
        for tp in types:
            if plat != tp['platform']:
                if plat != '':
                    f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
                plat = tp['platform']
                f.write(f"""
#if defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
            if tp.get('category') in ('struct', 'union'):
                write_copyfunc(f, tp)
        # finish the final platform
        f.write(f"""#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")

        plat = ''
        for tp in types:
            if plat != tp['platform']:
                if plat != '':
                    f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
                plat = tp['platform']
                f.write(f"""
#if defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
            if tp.get('category') in ('struct', 'union'):
                write_structwrapper(f, tp)
            elif tp.get('category') == 'handle':
                write_handlewrapper(f, tp)
        # finish the final platform
        f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")

        f.write("""
inline void* copyNextPtr(const void* pNext)
{
    if (pNext == nullptr) {
        return nullptr;
    }
    const VkStructureType sType = *static_cast<const VkStructureType*>(pNext);
    void* pCopy = nullptr;
    switch (sType)
    {""")
        plat = ''
        for tp in types:
            if plat != tp['platform']:
                if plat != '':
                    f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
                plat = tp['platform']
                f.write(f"""
#if defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
            if tp.get('category') in ('struct', 'union'):
                write_nextptrcase(f, tp)
        # finish the final platform
        f.write(f"""
#endif // defined({vulkan_hpp['platforms'][plat]['protect']})\n""")
        f.write("""
        default: break;
    }
    return pCopy;
}""")

        write_epilogue(f, header_guard=HEADER_GUARD)