#options GAME(IW3, IW4, IW5, T5, T6)

#filename "Game/" + GAME + "/Techset/TechniqueCompiler" + GAME + ".cpp"

#set COMPILER_HEADER "\"TechniqueCompiler" + GAME + ".h\""
#set GAME_HEADER "\"Game/" + GAME + "/" + GAME + ".h\""
#set TECHSET_CONSTANTS_HEADER "\"Game/" + GAME + "/Techset/TechsetConstants" + GAME + ".h\""

#if GAME == "IW3"
#define FEATURE_IW3
#define IS_DX9
#define SHADERS_ARE_SUBASSETS
#elif GAME == "IW4"
#define FEATURE_IW4
#define IS_DX9
#elif GAME == "IW5"
#define FEATURE_IW5
#define IS_DX9
#elif GAME == "T5"
#define FEATURE_T5
#define IS_DX9
#define SHADERS_ARE_SUBASSETS
#elif GAME == "T6"
#define FEATURE_T6
#define IS_DX11
#define SHADERS_ARE_SUBASSETS
#endif

// This file was templated.
// See TechniqueCompiler.cpp.template.
// Do not modify, changes will be lost.

#include COMPILER_HEADER

#include GAME_HEADER
#include TECHSET_CONSTANTS_HEADER
#include "Techset/CommonShaderArgCreator.h"
#include "Techset/CommonTechniqueLoader.h"
#include "Techset/CommonVertexDeclCreator.h"
#include "Techset/LiteralConstsZoneState.h"
#include "Utils/StringUtils.h"

#if defined(FEATURE_T6)
#set PRECOMPILED_INDEX_HEADER "\"Game/" + GAME + "/Techset/PrecompiledIndex" + GAME + ".h\""
#include PRECOMPILED_INDEX_HEADER
#endif

#include <cassert>
#include <optional>
#include <sstream>
#include <vector>

using namespace GAME;

#set SHADER_LOADER_CLASS_NAME "TechniqueShaderLoader" + GAME
#set COMPILER_CLASS_NAME "TechniqueCompiler" + GAME

namespace
{
    unsigned ConvertArgumentType(const techset::CommonShaderArgumentType& type)
    {
        if (type.m_shader_type == techset::CommonTechniqueShaderType::VERTEX)
        {
            switch (type.m_value_type)
            {
            case techset::CommonShaderValueType::LITERAL_CONST:
                return MTL_ARG_LITERAL_VERTEX_CONST;
            case techset::CommonShaderValueType::MATERIAL_CONST:
                return MTL_ARG_MATERIAL_VERTEX_CONST;
            case techset::CommonShaderValueType::CODE_CONST:
                return MTL_ARG_CODE_VERTEX_CONST;
            case techset::CommonShaderValueType::MATERIAL_SAMPLER:
#if defined(FEATURE_IW5)
                return MTL_ARG_MATERIAL_VERTEX_SAMPLER;
#endif
            case techset::CommonShaderValueType::CODE_SAMPLER:
            default:
                assert(false);
                return 0;
            }
        }

        assert(type.m_shader_type == techset::CommonTechniqueShaderType::PIXEL);

        switch (type.m_value_type)
        {
        case techset::CommonShaderValueType::LITERAL_CONST:
            return MTL_ARG_LITERAL_PIXEL_CONST;
        case techset::CommonShaderValueType::MATERIAL_CONST:
            return MTL_ARG_MATERIAL_PIXEL_CONST;
        case techset::CommonShaderValueType::CODE_CONST:
            return MTL_ARG_CODE_PIXEL_CONST;
        case techset::CommonShaderValueType::MATERIAL_SAMPLER:
            return MTL_ARG_MATERIAL_PIXEL_SAMPLER;
        case techset::CommonShaderValueType::CODE_SAMPLER:
            return MTL_ARG_CODE_PIXEL_SAMPLER;
        default:
            assert(false);
            return 0;
        }
    }

    void ConvertArgumentValue(MaterialArgumentDef& argValue, const techset::CommonShaderArg& commonArg, MemoryManager& memory, AssetCreationContext& context)
    {
        switch (commonArg.m_type.m_value_type)
        {
        case techset::CommonShaderValueType::LITERAL_CONST:
        {
            const techset::LiteralConst literal(commonArg.m_value.literal_value);
            argValue.literalConst = context.GetZoneAssetCreationState<techset::LiteralConstsZoneState<float>>().GetAllocatedLiteral(literal)->GamePtr();
            break;
        }
        case techset::CommonShaderValueType::CODE_CONST:
            argValue.codeConst.index = static_cast<decltype(MaterialArgumentCodeConst::index)>(commonArg.m_value.code_const_source.m_index);
            argValue.codeConst.firstRow = static_cast<decltype(MaterialArgumentCodeConst::firstRow)>(commonArg.m_value.code_const_source.m_first_row);
            argValue.codeConst.rowCount = static_cast<decltype(MaterialArgumentCodeConst::rowCount)>(commonArg.m_value.code_const_source.m_row_count);
            break;
        case techset::CommonShaderValueType::CODE_SAMPLER:
            argValue.codeSampler = static_cast<decltype(MaterialArgumentDef::codeSampler)>(commonArg.m_value.code_sampler_source);
            break;
        case techset::CommonShaderValueType::MATERIAL_CONST:
        case techset::CommonShaderValueType::MATERIAL_SAMPLER:
            argValue.nameHash = static_cast<decltype(MaterialArgumentDef::nameHash)>(commonArg.m_value.name_hash);
            break;
        case techset::CommonShaderValueType::COUNT:
            assert(false);
            break;
        }
    }

    void ConvertMaterialArgs(MaterialPass& pass, const techset::CommonPass& commonPass, MemoryManager& memory, AssetCreationContext& context)
    {
        pass.args = memory.Alloc<MaterialShaderArgument>(commonPass.m_args.size());

        const auto frequencyCount = commonPass.GetFrequencyCounts(commonCodeSourceInfos);

        pass.perObjArgCount = static_cast<unsigned char>(frequencyCount[std::to_underlying(techset::CommonCodeSourceUpdateFrequency::PER_OBJECT)]);
        pass.perPrimArgCount = static_cast<unsigned char>(frequencyCount[std::to_underlying(techset::CommonCodeSourceUpdateFrequency::PER_PRIM)]);
        pass.stableArgCount = static_cast<unsigned char>(frequencyCount[std::to_underlying(techset::CommonCodeSourceUpdateFrequency::RARELY)]);

        const auto commonArgCount = commonPass.m_args.size();
        for (auto argIndex = 0u; argIndex < commonArgCount; argIndex++)
        {
            auto& arg = pass.args[argIndex];
            const auto& commonArg = commonPass.m_args[argIndex];

            arg.type = static_cast<decltype(MaterialShaderArgument::type)>(ConvertArgumentType(commonArg.m_type));

#if defined(IS_DX9)
            arg.dest = static_cast<decltype(MaterialShaderArgument::dest)>(commonArg.m_destination.dx9.m_destination_register);
#else
            arg.size = static_cast<decltype(MaterialShaderArgument::size)>(commonArg.m_destination.dx11.m_size);
            arg.buffer = static_cast<decltype(MaterialShaderArgument::buffer)>(commonArg.m_destination.dx11.m_buffer);

            if (techset::IsConstValueType(commonArg.m_type.m_value_type))
            {
                arg.location.offset = static_cast<decltype(MaterialArgumentLocation::offset)>(commonArg.m_destination.dx11.m_location.constant_buffer_offset);
            }
            else
            {
                assert(techset::IsSamplerValueType(commonArg.m_type.m_value_type));
                arg.location.textureIndex =
                    static_cast<decltype(MaterialArgumentLocation::textureIndex)>(commonArg.m_destination.dx11.m_location.texture_index);
                arg.location.samplerIndex =
                    static_cast<decltype(MaterialArgumentLocation::samplerIndex)>(commonArg.m_destination.dx11.m_location.sampler_index);
            }
#endif

            ConvertArgumentValue(arg.u, commonArg, memory, context);
        }
    }

    void ConvertVertexDecl(MaterialPass& pass, const techset::CommonVertexDeclaration& commonDecl, AssetCreationContext& context)
    {
        std::ostringstream nameStream;
        for (const auto& entry : commonDecl.m_routing)
        {
            nameStream << commonRoutingInfos.GetSourceAbbreviation(entry.m_source);
            nameStream << commonRoutingInfos.GetDestinationAbbreviation(entry.m_destination);
        }

        const std::string declName(nameStream.str());
#if defined(SHADERS_ARE_SUBASSETS)
        auto* vertexDeclAsset = context.LoadSubAsset<SubAssetVertexDecl>(declName);
#else
        auto* vertexDeclAsset = context.LoadDependency<AssetVertexDecl>(declName);
#endif
        assert(vertexDeclAsset);
        pass.vertexDecl = vertexDeclAsset ? vertexDeclAsset->Asset() : nullptr;
    }

    void ConvertMaterialPass(MaterialPass& pass, const techset::CommonPass& commonPass, AssetCreationContext& context, MemoryManager& memory)
    {
        ConvertVertexDecl(pass, commonPass.m_vertex_declaration, context);

        if (!commonPass.m_vertex_shader.m_name.empty())
        {
#if defined(SHADERS_ARE_SUBASSETS)
            auto* vertexShaderAsset = context.LoadSubAsset<SubAssetVertexShader>(commonPass.m_vertex_shader.m_name);
#else
            auto* vertexShaderAsset = context.LoadDependency<AssetVertexShader>(commonPass.m_vertex_shader.m_name);
#endif
            assert(vertexShaderAsset);
            pass.vertexShader = vertexShaderAsset ? vertexShaderAsset->Asset() : nullptr;
        }

        if (!commonPass.m_pixel_shader.m_name.empty())
        {
#if defined(SHADERS_ARE_SUBASSETS)
            auto* pixelShaderAsset = context.LoadSubAsset<SubAssetPixelShader>(commonPass.m_pixel_shader.m_name);
#else
            auto* pixelShaderAsset = context.LoadDependency<AssetPixelShader>(commonPass.m_pixel_shader.m_name);
#endif
            assert(pixelShaderAsset);
            pass.pixelShader = pixelShaderAsset ? pixelShaderAsset->Asset() : nullptr;
        }

        ConvertMaterialArgs(pass, commonPass, memory, context);
        pass.customSamplerFlags = static_cast<decltype(MaterialPass::customSamplerFlags)>(commonPass.m_sampler_flags);
    }

#if defined(FEATURE_IW4) || defined(FEATURE_IW5)
    // Not sure if this is actually how this is calculated.
    // It produces identical results at least though.
    constexpr MaterialConstantSource ALLOWED_PIXEL_CONSTANTS_FOR_FLAG_200[]{
        CONST_SRC_CODE_RENDER_TARGET_SIZE,
        CONST_SRC_CODE_VIEWPORT_DIMENSIONS,
    };

    bool ShouldApplyFlag200(const MaterialTechnique& technique)
    {
        for (auto passIndex = 0u; passIndex < technique.passCount; passIndex++)
        {
            const auto& pass = technique.passArray[passIndex];
            if (!pass.args)
                continue;

            const unsigned argCount = pass.perPrimArgCount + pass.perObjArgCount + pass.stableArgCount;
            for (auto argIndex = 0u; argIndex < argCount; argIndex++)
            {
                const auto& arg = pass.args[argIndex];
                if (arg.type == MTL_ARG_MATERIAL_VERTEX_CONST || arg.type == MTL_ARG_MATERIAL_PIXEL_SAMPLER || arg.type == MTL_ARG_MATERIAL_PIXEL_CONST)
                    return false;

                if (arg.type == MTL_ARG_CODE_PIXEL_CONST)
                {
                    const auto foundAllowedConstant = std::ranges::find(ALLOWED_PIXEL_CONSTANTS_FOR_FLAG_200, arg.u.codeConst.index);
                    if (foundAllowedConstant == std::end(ALLOWED_PIXEL_CONSTANTS_FOR_FLAG_200))
                        return false;
                }
            }
        }

        return true;
    }
#endif

    bool AnyDeclHasOptionalSource(const MaterialTechnique& technique, AssetCreationContext& context)
    {
        for (auto passIndex = 0u; passIndex < technique.passCount; passIndex++)
        {
            const auto& pass = technique.passArray[passIndex];
            if (!pass.vertexDecl)
                continue;

#if defined(SHADERS_ARE_SUBASSETS)
            if (pass.vertexDecl->hasOptionalSource)
                return true;
#else
            if (pass.vertexDecl->name && pass.vertexDecl->name[0] == ',')
            {
                if (techset::HasOptionalSourceByName(&pass.vertexDecl->name[1], commonRoutingInfos).value_or(false))
                    return true;
            }
            else if (pass.vertexDecl->hasOptionalSource)
                return true;
#endif
        }

        return false;
    }

    void UpdateTechniqueFlags(MaterialTechnique& technique, const techset::CommonTechnique& commonTechnique, AssetCreationContext& context)
    {
        std::string lowerTechniqueName(commonTechnique.m_name);
        utils::MakeStringLowerCase(lowerTechniqueName);

#if defined(FEATURE_IW3)
        if (lowerTechniqueName == "zprepass")
            technique.flags |= MTL_TECHFLAG_ZPREPASS;
#elif defined(FEATURE_IW4) || defined(FEATURE_IW5)
        // Not a particularly cool way to do this but...
        // the game actually does this :shrug:
        if (lowerTechniqueName == "zprepass")
            technique.flags |= MTL_TECHFLAG_ZPREPASS;
        else if (lowerTechniqueName == "build_floatz")
            technique.flags |= MTL_TECHFLAG_BUILD_FLOATZ;
        else if (lowerTechniqueName == "build_shadowmap_depth" || lowerTechniqueName == "build_shadowmap_model")
            technique.flags |= MTL_TECHFLAG_BUILD_SHADOW_MAP_DEPTH_OR_MODEL;

        if (technique.flags & MTL_TECHFLAG_USES_FLOATZ && lowerTechniqueName.starts_with("distortion_"))
            technique.flags = (technique.flags & ~MTL_TECHFLAG_USES_FLOATZ) | MTL_TECHFLAG_USES_DISTORTION_FLOATZ;

        if (ShouldApplyFlag200(technique))
            technique.flags |= TECHNIQUE_FLAG_200;
#elif defined(FEATURE_T5)
        // Not a particularly cool way to do this but...
        // the game actually does this :shrug:
        if (lowerTechniqueName == "zprepass" || lowerTechniqueName.starts_with("pimp_technique_zprepass_")
            || lowerTechniqueName.starts_with("pimp_technique_layer_zprepass_"))
        {
            technique.flags |= MTL_TECHFLAG_ZPREPASS;
        }
#elif defined(FEATURE_T6)
        // Not a particularly cool way to do this but...
        // the game actually does this :shrug:
        if (lowerTechniqueName == "zprepass" || lowerTechniqueName.starts_with("pimp_technique_zprepass_")
            || lowerTechniqueName.starts_with("pimp_technique_layer_zprepass_") || lowerTechniqueName.starts_with("pimp_technique_buildshadowmap_"))
        {
            technique.flags |= MTL_TECHFLAG_ZPREPASS;
        }
#endif

        if (AnyDeclHasOptionalSource(technique, context))
            technique.flags |= MTL_TECHFLAG_DECL_HAS_OPTIONAL_SOURCE;
    }

    MaterialTechnique* ConvertTechnique(const techset::CommonTechnique& commonTechnique, AssetCreationContext& context, MemoryManager& memory)
    {
        const auto additionalPassCount = std::max(commonTechnique.m_passes.size(), 1uz) - 1uz;
        auto* technique = static_cast<MaterialTechnique*>(memory.AllocRaw(sizeof(MaterialTechnique) + additionalPassCount * sizeof(MaterialPass)));

        const auto passCount = static_cast<decltype(MaterialTechnique::passCount)>(commonTechnique.m_passes.size());

        technique->name = memory.Dup(commonTechnique.m_name.c_str());

        technique->passCount = passCount;

        for (auto passIndex = 0u; passIndex < passCount; passIndex++)
            ConvertMaterialPass(technique->passArray[passIndex], commonTechnique.m_passes[passIndex], context, memory);

        // Take common flags and apply further logic
        technique->flags = static_cast<decltype(MaterialTechnique::flags)>(commonTechnique.m_flags);
        UpdateTechniqueFlags(*technique, commonTechnique, context);

        return technique;
    }

#if defined(FEATURE_T5) || defined(FEATURE_T6)
    void ApplyTechFlagsFromMaterial(const Material& material, const Zone& zone)
    {
        if (!material.techniqueSet || !material.techniqueSet->name || !material.stateBitsTable)
            return;

        // Find the techniqueset asset from our zone since the material may link to a different one
        const auto techniqueSetAsset = zone.m_pools.GetAsset<AssetTechniqueSet>(material.techniqueSet->name);
        if (!techniqueSetAsset)
            return;

        for (auto techType = 0u; techType < TECHNIQUE_COUNT; techType++)
        {
            auto* technique = techniqueSetAsset->Asset()->techniques[techType];
            const auto stateBitsEntry = material.stateBitsEntry[techType];

            if (!technique || stateBitsEntry < 0 || static_cast<decltype(Material::stateBitsCount)>(stateBitsEntry) >= material.stateBitsCount)
                continue;

            const auto& stateBits = material.stateBitsTable[static_cast<decltype(Material::stateBitsCount)>(stateBitsEntry)].loadBits;
            const bool shouldSetFlag80 = stateBits.structured.depthTestDisabled == 0 || stateBits.structured.depthWrite > 0
                                         || stateBits.structured.depthTest > 0 || stateBits.structured.polygonOffset > 0;
            if (shouldSetFlag80)
                technique->flags |= TECHNIQUE_FLAG_80;
        }
    }
#endif

    class SHADER_LOADER_CLASS_NAME final : public techset::ITechniqueShaderLoader
    {
    public:
        explicit SHADER_LOADER_CLASS_NAME(AssetCreationContext& context)
            : m_context(context)
        {
        }

        std::optional<techset::CommonTechniqueShaderBin> LoadVertexShader(const std::string& name) override
        {
#if defined(SHADERS_ARE_SUBASSETS)
            auto* shaderAsset = m_context.LoadSubAsset<SubAssetVertexShader>(name);
#else
            auto* shaderAsset = m_context.ForceLoadDependency<AssetVertexShader>(name);
#endif
            if (!shaderAsset)
                return std::nullopt;

            const auto* shader = shaderAsset->Asset();
            assert(shader->prog.loadDef.program && shader->prog.loadDef.programSize > 0);
            if (!shader->prog.loadDef.program || shader->prog.loadDef.programSize == 0)
                return std::nullopt;

            return techset::CommonTechniqueShaderBin{
                .m_shader_bin = shader->prog.loadDef.program,
#if defined(IS_DX9)
                .m_shader_bin_size =
                    static_cast<size_t>(shader->prog.loadDef.programSize) * sizeof(std::remove_pointer_t<decltype(GfxVertexShaderLoadDef::program)>),
#else
                .m_shader_bin_size = shader->prog.loadDef.programSize,
#endif
            };
        }

        std::optional<techset::CommonTechniqueShaderBin> LoadPixelShader(const std::string& name) override
        {
#if defined(SHADERS_ARE_SUBASSETS)
            auto* shaderAsset = m_context.LoadSubAsset<SubAssetPixelShader>(name);
#else
            auto* shaderAsset = m_context.ForceLoadDependency<AssetPixelShader>(name);
#endif
            if (!shaderAsset)
                return std::nullopt;

            const auto* shader = shaderAsset->Asset();
            assert(shader->prog.loadDef.program && shader->prog.loadDef.programSize > 0);
            if (!shader->prog.loadDef.program || shader->prog.loadDef.programSize == 0)
                return std::nullopt;

            return techset::CommonTechniqueShaderBin{
                .m_shader_bin = shader->prog.loadDef.program,
#if defined(IS_DX9)
                .m_shader_bin_size =
                    static_cast<size_t>(shader->prog.loadDef.programSize) * sizeof(std::remove_pointer_t<decltype(GfxPixelShaderLoadDef::program)>),
#else
                .m_shader_bin_size = shader->prog.loadDef.programSize,
#endif
            };
        }

    private:
        AssetCreationContext& m_context;
    };

    class COMPILER_CLASS_NAME final : public SubAssetCreator<SubAssetTechnique>
    {
    public:
        COMPILER_CLASS_NAME(MemoryManager& memory, Zone& zone, ISearchPath& searchPath)
            : m_memory(memory),
              m_zone(zone),
              m_search_path(searchPath)
        {
        }

        AssetCreationResult CreateSubAsset(const std::string& subAssetName, AssetCreationContext& context) override
        {
            bool failure = false;
            SHADER_LOADER_CLASS_NAME shaderLoader(context);
#if defined(IS_DX9)
            const auto commonShaderArgCreator = techset::CommonShaderArgCreator::CreateDx9(shaderLoader, context, commonCodeSourceInfos);
#else
            const auto commonShaderArgCreator = techset::CommonShaderArgCreator::CreateDx11(shaderLoader, context, commonCodeSourceInfos);
#endif

            const auto commonTechnique =
                techset::LoadCommonTechnique(subAssetName, commonCodeSourceInfos, commonRoutingInfos, *commonShaderArgCreator, m_search_path, failure);

            if (!commonTechnique)
                return failure ? AssetCreationResult::Failure() : AssetCreationResult::NoAction();

            auto* convertedTechnique = ConvertTechnique(*commonTechnique, context, m_memory);
            assert(convertedTechnique);

            return AssetCreationResult::Success(context.AddSubAsset(AssetRegistration<SubAssetTechnique>(subAssetName, convertedTechnique)));
        }

        void FinalizeZone(AssetCreationContext& context) override
        {
#if defined(FEATURE_T5) || defined(FEATURE_T6)
            const auto materials = m_zone.m_pools.PoolAssets<AssetMaterial>();
            for (auto* materialAsset : materials)
            {
                ApplyTechFlagsFromMaterial(*materialAsset->Asset(), m_zone);
            }
#endif

#if defined(FEATURE_T6)
            const auto techniques = context.PoolSubAssets<SubAssetTechnique>();
            for (auto* techniqueSubAsset : techniques)
            {
                auto& technique = *techniqueSubAsset->Asset();
                for (auto passIndex = 0u; passIndex < technique.passCount; passIndex++)
                {
                    ApplyPrecompiledIndex(technique.passArray[passIndex]);
                }
            }
#endif
        }

    private:
        MemoryManager& m_memory;
        Zone& m_zone;
        ISearchPath& m_search_path;
    };
} // namespace

#set CREATE_COMPILER_METHOD "CreateTechniqueCompiler" + GAME

namespace techset
{
    std::unique_ptr<ISubAssetCreator> CREATE_COMPILER_METHOD(MemoryManager& memory, Zone& zone, ISearchPath& searchPath)
    {
        return std::make_unique<COMPILER_CLASS_NAME>(memory, zone, searchPath);
    }
} // namespace techset