#include "RenderDevice.h"

#include "MacroUtils.h"

#include <SDL3/SDL_log.h>

#include <array>
#include <optional>
#include <span>

#include "Frame.h"
#include "GlobalMemory.h"
#include "MathUtil.h"

RenderDevice::~RenderDevice()
{
  ASSERT( !isInit() );
}

// TODO: Failure Handling
RenderDevice* CreateRenderDevice( GlobalMemory* mem, RenderDevice::CreateInfo const& createInfo )
{
  ASSERT( createInfo.window );

  volkInitialize();

  VkInstance instance;
  // Create Instance
  {
    VkApplicationInfo constexpr applicationInfo = {
      .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
      .pNext = nullptr,
      .pApplicationName = "Test",
      .applicationVersion = VK_MAKE_API_VERSION( 0, 0, 1, 0 ),
      .pEngineName = "Blaze",
      .engineVersion = VK_MAKE_API_VERSION( 0, 0, 1, 0 ),
      .apiVersion = VK_API_VERSION_1_3,
    };

    uint32_t           instanceExtensionCount;
    char const* const* instanceExtensions = SDL_Vulkan_GetInstanceExtensions( &instanceExtensionCount );

    VkInstanceCreateInfo const instanceCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .pApplicationInfo = &applicationInfo,
      .enabledLayerCount = 0,
      .ppEnabledLayerNames = nullptr,
      .enabledExtensionCount = instanceExtensionCount,
      .ppEnabledExtensionNames = instanceExtensions,
    };

    VK_CHECK( vkCreateInstance(&instanceCreateInfo, nullptr, &instance) );
    volkLoadInstance( instance );
  }

  VkSurfaceKHR surface;
  // Create Surface
  ASSERT( SDL_Vulkan_CreateSurface(createInfo.window, instance, nullptr, &surface) );

  VkPhysicalDevice        physicalDeviceInUse    = nullptr;
  VkDevice                device                 = nullptr;
  VmaAllocator            gpuAllocator           = nullptr;
  std::optional<uint32_t> directQueueFamilyIndex = std::nullopt;
  VkQueue                 directQueue            = nullptr;
  // Create Device and Queue
  {
    auto tempAllocStart = mem->getState();

    uint32_t physicalDeviceCount;
    VK_CHECK( vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, nullptr) );
    SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "Found %u GPUs", physicalDeviceCount );

    VkPhysicalDevice* physicalDevices = reinterpret_cast<VkPhysicalDevice*>(mem->allocate(
      sizeof( VkPhysicalDevice ) * physicalDeviceCount ));
    VK_CHECK( vkEnumeratePhysicalDevices(instance, &physicalDeviceCount, physicalDevices) );

    for ( VkPhysicalDevice const physicalDevice : std::span{ physicalDevices, physicalDeviceCount } )
    {
      auto tempAllocQueueProperties = mem->getState();

      VkPhysicalDeviceProperties properties;
      vkGetPhysicalDeviceProperties( physicalDevice, &properties );

      SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "GPU: %s", properties.deviceName );

      SDL_LogInfo(
        SDL_LOG_CATEGORY_GPU,
        "- API Version %d.%d.%d",
        VK_API_VERSION_MAJOR( properties.apiVersion ),
        VK_API_VERSION_MINOR( properties.apiVersion ),
        VK_API_VERSION_PATCH( properties.apiVersion ) );

      constexpr static uint32_t API_PATCH_BITS = 0xFFF;
      if ( (properties.apiVersion & (~API_PATCH_BITS)) < VK_API_VERSION_1_3 )
      {
        continue;
      }

      if ( properties.deviceType == VK_PHYSICAL_DEVICE_TYPE_CPU )
      {
        continue;
      }

      uint32_t queueFamilyCount;
      vkGetPhysicalDeviceQueueFamilyProperties( physicalDevice, &queueFamilyCount, nullptr );
      VkQueueFamilyProperties* queueFamilyProperties = reinterpret_cast<VkQueueFamilyProperties*>(mem->allocate(
        sizeof( VkQueueFamilyProperties ) * queueFamilyCount ));
      vkGetPhysicalDeviceQueueFamilyProperties( physicalDevice, &queueFamilyCount, queueFamilyProperties );

      for ( uint32_t queueFamilyIndex = 0; queueFamilyIndex != queueFamilyCount;
            ++queueFamilyIndex )
      {
        VkQueueFamilyProperties const& qfp = queueFamilyProperties[queueFamilyIndex];

        bool hasGraphicsSupport = false;
        bool hasComputeSupport  = false;
        bool hasTransferSupport = false;
        bool hasPresentSupport  = false;

        SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "- Queue [%d]", queueFamilyIndex );
        if ( qfp.queueFlags & VK_QUEUE_GRAPHICS_BIT )
        {
          hasGraphicsSupport = true;
          SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "-- Graphic" );
        }
        if ( qfp.queueFlags & VK_QUEUE_COMPUTE_BIT )
        {
          hasComputeSupport = true;
          SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "-- Compute" );
        }
        if ( qfp.queueFlags & VK_QUEUE_TRANSFER_BIT )
        {
          hasTransferSupport = true;
          SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "-- Transfer" );
        }

        VkBool32 isSurfaceSupported;
        VK_CHECK(
          vkGetPhysicalDeviceSurfaceSupportKHR(physicalDevice, queueFamilyIndex, surface, &isSurfaceSupported) );

        if ( isSurfaceSupported )
        {
          hasPresentSupport = true;
          SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "-- Present" );
        }

        if ( hasGraphicsSupport and hasComputeSupport and hasTransferSupport and hasPresentSupport )
        {
          physicalDeviceInUse    = physicalDevice;
          directQueueFamilyIndex = queueFamilyIndex;
          break;
        }
      }

      mem->restoreState( tempAllocQueueProperties );
    }

    ASSERT( physicalDeviceInUse );
    ASSERT( directQueueFamilyIndex.has_value() );

    float                   priority        = 1.0f;
    VkDeviceQueueCreateInfo queueCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
      .pNext = nullptr,
      .flags = 0,
      .queueFamilyIndex = directQueueFamilyIndex.value(),
      .queueCount = 1,
      .pQueuePriorities = &priority,
    };

    VkPhysicalDeviceVulkan13Features constexpr features13 = {
      .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES,
      .pNext = nullptr,
      .synchronization2 = true,
      .dynamicRendering = true,
    };

    VkPhysicalDeviceFeatures features = {
      .depthClamp = true,
      .samplerAnisotropy = true,
    };

    std::array enabledDeviceExtensions = {
      VK_KHR_SWAPCHAIN_EXTENSION_NAME
    };

    VkDeviceCreateInfo const deviceCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
      .pNext = &features13,
      .flags = 0,
      .queueCreateInfoCount = 1,
      .pQueueCreateInfos = &queueCreateInfo,
      .enabledLayerCount = 0,
      .ppEnabledLayerNames = nullptr,
      .enabledExtensionCount = static_cast<uint32_t>(enabledDeviceExtensions.size()),
      .ppEnabledExtensionNames = enabledDeviceExtensions.data(),
      .pEnabledFeatures = &features,
    };

    VK_CHECK( vkCreateDevice(physicalDeviceInUse, &deviceCreateInfo, nullptr, &device) );
    volkLoadDevice( device );

    VmaAllocatorCreateInfo allocatorCreateInfo = {
      .flags = 0,
      .physicalDevice = physicalDeviceInUse,
      .device = device,
      .preferredLargeHeapBlockSize = 0,
      .pAllocationCallbacks = nullptr,
      .pDeviceMemoryCallbacks = nullptr,
      .pHeapSizeLimit = nullptr,
      .pVulkanFunctions = nullptr,
      .instance = instance,
      .vulkanApiVersion = VK_API_VERSION_1_3,
      .pTypeExternalMemoryHandleTypes = nullptr,
    };

    VmaVulkanFunctions vkFunctions;
    VK_CHECK( vmaImportVulkanFunctionsFromVolk(&allocatorCreateInfo, &vkFunctions) );
    allocatorCreateInfo.pVulkanFunctions = &vkFunctions;

    VK_CHECK( vmaCreateAllocator(&allocatorCreateInfo, &gpuAllocator) );

    vkGetDeviceQueue( device, directQueueFamilyIndex.value(), 0, &directQueue );

    mem->restoreState( tempAllocStart );
  }

  // Swapchain creation
  VkExtent2D     swapchainExtent = { createInfo.width, createInfo.height };
  VkFormat       swapchainFormat = VK_FORMAT_UNDEFINED;
  VkSwapchainKHR swapchain;
  VkImage*       swapchainImages;
  VkImageView*   swapchainViews;
  uint32_t       swapchainImageCount;
  {
    auto tempAllocStart = mem->getState();

    VkSurfaceCapabilitiesKHR capabilities;
    VK_CHECK( vkGetPhysicalDeviceSurfaceCapabilitiesKHR(physicalDeviceInUse, surface, &capabilities) );

    // Image Count Calculation
    swapchainImageCount = 3;
    if ( capabilities.maxImageCount > 0 )
    {
      swapchainImageCount = std::min( swapchainImageCount, capabilities.maxImageCount );
    }
    swapchainImageCount = std::max( swapchainImageCount, capabilities.minImageCount + 1 );

    // Image Size calculation
    {
      auto [minWidth, minHeight] = capabilities.minImageExtent;
      auto [maxWidth, maxHeight] = capabilities.maxImageExtent;
      swapchainExtent.width      = Clamp( swapchainExtent.width, minWidth, maxWidth );
      swapchainExtent.height     = Clamp( swapchainExtent.height, minHeight, maxHeight );
    }

    uint32_t surfaceFormatCount;
    vkGetPhysicalDeviceSurfaceFormatsKHR( physicalDeviceInUse, surface, &surfaceFormatCount, nullptr );
    VkSurfaceFormatKHR* surfaceFormats = reinterpret_cast<VkSurfaceFormatKHR*>(mem->allocate(
      sizeof( VkSurfaceFormatKHR ) * surfaceFormatCount ));
    vkGetPhysicalDeviceSurfaceFormatsKHR( physicalDeviceInUse, surface, &surfaceFormatCount, surfaceFormats );

    VkSurfaceFormatKHR format = {
      .format = VK_FORMAT_UNDEFINED,
      .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR,
    };
    for ( auto& surfaceFormat : std::span{ surfaceFormats, surfaceFormatCount } )
    {
      if ( surfaceFormat.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR )
      {
        SDL_LogInfo( SDL_LOG_CATEGORY_GPU, "Color Space SRGB Found %d", surfaceFormat.format );
        if ( surfaceFormat.format == VK_FORMAT_R8G8B8A8_SRGB )
        {
          format = surfaceFormat;
          break;
        }
        if ( surfaceFormat.format == VK_FORMAT_B8G8R8A8_SRGB )
        {
          format = surfaceFormat;
          break;
        }
        if ( surfaceFormat.format == VK_FORMAT_R8G8B8A8_UNORM )
        {
          format = surfaceFormat;
        }
      }
    }
    ASSERT( format.format != VK_FORMAT_UNDEFINED );
    swapchainFormat = format.format;

    uint32_t presentModeCount;
    vkGetPhysicalDeviceSurfacePresentModesKHR( physicalDeviceInUse, surface, &presentModeCount, nullptr );
    VkPresentModeKHR* presentModes = reinterpret_cast<VkPresentModeKHR*>(mem->allocate(
      sizeof( VkPresentModeKHR ) * presentModeCount ));
    vkGetPhysicalDeviceSurfacePresentModesKHR( physicalDeviceInUse, surface, &presentModeCount, presentModes );

    VkPresentModeKHR presentMode = VK_PRESENT_MODE_FIFO_KHR;
    for ( VkPresentModeKHR presentModeIter : std::span{ presentModes, presentModeCount } )
    {
      if ( presentModeIter == VK_PRESENT_MODE_FIFO_RELAXED_KHR )
      {
        presentMode = presentModeIter;
        break;
      }

      if ( presentModeIter == VK_PRESENT_MODE_MAILBOX_KHR )
      {
        presentMode = presentModeIter;
      }
    }

    mem->restoreState( tempAllocStart );

    VkSwapchainCreateInfoKHR const swapchainCreateInfo = {
      .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
      .pNext = nullptr,
      .flags = 0,
      .surface = surface,
      .minImageCount = swapchainImageCount,
      .imageFormat = format.format,
      .imageColorSpace = format.colorSpace,
      .imageExtent = swapchainExtent,
      .imageArrayLayers = 1,
      .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
      .imageSharingMode = VK_SHARING_MODE_EXCLUSIVE,
      .queueFamilyIndexCount = 0,
      .pQueueFamilyIndices = nullptr,
      .preTransform = VK_SURFACE_TRANSFORM_IDENTITY_BIT_KHR,
      .compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
      .presentMode = presentMode,
      .clipped = false,
      .oldSwapchain = nullptr,
    };

    VK_CHECK( vkCreateSwapchainKHR(device, &swapchainCreateInfo, nullptr, &swapchain) );

    swapchainImageCount = 0;
    vkGetSwapchainImagesKHR( device, swapchain, &swapchainImageCount, nullptr );
    swapchainImages = reinterpret_cast<VkImage*>(mem->allocate( sizeof( VkImage ) * swapchainImageCount ));
    vkGetSwapchainImagesKHR( device, swapchain, &swapchainImageCount, swapchainImages );

    swapchainViews = reinterpret_cast<VkImageView*>(mem->allocate( sizeof( VkImageView ) * swapchainImageCount ));
    for ( uint32_t i = 0; i != swapchainImageCount; ++i )
    {
      VkImageViewCreateInfo const viewCreateInfo = {
        .sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
        .pNext = nullptr,
        .flags = 0,
        .image = swapchainImages[i],
        .viewType = VK_IMAGE_VIEW_TYPE_2D,
        .format = format.format,
        .components = {
          VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY,
          VK_COMPONENT_SWIZZLE_IDENTITY
        },
        .subresourceRange = {
          .aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
          .baseMipLevel = 0,
          .levelCount = 1,
          .baseArrayLayer = 0,
          .layerCount = 1,
        }
      };

      VK_CHECK( vkCreateImageView(device, &viewCreateInfo, nullptr, &swapchainViews[i]) );
    }
  }

  // Init frames.

  Frame* frames = reinterpret_cast<Frame*>(mem->allocate( sizeof( Frame ) * swapchainImageCount ));
  for ( uint32_t i = 0; i != swapchainImageCount; ++i )
  {
    new( frames + i ) Frame( device, directQueueFamilyIndex.value() );
  }

  Byte* allocation = mem->allocate( sizeof( RenderDevice ), alignof( RenderDevice ) );
  return new( allocation ) RenderDevice{
    instance,
    surface,
    physicalDeviceInUse,
    device,
    gpuAllocator,
    directQueue,
    directQueueFamilyIndex.value(),
    swapchainFormat,
    swapchainExtent,
    swapchain,
    swapchainImages,
    swapchainViews,
    frames,
    swapchainImageCount,
  };
}

inline bool RenderDevice::isInit() const
{
  return instance and device;
}

void RenderDevice::destroy()
{
  if ( not isInit() ) return;

  for ( Frame& frame : std::span{ frames, swapchainImageCount } )
  {
    frame.destroy( *this );
  }

  for ( auto const& view : std::span{ swapchainViews, swapchainImageCount } )
  {
    vkDestroyImageView( device, view, nullptr );
  }

  vkDestroySwapchainKHR( device, Take( swapchain ), nullptr );

  vmaDestroyAllocator( Take( gpuAllocator ) );
  vkDestroyDevice( Take( device ), nullptr );

  SDL_Vulkan_DestroySurface( instance, Take( surface ), nullptr );

  vkDestroyInstance( Take( instance ), nullptr );

  volkFinalize();
}

void RenderDevice::waitIdle() const
{
  VK_CHECK( vkDeviceWaitIdle(device) );
}

uint32_t RenderDevice::getNumFrames() const
{
  return swapchainImageCount;
}

RenderDevice::RenderDevice(
  VkInstance const       instance,
  VkSurfaceKHR const     surface,
  VkPhysicalDevice const physicalDeviceInUse,
  VkDevice const         device,
  VmaAllocator const     gpuAllocator,
  VkQueue const          directQueue,
  uint32_t const         directQueueFamilyIndex,
  VkFormat const         swapchainFormat,
  VkExtent2D const       swapchainExtent,
  VkSwapchainKHR const   swapchain,
  VkImage*               swapchainImages,
  VkImageView*           swapchainViews,
  Frame*                 frames,
  uint32_t const         swapchainImageCount )
  : instance{ instance }
  , surface{ surface }
  , physicalDeviceInUse{ physicalDeviceInUse }
  , device{ device }
  , gpuAllocator{ gpuAllocator }
  , directQueue{ directQueue }
  , directQueueFamilyIndex{ directQueueFamilyIndex }
  , swapchainFormat{ swapchainFormat }
  , swapchainExtent{ swapchainExtent }
  , swapchain{ swapchain }
  , swapchainImages{ swapchainImages }
  , swapchainViews{ swapchainViews }
  , frames{ frames }
  , swapchainImageCount{ swapchainImageCount } {}