#include "vulkan_backend.h"
#include <logger.h>
#include <format>

bool VulkanTextureBackend::Initialize(IUnityInterfaces* unityInterfaces)
{
    if (m_Initialized)
    {
        return true;
    }

    m_UnityVulkan = unityInterfaces->Get<IUnityGraphicsVulkan>();
    if (!m_UnityVulkan)
    {
        Logger::Log("VulkanTextureBackend: IUnityGraphicsVulkan not available");
        return false;
    }

    // Get Unity's Vulkan instance
    UnityVulkanInstance vkInstance = m_UnityVulkan->Instance();
    m_Instance = vkInstance.instance;
    m_PhysicalDevice = vkInstance.physicalDevice;
    m_Device = vkInstance.device;
    m_QueueFamilyIndex = vkInstance.queueFamilyIndex;

    // Get a queue from Unity
    vkGetDeviceQueue(m_Device, m_QueueFamilyIndex, 0, &m_Queue);

    // Load extension function
    vkGetAndroidHardwareBufferPropertiesANDROID =
        reinterpret_cast<PFN_vkGetAndroidHardwareBufferPropertiesANDROID>(
            vkGetInstanceProcAddr(m_Instance, "vkGetAndroidHardwareBufferPropertiesANDROID")
        );

    if (!vkGetAndroidHardwareBufferPropertiesANDROID)
    {
        Logger::Log("VulkanTextureBackend: vkGetAndroidHardwareBufferPropertiesANDROID not available");
        return false;
    }

    // Create command pool
    VkCommandPoolCreateInfo poolInfo{};
    poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
    poolInfo.queueFamilyIndex = m_QueueFamilyIndex;
    poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;

    if (vkCreateCommandPool(m_Device, &poolInfo, nullptr, &m_CommandPool) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to create command pool");
        return false;
    }

    // Allocate command buffer
    VkCommandBufferAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = m_CommandPool;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = 1;

    if (vkAllocateCommandBuffers(m_Device, &allocInfo, &m_CommandBuffer) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to allocate command buffer");
        vkDestroyCommandPool(m_Device, m_CommandPool, nullptr);
        m_CommandPool = VK_NULL_HANDLE;
        return false;
    }

    m_Initialized = true;
    Logger::Log("VulkanTextureBackend: Initialized successfully");
    return true;
}

bool VulkanTextureBackend::Create(AHardwareBuffer* buffer)
{
    if (!m_Initialized)
    {
        Logger::Log("VulkanTextureBackend: Not initialized");
        return false;
    }

    if (m_Created)
    {
        Destroy();
    }

    // Get buffer dimensions
    AHardwareBuffer_Desc desc{};
    AHardwareBuffer_describe(buffer, &desc);
    m_Width = static_cast<int>(desc.width);
    m_Height = static_cast<int>(desc.height);

    Logger::Log(std::format("VulkanTextureBackend: Creating {}x{} texture", m_Width, m_Height));

    if (!ImportHardwareBuffer(buffer))
    {
        Logger::Log("VulkanTextureBackend: Failed to import hardware buffer");
        return false;
    }

    if (!CreateLocalImage())
    {
        Logger::Log("VulkanTextureBackend: Failed to create local image");
        DestroyImportedResources();
        return false;
    }

    m_Created = true;

    // Initial copy
    CopyImage();

    Logger::Log("VulkanTextureBackend: Created successfully");
    return true;
}

bool VulkanTextureBackend::ImportHardwareBuffer(AHardwareBuffer* buffer)
{
    // Query hardware buffer properties
    VkAndroidHardwareBufferFormatPropertiesANDROID formatProps{};
    formatProps.sType = VK_STRUCTURE_TYPE_ANDROID_HARDWARE_BUFFER_FORMAT_PROPERTIES_ANDROID;

    VkAndroidHardwareBufferPropertiesANDROID props{};
    props.sType = VK_STRUCTURE_TYPE_ANDROID_HARDWARE_BUFFER_PROPERTIES_ANDROID;
    props.pNext = &formatProps;

    if (vkGetAndroidHardwareBufferPropertiesANDROID(m_Device, buffer, &props) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to get hardware buffer properties");
        return false;
    }

    m_Format = formatProps.format;
    Logger::Log(std::format("VulkanTextureBackend: Buffer format: {}", static_cast<int>(m_Format)));

    // Create VkImage with external memory
    VkExternalMemoryImageCreateInfo extMemImageInfo{};
    extMemImageInfo.sType = VK_STRUCTURE_TYPE_EXTERNAL_MEMORY_IMAGE_CREATE_INFO;
    extMemImageInfo.handleTypes = VK_EXTERNAL_MEMORY_HANDLE_TYPE_ANDROID_HARDWARE_BUFFER_BIT_ANDROID;

    VkImageCreateInfo imageInfo{};
    imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageInfo.pNext = &extMemImageInfo;
    imageInfo.imageType = VK_IMAGE_TYPE_2D;
    imageInfo.format = m_Format;
    imageInfo.extent = {static_cast<uint32_t>(m_Width), static_cast<uint32_t>(m_Height), 1};
    imageInfo.mipLevels = 1;
    imageInfo.arrayLayers = 1;
    imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
    imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageInfo.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT;
    imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;

    if (vkCreateImage(m_Device, &imageInfo, nullptr, &m_ImportedImage) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to create imported image");
        return false;
    }

    // Import memory from hardware buffer
    VkImportAndroidHardwareBufferInfoANDROID importInfo{};
    importInfo.sType = VK_STRUCTURE_TYPE_IMPORT_ANDROID_HARDWARE_BUFFER_INFO_ANDROID;
    importInfo.buffer = buffer;

    VkMemoryDedicatedAllocateInfo dedicatedInfo{};
    dedicatedInfo.sType = VK_STRUCTURE_TYPE_MEMORY_DEDICATED_ALLOCATE_INFO;
    dedicatedInfo.pNext = &importInfo;
    dedicatedInfo.image = m_ImportedImage;

    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.pNext = &dedicatedInfo;
    allocInfo.allocationSize = props.allocationSize;
    allocInfo.memoryTypeIndex = FindMemoryType(props.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

    if (vkAllocateMemory(m_Device, &allocInfo, nullptr, &m_ImportedMemory) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to allocate imported memory");
        vkDestroyImage(m_Device, m_ImportedImage, nullptr);
        m_ImportedImage = VK_NULL_HANDLE;
        return false;
    }

    if (vkBindImageMemory(m_Device, m_ImportedImage, m_ImportedMemory, 0) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to bind imported memory");
        DestroyImportedResources();
        return false;
    }

    Logger::Log("VulkanTextureBackend: Hardware buffer imported");
    return true;
}

bool VulkanTextureBackend::CreateLocalImage()
{
    // Create local VkImage for safe rendering
    VkImageCreateInfo imageInfo{};
    imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageInfo.imageType = VK_IMAGE_TYPE_2D;
    imageInfo.format = m_Format;
    imageInfo.extent = {static_cast<uint32_t>(m_Width), static_cast<uint32_t>(m_Height), 1};
    imageInfo.mipLevels = 1;
    imageInfo.arrayLayers = 1;
    imageInfo.samples = VK_SAMPLE_COUNT_1_BIT;
    imageInfo.tiling = VK_IMAGE_TILING_OPTIMAL;
    imageInfo.usage = VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_DST_BIT;
    imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;
    imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;

    if (vkCreateImage(m_Device, &imageInfo, nullptr, &m_LocalImage) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to create local image");
        return false;
    }

    // Allocate memory for local image
    VkMemoryRequirements memReqs;
    vkGetImageMemoryRequirements(m_Device, m_LocalImage, &memReqs);

    VkMemoryAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize = memReqs.size;
    allocInfo.memoryTypeIndex = FindMemoryType(memReqs.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);

    if (vkAllocateMemory(m_Device, &allocInfo, nullptr, &m_LocalMemory) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to allocate local memory");
        vkDestroyImage(m_Device, m_LocalImage, nullptr);
        m_LocalImage = VK_NULL_HANDLE;
        return false;
    }

    if (vkBindImageMemory(m_Device, m_LocalImage, m_LocalMemory, 0) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to bind local memory");
        DestroyLocalResources();
        return false;
    }

    // Create image view
    VkImageViewCreateInfo viewInfo{};
    viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    viewInfo.image = m_LocalImage;
    viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
    viewInfo.format = m_Format;
    viewInfo.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    viewInfo.subresourceRange.baseMipLevel = 0;
    viewInfo.subresourceRange.levelCount = 1;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount = 1;

    if (vkCreateImageView(m_Device, &viewInfo, nullptr, &m_LocalImageView) != VK_SUCCESS)
    {
        Logger::Log("VulkanTextureBackend: Failed to create image view");
        DestroyLocalResources();
        return false;
    }

    Logger::Log("VulkanTextureBackend: Local image created");
    return true;
}

void VulkanTextureBackend::Update()
{
    if (m_Created)
    {
        CopyImage();
    }
}

void VulkanTextureBackend::CopyImage()
{
    // Begin command buffer
    VkCommandBufferBeginInfo beginInfo{};
    beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
    beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

    vkResetCommandBuffer(m_CommandBuffer, 0);
    vkBeginCommandBuffer(m_CommandBuffer, &beginInfo);

    // Transition imported image to transfer src
    TransitionImageLayout(m_ImportedImage,
        VK_IMAGE_LAYOUT_UNDEFINED,
        VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);

    // Transition local image to transfer dst
    TransitionImageLayout(m_LocalImage,
        VK_IMAGE_LAYOUT_UNDEFINED,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

    // Copy image
    VkImageCopy copyRegion{};
    copyRegion.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    copyRegion.srcSubresource.layerCount = 1;
    copyRegion.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    copyRegion.dstSubresource.layerCount = 1;
    copyRegion.extent = {static_cast<uint32_t>(m_Width), static_cast<uint32_t>(m_Height), 1};

    vkCmdCopyImage(
        m_CommandBuffer,
        m_ImportedImage, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
        m_LocalImage, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        1, &copyRegion
    );

    // Transition local image to shader read
    TransitionImageLayout(m_LocalImage,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);

    vkEndCommandBuffer(m_CommandBuffer);

    // Submit and wait (CPU sync for simplicity)
    VkSubmitInfo submitInfo{};
    submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
    submitInfo.commandBufferCount = 1;
    submitInfo.pCommandBuffers = &m_CommandBuffer;

    vkQueueSubmit(m_Queue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(m_Queue);
}

void VulkanTextureBackend::TransitionImageLayout(VkImage image, VkImageLayout oldLayout, VkImageLayout newLayout)
{
    VkImageMemoryBarrier barrier{};
    barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.oldLayout = oldLayout;
    barrier.newLayout = newLayout;
    barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
    barrier.image = image;
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    barrier.subresourceRange.baseMipLevel = 0;
    barrier.subresourceRange.levelCount = 1;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount = 1;

    VkPipelineStageFlags srcStage;
    VkPipelineStageFlags dstStage;

    if (oldLayout == VK_IMAGE_LAYOUT_UNDEFINED)
    {
        barrier.srcAccessMask = 0;
        srcStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    }
    else if (oldLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
    {
        barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        srcStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    }
    else
    {
        barrier.srcAccessMask = 0;
        srcStage = VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT;
    }

    if (newLayout == VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL)
    {
        barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        dstStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    }
    else if (newLayout == VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL)
    {
        barrier.dstAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
        dstStage = VK_PIPELINE_STAGE_TRANSFER_BIT;
    }
    else if (newLayout == VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
    {
        barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
        dstStage = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
    }
    else
    {
        barrier.dstAccessMask = 0;
        dstStage = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
    }

    vkCmdPipelineBarrier(
        m_CommandBuffer,
        srcStage, dstStage,
        0,
        0, nullptr,
        0, nullptr,
        1, &barrier
    );
}

uint32_t VulkanTextureBackend::FindMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties)
{
    VkPhysicalDeviceMemoryProperties memProps;
    vkGetPhysicalDeviceMemoryProperties(m_PhysicalDevice, &memProps);

    for (uint32_t i = 0; i < memProps.memoryTypeCount; i++)
    {
        if ((typeFilter & (1 << i)) &&
            (memProps.memoryTypes[i].propertyFlags & properties) == properties)
        {
            return i;
        }
    }

    Logger::Log("VulkanTextureBackend: Failed to find suitable memory type");
    return 0;
}

void VulkanTextureBackend::Destroy()
{
    if (m_Device && m_Created)
    {
        vkDeviceWaitIdle(m_Device);
    }

    DestroyLocalResources();
    DestroyImportedResources();

    m_Width = 0;
    m_Height = 0;
    m_Created = false;
}

void VulkanTextureBackend::DestroyImportedResources()
{
    if (m_Device)
    {
        if (m_ImportedMemory != VK_NULL_HANDLE)
        {
            vkFreeMemory(m_Device, m_ImportedMemory, nullptr);
            m_ImportedMemory = VK_NULL_HANDLE;
        }
        if (m_ImportedImage != VK_NULL_HANDLE)
        {
            vkDestroyImage(m_Device, m_ImportedImage, nullptr);
            m_ImportedImage = VK_NULL_HANDLE;
        }
    }
}

void VulkanTextureBackend::DestroyLocalResources()
{
    if (m_Device)
    {
        if (m_LocalImageView != VK_NULL_HANDLE)
        {
            vkDestroyImageView(m_Device, m_LocalImageView, nullptr);
            m_LocalImageView = VK_NULL_HANDLE;
        }
        if (m_LocalMemory != VK_NULL_HANDLE)
        {
            vkFreeMemory(m_Device, m_LocalMemory, nullptr);
            m_LocalMemory = VK_NULL_HANDLE;
        }
        if (m_LocalImage != VK_NULL_HANDLE)
        {
            vkDestroyImage(m_Device, m_LocalImage, nullptr);
            m_LocalImage = VK_NULL_HANDLE;
        }
    }
}

void* VulkanTextureBackend::GetNativeTexturePtr()
{
    // For Vulkan, Unity expects the VkImage
    return reinterpret_cast<void*>(m_LocalImage);
}