import vulkan
, vkTypes
, re
, bitops
, ../utils/lets

type 
 GPU* = object
  handle*: VkPhysicalDevice
  features*: VkPhysicalDeviceFeatures 
  properties*: VkPhysicalDeviceProperties
  memory_properties*: VkPhysicalDeviceMemoryProperties
  queueFamilyProperties*: seq[VkQueueFamilyProperties]
  requestedFeatures*: VkPhysicalDeviceFeatures

#from sets import toHashSet, excl, len
from vulkan_utils import `<`, `==`, `>`
from ../utils/etc import toString, flatten2
from sequtils import anyIt, mapIt, filterIt, maxIndex, toSeq
from strutils import contains, isAlphaNumeric

proc getQProperties*(gpu: VkPhysicalDevice): seq[VkQueueFamilyProperties] =
 var
  queueCount : uint32
#   formatCount: uint32
#   supportsPresent: seq[VkBool32]
#   gqni: uint32 = uint32.high
#   pqni: uint32 = uint32.high
  
 vkGetPhysicalDeviceQueueFamilyProperties(gpu, queueCount.addr, nil)
 result = newSeq[VkQueueFamilyProperties](queueCount)
 vkGetPhysicalDeviceQueueFamilyProperties(gpu, queueCount.addr, result[0].addr)

proc getSupportedDepthFormat*( gpu: VkPhysicalDevice
                              , depthFormat: var VkFormat
                              )
                               =  
 var
  formatProps: VkFormatProperties
  formats = [ VK_FORMAT_D32_SFLOAT_S8_UINT
            , VK_FORMAT_D32_SFLOAT
            , VK_FORMAT_D24_UNORM_S8_UINT
            , VK_FORMAT_D16_UNORM_S8_UINT
            , VK_FORMAT_D16_UNORM
            ]

 for format in formats:
  vkGetPhysicalDeviceFormatProperties( gpu
                                     , format
                                     , addr formatProps
                                     )
  
  if bitand(formatProps.optimalTilingFeatures.int
           , VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT.int
           ) > 0:
   depthFormat = format
   #echo "FOUND DEPTH_FORMAT:" & $depthFormat

   break
  else: echo "nope!" & $bitand(formatProps.optimalTilingFeatures.int
                              , VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT.int
                              )
 
proc get_gpu*( instance: VkInstance
            #, surface: VkSurfaceKHR
            ): GPU =
 load_vulkan_functions_for_gpu(instance)
 
 var gpuCount: uint32 = 0
 
 discard vkEnumeratePhysicalDevices(instance, gpuCount.addr, nil)
 var gpus = newSeq[VkPhysicalDevice](gpuCount)
 discard vkEnumeratePhysicalDevices(instance, gpuCount.addr, gpus[0].addr)
 case gpus.len
 of 0: quit "ERROR: No GPUs found by `vkEnumeratePhysicalDevices`. This is probably a driver-related issue."
 of 1: 
  var 
   devProps0: VKPhysicalDeviceProperties
   features0: VkPhysicalDeviceFeatures
   memoryProperties0: VkPhysicalDeviceMemoryProperties
   #qFamilyProperties0: seq[VkQueueFamilyProperties]
  
  vkGetPhysicalDeviceProperties( gpus[0]
                               , addr devProps0
                               )

  if devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU or 
     devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU: 

   vkGetPhysicalDeviceFeatures(gpus[0], addr features0)
   vkGetPhysicalDeviceMemoryProperties(gpus[0], addr memoryProperties0)
   result = GPU( handle: gpus[0]
               , features:  features0
               , properties: devProps0
               , memory_properties: memoryProperties0
               , queueFamilyProperties: gpus[0].getQProperties
               )  
    
  else: quit  "ERROR: The only GPU/device found is not a Discrete or Integrated type, it is: " & $devProps0.deviceType & " \n       Which is currently not supported."
 
 of 2: 
  #  Handle:
  #   1 GPU with 2 different drivers  -> Use the first
  #   2 GPUs: 1 Discrete 1 integrated -> Use the Discrete 
  #   2 GPUs: 2 Integrated || 2 Discrete -> Use the better one 
 
  var
   devProps0: VKPhysicalDeviceProperties
   devProps1: VKPhysicalDeviceProperties
   features0: VkPhysicalDeviceFeatures
   memoryProperties0: VkPhysicalDeviceMemoryProperties
   #qFamilyProperties0: seq[VkQueueFamilyProperties]
   
  vkGetPhysicalDeviceProperties(gpus[0], addr devProps0)
  vkGetPhysicalDeviceProperties(gpus[1], addr devProps1)
  
  # 1 GPU with 2 different drivers -> Use the first
  if devProps0.deviceID == devProps1.deviceID:

   vkGetPhysicalDeviceFeatures(gpus[0], addr features0)
   vkGetPhysicalDeviceMemoryProperties(gpus[0], addr memoryProperties0)
   result = GPU( handle: gpus[0]
               , features:  features0
               , properties: devProps0
               , memory_properties: memoryProperties0
               , queueFamilyProperties: gpus[0].getQProperties
               )  
  
  #2 GPUs: 1 Discrete 1 integrated, Discrete is the first GPU
  if devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU and 
     devProps1.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
    
    vkGetPhysicalDeviceFeatures(gpus[0], addr features0)
    vkGetPhysicalDeviceMemoryProperties(gpus[0], addr memoryProperties0)
    result = GPU( handle: gpus[0]
                , features:  features0
                , properties: devProps0
                , memory_properties: memoryProperties0
                , queueFamilyProperties: gpus[0].getQProperties
                )   
   
  #2 GPUs: 1 Discrete 1 integrated, Integrated is the first GPU, so we choose the second (Discrete)
  elif devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU and
       devProps1.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
    
    var
     features1: VkPhysicalDeviceFeatures
     memoryProperties1: VkPhysicalDeviceMemoryProperties
    
    vkGetPhysicalDeviceFeatures(gpus[1], addr features1)
    vkGetPhysicalDeviceMemoryProperties(gpus[1], addr memoryProperties1)
    result = GPU( handle: gpus[1]
                , features:  features1
                , properties: devProps1
                , memory_properties: memoryProperties1
                , queueFamilyProperties: gpus[1].getQProperties
                )   
   
  #2 GPUs: 2 Integrated || 2 Discrete -> Find VRAM sizes, and use the GPU with the larger one, or if they're the same, use the first
  if devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU and 
     devProps1.deviceType == VK_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU or
     devProps0.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU and 
     devProps1.deviceType == VK_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
    
    var 
     memoryProperties0: VkPhysicalDeviceMemoryProperties
     memoryProperties1: VkPhysicalDeviceMemoryProperties
     features1: VkPhysicalDeviceFeatures

    vkGetPhysicalDeviceMemoryProperties(gpus[0], addr memoryProperties0)
    vkGetPhysicalDeviceMemoryProperties(gpus[1], addr memoryProperties1)
    
    if memoryProperties0.memoryHeaps[0].size > memoryProperties1.memoryHeaps[0].size or
       memoryProperties0.memoryHeaps[0].size == memoryProperties1.memoryHeaps[0].size: 
     
     result = GPU( handle: gpus[0]
                 , features:  features0
                 , properties: devProps0
                 , memory_properties: memoryProperties0
                 , queueFamilyProperties: gpus[0].getQProperties
                 )   
    else:
     vkGetPhysicalDeviceFeatures(gpus[0], addr features1)
     result = GPU( handle: gpus[1]
                 , features:  features1
                 , properties: devProps1
                 , memory_properties: memoryProperties1
                 , queueFamilyProperties: gpus[1].getQProperties
                 )   

 # 3 or more GPUs, -> Find Discrete, or Integrated, else error
 else: 
  var
   memoryProperties: seq[VkPhysicalDeviceMemoryProperties]
   devProps: seq[VKPhysicalDeviceProperties]
   features: seq[VkPhysicalDeviceFeatures]
   gpuIndex: Natural

  for i, d in gpus: 
   var
    someFeatures: VkPhysicalDeviceFeatures
    someMemoryProperties: VkPhysicalDeviceMemoryProperties   
    someDevProps: VKPhysicalDeviceProperties
   
   vkGetPhysicalDeviceFeatures(d, addr someFeatures)
   vkGetPhysicalDeviceMemoryProperties(d, addr someMemoryProperties)
   vkGetPhysicalDeviceProperties(d, addr someDevProps)
   features.add someFeatures
   memoryProperties.add someMemoryProperties
   devProps.add someDevProps
    
  gpuIndex = maxIndex toSeq memoryProperties.mapit(it.memoryHeaps[0].size)

  result = GPU( handle: gpus[gpuIndex]
              , features:  features[gpuIndex]
              , properties: devProps[gpuIndex]
              , memory_properties: memoryProperties[gpuIndex]
              , queueFamilyProperties: gpus[gpuIndex].getQProperties
              )   

# NOTE: All of this is probably not needed, because i'm pretty sure there's 
#       just some temp-bug or something with q'ing extensions.

# It's UB to have a device extension that is not supported by the GPU (?)
# so we check that all coreDeviceExts exist in the q'd extensions from the GPU
# and when an ext is found, we "check it off" by popping it off the seq
# if there's any that are not found, we quit, and the resulting exts are shown
proc check_gpu_extensions*( gpu: VkPhysicalDevice
                          , dump_extensions: bool = false 
                          ) = 
  var 
   extCount: uint32
   checkList: seq[string] = toSeq coreDeviceExts
   #extsNotFound: seq[string]
   checkCount = checkList.len
#[    doubleExt = re"VK.*VK"
   doubleExtToSplit = re"VK.*(?=VK)"
   corruptedExt = re"VK.*[0-9](?=.*)"
   allExts: seq[string]
 ]#
  # var fixedList = availableExts.mapIt(it.extensionName.toString)
  #     .mapIt(it.split anExt)
  #     .flatten
  #     #.filterIt(it == "")
  discard vkEnumerateDeviceExtensionProperties(gpu, nil, extCount.addr, nil)
  var availableExts = newSeq[VkExtensionProperties](extCount)
  discard vkEnumerateDeviceExtensionProperties(gpu, nil, extCount.addr, availableExts[0].addr)
  
  
  for ext in availableExts.mapIt(it.extensionName):
   var ne = ext.filterIt((it.isAlphaNumeric or it == '_') ).toString
   if dump_extensions: echo ne

  # for ext in allExts:
  #  echo ext
  
  # for chExt in coreDeviceExts:
  #  echo availableExts.mapIt($it).contains chExt
     
  if checkCount > 0:
   echo checkList
  #  quit "wew ladddddddddddddddddddddddd"