// VulkanWindow: render passes, pipeline layout, shader modules and the pipeline cache keyed by D3D8 state.
#include "vulkan_window.h"

#include "host_util.h"
#include "blend_state.h"

#include <cstdlib>
#include <cstring>
#include <stdexcept>

namespace mt_host {

namespace {

std::vector<std::uint32_t> read_spirv(const char* name) {
    std::size_t size = 0;
    void* bytes = SDL_LoadFile(bundled_file_path(name).c_str(), &size);
#ifdef __ANDROID__
    if (!bytes) bytes = SDL_LoadFile((std::string("assets://") + name).c_str(), &size);
    if (!bytes) bytes = SDL_LoadFile(name, &size);
#endif
    if (!bytes) throw std::runtime_error(std::string("cannot open bundled shader: ") + name);
    if (size == 0 || size % 4) {
        SDL_free(bytes);
        throw std::runtime_error(std::string("invalid SPIR-V size: ") + name);
    }
    std::vector<std::uint32_t> words(size / 4);
    std::memcpy(words.data(), bytes, size);
    SDL_free(bytes);
    return words;
}

} // namespace

void VulkanWindow::create_render_pass_and_layout() {
    const bool msaa = samples_ != VK_SAMPLE_COUNT_1_BIT;
    VkAttachmentDescription attachments[3]{};
    attachments[0].format = swapchain_format_; attachments[0].samples = samples_;
    attachments[0].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
    attachments[0].storeOp = msaa ? VK_ATTACHMENT_STORE_OP_DONT_CARE : VK_ATTACHMENT_STORE_OP_STORE;
    attachments[0].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    attachments[0].finalLayout = msaa ? VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL : VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
    // MSAA resolve target: the swapchain image.
    attachments[2].format = swapchain_format_; attachments[2].samples = VK_SAMPLE_COUNT_1_BIT;
    attachments[2].loadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE; attachments[2].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    attachments[2].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED; attachments[2].finalLayout = VK_IMAGE_LAYOUT_PRESENT_SRC_KHR;
    attachments[1].format = VK_FORMAT_D32_SFLOAT; attachments[1].samples = samples_;
    attachments[1].loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR; attachments[1].storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    attachments[1].initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    attachments[1].finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;

    VkAttachmentReference color_ref{0, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL};
    VkAttachmentReference depth_ref{1, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL};
    VkAttachmentReference resolve_ref{2, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL};
    VkSubpassDescription subpass{};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1; subpass.pColorAttachments = &color_ref;
    if (msaa) subpass.pResolveAttachments = &resolve_ref;
    subpass.pDepthStencilAttachment = &depth_ref;

    VkSubpassDependency dependency{};
    dependency.srcSubpass = VK_SUBPASS_EXTERNAL; dependency.dstSubpass = 0;
    dependency.srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
    dependency.dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT;
    dependency.dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;

    VkRenderPassCreateInfo pass_info{VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO};
    pass_info.attachmentCount = msaa ? 3 : 2; pass_info.pAttachments = attachments;
    pass_info.subpassCount = 1; pass_info.pSubpasses = &subpass;
    pass_info.dependencyCount = 1; pass_info.pDependencies = &dependency;
    check(vkCreateRenderPass(device_, &pass_info, nullptr, &pass_), "vkCreateRenderPass");
    create_framebuffers();

    const auto vert = read_spirv("native.vert.spv");
    const auto frag = read_spirv("native.frag.spv");
    auto module = [&](const std::vector<std::uint32_t>& code) {
        VkShaderModuleCreateInfo info{VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO};
        info.codeSize = code.size() * sizeof(std::uint32_t); info.pCode = code.data();
        VkShaderModule handle = VK_NULL_HANDLE;
        check(vkCreateShaderModule(device_, &info, nullptr, &handle), "vkCreateShaderModule");
        return handle;
    };
    vertex_module_ = module(vert);
    fragment_module_ = module(frag);

    VkDescriptorSetLayout set_layouts[2] = {descriptor_layout_, bone_descriptor_layout_};
    VkPushConstantRange push_constants{};
    push_constants.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
    push_constants.size = sizeof(PushConstants);
    VkPipelineLayoutCreateInfo layout_info{VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO};
    layout_info.setLayoutCount = 2; layout_info.pSetLayouts = set_layouts;
    layout_info.pushConstantRangeCount = 1; layout_info.pPushConstantRanges = &push_constants;
    check(vkCreatePipelineLayout(device_, &layout_info, nullptr, &layout_), "vkCreatePipelineLayout");

    get_pipeline(0, 2, 0);
    create_offscreen_pass();
    create_scene_pass();
}

void VulkanWindow::create_offscreen_pass() {
    VkFormatProperties properties{};
    vkGetPhysicalDeviceFormatProperties(physical_, VK_FORMAT_R5G6B5_UNORM_PACK16, &properties);
    const VkFormatFeatureFlags wanted =
        VK_FORMAT_FEATURE_COLOR_ATTACHMENT_BIT | VK_FORMAT_FEATURE_SAMPLED_IMAGE_BIT;
    offscreen_format_ = (properties.optimalTilingFeatures & wanted) == wanted
        ? VK_FORMAT_R5G6B5_UNORM_PACK16 : VK_FORMAT_R8G8B8A8_UNORM;
    VkAttachmentDescription attachments[2]{};
    attachments[0].format = offscreen_format_; attachments[0].samples = VK_SAMPLE_COUNT_1_BIT;
    attachments[0].loadOp = VK_ATTACHMENT_LOAD_OP_LOAD; attachments[0].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    attachments[0].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    attachments[0].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    attachments[0].initialLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    attachments[0].finalLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    attachments[1].format = VK_FORMAT_D32_SFLOAT; attachments[1].samples = VK_SAMPLE_COUNT_1_BIT;
    attachments[1].loadOp = VK_ATTACHMENT_LOAD_OP_LOAD; attachments[1].storeOp = VK_ATTACHMENT_STORE_OP_STORE;
    attachments[1].stencilLoadOp = VK_ATTACHMENT_LOAD_OP_DONT_CARE;
    attachments[1].stencilStoreOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
    attachments[1].initialLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
    attachments[1].finalLayout = VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL;
    VkAttachmentReference color_ref{0, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL};
    VkAttachmentReference depth_ref{1, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL};
    VkSubpassDescription subpass{};
    subpass.pipelineBindPoint = VK_PIPELINE_BIND_POINT_GRAPHICS;
    subpass.colorAttachmentCount = 1; subpass.pColorAttachments = &color_ref;
    subpass.pDepthStencilAttachment = &depth_ref;
    VkSubpassDependency dependencies[2]{};
    // Earlier work on the queue -- the previous frame sampling the texture, the previous
    // offscreen pass, the first-use clear -- completes before this pass writes it.
    dependencies[0].srcSubpass = VK_SUBPASS_EXTERNAL; dependencies[0].dstSubpass = 0;
    dependencies[0].srcStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT |
        VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT | VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT |
        VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_TRANSFER_BIT;
    dependencies[0].srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT |
        VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT | VK_ACCESS_TRANSFER_WRITE_BIT;
    dependencies[0].dstStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT |
        VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT | VK_PIPELINE_STAGE_LATE_FRAGMENT_TESTS_BIT;
    dependencies[0].dstAccessMask = VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT |
        VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
    // The back-buffer pass samples the result.
    dependencies[1].srcSubpass = 0; dependencies[1].dstSubpass = VK_SUBPASS_EXTERNAL;
    dependencies[1].srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
    dependencies[1].srcAccessMask = VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
    dependencies[1].dstStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    dependencies[1].dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
    VkRenderPassCreateInfo pass_info{VK_STRUCTURE_TYPE_RENDER_PASS_CREATE_INFO};
    pass_info.attachmentCount = 2; pass_info.pAttachments = attachments;
    pass_info.subpassCount = 1; pass_info.pSubpasses = &subpass;
    pass_info.dependencyCount = 2; pass_info.pDependencies = dependencies;
    check(vkCreateRenderPass(device_, &pass_info, nullptr, &offscreen_pass_), "vkCreateRenderPass offscreen");
}

VkPipeline VulkanWindow::get_pipeline(std::uint8_t cull, std::uint8_t depth, std::uint8_t blend,
                        bool lines, std::uint8_t z_func, bool offscreen, const Render3DDraw* draw) {
    std::array<std::uint32_t, 17> key{};
    key[0] = std::uint32_t(cull) | (std::uint32_t(depth) << 8) |
                              (std::uint32_t(blend) << 16) | (lines ? (1u << 24) : 0u) |
                              (std::uint32_t(z_func & 0xfu) << 25) | (offscreen ? (1u << 29) : 0u);
    static const bool specialize = std::getenv("MT_GENERIC_SHADERS") == nullptr;
    key[1] = specialize ? 1u : 0u;
    if (specialize && draw) {
        key[2] = 1u;
        for (int stage = 0; stage < 2; ++stage) {
            const int base = 3 + stage * 6;
            key[base] = draw->color_op[stage];
            key[base + 1] = draw->color_arg1[stage];
            key[base + 2] = draw->color_arg2[stage];
            key[base + 3] = draw->alpha_op[stage];
            key[base + 4] = draw->alpha_arg1[stage];
            key[base + 5] = draw->alpha_arg2[stage];
        }
        key[15] = draw->alpha_test;
        key[16] = draw->alpha_test ? draw->alpha_func : 0u;
    }
    if (auto it = pipelines_.find(key); it != pipelines_.end()) return it->second;

    VkPipelineShaderStageCreateInfo stages[2]{};
    stages[0].sType = stages[1].sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
    stages[0].stage = VK_SHADER_STAGE_VERTEX_BIT; stages[0].module = vertex_module_; stages[0].pName = "main";
    stages[1].stage = VK_SHADER_STAGE_FRAGMENT_BIT; stages[1].module = fragment_module_; stages[1].pName = "main";
    std::array<VkSpecializationMapEntry, 16> entries{};
    for (std::uint32_t i = 0; i < entries.size(); ++i)
        entries[i] = {i, i * std::uint32_t(sizeof(std::uint32_t)), sizeof(std::uint32_t)};
    VkSpecializationInfo specialization{};
    specialization.mapEntryCount = entries.size();
    specialization.pMapEntries = entries.data();
    specialization.dataSize = 16 * sizeof(std::uint32_t);
    specialization.pData = key.data() + 1;
    stages[1].pSpecializationInfo = &specialization;

    VkVertexInputBindingDescription binding{0, sizeof(Vertex), VK_VERTEX_INPUT_RATE_VERTEX};
    VkVertexInputAttributeDescription attributes[9] = {
        {0, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, position)},
        {1, 0, VK_FORMAT_R32G32B32_SFLOAT, offsetof(Vertex, normal)},
        {2, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(Vertex, uv)},
        {3, 0, VK_FORMAT_R32G32B32A32_SFLOAT, offsetof(Vertex, color)},
        {4, 0, VK_FORMAT_R8G8B8A8_UINT, offsetof(Vertex, joints)},
        {5, 0, VK_FORMAT_R32G32B32A32_SFLOAT, offsetof(Vertex, weights)},
        {6, 0, VK_FORMAT_R32G32_SFLOAT, offsetof(Vertex, mask_uv)},
        {7, 0, VK_FORMAT_R32_SFLOAT, offsetof(Vertex, rhw)},
        {8, 0, VK_FORMAT_R32_SFLOAT, offsetof(Vertex, vertex_fog)}};
    VkPipelineVertexInputStateCreateInfo input{VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO};
    input.vertexBindingDescriptionCount = 1; input.pVertexBindingDescriptions = &binding;
    input.vertexAttributeDescriptionCount = 9; input.pVertexAttributeDescriptions = attributes;

    VkPipelineInputAssemblyStateCreateInfo assembly{VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO};
    assembly.topology = lines ? VK_PRIMITIVE_TOPOLOGY_LINE_LIST : VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;

    VkViewport viewport{0, 0, float(extent_.width), float(extent_.height), 0, 1};
    // An offscreen target is at most the device's 4096 (GetDeviceCaps); the viewport clips inside it.
    VkRect2D scissor{{0, 0}, offscreen ? VkExtent2D{4096, 4096} : extent_};
    VkPipelineViewportStateCreateInfo viewport_info{VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO};
    viewport_info.viewportCount = 1; viewport_info.pViewports = &viewport;
    viewport_info.scissorCount = 1; viewport_info.pScissors = &scissor;
    // D3DVIEWPORT8 changes per draw (SetViewport), so the viewport is dynamic state.
    const VkDynamicState dynamic_states[] = {VK_DYNAMIC_STATE_VIEWPORT};
    VkPipelineDynamicStateCreateInfo dynamic_info{VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO};
    dynamic_info.dynamicStateCount = 1; dynamic_info.pDynamicStates = dynamic_states;

    VkPipelineRasterizationStateCreateInfo raster{VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO};
    raster.polygonMode = VK_POLYGON_MODE_FILL;
    raster.cullMode = cull == 1 ? VK_CULL_MODE_BACK_BIT : (cull == 2 ? VK_CULL_MODE_FRONT_BIT : VK_CULL_MODE_NONE);
    raster.frontFace = VK_FRONT_FACE_CLOCKWISE;
    raster.lineWidth = 1;

    VkPipelineMultisampleStateCreateInfo multisample{VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO};
    multisample.rasterizationSamples = offscreen ? VK_SAMPLE_COUNT_1_BIT : samples_;

    VkPipelineDepthStencilStateCreateInfo depth_stencil{VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO};
    depth_stencil.depthTestEnable = depth > 0 ? VK_TRUE : VK_FALSE;
    depth_stencil.depthWriteEnable = depth > 1 ? VK_TRUE : VK_FALSE;
    switch (z_func) {
        case 1: depth_stencil.depthCompareOp = VK_COMPARE_OP_NEVER; break;
        case 2: depth_stencil.depthCompareOp = VK_COMPARE_OP_LESS; break;
        case 3: depth_stencil.depthCompareOp = VK_COMPARE_OP_EQUAL; break;
        case 5: depth_stencil.depthCompareOp = VK_COMPARE_OP_GREATER; break;
        case 6: depth_stencil.depthCompareOp = VK_COMPARE_OP_NOT_EQUAL; break;
        case 7: depth_stencil.depthCompareOp = VK_COMPARE_OP_GREATER_OR_EQUAL; break;
        case 8: depth_stencil.depthCompareOp = VK_COMPARE_OP_ALWAYS; break;
        default: depth_stencil.depthCompareOp = VK_COMPARE_OP_LESS_OR_EQUAL; break;
    }

    auto d3d_to_vk_blend = [](std::uint8_t d3d, VkBlendFactor fallback) -> VkBlendFactor {
        switch (d3d_blend_for_rgb_target(d3d)) {
            case 1: return VK_BLEND_FACTOR_ZERO;
            case 2: return VK_BLEND_FACTOR_ONE;
            case 3: return VK_BLEND_FACTOR_SRC_COLOR;
            case 4: return VK_BLEND_FACTOR_ONE_MINUS_SRC_COLOR;
            case 5: return VK_BLEND_FACTOR_SRC_ALPHA;
            case 6: return VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
            case 7: return VK_BLEND_FACTOR_DST_ALPHA;
            case 8: return VK_BLEND_FACTOR_ONE_MINUS_DST_ALPHA;
            case 9: return VK_BLEND_FACTOR_DST_COLOR;
            case 10: return VK_BLEND_FACTOR_ONE_MINUS_DST_COLOR;
            case 11: return VK_BLEND_FACTOR_SRC_ALPHA_SATURATE;
            default: return fallback;
        }
    };

    VkPipelineColorBlendAttachmentState blend_attachment{};
    blend_attachment.colorWriteMask = 0xf;
    if (blend > 0) {
        blend_attachment.blendEnable = VK_TRUE;
        if (blend >= 16) {
            blend_attachment.srcColorBlendFactor =
                d3d_to_vk_blend(blend & 0xFu, VK_BLEND_FACTOR_SRC_ALPHA);
            blend_attachment.dstColorBlendFactor =
                d3d_to_vk_blend((blend >> 4) & 0xFu, VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA);
        } else {
            blend_attachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
            blend_attachment.dstColorBlendFactor =
                blend == 2 ? VK_BLEND_FACTOR_ONE : VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
        }
        blend_attachment.colorBlendOp = VK_BLEND_OP_ADD;
        // D3D8 has no separate alpha blend: SRCBLEND/DESTBLEND apply to alpha as well.
        auto alpha_factor = [](VkBlendFactor factor) {
            switch (factor) {
                case VK_BLEND_FACTOR_SRC_COLOR: return VK_BLEND_FACTOR_SRC_ALPHA;
                case VK_BLEND_FACTOR_ONE_MINUS_SRC_COLOR: return VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
                case VK_BLEND_FACTOR_DST_COLOR: return VK_BLEND_FACTOR_DST_ALPHA;
                case VK_BLEND_FACTOR_ONE_MINUS_DST_COLOR: return VK_BLEND_FACTOR_ONE_MINUS_DST_ALPHA;
                case VK_BLEND_FACTOR_SRC_ALPHA_SATURATE: return VK_BLEND_FACTOR_ONE;
                default: return factor;
            }
        };
        blend_attachment.srcAlphaBlendFactor = alpha_factor(blend_attachment.srcColorBlendFactor);
        blend_attachment.dstAlphaBlendFactor = alpha_factor(blend_attachment.dstColorBlendFactor);
        blend_attachment.alphaBlendOp = VK_BLEND_OP_ADD;
    }
    VkPipelineColorBlendStateCreateInfo blending{VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO};
    blending.attachmentCount = 1; blending.pAttachments = &blend_attachment;

    VkGraphicsPipelineCreateInfo pipeline_info{VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO};
    pipeline_info.stageCount = 2; pipeline_info.pStages = stages;
    pipeline_info.pVertexInputState = &input; pipeline_info.pInputAssemblyState = &assembly;
    pipeline_info.pViewportState = &viewport_info; pipeline_info.pRasterizationState = &raster;
    pipeline_info.pMultisampleState = &multisample; pipeline_info.pDepthStencilState = &depth_stencil;
    pipeline_info.pColorBlendState = &blending;
    pipeline_info.pDynamicState = &dynamic_info;
    pipeline_info.layout = layout_; pipeline_info.renderPass = offscreen ? offscreen_pass_ : pass_;
    VkPipeline handle = VK_NULL_HANDLE;
    check(vkCreateGraphicsPipelines(device_, VK_NULL_HANDLE, 1, &pipeline_info, nullptr, &handle), "vkCreateGraphicsPipelines");
    pipelines_.emplace(key, handle);
    return handle;
}

} // namespace mt_host