#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 "render_target.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;

void beginColorPass(vk::CommandBuffer commandBuffer, vk::ImageView view,
                    vk::Extent2D extent) {
  vk::ClearValue clearValue{};
  clearValue.color =
      vk::ClearColorValue{std::array<float, 4>{0.01f, 0.01f, 0.02f, 1.0f}};

  vk::RenderingAttachmentInfo colorAttachment{};
  colorAttachment.imageView = view;
  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);
}

}  // 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.resourceRegistry};
    this->application._ui.renderer = &imguiRenderer;

    this->application.consoleCommandRegistry.registerCommand(
        "reset_layout", "Restore the default panel layout",
        [&imguiRenderer](const std::vector<std::string>&) {
          imguiRenderer.resetLayout();
          return std::string{"Layout reset"};
        });

    this->application.onCreate();
    imguiRenderer.setDefaultLayout(
        this->application.panelRegistry.getDefaultLayout());

    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();

      // Resizing waits for the GPU, so it has to happen before any of this
      // frame's work (or ImGui's references to the targets) is recorded.
      this->application.resourceRegistry.applyPendingResizes();

      RenderTarget* sceneTarget =
          this->sceneTargetHandle.isValid()
              ? &this->application.resourceRegistry.getRenderTarget(
                    this->sceneTargetHandle)
              : nullptr;

      vk::Extent2D extent = renderContext.getExtent();
      vk::Extent2D sceneExtent =
          sceneTarget ? sceneTarget->getExtent() : extent;
      FrameInfo frameInfo{frameIndex, elapsedTime, deltaTime,
                          glm::vec2(static_cast<float>(sceneExtent.width),
                                    static_cast<float>(sceneExtent.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();

      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());
      this->drawCalls.clear();

      if (sceneTarget) {
        Image& targetImage = sceneTarget->getImage();
        targetImage.transitionLayout(commandBuffer,
                                     vk::ImageLayout::eColorAttachmentOptimal);
        beginColorPass(commandBuffer, targetImage.getView(), sceneExtent);
        modelRenderSystem.render(commandBuffer, instances, frameInfo);
        commandBuffer.endRendering();
        targetImage.transitionLayout(commandBuffer,
                                     vk::ImageLayout::eShaderReadOnlyOptimal);
      }

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

      if (!sceneTarget) {
        modelRenderSystem.render(commandBuffer, instances, frameInfo);
      }

      imguiRenderer.endFrame(commandBuffer);

      commandBuffer.endRendering();

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

      renderContext.endFrame();

      frameIndex++;
    }

    this->application.onDestroy();

    device.waitIdle();

    this->application._ui.renderer = nullptr;

    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});
  }

  void setSceneTarget(RenderTargetHandle handle) {
    this->sceneTargetHandle = handle;
  }

  u32 getColorFormat() const {
    return static_cast<u32>(this->renderContext.getImageFormat());
  }

  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;

  RenderTargetHandle sceneTargetHandle = RenderTargetHandle::invalid();
};

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)},
      resourceRegistry{impl->getDevice(), impl->getColorFormat()},
      _ui{resourceRegistry},
      _debugUI{this->panelRegistry, this->consoleCommandRegistry} {}

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::setSceneTarget(RenderTargetHandle handle) {
  impl->setSceneTarget(handle);
}

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

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

}  // namespace sophia