#include "ShaderArtifactLoader.h"

#include "ShaderFileUtils.h"
#include "ShaderRuntimeBuildUtils.h"

#include <XCEngine/Core/Asset/ArtifactFormats.h>

#include <cstring>
#include <memory>

namespace XCEngine {
namespace Resources {

namespace {

template<typename T>
bool ReadShaderArtifactValue(const Containers::Array<Core::uint8>& data, size_t& offset, T& outValue) {
    if (offset + sizeof(T) > data.Size()) {
        return false;
    }

    std::memcpy(&outValue, data.Data() + offset, sizeof(T));
    offset += sizeof(T);
    return true;
}

bool ReadShaderArtifactString(
    const Containers::Array<Core::uint8>& data,
    size_t& offset,
    Containers::String& outValue) {
    Core::uint32 length = 0;
    if (!ReadShaderArtifactValue(data, offset, length)) {
        return false;
    }

    if (length == 0) {
        outValue.Clear();
        return true;
    }

    if (offset + length > data.Size()) {
        return false;
    }

    outValue = Containers::String(
        std::string(reinterpret_cast<const char*>(data.Data() + offset), length).c_str());
    offset += length;
    return true;
}

} // namespace

LoadResult LoadShaderArtifact(const Containers::String& path) {
    const Containers::Array<Core::uint8> data = ReadShaderFileData(path);
    if (data.Empty()) {
        return LoadResult("Failed to read shader artifact: " + path);
    }

    size_t offset = 0;
    ShaderArtifactFileHeader fileHeader;
    if (!ReadShaderArtifactValue(data, offset, fileHeader)) {
        return LoadResult("Failed to parse shader artifact header: " + path);
    }

    const std::string magic(fileHeader.magic, fileHeader.magic + 7);
    const bool isCurrentSchema =
        magic == "XCSHD05" && fileHeader.schemaVersion == kShaderArtifactSchemaVersion;
    if (!isCurrentSchema) {
        return LoadResult("Invalid shader artifact header: " + path);
    }

    auto shader = std::make_unique<Shader>();

    Containers::String shaderName;
    Containers::String shaderSourcePath;
    Containers::String shaderFallback;
    if (!ReadShaderArtifactString(data, offset, shaderName) ||
        !ReadShaderArtifactString(data, offset, shaderSourcePath) ||
        !ReadShaderArtifactString(data, offset, shaderFallback)) {
        return LoadResult("Failed to parse shader artifact strings: " + path);
    }

    shader->m_name = shaderName.Empty() ? path : shaderName;
    shader->m_path = shaderSourcePath.Empty() ? path : shaderSourcePath;
    shader->m_guid = ResourceGUID::Generate(shader->m_path);
    shader->SetFallback(shaderFallback);

    ShaderArtifactHeader shaderHeader;
    if (!ReadShaderArtifactValue(data, offset, shaderHeader)) {
        return LoadResult("Failed to parse shader artifact body: " + path);
    }

    for (Core::uint32 propertyIndex = 0; propertyIndex < shaderHeader.propertyCount; ++propertyIndex) {
        ShaderPropertyDesc property = {};
        ShaderPropertyArtifact propertyArtifact;
        if (!ReadShaderArtifactString(data, offset, property.name) ||
            !ReadShaderArtifactString(data, offset, property.displayName) ||
            !ReadShaderArtifactString(data, offset, property.defaultValue) ||
            !ReadShaderArtifactString(data, offset, property.semantic) ||
            !ReadShaderArtifactValue(data, offset, propertyArtifact)) {
            return LoadResult("Failed to read shader artifact properties: " + path);
        }

        property.type = static_cast<ShaderPropertyType>(propertyArtifact.propertyType);
        shader->AddProperty(property);
    }

    for (Core::uint32 passIndex = 0; passIndex < shaderHeader.passCount; ++passIndex) {
        Containers::String passName;
        Core::uint32 tagCount = 0;
        Core::uint32 resourceCount = 0;
        Core::uint32 keywordDeclarationCount = 0;
        Core::uint32 variantCount = 0;
        Core::uint32 hasFixedFunctionState = 0;
        MaterialRenderState fixedFunctionState = {};
        if (!ReadShaderArtifactString(data, offset, passName)) {
            return LoadResult("Failed to read shader artifact passes: " + path);
        }

        ShaderPassArtifactHeaderV5 passHeader = {};
        if (!ReadShaderArtifactValue(data, offset, passHeader)) {
            return LoadResult("Failed to read shader artifact passes: " + path);
        }

        tagCount = passHeader.tagCount;
        resourceCount = passHeader.resourceCount;
        keywordDeclarationCount = passHeader.keywordDeclarationCount;
        variantCount = passHeader.variantCount;
        hasFixedFunctionState = passHeader.hasFixedFunctionState;
        fixedFunctionState = passHeader.fixedFunctionState;

        ShaderPass pass = {};
        pass.name = passName;
        pass.hasFixedFunctionState = hasFixedFunctionState != 0u;
        pass.fixedFunctionState = fixedFunctionState;
        shader->AddPass(pass);

        for (Core::uint32 tagIndex = 0; tagIndex < tagCount; ++tagIndex) {
            Containers::String tagName;
            Containers::String tagValue;
            if (!ReadShaderArtifactString(data, offset, tagName) ||
                !ReadShaderArtifactString(data, offset, tagValue)) {
                return LoadResult("Failed to read shader artifact pass tags: " + path);
            }

            shader->SetPassTag(passName, tagName, tagValue);
        }

        for (Core::uint32 resourceIndex = 0; resourceIndex < resourceCount; ++resourceIndex) {
            ShaderResourceBindingDesc binding = {};
            ShaderResourceArtifact resourceArtifact;
            if (!ReadShaderArtifactString(data, offset, binding.name) ||
                !ReadShaderArtifactString(data, offset, binding.semantic) ||
                !ReadShaderArtifactValue(data, offset, resourceArtifact)) {
                return LoadResult("Failed to read shader artifact pass resources: " + path);
            }

            binding.type = static_cast<ShaderResourceType>(resourceArtifact.resourceType);
            binding.set = resourceArtifact.set;
            binding.binding = resourceArtifact.binding;
            shader->AddPassResourceBinding(passName, binding);
        }

        for (Core::uint32 declarationIndex = 0; declarationIndex < keywordDeclarationCount; ++declarationIndex) {
            ShaderKeywordDeclaration declaration = {};
            ShaderKeywordDeclarationArtifactHeader declarationHeader = {};
            if (!ReadShaderArtifactValue(data, offset, declarationHeader)) {
                return LoadResult("Failed to read shader artifact pass keywords: " + path);
            }

            declaration.type =
                static_cast<ShaderKeywordDeclarationType>(declarationHeader.declarationType);
            declaration.options.Reserve(declarationHeader.optionCount);
            for (Core::uint32 optionIndex = 0; optionIndex < declarationHeader.optionCount; ++optionIndex) {
                Containers::String option;
                if (!ReadShaderArtifactString(data, offset, option)) {
                    return LoadResult("Failed to read shader artifact keyword options: " + path);
                }

                declaration.options.PushBack(option);
            }

            shader->AddPassKeywordDeclaration(passName, declaration);
        }

        for (Core::uint32 variantIndex = 0; variantIndex < variantCount; ++variantIndex) {
            ShaderStageVariant variant = {};
            Core::uint64 compiledBinarySize = 0;
            Core::uint32 keywordCount = 0;
            ShaderVariantArtifactHeader variantHeader = {};
            if (!ReadShaderArtifactValue(data, offset, variantHeader)) {
                return LoadResult("Failed to read shader artifact variants: " + path);
            }

            variant.stage = static_cast<ShaderType>(variantHeader.stage);
            variant.language = static_cast<ShaderLanguage>(variantHeader.language);
            variant.backend = static_cast<ShaderBackend>(variantHeader.backend);
            keywordCount = variantHeader.keywordCount;
            compiledBinarySize = variantHeader.compiledBinarySize;

            if (!ReadShaderArtifactString(data, offset, variant.entryPoint) ||
                !ReadShaderArtifactString(data, offset, variant.profile) ||
                !ReadShaderArtifactString(data, offset, variant.sourceCode)) {
                return LoadResult("Failed to read shader artifact variants: " + path);
            }

            for (Core::uint32 keywordIndex = 0; keywordIndex < keywordCount; ++keywordIndex) {
                Containers::String keyword;
                if (!ReadShaderArtifactString(data, offset, keyword)) {
                    return LoadResult("Failed to read shader artifact variant keywords: " + path);
                }

                variant.requiredKeywords.enabledKeywords.PushBack(keyword);
            }
            NormalizeShaderKeywordSetInPlace(variant.requiredKeywords);

            if (compiledBinarySize > 0) {
                if (offset + compiledBinarySize > data.Size()) {
                    return LoadResult("Shader artifact variant binary payload is truncated: " + path);
                }

                variant.compiledBinary.Resize(static_cast<size_t>(compiledBinarySize));
                std::memcpy(
                    variant.compiledBinary.Data(),
                    data.Data() + offset,
                    static_cast<size_t>(compiledBinarySize));
                offset += static_cast<size_t>(compiledBinarySize);
            }

            shader->AddPassVariant(passName, variant);
        }
    }

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

} // namespace Resources
} // namespace XCEngine