#ifndef _VULKAN_GAME_H
#define _VULKAN_GAME_H

#include <algorithm>
#include <chrono>
#include <map>
#include <vector>

#include <vulkan/vulkan.h>

#include <SDL2/SDL.h>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE // Since, in Vulkan, the depth range is 0 to 1 instead of -1 to 1
#define GLM_FORCE_RIGHT_HANDED

#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>

#include "IMGUI/imgui_impl_vulkan.h"

#include "consts.hpp"
#include "utils.hpp"
#include "game-gui-sdl.hpp"
#include "vulkan-utils.hpp"
#include "graphics-pipeline_vulkan.hpp"
#include "vulkan-buffer.hpp"

using namespace glm;
using namespace std::chrono;

#ifdef NDEBUG
   const bool ENABLE_VALIDATION_LAYERS = false;
#else
   const bool ENABLE_VALIDATION_LAYERS = true;
#endif

// TODO: Consider if there is a better way of dealing with all the vertex types and ssbo types, maybe
// by consolidating some and trying to keep new ones to a minimum

struct ModelVertex {
   vec3 pos;
   vec3 color;
   vec2 texCoord;
   vec3 normal;
   unsigned int objIndex;
};

struct LaserVertex {
   vec3 pos;
   vec2 texCoord;
   unsigned int objIndex;
};

struct ExplosionVertex {
   vec3 particleStartVelocity;
   float particleStartTime;
   unsigned int objIndex;
};

struct SSBO_ModelObject {
   alignas(16) mat4 model;
};

struct SSBO_Asteroid {
   alignas(16) mat4 model;
   alignas(4) float hp;
   alignas(4) unsigned int deleted;
};

struct SSBO_Laser {
   alignas(16) mat4 model;
   alignas(4) vec3 color;
   alignas(4) unsigned int deleted;
};

struct SSBO_Explosion {
   alignas(16) mat4 model;
   alignas(4) float explosionStartTime;
   alignas(4) float explosionDuration;
   alignas(4) unsigned int deleted;
};

struct UBO_VP_mats {
   alignas(16) mat4 view;
   alignas(16) mat4 proj;
};

struct UBO_Explosion {
   alignas(16) mat4 view;
   alignas(16) mat4 proj;
   alignas(4) float cur_time;
};

// TODO: Use this struct for uniform buffers as well and probably combine it with the VulkanBuffer class
// Also, probably better to make this a vector of structs where each struct
// has a VkBuffer, VkDeviceMemory, and VkDescriptorBufferInfo
// TODO: Maybe change the structure here since VkDescriptorBufferInfo already stores a reference to the VkBuffer
struct BufferSet {
   vector<VkBuffer> buffers;
   vector<VkDeviceMemory> memory;
   vector<VkDescriptorBufferInfo> infoSet;
   VkBufferUsageFlags usages;
   VkMemoryPropertyFlags properties;
};

// TODO: Change the index type to uint32_t and check the Vulkan Tutorial loading model section as a reference
// TODO: Create a typedef for index type so I can easily change uin16_t to something else later
// TODO: Maybe create a typedef for each of the templated SceneObject types
template<class VertexType>
struct SceneObject {
   vector<VertexType> vertices;
   vector<uint16_t> indices;

   mat4 model_base;
   mat4 model_transform;

   // TODO: Figure out if I should make child classes that have these fields instead of putting them in the
   // parent class
   vec3 center; // currently only matters for asteroids
   float radius; // currently only matters for asteroids

   // Move the targetAsteroid stuff out of this class since it is very specific to lasers
   // and makes moving SceneObject into its own header file more problematic
   SceneObject<ModelVertex>* targetAsteroid; // currently only used for lasers
};

// TODO: Figure out how to include an optional ssbo parameter for each object
// Could probably use the same approach to make indices optional
// Figure out if there are sufficient use cases to make either of these optional or is it fine to make
// them mandatory


// TODO: Look into using dynamic_cast to check types of SceneObject and EffectOverTime

struct BaseEffectOverTime {
   bool deleted;

   virtual void applyEffect(float curTime) = 0;

   BaseEffectOverTime() :
         deleted(false) {
   }

   virtual ~BaseEffectOverTime() {
   }
};

// TODO: Move this into its own hpp and cpp files
// TODO: Actually, since this is only used to make lasers deal damage to asteroids, review the logic
// and see if there is a more straightforward way of implementing this.
// If there is a simple and straightforward way to implement this in updateScene(), I should just do that instead of
// using this class. A general feature like this is useful, but, depending on the actual game, it might not be used
// that often, and using this class might actually make things more complicated if it doesn't quite implement the
// desired features
template<class SSBOType>
struct EffectOverTime : public BaseEffectOverTime {
   VulkanBuffer<SSBOType> &buffer;
   uint32_t objectIndex;
   size_t effectedFieldOffset;
   float startValue;
   float startTime;
   float changePerSecond;

   EffectOverTime(VulkanBuffer<SSBOType> &buffer, uint32_t objectIndex, size_t effectedFieldOffset, float startTime,
                  float changePerSecond)
               : buffer(buffer)
               , objectIndex(objectIndex)
               , effectedFieldOffset(effectedFieldOffset)
               , startTime(startTime)
               , changePerSecond(changePerSecond) {
      startValue = *reinterpret_cast<float*>(getEffectedFieldPtr());
   }

   unsigned char* getEffectedFieldPtr() {
      return reinterpret_cast<unsigned char*>(&buffer.get(objectIndex)) + effectedFieldOffset;
   }

   void applyEffect(float curTime) {
      if (buffer.get(objectIndex).deleted) {
         deleted = true;
         return;
      }

      *reinterpret_cast<float*>(getEffectedFieldPtr()) = startValue + (curTime - startTime) * changePerSecond;
   }
};

// TODO: Maybe move this to a different header

enum UIValueType {
   UIVALUE_INT,
   UIVALUE_DOUBLE,
};

struct UIValue {
   UIValueType type;
   string label;
   void* value;

   UIValue(UIValueType _type, string _label, void* _value) : type(_type), label(_label), value(_value) {}
};

/* TODO: The following syntax (note the const keyword) means the function will not modify
 * its params. I should use this where appropriate
 *
 * [return-type] [func-name](params...) const { ... }
 */

class VulkanGame {

   public:

      VulkanGame();
      ~VulkanGame();

      void run(int width, int height, unsigned char guiFlags);

   private:

      static VKAPI_ATTR VkBool32 VKAPI_CALL debugCallback(
         VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
         VkDebugUtilsMessageTypeFlagsEXT messageType,
         const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
         void* pUserData);

      // TODO: Maybe pass these in as parameters to some Camera class
      const float NEAR_CLIP = 0.1f;
      const float FAR_CLIP = 100.0f;
      const float FOV_ANGLE = 67.0f; // means the camera lens goes from -33 deg to 33 deg

      const int EXPLOSION_PARTICLE_COUNT = 300;
      const vec3 LASER_COLOR = vec3(0.2f, 1.0f, 0.2f);

      bool done;

      vec3 cam_pos;

      // TODO: Good place to start using smart pointers
      GameGui* gui;

      SDL_version sdlVersion;
      SDL_Window* window = nullptr;

      int drawableWidth, drawableHeight;

      VkInstance instance;
      VkDebugUtilsMessengerEXT debugMessenger;
      VkSurfaceKHR vulkanSurface;
      VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;
      VkDevice device;

      VkQueue graphicsQueue;
      VkQueue presentQueue;

      // TODO: Maybe make a swapchain struct for convenience
      VkSurfaceFormatKHR swapChainSurfaceFormat;
      VkPresentModeKHR swapChainPresentMode;
      VkExtent2D swapChainExtent;
      uint32_t swapChainMinImageCount;
      uint32_t swapChainImageCount;
      VkSwapchainKHR swapChain;
      vector<VkImage> swapChainImages;
      vector<VkImageView> swapChainImageViews;
      vector<VkFramebuffer> swapChainFramebuffers;

      VkRenderPass renderPass;

      VkCommandPool resourceCommandPool;

      vector<VkCommandPool> commandPools;
      vector<VkCommandBuffer> commandBuffers;

      VulkanImage depthImage;

      // These are per frame
      vector<VkSemaphore> imageAcquiredSemaphores;
      vector<VkSemaphore> renderCompleteSemaphores;

      // These are per swap chain image
      vector<VkFence> inFlightFences;

      uint32_t imageIndex;
      uint32_t currentFrame;

      bool shouldRecreateSwapChain;

      VkSampler textureSampler;

      VulkanImage floorTextureImage;
      VkDescriptorImageInfo floorTextureImageDescriptor;

      VulkanImage laserTextureImage;
      VkDescriptorImageInfo laserTextureImageDescriptor;

      mat4 viewMat, projMat;

      // Maybe at some point create an imgui pipeline class, but I don't think it makes sense right now
      VkDescriptorPool imguiDescriptorPool;

      // TODO: Probably restructure the GraphicsPipeline_Vulkan class based on what I learned about descriptors and textures
      // while working on graphics-library. Double-check exactly what this was and note it down here.
      // Basically, I think the point was that if I have several modesl that all use the same shaders and, therefore,
      // the same pipeline, but use different textures, the approach I took when initially creating GraphicsPipeline_Vulkan
      // wouldn't work since the whole pipeline couldn't have a common set of descriptors for the textures
      GraphicsPipeline_Vulkan<ModelVertex> modelPipeline;
      GraphicsPipeline_Vulkan<ModelVertex> shipPipeline;
      GraphicsPipeline_Vulkan<ModelVertex> asteroidPipeline;
      GraphicsPipeline_Vulkan<LaserVertex> laserPipeline;
      GraphicsPipeline_Vulkan<ExplosionVertex> explosionPipeline;

      BufferSet uniformBuffers_modelPipeline;
      BufferSet objectBuffers_modelPipeline;

      VulkanBuffer<UBO_VP_mats> uniforms_modelPipeline;
      VulkanBuffer<SSBO_ModelObject> objects_modelPipeline;

      BufferSet uniformBuffers_shipPipeline;
      BufferSet objectBuffers_shipPipeline;

      VulkanBuffer<UBO_VP_mats> uniforms_shipPipeline;
      VulkanBuffer<SSBO_ModelObject> objects_shipPipeline;

      BufferSet uniformBuffers_asteroidPipeline;
      BufferSet objectBuffers_asteroidPipeline;

      VulkanBuffer<UBO_VP_mats> uniforms_asteroidPipeline;
      VulkanBuffer<SSBO_Asteroid> objects_asteroidPipeline;

      BufferSet uniformBuffers_laserPipeline;
      BufferSet objectBuffers_laserPipeline;

      VulkanBuffer<UBO_VP_mats> uniforms_laserPipeline;
      VulkanBuffer<SSBO_Laser> objects_laserPipeline;

      BufferSet uniformBuffers_explosionPipeline;
      BufferSet objectBuffers_explosionPipeline;

      VulkanBuffer<UBO_Explosion> uniforms_explosionPipeline;
      VulkanBuffer<SSBO_Explosion> objects_explosionPipeline;

      // TODO: Maybe make the ubo objects part of the pipeline class since there's only one ubo
      // per pipeline.
      // Or maybe create a higher level wrapper around GraphicsPipeline_Vulkan to hold things like
      // the objects vector, the ubo, and the ssbo

      // TODO: Rename *_VP_mats to *_uniforms and possibly use different types for each one
      // if there is a need to add other uniform variables to one or more of the shaders

      vector<SceneObject<ModelVertex>> modelObjects;
      vector<SceneObject<ModelVertex>> shipObjects;
      vector<SceneObject<ModelVertex>> asteroidObjects;
      vector<SceneObject<LaserVertex>> laserObjects;
      vector<SceneObject<ExplosionVertex>> explosionObjects;

      //UBO_Explosion explosion_UBO;

      vector<BaseEffectOverTime*> effects;

      float shipSpeed = 0.5f;
      float asteroidSpeed = 2.0f;

      float spawnRate_asteroid = 0.5;
      float lastSpawn_asteroid;

      unsigned int leftLaserIdx = -1;
      EffectOverTime<SSBO_Asteroid>* leftLaserEffect = nullptr;

      unsigned int rightLaserIdx = -1;
      EffectOverTime<SSBO_Asteroid>* rightLaserEffect = nullptr;

      /*** High-level vars ***/

      // TODO: Just typedef the type of this function to RenderScreenFn or something since it's used in a few places
      void (VulkanGame::* currentRenderScreenFn)(int width, int height);

      map<string, vector<UIValue>> valueLists;

      int score;
      float fps;

      // TODO: Make a separate TImer class
      time_point<steady_clock> startTime;
      float fpsStartTime, curTime, prevTime, elapsedTime;

      int frameCount;

      /*** Functions ***/

      bool initUI(int width, int height, unsigned char guiFlags);
      void initVulkan();
      void initGraphicsPipelines();
      void initMatrices();
      void renderLoop();
      void updateScene();
      void cleanup();

      void createVulkanInstance(const vector<const char*>& validationLayers);
      void setupDebugMessenger();
      void populateDebugMessengerCreateInfo(VkDebugUtilsMessengerCreateInfoEXT& createInfo);
      void createVulkanSurface();
      void pickPhysicalDevice(const vector<const char*>& deviceExtensions);
      bool isDeviceSuitable(VkPhysicalDevice physicalDevice, const vector<const char*>& deviceExtensions);
      void createLogicalDevice(const vector<const char*>& validationLayers,
         const vector<const char*>& deviceExtensions);
      void chooseSwapChainProperties();
      void createSwapChain();
      void createImageViews();
      void createResourceCommandPool();
      void createImageResources();
      VkFormat findDepthFormat(); // TODO: Declare/define (in the cpp file) this function in some util functions section
      void createRenderPass();
      void createCommandPools();
      void createFramebuffers();
      void createCommandBuffers();
      void createSyncObjects();

      void createTextureSampler();

      void initImGuiOverlay();
      void cleanupImGuiOverlay();

      // TODO: Maybe move these to a different class, possibly VulkanBuffer or some new related class

      void createBufferSet(VkDeviceSize bufferSize, VkBufferUsageFlags usages, VkMemoryPropertyFlags properties,
                           BufferSet& set);

      void resizeBufferSet(BufferSet& set, VkDeviceSize newSize, VkCommandPool commandPool, VkQueue graphicsQueue,
                           bool copyData);

      // TODO: Since addObject() returns a reference to the new object now,
      // stop using objects.back() to access the object that was just created
      template<class VertexType, class SSBOType>
      SceneObject<VertexType>& addObject(vector<SceneObject<VertexType>>& objects,
                                         GraphicsPipeline_Vulkan<VertexType>& pipeline,
                                         const vector<VertexType>& vertices, vector<uint16_t> indices,
                                         VulkanBuffer<SSBOType>& objectBuffer, SSBOType ssbo);

      template<class VertexType>
      vector<VertexType> addObjectIndex(unsigned int objIndex, vector<VertexType> vertices);

      template<class VertexType>
      vector<VertexType> addVertexNormals(vector<VertexType> vertices);

      template<class VertexType>
      void centerObject(SceneObject<VertexType>& object);

      template<class VertexType>
      void updateObjectVertices(GraphicsPipeline_Vulkan<VertexType>& pipeline, SceneObject<VertexType>& obj,
                                size_t index);

      void addLaser(vec3 start, vec3 end, vec3 color, float width);
      void translateLaser(size_t index, const vec3& translation);
      void updateLaserTarget(size_t index);
      bool getLaserAndAsteroidIntersection(SceneObject<ModelVertex>& asteroid, vec3& start, vec3& end,
                                           vec3& intersection);

      void addExplosion(mat4 model_mat, float duration, float cur_time);

      void renderFrame(ImDrawData* draw_data);
      void presentFrame();

      void recreateSwapChain();

      void cleanupSwapChain();

      /*** High-level functions ***/

      void renderMainScreen(int width, int height);
      void renderGameScreen(int width, int height);

      void initGuiValueLists(map<string, vector<UIValue>>& valueLists);
      void renderGuiValueList(vector<UIValue>& values);

      void goToScreen(void (VulkanGame::* renderScreenFn)(int width, int height));
      void quitGame();
};

// Start of specialized no-op functions

template<>
inline void VulkanGame::centerObject(SceneObject<ExplosionVertex>& object) {
}

// End of specialized no-op functions

// TODO: Right now, it's basically necessary to pass the identity matrix in for ssbo.model and to change
// the model matrix later by setting model_transform and then calculating the new ssbo.model.
// Figure out a better way to allow the model matrix to be set during object creation
template<class VertexType, class SSBOType>
SceneObject<VertexType>& VulkanGame::addObject(vector<SceneObject<VertexType>>& objects,
                                               GraphicsPipeline_Vulkan<VertexType>& pipeline,
                                               const vector<VertexType>& vertices, vector<uint16_t> indices,
                                               VulkanBuffer<SSBOType>& objectBuffer, SSBOType ssbo) {
   // TODO: Use the model field of ssbo to set the object's model_base
   // currently, the passed-in model is useless since it gets overridden when ssbo.model is recalculated
   size_t numVertices = pipeline.getNumVertices();

   for (uint16_t& idx : indices) {
      idx += numVertices;
   }

   objects.push_back({ vertices, indices, mat4(1.0f), mat4(1.0f) });
   objectBuffer.add(ssbo);

   SceneObject<VertexType>& obj = objects.back();

   // TODO: Specify whether to center the object outside of this function or, worst case, maybe
   // with a boolean being passed in here, so that I don't have to rely on checking the specific object
   // type
   // TODO: Actually, I've already defined a no-op centerObject method for explosions
   // Maybe I should do the same for lasers and remove this conditional altogether
   if (!is_same_v<VertexType, LaserVertex> && !is_same_v<VertexType, ExplosionVertex>) {
      centerObject(obj);
   }

   pipeline.addObject(obj.vertices, obj.indices, resourceCommandPool, graphicsQueue);

   return obj;
}

template<class VertexType>
vector<VertexType> VulkanGame::addObjectIndex(unsigned int objIndex, vector<VertexType> vertices) {
   for (VertexType& vertex : vertices) {
      vertex.objIndex = objIndex;
   }

   return vertices;
}

// This function sets all the normals for a face to be parallel
// This is good for models that should have distinct faces, but bad for models that should appear smooth
// Maybe add an option to set all copies of a point to have the same normal and have the direction of
// that normal be the weighted average of all the faces it is a part of, where the weight from each face
// is its surface area.

// TODO: Since the current approach to normal calculation basicaly makes indexed drawing useless, see if it's
// feasible to automatically enable/disable indexed drawing based on which approach is used
template<class VertexType>
vector<VertexType> VulkanGame::addVertexNormals(vector<VertexType> vertices) {
   for (unsigned int i = 0; i < vertices.size(); i += 3) {
      vec3 p1 = vertices[i].pos;
      vec3 p2 = vertices[i + 1].pos;
      vec3 p3 = vertices[i + 2].pos;

      vec3 normal = normalize(cross(p2 - p1, p3 - p1));

      // Add the same normal for all 3 vertices
      vertices[i].normal = normal;
      vertices[i + 1].normal = normal;
      vertices[i + 2].normal = normal;
   }

   return vertices;
}

template<class VertexType>
void VulkanGame::centerObject(SceneObject<VertexType>& object) {
   vector<VertexType>& vertices = object.vertices;

   float min_x = vertices[0].pos.x;
   float max_x = vertices[0].pos.x;
   float min_y = vertices[0].pos.y;
   float max_y = vertices[0].pos.y;
   float min_z = vertices[0].pos.z;
   float max_z = vertices[0].pos.z;

   // start from the second point
   for (unsigned int i = 1; i < vertices.size(); i++) {
      vec3& pos = vertices[i].pos;

      if (min_x > pos.x) {
         min_x = pos.x;
      } else if (max_x < pos.x) {
         max_x = pos.x;
      }

      if (min_y > pos.y) {
         min_y = pos.y;
      } else if (max_y < pos.y) {
         max_y = pos.y;
      }

      if (min_z > pos.z) {
         min_z = pos.z;
      } else if (max_z < pos.z) {
         max_z = pos.z;
      }
   }

   vec3 center = vec3(min_x + max_x, min_y + max_y, min_z + max_z) / 2.0f;

   for (unsigned int i = 0; i < vertices.size(); i++) {
      vertices[i].pos -= center;
   }

   object.radius = std::max(max_x - center.x, max_y - center.y);
   object.radius = std::max(object.radius, max_z - center.z);

   object.center = vec3(0.0f, 0.0f, 0.0f);
}

template<class VertexType>
void VulkanGame::updateObjectVertices(GraphicsPipeline_Vulkan<VertexType>& pipeline, SceneObject<VertexType>& obj,
                                      size_t index) {
   pipeline.updateObjectVertices(index, obj.vertices, resourceCommandPool, graphicsQueue);
}

#endif // _VULKAN_GAME_H