#define GLFW_INCLUDE_VULKAN
#include <GLFW/glfw3.h>
#include <imgui.h>

#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/mat4x4.hpp>
#include <glm/vec4.hpp>
#include <iostream>
#include <map>
#include <memory>
#include <optional>
#include <set>
#include <sophia/application.hpp>
#include <sophia/canvas.hpp>
#include <sophia/console/console_command_registry.hpp>
#include <sophia/frame_info.hpp>
#include <sophia/math.hpp>
#include <sophia/resource_registry_ref.hpp>
#include <sophia/scene_registry.hpp>
#include <sophia/scene_view.hpp>
#include <sophia/scene_view_registry.hpp>
#include <sophia/types.hpp>
#include <sophia/ui/debug_ui.hpp>
#include <sophia/ui/panel_registry.hpp>
#include <sophia/ui/ui.hpp>
#include <sophia/util/logger.hpp>
#include <string>
#include <tracy/Tracy.hpp>
#include <utility>
#include <vector>
#include <vulkan/vulkan.hpp>

#include "renderer/imgui_renderer.hpp"
#include "renderer/model_render_system.hpp"
#include "renderer/render_context.hpp"
#include "renderer/render_graph.hpp"
#include "renderer/scene_renderer.hpp"
#include "resources/mesh.hpp"
#include "resources/render_target.hpp"
#include "resources/resource_registry.hpp"
#include "vulkan/buffer.hpp"
#include "vulkan/device.hpp"
#include "vulkan/pipelines/graphics_pipeline.hpp"
#include "vulkan/shader.hpp"
#include "window.hpp"

constexpr uint32_t WIDTH = 1280;
constexpr uint32_t HEIGHT = 720;
constexpr const char* NAME = "sandbox";

constexpr sophia::ClearColor CLEAR_COLOR = {0.01f, 0.01f, 0.02f, 1.0f};

bool hasFlag(sophia::ClearFlags flags, sophia::ClearFlags flag) {
  return (static_cast<sophia::u8>(flags) & static_cast<sophia::u8>(flag)) != 0;
}

int main(int argc, const char** argv) {
  std::vector<std::string> args;
  args.assign(argv, argv + argc);

  std::cout << __FILE__ << "::" << __LINE__ << '\n';

  sophia::Window window{WIDTH, HEIGHT, NAME};
  sophia::Device device{&window};
  sophia::RenderContext renderContext{window, device};
  sophia::ResourceRegistry resourceRegistry{
      device, renderContext.getImageFormat()};  // remove?

  // remove?
  std::vector<std::unique_ptr<sophia::Mesh>> meshes;

  // remove?
  std::vector<sophia::ModelRenderSystem::MeshInstance> drawCalls;

  // todo: figure out which to remove
  sophia::PanelRegistry panelRegistry;
  sophia::SceneRegistry sceneRegistry;
  sophia::SceneViewRegistry sceneViewRegistry;
  sophia::ConsoleCommandRegistry consoleCommandRegistry;

  sophia::DebugUI debugUI{panelRegistry, consoleCommandRegistry};
  debugUI.showFrameStats(true);

  sophia::ModelRenderSystem modelRenderSystem{device, renderContext};
  sophia::SceneRenderer sceneRenderer{sceneRegistry, resourceRegistry, device,
                                      renderContext.getImageFormat()};
  sophia::ImGuiRenderer imguiRenderer{
      window, device, renderContext.getImageFormat(),
      renderContext.getImageCount(), resourceRegistry};

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

  auto& meshPool = resourceRegistry.meshPool;
  auto& renderTargetPool = resourceRegistry.renderTargetPool;

  imguiRenderer.setDefaultLayout(panelRegistry.getDefaultLayout());

  // Views already warned about, so each warns once.
  std::set<sophia::SceneViewHandle> viewsWithoutCamera;
  std::set<sophia::SceneViewHandle> viewsOverwritingOthers;

  auto startTime = std::chrono::high_resolution_clock::now();
  auto currentTime = startTime;
  sophia::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.
    renderTargetPool.applyPendingResizes();

    vk::Extent2D extent = renderContext.getExtent();
    sophia::FrameInfo frameInfo{frameIndex, elapsedTime, deltaTime,
                                glm::vec2(static_cast<float>(extent.width),
                                          static_cast<float>(extent.height))};

    // Canvas draws in window pixels, origin at the center.
    // TODO: replace with Debug/Gizmos drawn into scene views through their
    // cameras.
    glm::mat4 canvasProjection =
        glm::ortho(-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();
    sceneRenderer.beginFrame(renderContext.getFrameIndex());

    panelRegistry.update(frameInfo);

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

    sophia::RenderGraph graph{renderContext.getTracyContext()};
    sophia::ImageRef backbuffer = graph.importImage(
        renderContext.getCurrentImage(), sophia::Access::Present);

    struct ViewTarget {
      sophia::ImageRef color;
      sophia::ImageRef depth;
    };
    std::map<sophia::RenderTargetHandle, ViewTarget> viewTargets;
    // Views without a target render straight into the backbuffer.
    std::optional<ViewTarget> swapchainTarget;
    for (sophia::SceneViewHandle handle : sceneViewRegistry.GetRenderQueue()) {
      const sophia::SceneView& view = sceneViewRegistry.peek(handle);
      bool toSwapchain = !view.target.isValid();
      if (!sceneRegistry.isAlive(view.scene) ||
          !sceneRegistry.get(view.scene).isVisible() ||
          (!toSwapchain && !renderTargetPool.isAlive(view.target))) {
        continue;
      }

      ViewTarget target;
      vk::Extent2D targetExtent;
      // Whether this is the first view this frame to render to its target.
      bool firstOnTarget = false;
      if (toSwapchain) {
        firstOnTarget = !swapchainTarget;
        if (!swapchainTarget) {
          swapchainTarget = {
              backbuffer,
              graph.importImage(renderContext.getCurrentDepthImage(),
                                sophia::Access::DepthWrite)};
        }
        target = *swapchainTarget;
        targetExtent = renderContext.getExtent();
      } else {
        sophia::RenderTarget& renderTarget = renderTargetPool.get(view.target);
        auto [it, inserted] = viewTargets.try_emplace(view.target);
        firstOnTarget = inserted;
        if (inserted) {
          it->second = {graph.importImage(renderTarget.getImage(),
                                          sophia::Access::SampledRead),
                        graph.importImage(renderTarget.getDepthImage(),
                                          sophia::Access::DepthWrite)};
        }
        target = it->second;
        targetExtent = renderTarget.getExtent();
      }

      if (!firstOnTarget &&
          hasFlag(view.clearFlags, sophia::ClearFlags::Color) &&
          viewsOverwritingOthers.insert(handle).second) {
        sophia::log::warning(
            "scene view clears its target after another view already "
            "rendered to it this frame, hiding that view (give it clearFlags "
            "without Color, or render it first with a lower priority)");
      }

      const sophia::vec4& color = view.clearColor;
      graph.addPass(
          "scene view " + std::to_string(handle.index),
          [&](sophia::RenderGraph::PassBuilder& pass) {
            pass.writeColor(target.color,
                            hasFlag(view.clearFlags, sophia::ClearFlags::Color)
                                ? vk::AttachmentLoadOp::eClear
                                : vk::AttachmentLoadOp::eLoad,
                            {color.r, color.g, color.b, color.a});
            pass.writeDepth(target.depth,
                            hasFlag(view.clearFlags, sophia::ClearFlags::Depth)
                                ? vk::AttachmentLoadOp::eClear
                                : vk::AttachmentLoadOp::eLoad);
          },
          [&sceneRenderer, &sceneRegistry, &sceneViewRegistry,
           &viewsWithoutCamera, handle,
           targetExtent](vk::CommandBuffer commandBuffer) {
            ZoneScopedN("scene view pass");
            const sophia::SceneView& view = sceneViewRegistry.peek(handle);
            sophia::Entity camera =
                view.camera ? view.camera
                            : sceneRegistry.get(view.scene).getActiveCamera();
            // No camera to look through: the target is just cleared.
            if (!camera) {
              if (viewsWithoutCamera.insert(handle).second) {
                sophia::log::warning(
                    "scene view has no camera and its scene has no active "
                    "camera; it will only be cleared (see "
                    "Scene::createCamera)");
              }
              return;
            }
            sceneRenderer.render(commandBuffer, view, camera, targetExtent);
          });
    }

    graph.addPass(
        "main",
        [&](sophia::RenderGraph::PassBuilder& pass) {
          // Keep what scene views drew into the backbuffer.
          pass.writeColor(backbuffer,
                          swapchainTarget ? vk::AttachmentLoadOp::eLoad
                                          : vk::AttachmentLoadOp::eClear,
                          CLEAR_COLOR);
          for (const auto& [handle, target] : viewTargets) {
            pass.sample(target.color);
          }
        },
        [&](vk::CommandBuffer commandBuffer) {
          ZoneScopedN("main pass");
          modelRenderSystem.render(commandBuffer, instances, canvasProjection);

          imguiRenderer.endFrame(commandBuffer);
        });

    graph.execute(commandBuffer);

    renderContext.endFrame();
    FrameMark;

    frameIndex++;
  }

  device.waitIdle();

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

  return EXIT_SUCCESS;
}