#include <vulkan/vulkan.h>
#include <vulkan/vk_layer.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <errno.h>
#include <new>

#include "overlay.h"
#include "vulkan.h"

#define FH_EXPORT __attribute__(( visibility( "default" ) ))
#define LAYER_NAME "VK_LAYER_fhvkov_overlay"

VkLayerInstanceCreateInfo* find_instance_link( const VkInstanceCreateInfo* ci ) {
  VkLayerInstanceCreateInfo* li = (VkLayerInstanceCreateInfo*)ci->pNext;
  while( li && !(li->sType == VK_STRUCTURE_TYPE_LOADER_INSTANCE_CREATE_INFO && li->function == VK_LAYER_LINK_INFO) )
    li = (VkLayerInstanceCreateInfo*)li->pNext;
  return li;
}

VkLayerDeviceCreateInfo* find_device_link( const VkDeviceCreateInfo* ci ) {
  VkLayerDeviceCreateInfo* li = (VkLayerDeviceCreateInfo*)ci->pNext;
  while( li && !(li->sType == VK_STRUCTURE_TYPE_LOADER_DEVICE_CREATE_INFO && li->function == VK_LAYER_LINK_INFO) )
    li = (VkLayerDeviceCreateInfo*)li->pNext;
  return li;
}

VKAPI_ATTR VkResult VKAPI_CALL hook_create_instance(
  const VkInstanceCreateInfo* ci,
  const VkAllocationCallbacks* ac,
  VkInstance* out
) {
  VkLayerInstanceCreateInfo* li = find_instance_link( ci );
  if( !li )
    return VK_ERROR_INITIALIZATION_FAILED;

  PFN_vkGetInstanceProcAddr gipa = li->u.pLayerInfo->pfnNextGetInstanceProcAddr;
  li->u.pLayerInfo = li->u.pLayerInfo->pNext;

  PFN_vkCreateInstance create = (PFN_vkCreateInstance)gipa( 0, "vkCreateInstance" );
  if( !create )
    return VK_ERROR_INITIALIZATION_FAILED;

  VkResult r = create( ci, ac, out );
  if( r != VK_SUCCESS )
    return r;

  vk_instance_register(* out, gipa );
  return r;
}

VKAPI_ATTR void VKAPI_CALL hook_destroy_instance( VkInstance inst, const VkAllocationCallbacks* ac ) {
  INSTANCE_DATA* id = vk_find_instance( GET_KEY( inst ) );
  PFN_vkDestroyInstance down = id ? id->destroy_instance : 0;
  vk_instance_forget( inst );
  if( down )
    down( inst, ac );
}

VKAPI_ATTR VkResult VKAPI_CALL hook_create_device(
  VkPhysicalDevice phys,
  const VkDeviceCreateInfo* ci,
  const VkAllocationCallbacks* ac,
  VkDevice* out
) {
  VkLayerDeviceCreateInfo* li = find_device_link( ci );
  if( !li )
    return VK_ERROR_INITIALIZATION_FAILED;

  PFN_vkGetInstanceProcAddr gipa = li->u.pLayerInfo->pfnNextGetInstanceProcAddr;
  PFN_vkGetDeviceProcAddr gdpa = li->u.pLayerInfo->pfnNextGetDeviceProcAddr;
  li->u.pLayerInfo = li->u.pLayerInfo->pNext;

  PFN_vkCreateDevice create = (PFN_vkCreateDevice)gipa( 0, "vkCreateDevice" );
  if( !create )
    return VK_ERROR_INITIALIZATION_FAILED;

  VkResult r = create( phys, ci, ac, out );
  if( r != VK_SUCCESS )
    return r;

  DEVICE_DATA* d = (DEVICE_DATA*)calloc( 1, sizeof(DEVICE_DATA) );
  OVERLAY_DATA* od = overlay_inst();
  d->overlay = od;
  new( &d->udp )   UDP_LISTENER();
  new( &d->gauge ) GAUGE_RENDERER();

  d->key    = GET_KEY(* out );
  d->dev    =* out;
  d->phys   = phys;
  d->gdpa   = gdpa;
  d->inst   = vk_find_instance( GET_KEY( phys ) );

  od->t_prev = t_nows();
  vk_resolve_device_procs( d,* out, gdpa );

  config_init( &od->cfg );
  const char* e = getenv( "FHVKOV_PORT" );
  if( e )
    od->cfg.port = atoi( e );

  if( !OK( udp_open( &d->udp, od->cfg.port ) ) )
    dlog( "fhvkov: cannot bind udp/%d (%s), overlay will stay hidden\n", od->cfg.port, strerror(errno) );
  else
    printf( "fhvkov: active, udp/%d, config %s\n", od->cfg.port, od->cfg.path );

  vk_device_register( d );
  return r;
}

VKAPI_ATTR void VKAPI_CALL hook_destroy_device(
  VkDevice dev,
  const VkAllocationCallbacks* ac
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( dev ) );
  if( !d )
    return;

  for( I32 i = 0; i < MAX_SWAPS; i++ ) {
    if( d->swaps[i] ) {
      vk_swap_destroy( d, d->swaps[i] );
      d->swaps[i] = 0;
    }
  }
  vk_device_unregister( d );

  PFN_vkDestroyDevice down = d->destroy_device;
  d->udp.~UDP_LISTENER();
  free( d );
  down( dev, ac );
}

VKAPI_ATTR void VKAPI_CALL hook_get_device_queue(
  VkDevice dev,
  U32 family,
  U32 index,
  VkQueue* q
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( dev ) );
  d->get_device_queue( dev, family, index, q );
  vk_device_add_queue( d,* q, family );
}

VKAPI_ATTR void VKAPI_CALL hook_get_device_queue2(
  VkDevice dev,
  const VkDeviceQueueInfo2* qi,
  VkQueue* q
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( dev ) );
  d->get_device_queue2( dev, qi, q );
  vk_device_add_queue( d,* q, qi->queueFamilyIndex );
}

VKAPI_ATTR VkResult VKAPI_CALL hook_create_swapchain(
  VkDevice dev,
  const VkSwapchainCreateInfoKHR* ci,
  const VkAllocationCallbacks* ac,
  VkSwapchainKHR* out
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( dev ) );
  if( !d )
    return VK_ERROR_INITIALIZATION_FAILED;

  VkSwapchainCreateInfoKHR ci2 =* ci;
  ci2.imageUsage |= VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

  VkResult r = d->create_swapchain( dev, &ci2, ac, out );
  if( r != VK_SUCCESS )
    return r;

  vk_swap_create( d,* out, ci );
  return r;
}

VKAPI_ATTR void VKAPI_CALL hook_destroy_swapchain(
  VkDevice dev,
  VkSwapchainKHR sc,
  const VkAllocationCallbacks* ac
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( dev ) );
  if( !d )
    return;

  vk_swap_forget( d, sc );
  d->destroy_swapchain( dev, sc, ac );
}

VKAPI_ATTR VkResult VKAPI_CALL hook_queue_present(
  VkQueue queue,
  const VkPresentInfoKHR* pi
) {
  DEVICE_DATA* d = vk_find_device( GET_KEY( queue ) );
  if( !d )
    return VK_ERROR_DEVICE_LOST;
  OVERLAY_DATA* od = d->overlay;

  udp_parse( &d->udp, &od->telem );
  if( ++od->frame % 60 == 0 ) {
    if( config_poll( &od->cfg ) )
      od->valid_redline = od->target_redline = 0.86;
  }

  F64 now = t_nows();
  F64 dt  = now - od->t_prev;
  od->t_prev = now;
  if( dt < 0.0 ) dt = 0.0;
  if( dt > 0.1 ) dt = 0.1;
  overlay_update( od, now, dt );

  if( !overlay_visible( od, now ) || pi->swapchainCount != 1 )
    return d->queue_present( queue, pi );

  pthread_mutex_lock( &g_lock );
  SWAP_DATA* s = vk_find_swap( d, pi->pSwapchains[0] );
  bool ok = s && !s->broken;
  if( ok && !s->frames_ready && !OK( vk_swap_setup_frames( d, s, queue ) ) ) {
    s->broken = true;
    ok = false;
  }
  pthread_mutex_unlock( &g_lock );

  if( !ok )
    return d->queue_present( queue, pi );
  return vk_present( d, s, queue, pi );
}

#define HOOK( vk_name, fn_name ) \
  if( !strcmp( name, #vk_name ) ) \
    return (PFN_vkVoidFunction)fn_name

PFN_vkVoidFunction hooked( const char* name ) {
  HOOK( vkCreateInstance,       hook_create_instance );
  HOOK( vkDestroyInstance,      hook_destroy_instance );
  HOOK( vkCreateDevice,         hook_create_device );
  HOOK( vkDestroyDevice,        hook_destroy_device );
  HOOK( vkGetDeviceQueue,       hook_get_device_queue );
  HOOK( vkGetDeviceQueue2,      hook_get_device_queue2 );
  HOOK( vkCreateSwapchainKHR,   hook_create_swapchain );
  HOOK( vkDestroySwapchainKHR,  hook_destroy_swapchain );
  HOOK( vkQueuePresentKHR,      hook_queue_present );
  HOOK( vkGetInstanceProcAddr,  vkGetInstanceProcAddr );
  HOOK( vkGetDeviceProcAddr,    vkGetDeviceProcAddr );
  return 0;
}

extern "C" {
  FH_EXPORT VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
  vkGetDeviceProcAddr( VkDevice dev, const char* name );

  FH_EXPORT VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
  vkGetInstanceProcAddr( VkInstance inst, const char* name );


  FH_EXPORT VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
  vkGetDeviceProcAddr( VkDevice dev, const char* name ) {
    PFN_vkVoidFunction f = hooked( name );
    if( f )
      return f;

    DEVICE_DATA* d = dev ? vk_find_device( GET_KEY( dev ) ) : 0;
    return d ? d->gdpa( dev, name ) : 0;
  }

  FH_EXPORT VKAPI_ATTR PFN_vkVoidFunction VKAPI_CALL
  vkGetInstanceProcAddr( VkInstance inst, const char* name ) {
    PFN_vkVoidFunction f = hooked( name );
    if( f )
      return f;

    INSTANCE_DATA* id = inst ? vk_find_instance( GET_KEY( inst ) ) : 0;
    return id ? id->gipa( inst, name ) : 0;
  }
}