#include "viewport_panel.hpp"

#include <imgui.h>

#include <array>
#include <optional>
#include <sophia/components.hpp>
#include <tuple>

namespace {

// Corners of the unit quad a SpriteComponent's model matrix maps to world.
constexpr std::array<glm::vec2, 4> UNIT_QUAD{
    glm::vec2{0.f, 0.f}, glm::vec2{1.f, 0.f}, glm::vec2{1.f, 1.f},
    glm::vec2{0.f, 1.f}};

const ImU32 SELECTION_COLOR = IM_COL32(255, 170, 40, 255);

}  // namespace

ImVec2 ViewportPanel::ViewportTransform::worldToScreen(
    const glm::vec3& world) const {
  glm::vec4 clip = viewProjection * glm::vec4{world, 1.f};
  glm::vec2 ndc = glm::vec2{clip} / clip.w;
  // Vulkan NDC has +y pointing down, same as the screen.
  return {origin.x + (ndc.x * 0.5f + 0.5f) * size.x,
          origin.y + (ndc.y * 0.5f + 0.5f) * size.y};
}

glm::vec2 ViewportPanel::ViewportTransform::screenToWorld(
    const ImVec2& screen) const {
  glm::vec2 ndc{(screen.x - origin.x) / size.x * 2.f - 1.f,
                (screen.y - origin.y) / size.y * 2.f - 1.f};
  glm::vec4 world = inverseViewProjection * glm::vec4{ndc, 0.f, 1.f};
  return glm::vec2{world} / world.w;
}

void ViewportPanel::onUpdate(const sophia::FrameInfo&) {
  drawToolbar();
  ImGui::Separator();

  const ImVec2 origin = ImGui::GetCursorScreenPos();
  const ImVec2 size = ImGui::GetContentRegionAvail();
  if (size.x < 1.0f || size.y < 1.0f) {
    return;
  }

  ui.image(target);

  // The image doesn't take mouse input, so without this, dragging in it
  // would drag the window instead.
  ImGui::SetCursorScreenPos(origin);
  ImGui::InvisibleButton("##viewport", size);

  ImDrawList* drawList = ImGui::GetWindowDrawList();
  drawList->PushClipRect(origin, ImVec2(origin.x + size.x, origin.y + size.y),
                         true);

  if (showGrid) {
    drawGrid(origin, size);
  }

  sophia::Entity camera = scene.getActiveCamera();
  if (camera && camera.hasComponent<sophia::CameraComponent>()) {
    // Use the panel's size rather than the target's: the target is resized
    // a frame late, and the mouse is over the panel.
    float aspect = size.x / size.y;
    glm::mat4 viewProjection =
        camera.getComponent<sophia::CameraComponent>().getViewProjection(
            camera.getComponent<sophia::TransformComponent>(), aspect);

    ViewportTransform viewport{origin, size, viewProjection,
                               glm::inverse(viewProjection)};
    handleInput(viewport);
    drawSelection(viewport);
  }

  drawList->PopClipRect();

  ImGui::SetCursorScreenPos(ImVec2(origin.x + 8.0f, origin.y + 8.0f));
  ImGui::TextDisabled("%.0f x %.0f", size.x, size.y);
}

void ViewportPanel::drawToolbar() {
  const struct {
    const char* label;
    GizmoMode mode;
  } modes[] = {{"Move", GizmoMode::Move},
               {"Rotate", GizmoMode::Rotate},
               {"Scale", GizmoMode::Scale}};

  for (const auto& entry : modes) {
    if (ImGui::RadioButton(entry.label, gizmoMode == entry.mode)) {
      gizmoMode = entry.mode;
    }
    ImGui::SameLine();
  }
  ImGui::Checkbox("Grid", &showGrid);
}

void ViewportPanel::drawGrid(ImVec2 origin, ImVec2 size) {
  constexpr float GRID_STEP = 40.0f;
  const ImU32 gridColor = IM_COL32(255, 255, 255, 18);
  const ImVec2 end(origin.x + size.x, origin.y + size.y);

  ImDrawList* drawList = ImGui::GetWindowDrawList();
  for (float x = GRID_STEP; x < size.x; x += GRID_STEP) {
    drawList->AddLine(ImVec2(origin.x + x, origin.y),
                      ImVec2(origin.x + x, end.y), gridColor);
  }
  for (float y = GRID_STEP; y < size.y; y += GRID_STEP) {
    drawList->AddLine(ImVec2(origin.x, origin.y + y),
                      ImVec2(end.x, origin.y + y), gridColor);
  }
}

void ViewportPanel::handleInput(const ViewportTransform& viewport) {
  const ImGuiIO& io = ImGui::GetIO();

  if (ImGui::IsItemActivated()) {
    selected = pick(viewport.screenToWorld(io.MousePos));
    dragging = selected && gizmoMode == GizmoMode::Move;
  }

  if (!ImGui::IsItemActive()) {
    dragging = false;
    return;
  }

  if (dragging && (io.MouseDelta.x != 0.f || io.MouseDelta.y != 0.f)) {
    ImVec2 previous(io.MousePos.x - io.MouseDelta.x,
                    io.MousePos.y - io.MouseDelta.y);
    glm::vec2 delta =
        viewport.screenToWorld(io.MousePos) - viewport.screenToWorld(previous);

    auto& transform = selected.getComponent<sophia::TransformComponent>();
    transform.position.x += delta.x;
    transform.position.y += delta.y;
  }
}

void ViewportPanel::drawSelection(const ViewportTransform& viewport) {
  if (!selected) {
    return;
  }

  ImDrawList* drawList = ImGui::GetWindowDrawList();
  const auto& transform = selected.getComponent<sophia::TransformComponent>();

  if (selected.hasComponent<sophia::SpriteComponent>()) {
    glm::mat4 model =
        selected.getComponent<sophia::SpriteComponent>().getModel(transform);

    std::array<ImVec2, 4> corners;
    for (size_t i = 0; i < UNIT_QUAD.size(); i++) {
      corners[i] = viewport.worldToScreen(
          glm::vec3{model * glm::vec4{UNIT_QUAD[i], 0.f, 1.f}});
    }
    drawList->AddQuad(corners[0], corners[1], corners[2], corners[3],
                      SELECTION_COLOR, 2.0f);
  }

  // Mark the transform's position (the sprite's pivot), which is also all
  // there is to show for entities without a sprite, such as the camera.
  ImVec2 center = viewport.worldToScreen(transform.position);
  drawList->AddCircle(center, 4.0f, SELECTION_COLOR, 0, 2.0f);
}

sophia::Entity ViewportPanel::pick(glm::vec2 world) {
  sophia::Entity best;
  std::optional<std::tuple<sophia::s32, float>> bestOrder;

  scene.forEachEntity([&](sophia::Entity entity) {
    if (!entity.hasComponent<sophia::SpriteComponent>()) {
      return;
    }

    const auto& transform = entity.getComponent<sophia::TransformComponent>();
    const auto& sprite = entity.getComponent<sophia::SpriteComponent>();

    glm::mat4 model = sprite.getModel(transform);
    // Zero size or scale; nothing to hit.
    if (glm::determinant(model) == 0.f) {
      return;
    }

    glm::vec4 local =
        glm::inverse(model) * glm::vec4{world, transform.position.z, 1.f};
    if (local.x < 0.f || local.x > 1.f || local.y < 0.f || local.y > 1.f) {
      return;
    }

    // Same ordering the renderer draws in, so this is the one on top.
    std::tuple order{sprite.sortOrder, transform.position.z};
    if (!bestOrder || order >= *bestOrder) {
      best = entity;
      bestOrder = order;
    }
  });

  return best;
}