#include "XCEngine/RHI/Vulkan/VulkanPipelineState.h"

#include "XCEngine/RHI/Vulkan/VulkanDevice.h"

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <map>
#include <vector>

namespace XCEngine {
namespace RHI {

namespace {

std::string NarrowAscii(const std::wstring& value) {
    std::string result;
    result.reserve(value.size());
    for (wchar_t ch : value) {
        result.push_back(static_cast<char>(ch));
    }
    return result;
}

bool LoadSpirvBytes(const ShaderCompileDesc& desc, std::vector<uint32_t>& words, std::string& entryPoint) {
    entryPoint = NarrowAscii(desc.entryPoint);
    if (entryPoint.empty()) {
        entryPoint = "main";
    }

    if (desc.sourceLanguage != ShaderLanguage::SPIRV) {
        return false;
    }

    std::vector<char> bytes;
    if (!desc.source.empty()) {
        bytes.assign(desc.source.begin(), desc.source.end());
    } else if (!desc.fileName.empty()) {
        std::ifstream file(std::filesystem::path(desc.fileName), std::ios::binary | std::ios::ate);
        if (!file.is_open()) {
            return false;
        }

        const std::streamsize fileSize = file.tellg();
        if (fileSize <= 0 || (fileSize % 4) != 0) {
            return false;
        }

        bytes.resize(static_cast<size_t>(fileSize));
        file.seekg(0, std::ios::beg);
        if (!file.read(bytes.data(), fileSize)) {
            return false;
        }
    } else {
        return false;
    }

    if ((bytes.size() % sizeof(uint32_t)) != 0) {
        return false;
    }

    words.resize(bytes.size() / sizeof(uint32_t));
    std::memcpy(words.data(), bytes.data(), bytes.size());
    return !words.empty();
}

VkShaderModule CreateShaderModule(VkDevice device, const std::vector<uint32_t>& words) {
    VkShaderModuleCreateInfo createInfo = {};
    createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
    createInfo.codeSize = words.size() * sizeof(uint32_t);
    createInfo.pCode = words.data();

    VkShaderModule module = VK_NULL_HANDLE;
    if (vkCreateShaderModule(device, &createInfo, nullptr, &module) != VK_SUCCESS) {
        return VK_NULL_HANDLE;
    }

    return module;
}

} // namespace

VulkanPipelineState::~VulkanPipelineState() {
    Shutdown();
}

bool VulkanPipelineState::Initialize(VulkanDevice* device, const GraphicsPipelineDesc& desc) {
    if (device == nullptr || device->GetDevice() == VK_NULL_HANDLE) {
        return false;
    }

    m_deviceOwner = device;
    m_device = device->GetDevice();
    m_inputLayoutDesc = desc.inputLayout;
    m_rasterizerDesc = desc.rasterizerState;
    m_blendDesc = desc.blendState;
    m_depthStencilDesc = desc.depthStencilState;
    m_topologyType = desc.topologyType;
    m_renderTargetCount = desc.renderTargetCount;
    m_depthStencilFormat = desc.depthStencilFormat;
    m_sampleCount = desc.sampleCount > 0 ? desc.sampleCount : 1;
    for (uint32_t i = 0; i < 8; ++i) {
        m_renderTargetFormats[i] = desc.renderTargetFormats[i];
    }

    if (m_renderTargetCount != 1 || m_renderTargetFormats[0] == 0) {
        return false;
    }

    std::vector<uint32_t> vertexWords;
    std::vector<uint32_t> fragmentWords;
    std::string vertexEntryPoint;
    std::string fragmentEntryPoint;
    if (!LoadSpirvBytes(desc.vertexShader, vertexWords, vertexEntryPoint) ||
        !LoadSpirvBytes(desc.fragmentShader, fragmentWords, fragmentEntryPoint)) {
        return false;
    }

    const VkShaderModule vertexModule = CreateShaderModule(m_device, vertexWords);
    const VkShaderModule fragmentModule = CreateShaderModule(m_device, fragmentWords);
    if (vertexModule == VK_NULL_HANDLE || fragmentModule == VK_NULL_HANDLE) {
        if (vertexModule != VK_NULL_HANDLE) {
            vkDestroyShaderModule(m_device, vertexModule, nullptr);
        }
        if (fragmentModule != VK_NULL_HANDLE) {
            vkDestroyShaderModule(m_device, fragmentModule, nullptr);
        }
        return false;
    }

    VkPipelineLayoutCreateInfo pipelineLayoutInfo = {};
    pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
    if (vkCreatePipelineLayout(m_device, &pipelineLayoutInfo, nullptr, &m_pipelineLayout) != VK_SUCCESS) {
        vkDestroyShaderModule(m_device, fragmentModule, nullptr);
        vkDestroyShaderModule(m_device, vertexModule, nullptr);
        return false;
    }

    VkAttachmentDescription colorAttachment = {};
    colorAttachment.format = ToVulkanFormat(static_cast<Format>(m_renderTargetFormats[0]));
    colorAttachment.samples = ToVulkanSampleCount(m_sampleCount);
    colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_LOAD;
    colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    colorAttachment.stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    colorAttachment.stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    colorAttachment.initialLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
    colorAttachment.finalLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    VkAttachmentReference colorAttachmentRef = {};
    colorAttachmentRef.attachment = 0;
    colorAttachmentRef.layout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

    VkSubpassDescription subpass = {};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1;
    subpass.pColorAttachments = &colorAttachmentRef;

    VkRenderPassCreateInfo renderPassInfo = {};
    renderPassInfo.sType = VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO;
    renderPassInfo.attachmentCount = 1;
    renderPassInfo.pAttachments = &colorAttachment;
    renderPassInfo.subpassCount = 1;
    renderPassInfo.pSubpasses = &subpass;

    if (vkCreateRenderPass(m_device, &renderPassInfo, nullptr, &m_renderPass) != VK_SUCCESS) {
        vkDestroyPipelineLayout(m_device, m_pipelineLayout, nullptr);
        m_pipelineLayout = VK_NULL_HANDLE;
        vkDestroyShaderModule(m_device, fragmentModule, nullptr);
        vkDestroyShaderModule(m_device, vertexModule, nullptr);
        return false;
    }

    std::map<uint32_t, uint32_t> strideBySlot;
    for (const InputElementDesc& element : m_inputLayoutDesc.elements) {
        const uint32_t attributeSize = GetFormatSize(static_cast<Format>(element.format));
        strideBySlot[element.inputSlot] = (std::max)(strideBySlot[element.inputSlot], element.alignedByteOffset + attributeSize);
    }

    std::vector<VkVertexInputBindingDescription> bindings;
    bindings.reserve(strideBySlot.size());
    for (const auto& entry : strideBySlot) {
        VkVertexInputBindingDescription binding = {};
        binding.binding = entry.first;
        binding.stride = entry.second;
        binding.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;
        bindings.push_back(binding);
    }

    std::vector<VkVertexInputAttributeDescription> attributes;
    attributes.reserve(m_inputLayoutDesc.elements.size());
    for (uint32_t location = 0; location < m_inputLayoutDesc.elements.size(); ++location) {
        const InputElementDesc& element = m_inputLayoutDesc.elements[location];
        VkVertexInputAttributeDescription attribute = {};
        attribute.location = location;
        attribute.binding = element.inputSlot;
        attribute.format = ToVulkanFormat(static_cast<Format>(element.format));
        attribute.offset = element.alignedByteOffset;
        attributes.push_back(attribute);
    }

    VkPipelineShaderStageCreateInfo shaderStages[2] = {};
    shaderStages[0].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    shaderStages[0].stage = VK_SHADER_STAGE_VERTEX_BIT;
    shaderStages[0].module = vertexModule;
    shaderStages[0].pName = vertexEntryPoint.c_str();

    shaderStages[1].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    shaderStages[1].stage = VK_SHADER_STAGE_FRAGMENT_BIT;
    shaderStages[1].module = fragmentModule;
    shaderStages[1].pName = fragmentEntryPoint.c_str();

    VkPipelineVertexInputStateCreateInfo vertexInputInfo = {};
    vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
    vertexInputInfo.vertexBindingDescriptionCount = static_cast<uint32_t>(bindings.size());
    vertexInputInfo.pVertexBindingDescriptions = bindings.data();
    vertexInputInfo.vertexAttributeDescriptionCount = static_cast<uint32_t>(attributes.size());
    vertexInputInfo.pVertexAttributeDescriptions = attributes.data();

    VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
    inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
    inputAssembly.topology = ToVulkanPrimitiveTopology(static_cast<PrimitiveTopologyType>(m_topologyType));
    inputAssembly.primitiveRestartEnable = VK_FALSE;

    VkPipelineViewportStateCreateInfo viewportState = {};
    viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
    viewportState.viewportCount = 1;
    viewportState.scissorCount = 1;

    VkPipelineRasterizationStateCreateInfo rasterizer = {};
    rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
    rasterizer.depthClampEnable = m_rasterizerDesc.depthClipEnable ? VK_FALSE : VK_TRUE;
    rasterizer.rasterizerDiscardEnable = VK_FALSE;
    rasterizer.polygonMode = ToVulkanPolygonMode(static_cast<FillMode>(m_rasterizerDesc.fillMode));
    rasterizer.lineWidth = 1.0f;
    rasterizer.cullMode = ToVulkanCullMode(static_cast<CullMode>(m_rasterizerDesc.cullMode));
    rasterizer.frontFace = ToVulkanFrontFace(static_cast<FrontFace>(m_rasterizerDesc.frontFace));
    rasterizer.depthBiasEnable = m_rasterizerDesc.depthBias != 0 || m_rasterizerDesc.slopeScaledDepthBias != 0.0f;
    rasterizer.depthBiasConstantFactor = static_cast<float>(m_rasterizerDesc.depthBias);
    rasterizer.depthBiasClamp = m_rasterizerDesc.depthBiasClamp;
    rasterizer.depthBiasSlopeFactor = m_rasterizerDesc.slopeScaledDepthBias;

    VkPipelineMultisampleStateCreateInfo multisampling = {};
    multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
    multisampling.rasterizationSamples = ToVulkanSampleCount(m_sampleCount);
    multisampling.sampleShadingEnable = VK_FALSE;

    VkPipelineColorBlendAttachmentState colorBlendAttachment = {};
    colorBlendAttachment.colorWriteMask =
        VK_COLOR_COMPONENT_R_BIT |
        VK_COLOR_COMPONENT_G_BIT |
        VK_COLOR_COMPONENT_B_BIT |
        VK_COLOR_COMPONENT_A_BIT;
    colorBlendAttachment.blendEnable = m_blendDesc.blendEnable ? VK_TRUE : VK_FALSE;
    colorBlendAttachment.srcColorBlendFactor = ToVulkanBlendFactor(static_cast<BlendFactor>(m_blendDesc.srcBlend));
    colorBlendAttachment.dstColorBlendFactor = ToVulkanBlendFactor(static_cast<BlendFactor>(m_blendDesc.dstBlend));
    colorBlendAttachment.colorBlendOp = ToVulkanBlendOp(static_cast<BlendOp>(m_blendDesc.blendOp));
    colorBlendAttachment.srcAlphaBlendFactor = ToVulkanBlendFactor(static_cast<BlendFactor>(m_blendDesc.srcBlendAlpha));
    colorBlendAttachment.dstAlphaBlendFactor = ToVulkanBlendFactor(static_cast<BlendFactor>(m_blendDesc.dstBlendAlpha));
    colorBlendAttachment.alphaBlendOp = ToVulkanBlendOp(static_cast<BlendOp>(m_blendDesc.blendOpAlpha));

    VkPipelineColorBlendStateCreateInfo colorBlending = {};
    colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
    colorBlending.logicOpEnable = VK_FALSE;
    colorBlending.attachmentCount = 1;
    colorBlending.pAttachments = &colorBlendAttachment;

    VkPipelineDepthStencilStateCreateInfo depthStencil = {};
    depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
    depthStencil.depthTestEnable = m_depthStencilDesc.depthTestEnable ? VK_TRUE : VK_FALSE;
    depthStencil.depthWriteEnable = m_depthStencilDesc.depthWriteEnable ? VK_TRUE : VK_FALSE;
    depthStencil.depthCompareOp = ToVulkanCompareOp(static_cast<ComparisonFunc>(m_depthStencilDesc.depthFunc));
    depthStencil.depthBoundsTestEnable = m_depthStencilDesc.depthBoundsEnable ? VK_TRUE : VK_FALSE;
    depthStencil.stencilTestEnable = m_depthStencilDesc.stencilEnable ? VK_TRUE : VK_FALSE;

    VkDynamicState dynamicStates[] = {
        VK_DYNAMIC_STATE_VIEWPORT,
        VK_DYNAMIC_STATE_SCISSOR
    };
    VkPipelineDynamicStateCreateInfo dynamicState = {};
    dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
    dynamicState.dynamicStateCount = static_cast<uint32_t>(std::size(dynamicStates));
    dynamicState.pDynamicStates = dynamicStates;

    VkGraphicsPipelineCreateInfo pipelineInfo = {};
    pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
    pipelineInfo.stageCount = 2;
    pipelineInfo.pStages = shaderStages;
    pipelineInfo.pVertexInputState = &vertexInputInfo;
    pipelineInfo.pInputAssemblyState = &inputAssembly;
    pipelineInfo.pViewportState = &viewportState;
    pipelineInfo.pRasterizationState = &rasterizer;
    pipelineInfo.pMultisampleState = &multisampling;
    pipelineInfo.pDepthStencilState = &depthStencil;
    pipelineInfo.pColorBlendState = &colorBlending;
    pipelineInfo.pDynamicState = &dynamicState;
    pipelineInfo.layout = m_pipelineLayout;
    pipelineInfo.renderPass = m_renderPass;
    pipelineInfo.subpass = 0;

    const bool success = vkCreateGraphicsPipelines(
        m_device,
        VK_NULL_HANDLE,
        1,
        &pipelineInfo,
        nullptr,
        &m_pipeline) == VK_SUCCESS;

    vkDestroyShaderModule(m_device, fragmentModule, nullptr);
    vkDestroyShaderModule(m_device, vertexModule, nullptr);

    if (!success) {
        Shutdown();
        return false;
    }

    return true;
}

void VulkanPipelineState::SetInputLayout(const InputLayoutDesc& layout) {
    m_inputLayoutDesc = layout;
}

void VulkanPipelineState::SetRasterizerState(const RasterizerDesc& state) {
    m_rasterizerDesc = state;
}

void VulkanPipelineState::SetBlendState(const BlendDesc& state) {
    m_blendDesc = state;
}

void VulkanPipelineState::SetDepthStencilState(const DepthStencilStateDesc& state) {
    m_depthStencilDesc = state;
}

void VulkanPipelineState::SetTopology(uint32_t topologyType) {
    m_topologyType = topologyType;
}

void VulkanPipelineState::SetRenderTargetFormats(uint32_t count, const uint32_t* formats, uint32_t depthFormat) {
    m_renderTargetCount = count;
    m_depthStencilFormat = depthFormat;
    for (uint32_t i = 0; i < count && i < 8; ++i) {
        m_renderTargetFormats[i] = formats[i];
    }
}

void VulkanPipelineState::SetSampleCount(uint32_t count) {
    m_sampleCount = count > 0 ? count : 1;
}

void VulkanPipelineState::SetComputeShader(RHIShader* shader) {
    (void)shader;
}

PipelineStateHash VulkanPipelineState::GetHash() const {
    PipelineStateHash hash = {};
    hash.topologyHash = m_topologyType;
    hash.renderTargetHash = m_renderTargetCount ^ (m_renderTargetFormats[0] << 8);
    return hash;
}

void VulkanPipelineState::Shutdown() {
    if (m_pipeline != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkDestroyPipeline(m_device, m_pipeline, nullptr);
        m_pipeline = VK_NULL_HANDLE;
    }

    if (m_renderPass != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkDestroyRenderPass(m_device, m_renderPass, nullptr);
        m_renderPass = VK_NULL_HANDLE;
    }

    if (m_pipelineLayout != VK_NULL_HANDLE && m_device != VK_NULL_HANDLE) {
        vkDestroyPipelineLayout(m_device, m_pipelineLayout, nullptr);
        m_pipelineLayout = VK_NULL_HANDLE;
    }

    m_deviceOwner = nullptr;
    m_device = VK_NULL_HANDLE;
}

} // namespace RHI
} // namespace XCEngine