#include <array>
#include <chrono>
#include <cstddef>
#include <sophia/application.hpp>
#include <sophia/util/logger.hpp>
#include <string>
#include <utility>
#include <vector>
#include <vulkan/vulkan.hpp>

#define GLM_FORCE_RADIANS
#define GLM_FORCE_DEPTH_ZERO_TO_ONE
#include <imgui.h>

#include <glm/glm.hpp>

#include "buffer.hpp"
#include "device.hpp"
#include "imgui_renderer.hpp"
#include "mesh.hpp"
#include "pipelines/graphics_pipeline.hpp"
#include "render_context.hpp"
#include "shader.hpp"
#include "systems/model_render_system.hpp"
#include "window.hpp"

namespace sophia {

namespace {

struct Vertex {
  glm::vec2 position;
  glm::vec3 color;
};

struct PushConstants {
  glm::mat4 transform;
  glm::vec4 color;
  glm::vec4 data;
};

std::vector<std::string> commandLineArgs;

}  // namespace

class Application::Impl {
 public:
  Impl(Application& application, const Application::Config& config)
      : application{application},
        config{config},
        window{config.width, config.height, config.name},
        device{&window},
        renderContext{window, device} {}

  void run() {
    ModelRenderSystem modelRenderSystem{device, renderContext};
    ImGuiRenderer imguiRenderer{window, device, renderContext.getImageFormat(),
                                renderContext.getImageCount()};

    this->application.onCreate();

    auto startTime = std::chrono::high_resolution_clock::now();
    auto currentTime = startTime;
    u32 frameIndex = 0;

    while (!window.shouldClose()) {
      glfwPollEvents();

      auto newTime = std::chrono::high_resolution_clock::now();
      double deltaTime =
          std::chrono::duration<double>(newTime - currentTime).count();
      currentTime = newTime;

      double elapsedTime =
          std::chrono::duration<double>(currentTime - startTime).count();

      vk::Extent2D extent = renderContext.getExtent();
      FrameInfo frameInfo{frameIndex, elapsedTime, deltaTime,
                          glm::vec2(static_cast<float>(extent.width),
                                    static_cast<float>(extent.height)),
                          this->application.camera};
      this->application.onUpdate(frameInfo);
      this->application.camera.setProjection(
          -frameInfo.currentFramebufferExtent.x / 2.f,
          frameInfo.currentFramebufferExtent.x / 2.f,
          -frameInfo.currentFramebufferExtent.y / 2.f,
          frameInfo.currentFramebufferExtent.y / 2.f);

      vk::CommandBuffer commandBuffer = renderContext.beginFrame();
      if (!commandBuffer) {
        continue;
      }

      imguiRenderer.beginFrame();

      this->application.panelRegistry.update(frameInfo);
      this->application.onRender();

      renderContext.transitionCurrentImage(
          commandBuffer, vk::ImageLayout::eColorAttachmentOptimal);

      vk::ClearValue clearValue{};
      clearValue.color =
          vk::ClearColorValue{std::array<float, 4>{0.01f, 0.01f, 0.02f, 1.0f}};

      vk::RenderingAttachmentInfo colorAttachment{};
      colorAttachment.imageView = renderContext.getCurrentImageView();
      colorAttachment.imageLayout = vk::ImageLayout::eColorAttachmentOptimal;
      colorAttachment.loadOp = vk::AttachmentLoadOp::eClear;
      colorAttachment.storeOp = vk::AttachmentStoreOp::eStore;
      colorAttachment.clearValue = clearValue;

      vk::RenderingInfo renderingInfo{};
      renderingInfo.renderArea = vk::Rect2D{{0, 0}, extent};
      renderingInfo.layerCount = 1;
      renderingInfo.colorAttachmentCount = 1;
      renderingInfo.pColorAttachments = &colorAttachment;

      commandBuffer.beginRendering(renderingInfo);

      vk::Viewport viewport{0.0f,
                            0.0f,
                            static_cast<float>(extent.width),
                            static_cast<float>(extent.height),
                            0.0f,
                            1.0f};
      vk::Rect2D scissor{{0, 0}, extent};
      commandBuffer.setViewport(0, viewport);
      commandBuffer.setScissor(0, scissor);

      std::vector<ModelRenderSystem::MeshInstance> instances;
      instances.reserve(this->meshes.size() + this->drawCalls.size());
      for (const std::unique_ptr<Mesh>& mesh : this->meshes) {
        instances.push_back({mesh.get(), glm::mat4{1.f}});
      }
      instances.insert(instances.end(), this->drawCalls.begin(),
                       this->drawCalls.end());

      modelRenderSystem.render(commandBuffer, instances, frameInfo);
      this->drawCalls.clear();

      imguiRenderer.endFrame(commandBuffer);

      commandBuffer.endRendering();

      renderContext.transitionCurrentImage(commandBuffer,
                                           vk::ImageLayout::ePresentSrcKHR);

      renderContext.endFrame();

      frameIndex++;
    }

    this->application.onDestroy();

    device.waitIdle();

    auto endTime = std::chrono::high_resolution_clock::now();
    double totalRuntime =
        std::chrono::duration<double>(endTime - startTime).count();
    log::info("Application ran for", totalRuntime, "s");
  }

  void pushMesh(std::vector<CanvasVertex> vertices, std::vector<u32> indices) {
    Mesh::Builder builder{};
    builder.vertices.reserve(vertices.size());
    for (const CanvasVertex& vertex : vertices) {
      builder.vertices.push_back({vertex.position, vertex.color});
    }
    builder.indices = std::move(indices);

    this->meshes.push_back(std::make_unique<Mesh>(this->device, builder));
  }

  void clearMeshes() {
    // Meshes back live GPU vertex buffers and up to MAX_FRAMES_IN_FLIGHT
    // command buffers may still be executing against them, so wait for the
    // device to go idle before freeing anything.
    this->device.waitIdle();
    this->meshes.clear();
  }

  void drawMesh(const Mesh& mesh, const glm::mat4& transform) {
    this->drawCalls.push_back({&mesh, transform});
  }

  Device& getDevice() { return device; }

 private:
  Application& application;
  const Application::Config config;
  Window window;
  Device device;
  RenderContext renderContext;

  std::vector<std::unique_ptr<Mesh>> meshes;
  std::vector<ModelRenderSystem::MeshInstance> drawCalls;
};

void Application::setCommandLineArgs(int argc, const char* const* argv) {
  commandLineArgs.assign(argv, argv + argc);
}

Application::Application(const Config& config)
    : impl{std::make_unique<Impl>(*this, config)},
      canvas{this},
      args{std::move(commandLineArgs)},
      assetRegistry{impl->getDevice()},
      _debugUI{this->panelRegistry} {}

Application::~Application() = default;

void Application::run() { impl->run(); }

void Application::pushMesh(std::vector<CanvasVertex> vertices,
                           std::vector<u32> indices) {
  impl->pushMesh(std::move(vertices), std::move(indices));
}

void Application::clearMeshes() { impl->clearMeshes(); }

void Application::drawMesh(MeshHandle handle, const glm::mat4& transform) {
  impl->drawMesh(assetRegistry.get<Mesh>(handle), transform);
}

}  // namespace sophia