#include "ShaderRuntimeBuildUtils.h"

#include "../ShaderAuthoringParser.h"
#include "ShaderFileUtils.h"
#include "../ShaderSourceUtils.h"

#include <memory>

namespace XCEngine {
namespace Resources {

Containers::String GetDefaultEntryPoint(ShaderLanguage language, ShaderType stage) {
    if (language == ShaderLanguage::HLSL) {
        switch (stage) {
        case ShaderType::Vertex: return "VSMain";
        case ShaderType::Fragment: return "PSMain";
        case ShaderType::Geometry: return "GSMain";
        case ShaderType::Compute: return "CSMain";
        case ShaderType::Hull: return "HSMain";
        case ShaderType::Domain: return "DSMain";
        default: return "main";
        }
    }

    return "main";
}

Containers::String GetDefaultProfile(
    ShaderLanguage language,
    ShaderBackend backend,
    ShaderType stage) {
    if (language == ShaderLanguage::HLSL) {
        switch (stage) {
        case ShaderType::Vertex: return "vs_5_0";
        case ShaderType::Fragment: return "ps_5_0";
        case ShaderType::Geometry: return "gs_5_0";
        case ShaderType::Compute: return "cs_5_0";
        case ShaderType::Hull: return "hs_5_0";
        case ShaderType::Domain: return "ds_5_0";
        default: return Containers::String();
        }
    }

    const bool isVulkan = backend == ShaderBackend::Vulkan;
    switch (stage) {
    case ShaderType::Vertex:
        return isVulkan ? "vs_4_50" : "vs_4_30";
    case ShaderType::Fragment:
        return isVulkan ? "fs_4_50" : "fs_4_30";
    case ShaderType::Geometry:
        return isVulkan ? "gs_4_50" : "gs_4_30";
    case ShaderType::Compute:
        return isVulkan ? "cs_4_50" : "cs_4_30";
    case ShaderType::Hull:
        return isVulkan ? "hs_4_50" : "hs_4_30";
    case ShaderType::Domain:
        return isVulkan ? "ds_4_50" : "ds_4_30";
    default:
        return Containers::String();
    }
}

size_t CalculateShaderMemorySize(const Shader& shader) {
    size_t memorySize =
        sizeof(Shader) + shader.GetName().Length() + shader.GetPath().Length() + shader.GetFallback().Length();
    for (const ShaderPropertyDesc& property : shader.GetProperties()) {
        memorySize += property.name.Length();
        memorySize += property.displayName.Length();
        memorySize += property.defaultValue.Length();
        memorySize += property.semantic.Length();
    }
    for (const ShaderPass& pass : shader.GetPasses()) {
        memorySize += pass.name.Length();
        for (const ShaderPassTagEntry& tag : pass.tags) {
            memorySize += tag.name.Length();
            memorySize += tag.value.Length();
        }
        for (const ShaderResourceBindingDesc& binding : pass.resources) {
            memorySize += binding.name.Length();
            memorySize += binding.semantic.Length();
        }
        for (const ShaderKeywordDeclaration& declaration : pass.keywordDeclarations) {
            for (const Containers::String& option : declaration.options) {
                memorySize += option.Length();
            }
        }
        for (const ShaderStageVariant& variant : pass.variants) {
            for (const Containers::String& keyword : variant.requiredKeywords.enabledKeywords) {
                memorySize += keyword.Length();
            }
            memorySize += variant.entryPoint.Length();
            memorySize += variant.profile.Length();
            memorySize += variant.sourceCode.Length();
            memorySize += variant.compiledBinary.Size();
        }
    }

    return memorySize;
}

LoadResult BuildShaderFromIR(
    const Containers::String& path,
    const ShaderIR& shaderIR) {
    auto shader = std::make_unique<Shader>();
    IResource::ConstructParams params;
    params.path = path;
    params.guid = ResourceGUID::Generate(path);
    params.name = shaderIR.name;
    shader->Initialize(params);
    shader->SetFallback(shaderIR.fallback);

    for (const ShaderPropertyDesc& property : shaderIR.properties) {
        shader->AddProperty(property);
    }

    for (const ShaderSubShaderIR& subShader : shaderIR.subShaders) {
        for (const ShaderPassIR& pass : subShader.passes) {
            ShaderPass shaderPass = {};
            shaderPass.name = pass.name;
            shaderPass.hasFixedFunctionState = pass.hasFixedFunctionState;
            shaderPass.fixedFunctionState = pass.fixedFunctionState;
            shader->AddPass(shaderPass);

            for (const ShaderTagIR& subShaderTag : subShader.tags) {
                shader->SetPassTag(pass.name, subShaderTag.name, subShaderTag.value);
            }
            for (const ShaderTagIR& passTag : pass.tags) {
                shader->SetPassTag(pass.name, passTag.name, passTag.value);
            }

            for (const ShaderResourceBindingDesc& resourceBinding : pass.resources) {
                shader->AddPassResourceBinding(pass.name, resourceBinding);
            }
            for (const ShaderKeywordDeclaration& keywordDeclaration : pass.keywordDeclarations) {
                shader->AddPassKeywordDeclaration(pass.name, keywordDeclaration);
            }

            if (!pass.programSource.Empty()) {
                Containers::String combinedSource;
                AppendAuthoringSourceBlock(combinedSource, shaderIR.sharedProgramSource);
                AppendAuthoringSourceBlock(combinedSource, subShader.sharedProgramSource);
                AppendAuthoringSourceBlock(combinedSource, pass.sharedProgramSource);
                AppendAuthoringSourceBlock(combinedSource, pass.programSource);
                const Containers::String strippedCombinedSource =
                    StripShaderAuthoringPragmas(combinedSource);
                const std::vector<ShaderKeywordSet> keywordSets =
                    BuildShaderKeywordVariantSets(pass.keywordDeclarations);

                for (const ShaderKeywordSet& keywordSet : keywordSets) {
                    const Containers::String variantSource =
                        BuildKeywordVariantSource(strippedCombinedSource, keywordSet);

                    ShaderStageVariant vertexVariant = {};
                    vertexVariant.stage = ShaderType::Vertex;
                    vertexVariant.backend = ShaderBackend::Generic;
                    vertexVariant.language = ShaderLanguage::HLSL;
                    vertexVariant.requiredKeywords = keywordSet;
                    vertexVariant.entryPoint =
                        !pass.vertexEntryPoint.Empty()
                            ? pass.vertexEntryPoint
                            : GetDefaultEntryPoint(ShaderLanguage::HLSL, ShaderType::Vertex);
                    vertexVariant.profile = GetDefaultProfile(
                        ShaderLanguage::HLSL,
                        ShaderBackend::Generic,
                        ShaderType::Vertex);
                    vertexVariant.sourceCode = variantSource;
                    shader->AddPassVariant(pass.name, vertexVariant);

                    ShaderStageVariant fragmentVariant = {};
                    fragmentVariant.stage = ShaderType::Fragment;
                    fragmentVariant.backend = ShaderBackend::Generic;
                    fragmentVariant.language = ShaderLanguage::HLSL;
                    fragmentVariant.requiredKeywords = keywordSet;
                    fragmentVariant.entryPoint =
                        !pass.fragmentEntryPoint.Empty()
                            ? pass.fragmentEntryPoint
                            : GetDefaultEntryPoint(ShaderLanguage::HLSL, ShaderType::Fragment);
                    fragmentVariant.profile = GetDefaultProfile(
                        ShaderLanguage::HLSL,
                        ShaderBackend::Generic,
                        ShaderType::Fragment);
                    fragmentVariant.sourceCode = variantSource;
                    shader->AddPassVariant(pass.name, fragmentVariant);
                }
            }
        }
    }

    shader->m_memorySize = CalculateShaderMemorySize(*shader);
    return LoadResult(shader.release());
}

LoadResult LoadLegacySingleStageShader(
    const Containers::String& path,
    const std::string& sourceText) {
    auto shader = std::make_unique<Shader>();
    shader->m_path = path;
    shader->m_name = path;
    shader->m_guid = ResourceGUID::Generate(path);

    const Containers::String ext = GetShaderPathExtension(path).ToLower();
    if (ext == "hlsl") {
        shader->SetShaderLanguage(ShaderLanguage::HLSL);
    } else {
        shader->SetShaderLanguage(ShaderLanguage::GLSL);
    }

    shader->SetShaderType(DetectShaderTypeFromPath(path));
    shader->SetSourceCode(sourceText.c_str());
    shader->m_isValid = true;
    shader->m_memorySize =
        sizeof(Shader) +
        shader->m_name.Length() +
        shader->m_path.Length() +
        shader->GetSourceCode().Length();

    return LoadResult(shader.release());
}

} // namespace Resources
} // namespace XCEngine