import std/enumerate

import ./vulkan
import ./vulkan_helpers
import ./xlib_helpers

import ./glslang/glslang

var vertexShaderCode: string = """#version 450
layout(location = 0) out vec3 fragColor;
vec3 colors[3] = vec3[](
    vec3(1.0, 0.0, 0.0),
    vec3(0.0, 1.0, 0.0),
    vec3(0.0, 0.0, 1.0)
);
vec2 positions[3] = vec2[](
  vec2(0.0, -0.5),
  vec2(0.5, 0.5),
  vec2(-0.5, 0.5)
);
void main() {
  gl_Position = vec4(positions[gl_VertexIndex], 0.0, 1.0);
  fragColor = colors[gl_VertexIndex];
}"""

var fragmentShaderCode: string = """#version 450
layout(location = 0) out vec4 outColor;
layout(location = 0) in vec3 fragColor;
void main() {
  outColor = vec4(fragColor, 1.0);
}"""

import
  x11/xlib,
  x11/x

const VULKAN_VERSION = VK_MAKE_API_VERSION(0'u32, 1'u32, 2'u32, 0'u32)

type
  GraphicsPipeline = object
    shaderStages*: seq[VkPipelineShaderStageCreateInfo]
  QueueFamily = object
    properties*: VkQueueFamilyProperties
    hasSurfaceSupport*: bool
  PhyscialDevice = object
    device*: VkPhysicalDevice
    extensions*: seq[string]
    properties*: VkPhysicalDeviceProperties
    features*: VkPhysicalDeviceFeatures
    queueFamilies*: seq[QueueFamily]
    surfaceCapabilities*: VkSurfaceCapabilitiesKHR
    surfaceFormats: seq[VkSurfaceFormatKHR]
    presentModes: seq[VkPresentModeKHR]
  Vulkan* = object
    instance*: VkInstance
    deviceList*: seq[PhyscialDevice]
    activePhysicalDevice*: PhyscialDevice
    activeQueueFamily*: uint32
    device*: VkDevice
    presentationQueue*: VkQueue
    surface*: VkSurfaceKHR
    selectedSurfaceFormat: VkSurfaceFormatKHR
    selectedPresentationMode: VkPresentModeKHR
    selectedExtent: VkExtent2D
    swapChain: VkSwapchainKHR
    swapImages: seq[VkImage]
    swapImageViews: seq[VkImageView]
  Engine* = object
    display*: PDisplay
    window*: x.Window
    vulkan*: Vulkan
    pipeline*: GraphicsPipeline


proc getAllPhysicalDevices(instance: VkInstance, surface: VkSurfaceKHR): seq[PhyscialDevice] =
  for vulkanPhysicalDevice in getVulkanPhysicalDevices(instance):
    var device = PhyscialDevice(device: vulkanPhysicalDevice, extensions: getDeviceExtensions(vulkanPhysicalDevice))
    vkGetPhysicalDeviceProperties(vulkanPhysicalDevice, addr(device.properties))
    vkGetPhysicalDeviceFeatures(vulkanPhysicalDevice, addr(device.features))
    checkVkResult vkGetPhysicalDeviceSurfaceCapabilitiesKHR(vulkanPhysicalDevice, surface, addr(device.surfaceCapabilities))
    device.surfaceFormats = getDeviceSurfaceFormats(vulkanPhysicalDevice, surface)
    device.presentModes = getDeviceSurfacePresentModes(vulkanPhysicalDevice, surface)

    for i, queueFamilyProperty in enumerate(getQueueFamilies(vulkanPhysicalDevice)):
      var hasSurfaceSupport: VkBool32 = VK_FALSE
      checkVkResult vkGetPhysicalDeviceSurfaceSupportKHR(vulkanPhysicalDevice, uint32(i), surface, addr(hasSurfaceSupport))
      device.queueFamilies.add(QueueFamily(properties: queueFamilyProperty, hasSurfaceSupport: bool(hasSurfaceSupport)))

    result.add(device)

proc filterForDevice(devices: seq[PhyscialDevice]): seq[(PhyscialDevice, uint32)] =
  for device in devices:
    if "VK_KHR_swapchain" in device.extensions:
      for i, queueFamily in enumerate(device.queueFamilies):
        let hasGraphics = bool(uint32(queueFamily.properties.queueFlags) and ord(VK_QUEUE_GRAPHICS_BIT))
        if (
          queueFamily.hasSurfaceSupport and
          hasGraphics and
          device.surfaceFormats.len > 0 and
          device.presentModes.len > 0
        ):
          result.add((device, uint32(i)))

proc filterForSurfaceFormat(formats: seq[VkSurfaceFormatKHR]): seq[VkSurfaceFormatKHR] =
  for format in formats:
    if format.format == VK_FORMAT_B8G8R8A8_SRGB and format.colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR:
      result.add(format)

proc getSwapExtent(display: PDisplay, window: Window, capabilities: VkSurfaceCapabilitiesKHR): VkExtent2D =
  if capabilities.currentExtent.width != high(uint32):
    return capabilities.currentExtent
  else:
    let (width, height) = xlibFramebufferSize(display, window)
    return VkExtent2D(
      width: min(max(uint32(width), capabilities.minImageExtent.width), capabilities.maxImageExtent.width),
      height: min(max(uint32(height), capabilities.minImageExtent.height), capabilities.maxImageExtent.height),
    )

proc igniteEngine*(): Engine =
  vkLoad1_0()
  vkLoad1_1()
  vkLoad1_2()

  # init X11 window
  (result.display, result.window) = xlibInit()

  # create vulkan instance
  result.vulkan.instance = createVulkanInstance(VULKAN_VERSION)

  # create vulkan-X11 surface
  var surfaceCreateInfo = VkXlibSurfaceCreateInfoKHR(
    sType: VK_STRUCTURE_TYPE_XLIB_SURFACE_CREATE_INFO_KHR,
    dpy: result.display,
    window: result.window,
  )
  checkVkResult vkCreateXlibSurfaceKHR(result.vulkan.instance, addr(surfaceCreateInfo), nil, addr(result.vulkan.surface))

  # determine device and queue to use and instantiate
  result.vulkan.deviceList = result.vulkan.instance.getAllPhysicalDevices(result.vulkan.surface)
  let usableDevices = result.vulkan.deviceList.filterForDevice()
  if len(usableDevices) == 0:
    raise newException(Exception, "No suitable graphics device found")
  (result.vulkan.activePhysicalDevice, result.vulkan.activeQueueFamily) = usableDevices[0]
  
  (result.vulkan.device, result.vulkan.presentationQueue) = getVulcanDevice(
    result.vulkan.activePhysicalDevice.device,
    result.vulkan.activePhysicalDevice.features,
    result.vulkan.activeQueueFamily
  )
  
  # determine surface format for swapchain
  let usableSurfaceFormats = filterForSurfaceFormat(result.vulkan.activePhysicalDevice.surfaceFormats)
  if len(usableSurfaceFormats) == 0:
    raise newException(Exception, "No suitable surface formats found")
  result.vulkan.selectedSurfaceFormat = usableSurfaceFormats[0]
  result.vulkan.selectedPresentationMode = getPresentMode(result.vulkan.activePhysicalDevice.presentModes)
  result.vulkan.selectedExtent = getSwapExtent(result.display, result.window, result.vulkan.activePhysicalDevice.surfaceCapabilities)

  # setup swapchain
  var swapchainCreateInfo = VkSwapchainCreateInfoKHR(
    sType: VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
    surface: result.vulkan.surface,
    minImageCount: max(result.vulkan.activePhysicalDevice.surfaceCapabilities.minImageCount + 1, result.vulkan.activePhysicalDevice.surfaceCapabilities.maxImageCount),
    imageFormat: result.vulkan.selectedSurfaceFormat.format,
    imageColorSpace: result.vulkan.selectedSurfaceFormat.colorSpace,
    imageExtent: result.vulkan.selectedExtent,
    imageArrayLayers: 1,
    imageUsage: VkImageUsageFlags(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT),
    # VK_SHARING_MODE_CONCURRENT no supported (i.e cannot use different queue families for  drawing to swap surface?)
    imageSharingMode: VK_SHARING_MODE_EXCLUSIVE,
    preTransform: result.vulkan.activePhysicalDevice.surfaceCapabilities.currentTransform,
    compositeAlpha: VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR,
    presentMode: result.vulkan.selectedPresentationMode,
    clipped: VK_TRUE,
    oldSwapchain: VkSwapchainKHR(0),
  )
  checkVkResult vkCreateSwapchainKHR(result.vulkan.device, addr(swapchainCreateInfo), nil, addr(result.vulkan.swapChain))
  result.vulkan.swapImages = getSwapChainImages(result.vulkan.device, result.vulkan.swapChain)

  # setup swapchian image views
  result.vulkan.swapImageViews = newSeq[VkImageView](result.vulkan.swapImages.len)
  for i, image in enumerate(result.vulkan.swapImages):
    var imageViewCreateInfo = VkImageViewCreateInfo(
      sType: VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO,
      image: image,
      viewType: VK_IMAGE_VIEW_TYPE_2D,
      format: result.vulkan.selectedSurfaceFormat.format,
      components: VkComponentMapping(
        r: VK_COMPONENT_SWIZZLE_IDENTITY,
        g: VK_COMPONENT_SWIZZLE_IDENTITY,
        b: VK_COMPONENT_SWIZZLE_IDENTITY,
        a: VK_COMPONENT_SWIZZLE_IDENTITY,
      ),
      subresourceRange: VkImageSubresourceRange(
        aspectMask: VkImageAspectFlags(VK_IMAGE_ASPECT_COLOR_BIT),
        baseMipLevel: 0,
        levelCount: 1,
        baseArrayLayer: 0,
        layerCount: 1,
      ),
    )
    checkVkResult vkCreateImageView(result.vulkan.device, addr(imageViewCreateInfo), nil, addr(result.vulkan.swapImageViews[i]))

  # init shader system
  checkGlslangResult glslang_initialize_process()

  # load shaders
  result.pipeline.shaderStages.add(createShaderStage(result.vulkan.device, VK_SHADER_STAGE_VERTEX_BIT, vertexShaderCode))
  result.pipeline.shaderStages.add(createShaderStage(result.vulkan.device, VK_SHADER_STAGE_FRAGMENT_BIT, fragmentShaderCode))

  # create graphis pipeline
  var dynamicStates = [VK_DYNAMIC_STATE_VIEWPORT, VK_DYNAMIC_STATE_SCISSOR]
  var dynamicState = VkPipelineDynamicStateCreateInfo(
    sType: VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO,
    dynamicStateCount: uint32(dynamicStates.len),
    pDynamicStates: addr(dynamicStates[0]),
  )
  var vertexInputInfo = VkPipelineVertexInputStateCreateInfo(
    sType: VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO,
    vertexBindingDescriptionCount: 0,
    pVertexBindingDescriptions: nil,
    vertexAttributeDescriptionCount: 0,
    pVertexAttributeDescriptions: nil,
  )
  var inputAssembly = VkPipelineInputAssemblyStateCreateInfo(
    sType: VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO,
    topology: VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST,
    primitiveRestartEnable: VK_FALSE,
  )

  # setup viewport
  var viewport = VkViewport(
    x: 0.0,
    y: 0.0,
    width: (float) result.vulkan.selectedExtent.width,
    height: (float) result.vulkan.selectedExtent.height,
    minDepth: 0.0,
    maxDepth: 1.0,
  )
  var scissor = VkRect2D(
    offset: VkOffset2D(x: 0, y: 0),
    extent: result.vulkan.selectedExtent
  )


proc fullThrottle*(engine: Engine) =
  var event: XEvent
  while true:
    discard XNextEvent(engine.display, addr(event))
    case event.theType
    of Expose:
      discard
    of ClientMessage:
      if cast[Atom](event.xclient.data.l[0]) == deleteMessage:
        break
    of KeyPress:
      let key = XLookupKeysym(cast[PXKeyEvent](addr(event)), 0)
      if key != 0:
        echo "Key ", key, " pressed"
    of ButtonPressMask:
      echo "Mouse button ", event.xbutton.button, " pressed at ",
          event.xbutton.x, ",", event.xbutton.y
    else:
      discard

proc trash*(engine: Engine) =
  for shaderStage in engine.pipeline.shaderStages:
    vkDestroyShaderModule(engine.vulkan.device, shaderStage.module, nil);
  glslang_finalize_process()
  vkDestroySwapchainKHR(engine.vulkan.device, engine.vulkan.swapChain, nil);
  vkDestroySurfaceKHR(engine.vulkan.instance, engine.vulkan.surface, nil);
  vkDestroyDevice(engine.vulkan.device, nil)
  vkDestroyInstance(engine.vulkan.instance, nil)
  checkXlibResult engine.display.XDestroyWindow(engine.window)
  discard engine.display.XCloseDisplay() # always returns 0