#include "scene_panel.hpp"

#include <imgui.h>
// After imgui.h, which it relies on.
#include <ImGuizmo.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <glm/gtc/type_ptr.hpp>
#include <limits>
#include <optional>
#include <sophia/aabb.hpp>
#include <sophia/components.hpp>
#include <tuple>
#include <utility>
#include <vector>

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

// Holding Ctrl snaps the gizmo to these steps.
constexpr float TRANSLATE_SNAP = 0.5f;
constexpr float ROTATE_SNAP_DEGREES = 15.f;
constexpr float SCALE_SNAP = 0.1f;

// Writes a matrix built from a translation, rotation and scale (no shear)
// back into `transform`. A mirroring matrix comes back with a negative x
// scale.
void decompose(const glm::mat4& matrix, sophia::TransformComponent& transform) {
  transform.position = glm::vec3{matrix[3]};

  glm::vec3 scale{glm::length(glm::vec3{matrix[0]}),
                  glm::length(glm::vec3{matrix[1]}),
                  glm::length(glm::vec3{matrix[2]})};
  if (glm::determinant(glm::mat3{matrix}) < 0.f) {
    scale.x = -scale.x;
  }
  transform.scale = scale;

  // A zero scale loses the axis, and with it the rotation; keep the old one.
  if (scale.x == 0.f || scale.y == 0.f || scale.z == 0.f) {
    return;
  }
  glm::mat3 rotation{glm::vec3{matrix[0]} / scale.x,
                     glm::vec3{matrix[1]} / scale.y,
                     glm::vec3{matrix[2]} / scale.z};
  transform.rotation = glm::normalize(glm::quat_cast(rotation));
}

// How far along a ray (from `origin`, along the normalized `direction`) it
// enters `bounds` placed by `model`, or null if it misses. A ray starting
// inside hits at 0.
std::optional<float> intersectBox(const glm::vec3& origin,
                                  const glm::vec3& direction,
                                  const glm::mat4& model,
                                  const sophia::AABB& bounds) {
  if (glm::determinant(model) == 0.f) {
    return std::nullopt;
  }

  // Test in the box's space, where it's axis-aligned. The direction isn't
  // renormalized, so distances along it are still world distances.
  glm::mat4 inverse = glm::inverse(model);
  glm::vec3 localOrigin = glm::vec3{inverse * glm::vec4{origin, 1.f}};
  glm::vec3 localDirection = glm::vec3{inverse * glm::vec4{direction, 0.f}};

  float enter = 0.f;
  float exit = std::numeric_limits<float>::infinity();
  for (int axis = 0; axis < 3; axis++) {
    if (std::abs(localDirection[axis]) < 1e-8f) {
      // Parallel to this pair of faces: hits only if already between them.
      if (localOrigin[axis] < bounds.min[axis] ||
          localOrigin[axis] > bounds.max[axis]) {
        return std::nullopt;
      }
      continue;
    }

    float toMin = (bounds.min[axis] - localOrigin[axis]) / localDirection[axis];
    float toMax = (bounds.max[axis] - localOrigin[axis]) / localDirection[axis];
    if (toMin > toMax) {
      std::swap(toMin, toMax);
    }
    enter = std::max(enter, toMin);
    exit = std::min(exit, toMax);
    if (enter > exit) {
      return std::nullopt;
    }
  }
  return enter;
}

// How far along a ray it hits the unit quad (0, 0)-(1, 1) in the XY plane,
// placed by `model`, or null if it misses.
std::optional<float> intersectQuad(const glm::vec3& origin,
                                   const glm::vec3& direction,
                                   const glm::mat4& model) {
  if (glm::determinant(model) == 0.f) {
    return std::nullopt;
  }

  glm::mat4 inverse = glm::inverse(model);
  glm::vec3 localOrigin = glm::vec3{inverse * glm::vec4{origin, 1.f}};
  glm::vec3 localDirection = glm::vec3{inverse * glm::vec4{direction, 0.f}};
  // Edge on to the quad.
  if (std::abs(localDirection.z) < 1e-8f) {
    return std::nullopt;
  }

  float distance = -localOrigin.z / localDirection.z;
  if (distance < 0.f) {
    return std::nullopt;
  }
  glm::vec3 hit = localOrigin + distance * localDirection;
  if (hit.x < 0.f || hit.x > 1.f || hit.y < 0.f || hit.y > 1.f) {
    return std::nullopt;
  }
  return distance;
}

}  // namespace

ImVec2 ScenePanel::ViewportTransform::clipToScreen(
    const glm::vec4& clip) const {
  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};
}

std::optional<ImVec2> ScenePanel::ViewportTransform::worldToScreen(
    const glm::vec3& world) const {
  glm::vec4 clip = viewProjection * glm::vec4{world, 1.f};
  if (clip.w <= 0.f) {
    return std::nullopt;
  }
  return clipToScreen(clip);
}

glm::vec2 ScenePanel::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;
}

ScenePanel::Ray ScenePanel::ViewportTransform::screenToRay(
    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};
  // Depth runs 0 at the near plane to 1 at the far one.
  glm::vec4 nearPoint = inverseViewProjection * glm::vec4{ndc, 0.f, 1.f};
  glm::vec4 farPoint = inverseViewProjection * glm::vec4{ndc, 1.f, 1.f};
  glm::vec3 from = glm::vec3{nearPoint} / nearPoint.w;
  glm::vec3 to = glm::vec3{farPoint} / farPoint.w;
  return {from, glm::normalize(to - from)};
}

void ScenePanel::ViewportTransform::drawLine(ImDrawList* drawList,
                                             const glm::vec3& a,
                                             const glm::vec3& b, ImU32 color,
                                             float thickness) const {
  // Just in front of the camera. Points behind it project mirrored, so the
  // part of the line there is cut off rather than drawn.
  constexpr float MIN_W = 1e-4f;

  glm::vec4 clipA = viewProjection * glm::vec4{a, 1.f};
  glm::vec4 clipB = viewProjection * glm::vec4{b, 1.f};
  if (clipA.w < MIN_W && clipB.w < MIN_W) {
    return;
  }
  if (clipA.w < MIN_W) {
    clipA = glm::mix(clipA, clipB, (MIN_W - clipA.w) / (clipB.w - clipA.w));
  } else if (clipB.w < MIN_W) {
    clipB = glm::mix(clipB, clipA, (MIN_W - clipB.w) / (clipA.w - clipB.w));
  }
  drawList->AddLine(clipToScreen(clipA), clipToScreen(clipB), color, thickness);
}

void ScenePanel::onUpdate(const sophia::FrameInfo& frameInfo) {
  syncSceneView(context, sceneViews, view, isShown());

  drawToolbar();
  ImGui::Separator();

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

  if (!context.hasScene()) {
    ImGui::TextDisabled("No scene open");
    return;
  }

  ui.image(target);

  // Rolls over ImGuizmo's "mouse is over a handle" state, so IsOver() below
  // reports last frame's.
  ImGuizmo::BeginFrame();
  // A handle can't be grabbed while any ImGui item is hovered, so while the
  // mouse is over one, leave out the viewport button, and keep the window
  // still so dragging the handle doesn't drag it. Unless a camera drag is
  // already under way: dropping the button would end it.
  bool overGizmo =
      (ImGuizmo::IsOver() || ImGuizmo::IsUsing()) && !viewportActive;
  getConfig().windowFlags =
      overGizmo ? WINDOW_FLAGS | ImGuiWindowFlags_NoMove : WINDOW_FLAGS;

  // The image doesn't take mouse input, so without this, dragging in it
  // would drag the window instead.
  ImGui::SetCursorScreenPos(origin);
  if (overGizmo) {
    ImGui::Dummy(size);
  } else {
    ImGui::InvisibleButton("##viewport", size,
                           ImGuiButtonFlags_MouseButtonLeft |
                               ImGuiButtonFlags_MouseButtonRight |
                               ImGuiButtonFlags_MouseButtonMiddle);
  }
  viewportActive = !overGizmo && ImGui::IsItemActive();

  handleCameraInput(frameInfo, size);
  handleGizmoHotkeys();

  // Drive the view's camera from the editor camera. The scene view renders
  // after the panels update, so this frame already uses it.
  camera.getComponent<sophia::TransformComponent>() =
      editorCamera.getTransform();
  camera.getComponent<sophia::CameraComponent>() = editorCamera.getCamera();

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

  // Screen-space, so it only lines up with anything in 2D.
  // TODO: a world-space grid on the XZ plane for 3D, once there's debug
  // drawing in scene views.
  if (showGrid && editorCamera.mode == EditorCamera::Mode::Scene2D) {
    drawGrid(origin, size);
  }

  // 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 view = sophia::CameraComponent::getView(
      camera.getComponent<sophia::TransformComponent>());
  glm::mat4 viewProjection =
      camera.getComponent<sophia::CameraComponent>().getProjection(aspect) *
      view;
  ViewportTransform viewport{origin, size, viewProjection,
                             glm::inverse(viewProjection)};

  handleSelectionInput(viewport);
  drawSelection(viewport);

  ImGuizmo::SetDrawlist();
  ImGuizmo::SetRect(origin.x, origin.y, size.x, size.y);
  drawGizmo(view, aspect);

  drawList->PopClipRect();

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

void ScenePanel::drawToolbar() {
  bool is2D = editorCamera.mode == EditorCamera::Mode::Scene2D;
  if (ImGui::Checkbox("2D", &is2D)) {
    editorCamera.mode =
        is2D ? EditorCamera::Mode::Scene2D : EditorCamera::Mode::Scene3D;
    dragging = false;
  }
  ImGui::SameLine();
  ImGui::BeginDisabled(!is2D);
  ImGui::Checkbox("Grid", &showGrid);
  ImGui::EndDisabled();

  ImGui::SameLine();
  ImGui::TextDisabled("|");
  ImGui::SameLine();
  const struct {
    const char* label;
    GizmoOperation operation;
  } operations[] = {{"Move (W)", GizmoOperation::Translate},
                    {"Rotate (E)", GizmoOperation::Rotate},
                    {"Scale (R)", GizmoOperation::Scale}};
  for (const auto& entry : operations) {
    if (ImGui::RadioButton(entry.label, gizmoOperation == entry.operation)) {
      gizmoOperation = entry.operation;
    }
    ImGui::SameLine();
  }
  if (ImGui::Button(gizmoLocal ? "Local" : "World")) {
    gizmoLocal = !gizmoLocal;
  }
}

void ScenePanel::handleGizmoHotkeys() {
  const ImGuiIO& io = ImGui::GetIO();
  // W/E/Q fly the camera while the right button is held.
  if (!ImGui::IsWindowFocused() || io.WantTextInput ||
      ImGui::IsMouseDown(ImGuiMouseButton_Right)) {
    return;
  }
  if (ImGui::IsKeyPressed(ImGuiKey_W, false)) {
    gizmoOperation = GizmoOperation::Translate;
  } else if (ImGui::IsKeyPressed(ImGuiKey_E, false)) {
    gizmoOperation = GizmoOperation::Rotate;
  } else if (ImGui::IsKeyPressed(ImGuiKey_R, false)) {
    gizmoOperation = GizmoOperation::Scale;
  }
}

void ScenePanel::drawGizmo(const glm::mat4& view, float aspect) {
  sophia::Entity selected = context.selected;
  if (!selected) {
    return;
  }

  bool is2D = editorCamera.mode == EditorCamera::Mode::Scene2D;

  // In 2D, only the axes that keep things in the XY plane.
  ImGuizmo::OPERATION operation{};
  float snapStep = 0.f;
  switch (gizmoOperation) {
    case GizmoOperation::Translate:
      operation = is2D ? ImGuizmo::TRANSLATE_X | ImGuizmo::TRANSLATE_Y
                       : ImGuizmo::TRANSLATE;
      snapStep = TRANSLATE_SNAP;
      break;
    case GizmoOperation::Rotate:
      operation = is2D ? ImGuizmo::ROTATE_Z : ImGuizmo::ROTATE;
      snapStep = ROTATE_SNAP_DEGREES;
      break;
    case GizmoOperation::Scale:
      operation =
          is2D ? ImGuizmo::SCALE_X | ImGuizmo::SCALE_Y : ImGuizmo::SCALE;
      snapStep = SCALE_SNAP;
      break;
  }
  std::array<float, 3> snap{snapStep, snapStep, snapStep};
  bool snapping = ImGui::GetIO().KeyCtrl;

  // ImGuizmo maps clip space to the screen the OpenGL way, with +y up, so it
  // needs the projection without the Y flip CameraComponent adds for Vulkan.
  glm::mat4 projection =
      camera.getComponent<sophia::CameraComponent>().getProjection(aspect);
  projection[1][1] *= -1.f;

  auto& transform = selected.getComponent<sophia::TransformComponent>();
  glm::mat4 model = transform.mat4();

  ImGuizmo::SetOrthographic(is2D);
  // Per entity, so a drag in progress doesn't carry over to a new selection.
  ImGuizmo::PushID(static_cast<int>(static_cast<sophia::u32>(selected)));
  bool changed = ImGuizmo::Manipulate(
      glm::value_ptr(view), glm::value_ptr(projection), operation,
      gizmoLocal ? ImGuizmo::LOCAL : ImGuizmo::WORLD, glm::value_ptr(model),
      nullptr, snapping ? snap.data() : nullptr);
  ImGuizmo::PopID();

  if (changed) {
    decompose(model, transform);
  }
}

void ScenePanel::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 ScenePanel::handleCameraInput(const sophia::FrameInfo& frameInfo,
                                   ImVec2 size) {
  const ImGuiIO& io = ImGui::GetIO();
  bool is3D = editorCamera.mode == EditorCamera::Mode::Scene3D;
  glm::vec2 mouseDelta{io.MouseDelta.x, io.MouseDelta.y};

  // The button stays active while any of its mouse buttons is held, even
  // once the mouse leaves the panel.
  if (ImGui::IsItemActive()) {
    bool rightDown = ImGui::IsMouseDown(ImGuiMouseButton_Right);
    bool middleDown = ImGui::IsMouseDown(ImGuiMouseButton_Middle);

    if (is3D && rightDown) {
      editorCamera.orbit(mouseDelta);

      glm::vec3 direction{0.f};
      if (ImGui::IsKeyDown(ImGuiKey_D)) direction.x += 1.f;
      if (ImGui::IsKeyDown(ImGuiKey_A)) direction.x -= 1.f;
      if (ImGui::IsKeyDown(ImGuiKey_E)) direction.y += 1.f;
      if (ImGui::IsKeyDown(ImGuiKey_Q)) direction.y -= 1.f;
      if (ImGui::IsKeyDown(ImGuiKey_S)) direction.z += 1.f;
      if (ImGui::IsKeyDown(ImGuiKey_W)) direction.z -= 1.f;
      editorCamera.fly(direction, static_cast<float>(frameInfo.deltaTime));
    } else if (middleDown || (!is3D && rightDown)) {
      editorCamera.pan(mouseDelta, size.y);
    }
  }

  if (ImGui::IsItemHovered() && io.MouseWheel != 0.f) {
    editorCamera.zoom(io.MouseWheel);
  }

  if (ImGui::IsWindowFocused() && !io.WantTextInput &&
      ImGui::IsKeyPressed(ImGuiKey_F, false)) {
    frameSelection();
  }
}

void ScenePanel::handleSelectionInput(const ViewportTransform& viewport) {
  const ImGuiIO& io = ImGui::GetIO();

  if (ImGui::IsItemClicked(ImGuiMouseButton_Left)) {
    context.selected = pick(viewport.screenToRay(io.MousePos));
    // Dragging the selection around is 2D only; in 3D, use the gizmo.
    dragging =
        context.selected && editorCamera.mode == EditorCamera::Mode::Scene2D;
  }

  if (!ImGui::IsItemActive() || !ImGui::IsMouseDown(ImGuiMouseButton_Left)) {
    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 =
        context.selected.getComponent<sophia::TransformComponent>();
    transform.position.x += delta.x;
    transform.position.y += delta.y;
  }
}

void ScenePanel::drawSelection(const ViewportTransform& viewport) {
  sophia::Entity selected = context.selected;
  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);
    for (size_t i = 0; i < UNIT_QUAD.size(); i++) {
      const glm::vec2& from = UNIT_QUAD[i];
      const glm::vec2& to = UNIT_QUAD[(i + 1) % UNIT_QUAD.size()];
      viewport.drawLine(drawList, glm::vec3{model * glm::vec4{from, 0.f, 1.f}},
                        glm::vec3{model * glm::vec4{to, 0.f, 1.f}},
                        SELECTION_COLOR, 2.0f);
    }
  }

  if (std::optional<sophia::AABB> bounds = getMeshBounds(selected)) {
    glm::mat4 model = transform.mat4();
    std::array<glm::vec3, 8> corners = bounds->getCorners();
    for (glm::vec3& corner : corners) {
      corner = glm::vec3{model * glm::vec4{corner, 1.f}};
    }
    // Corners whose indices differ in one bit share an edge.
    for (int i = 0; i < 8; i++) {
      for (int bit : {1, 2, 4}) {
        if ((i & bit) == 0) {
          viewport.drawLine(drawList, corners[i], corners[i | bit],
                            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 or mesh, such as the
  // camera.
  if (std::optional<ImVec2> center =
          viewport.worldToScreen(transform.position)) {
    drawList->AddCircle(*center, 4.0f, SELECTION_COLOR, 0, 2.0f);
  }
}

void ScenePanel::frameSelection() {
  sophia::Entity selected = context.selected;
  if (!selected) {
    return;
  }

  const auto& transform = selected.getComponent<sophia::TransformComponent>();

  // The selection's corners in world space, from its mesh or sprite.
  std::vector<glm::vec3> points;
  if (std::optional<sophia::AABB> bounds = getMeshBounds(selected)) {
    glm::mat4 model = transform.mat4();
    for (const glm::vec3& corner : bounds->getCorners()) {
      points.push_back(glm::vec3{model * glm::vec4{corner, 1.f}});
    }
  }
  if (selected.hasComponent<sophia::SpriteComponent>()) {
    glm::mat4 model =
        selected.getComponent<sophia::SpriteComponent>().getModel(transform);
    for (const glm::vec2& corner : UNIT_QUAD) {
      points.push_back(glm::vec3{model * glm::vec4{corner, 0.f, 1.f}});
    }
  }

  if (points.empty()) {
    // Nothing with a size: frame a unit around its position.
    editorCamera.frame(transform.position, 0.5f);
    return;
  }

  glm::vec3 min = points.front();
  glm::vec3 max = points.front();
  for (const glm::vec3& point : points) {
    min = glm::min(min, point);
    max = glm::max(max, point);
  }
  editorCamera.frame((min + max) * 0.5f, glm::length(max - min) * 0.5f);
}

std::optional<sophia::AABB> ScenePanel::getMeshBounds(sophia::Entity entity) {
  if (!entity.hasComponent<sophia::MeshComponent>()) {
    return std::nullopt;
  }
  sophia::MeshHandle mesh = entity.getComponent<sophia::MeshComponent>().handle;
  if (!resources.isAlive(mesh)) {
    return std::nullopt;
  }
  return resources.getBounds(mesh);
}

sophia::Entity ScenePanel::pick(const Ray& ray) {
  sophia::Scene& scene = context.getScene();

  // Meshes are opaque and depth tested, so the nearest one hides anything
  // behind it.
  sophia::Entity nearestMesh;
  float meshDistance = std::numeric_limits<float>::infinity();
  scene.forEachEntity([&](sophia::Entity entity) {
    std::optional<sophia::AABB> bounds = getMeshBounds(entity);
    if (!bounds) {
      return;
    }
    glm::mat4 model = entity.getComponent<sophia::TransformComponent>().mat4();
    std::optional<float> distance =
        intersectBox(ray.origin, ray.direction, model, *bounds);
    if (distance && *distance < meshDistance) {
      nearestMesh = entity;
      meshDistance = *distance;
    }
  });

  // Sprites in front of that mesh draw over it, in the renderer's order:
  // higher sortOrder on top, then nearer.
  sophia::Entity topSprite;
  std::optional<std::tuple<sophia::s32, float>> topOrder;
  scene.forEachEntity([&](sophia::Entity entity) {
    if (!entity.hasComponent<sophia::SpriteComponent>()) {
      return;
    }
    const auto& sprite = entity.getComponent<sophia::SpriteComponent>();
    glm::mat4 model =
        sprite.getModel(entity.getComponent<sophia::TransformComponent>());
    std::optional<float> distance =
        intersectQuad(ray.origin, ray.direction, model);
    if (!distance || *distance > meshDistance) {
      return;
    }

    // Negated distance, so nearer compares greater.
    std::tuple order{sprite.sortOrder, -*distance};
    if (!topOrder || order > *topOrder) {
      topSprite = entity;
      topOrder = order;
    }
  });

  return topSprite ? topSprite : nearestMesh;
}