#include "Features/Scene/SceneViewportTransformGizmo.h"

#include "Features/Scene/SceneViewportTransformGizmoSupport.h"
#include "Scene/EditorSceneRuntime.h"
#include "Scene/SceneToolState.h"

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Components/MeshFilterComponent.h>
#include <XCEngine/Components/TransformComponent.h>
#include <XCEngine/Resources/Mesh/Mesh.h>
#include <XCEngine/Scene/Scene.h>

#include <utility>

namespace XCEngine::UI::Editor::App {

namespace {

using ::XCEngine::Components::CameraComponent;
using ::XCEngine::Components::GameObject;
using ::XCEngine::Components::MeshFilterComponent;
using ::XCEngine::Components::TransformComponent;
using ::XCEngine::Math::Color;
using ::XCEngine::Math::Quaternion;
using ::XCEngine::Math::Vector2;
using ::XCEngine::Math::Vector3;
using ::XCEngine::UI::UIColor;
using ::XCEngine::UI::UIPoint;
using ::XCEngine::UI::UIRect;
namespace SceneViewportGizmoInternal = ::XCEngine::UI::Editor::App::SceneViewportGizmoInternal;
using SceneViewportGizmoInternal::BuildSceneViewportTransformGizmoOverlayFrameData;
using SceneViewportGizmoInternal::IUndoManager;
using SceneViewportGizmoInternal::SceneViewportMoveGizmo;
using SceneViewportGizmoInternal::SceneViewportMoveGizmoContext;
using SceneViewportGizmoInternal::SceneViewportOverlayData;
using SceneViewportGizmoInternal::SceneViewportOverlayFrameData;
using SceneViewportGizmoInternal::SceneViewportRotateGizmo;
using SceneViewportGizmoInternal::SceneViewportRotateGizmoContext;
using SceneViewportGizmoInternal::SceneViewportScaleGizmo;
using SceneViewportGizmoInternal::SceneViewportScaleGizmoContext;
using SceneViewportGizmoInternal::SceneViewportTransformGizmoHandleBuildInputs;

enum class ActiveTransformGizmoKind : std::uint8_t {
    None = 0,
    Move,
    Rotate,
    Scale
};

struct TransformGizmoSelectionState {
    GameObject* primaryObject = nullptr;
    std::vector<GameObject*> selectedObjects = {};
    Vector3 pivotWorldPosition = Vector3::Zero();
    Quaternion primaryWorldRotation = Quaternion::Identity();
};

UIColor ToUIColor(const Color& color) {
    return UIColor(color.r, color.g, color.b, color.a);
}

UIPoint ToScreenPoint(const Vector2& point, const UIRect& viewportRect) {
    return UIPoint(viewportRect.x + point.x, viewportRect.y + point.y);
}

Vector2 ToLocalPoint(const UIRect& viewportRect, const UIPoint& point) {
    return Vector2(point.x - viewportRect.x, point.y - viewportRect.y);
}

Quaternion ComputeStableWorldRotation(const GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Quaternion::Identity();
    }

    return gameObject->GetTransform()->GetRotation().Normalized();
}

Vector3 ResolvePivotWorldPosition(const GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Vector3::Zero();
    }

    return gameObject->GetTransform()->GetPosition();
}

Vector3 ResolveCenterWorldPosition(const GameObject* gameObject) {
    if (gameObject == nullptr || gameObject->GetTransform() == nullptr) {
        return Vector3::Zero();
    }

    if (const MeshFilterComponent* meshFilter =
            gameObject->GetComponent<MeshFilterComponent>();
        meshFilter != nullptr) {
        if (::XCEngine::Resources::Mesh* mesh = meshFilter->GetMesh();
            mesh != nullptr && mesh->IsValid()) {
            return gameObject->GetTransform()->TransformPoint(mesh->GetBounds().center);
        }
    }

    return gameObject->GetTransform()->GetPosition();
}

TransformGizmoSelectionState BuildSelectionState(
    EditorSceneRuntime& sceneRuntime,
    bool useCenterPivot) {
    TransformGizmoSelectionState state = {};
    const auto selectedId = sceneRuntime.GetSelectedGameObjectId();
    if (!selectedId.has_value()) {
        return state;
    }

    GameObject* selectedObject =
        sceneRuntime.GetActiveScene() != nullptr
            ? sceneRuntime.GetActiveScene()->FindByID(selectedId.value())
            : nullptr;
    if (selectedObject == nullptr || selectedObject->GetTransform() == nullptr) {
        return state;
    }

    state.primaryObject = selectedObject;
    state.selectedObjects.push_back(selectedObject);
    state.primaryWorldRotation = ComputeStableWorldRotation(selectedObject);
    state.pivotWorldPosition = useCenterPivot
        ? ResolveCenterWorldPosition(selectedObject)
        : ResolvePivotWorldPosition(selectedObject);
    return state;
}

SceneViewportOverlayData BuildOverlayData(const EditorSceneRuntime& sceneRuntime) {
    SceneViewportOverlayData overlay = {};
    const CameraComponent* camera = sceneRuntime.GetSceneViewCamera();
    if (camera == nullptr || camera->GetGameObject() == nullptr) {
        return overlay;
    }

    const TransformComponent* transform = camera->GetGameObject()->GetTransform();
    if (transform == nullptr) {
        return overlay;
    }

    overlay.valid = true;
    overlay.cameraPosition = transform->GetPosition();
    overlay.cameraForward = transform->GetForward();
    overlay.cameraRight = transform->GetRight();
    overlay.cameraUp = transform->GetUp();
    overlay.verticalFovDegrees = camera->GetFieldOfView();
    overlay.nearClipPlane = camera->GetNearClipPlane();
    overlay.farClipPlane = camera->GetFarClipPlane();
    overlay.orbitDistance = sceneRuntime.GetSceneViewOrbitDistance();
    return overlay;
}

SceneViewportMoveGizmoContext BuildMoveContext(
    const TransformGizmoSelectionState& selection,
    const SceneViewportOverlayData& overlay,
    const UIRect& viewportRect,
    const Vector2& localPointer,
    bool localSpace) {
    SceneViewportMoveGizmoContext context = {};
    context.overlay = overlay;
    context.viewportSize = Vector2(viewportRect.width, viewportRect.height);
    context.mousePosition = localPointer;
    context.selectedObject = selection.primaryObject;
    context.selectedObjects = selection.selectedObjects;
    context.pivotWorldPosition = selection.pivotWorldPosition;
    context.axisOrientation = localSpace
        ? selection.primaryWorldRotation
        : Quaternion::Identity();
    return context;
}

SceneViewportRotateGizmoContext BuildRotateContext(
    const TransformGizmoSelectionState& selection,
    const SceneViewportOverlayData& overlay,
    const UIRect& viewportRect,
    const Vector2& localPointer,
    bool localSpace,
    bool useCenterPivot) {
    SceneViewportRotateGizmoContext context = {};
    context.overlay = overlay;
    context.viewportSize = Vector2(viewportRect.width, viewportRect.height);
    context.mousePosition = localPointer;
    context.selectedObject = selection.primaryObject;
    context.selectedObjects = selection.selectedObjects;
    context.pivotWorldPosition = selection.pivotWorldPosition;
    context.axisOrientation = localSpace
        ? selection.primaryWorldRotation
        : Quaternion::Identity();
    context.localSpace = localSpace;
    context.rotateAroundSharedPivot = useCenterPivot;
    return context;
}

SceneViewportScaleGizmoContext BuildScaleContext(
    const TransformGizmoSelectionState& selection,
    const SceneViewportOverlayData& overlay,
    const UIRect& viewportRect,
    const Vector2& localPointer,
    bool localSpace) {
    SceneViewportScaleGizmoContext context = {};
    context.overlay = overlay;
    context.viewportSize = Vector2(viewportRect.width, viewportRect.height);
    context.mousePosition = localPointer;
    context.selectedObject = selection.primaryObject;
    context.pivotWorldPosition = selection.pivotWorldPosition;
    context.axisOrientation = localSpace
        ? selection.primaryWorldRotation
        : Quaternion::Identity();
    return context;
}

ActiveTransformGizmoKind ResolveActiveGizmoKind(
    const SceneViewportMoveGizmo& moveGizmo,
    const SceneViewportRotateGizmo& rotateGizmo,
    const SceneViewportScaleGizmo& scaleGizmo) {
    if (moveGizmo.IsActive()) {
        return ActiveTransformGizmoKind::Move;
    }
    if (rotateGizmo.IsActive()) {
        return ActiveTransformGizmoKind::Rotate;
    }
    if (scaleGizmo.IsActive()) {
        return ActiveTransformGizmoKind::Scale;
    }

    return ActiveTransformGizmoKind::None;
}

bool IsTransformToolMode(SceneToolMode mode) {
    return mode == SceneToolMode::Transform;
}

bool ShouldShowMoveGizmo(SceneToolMode mode) {
    return mode == SceneToolMode::Translate || IsTransformToolMode(mode);
}

bool ShouldShowRotateGizmo(SceneToolMode mode) {
    return mode == SceneToolMode::Rotate || IsTransformToolMode(mode);
}

bool ShouldShowScaleGizmo(SceneToolMode mode) {
    return mode == SceneToolMode::Scale || IsTransformToolMode(mode);
}

Vector2 ResolveUpdatePointerPosition(
    const Vector2& pointerPosition,
    bool hoverEnabled,
    ActiveTransformGizmoKind activeKind,
    ActiveTransformGizmoKind gizmoKind) {
    if (!hoverEnabled && activeKind == ActiveTransformGizmoKind::None) {
        return Vector2(-1.0f, -1.0f);
    }

    if (activeKind != ActiveTransformGizmoKind::None && activeKind != gizmoKind) {
        return Vector2(-1.0f, -1.0f);
    }

    return pointerPosition;
}

class SceneGizmoUndoBridge final : public IUndoManager {
public:
    void Bind(EditorSceneRuntime& sceneRuntime) {
        m_sceneRuntime = &sceneRuntime;
    }

    void BeginInteractiveChange(const std::string& label) override {
        if (m_sceneRuntime == nullptr || HasPendingInteractiveChange()) {
            return;
        }

        SceneTransformSnapshot snapshot = {};
        if (!m_sceneRuntime->CaptureSelectedTransformSnapshot(snapshot)) {
            return;
        }

        m_pendingLabel = label;
        m_beforeSnapshot = snapshot;
    }

    bool HasPendingInteractiveChange() const override {
        return m_beforeSnapshot.IsValid();
    }

    void FinalizeInteractiveChange() override {
        if (m_sceneRuntime == nullptr || !HasPendingInteractiveChange()) {
            return;
        }

        SceneTransformSnapshot afterSnapshot = {};
        m_sceneRuntime->CaptureSelectedTransformSnapshot(afterSnapshot);
        m_sceneRuntime->RecordTransformEdit(m_beforeSnapshot, afterSnapshot);
        m_pendingLabel.clear();
        m_beforeSnapshot = {};
    }

    void CancelInteractiveChange() override {
        if (m_sceneRuntime == nullptr || !HasPendingInteractiveChange()) {
            return;
        }

        m_sceneRuntime->ApplyTransformSnapshot(m_beforeSnapshot);
        m_pendingLabel.clear();
        m_beforeSnapshot = {};
    }

private:
    EditorSceneRuntime* m_sceneRuntime = nullptr;
    std::string m_pendingLabel = {};
    SceneTransformSnapshot m_beforeSnapshot = {};
};

} // namespace

struct SceneViewportTransformGizmo::State {
    SceneGizmoUndoBridge undoBridge = {};
    SceneViewportMoveGizmo moveGizmo = {};
    SceneViewportRotateGizmo rotateGizmo = {};
    SceneViewportScaleGizmo scaleGizmo = {};
    SceneViewportMoveGizmoContext moveContext = {};
    SceneViewportRotateGizmoContext rotateContext = {};
    SceneViewportScaleGizmoContext scaleContext = {};
    SceneViewportTransformGizmoFrame frame = {};
};

SceneViewportTransformGizmo::SceneViewportTransformGizmo()
    : m_state(std::make_unique<State>()) {
}

SceneViewportTransformGizmo::~SceneViewportTransformGizmo() = default;

SceneViewportTransformGizmo::SceneViewportTransformGizmo(SceneViewportTransformGizmo&&) noexcept =
    default;

SceneViewportTransformGizmo& SceneViewportTransformGizmo::operator=(
    SceneViewportTransformGizmo&&) noexcept = default;

void SceneViewportTransformGizmo::Refresh(
    EditorSceneRuntime& sceneRuntime,
    const UIRect& viewportRect,
    const UIPoint& pointerScreen,
    bool hoverEnabled) {
    State& state = *m_state;
    state.undoBridge.Bind(sceneRuntime);
    state.frame = {};
    state.frame.clipRect = viewportRect;

    if (viewportRect.width <= 1.0f || viewportRect.height <= 1.0f) {
        return;
    }

    const SceneViewportOverlayData overlay = BuildOverlayData(sceneRuntime);
    if (!overlay.valid) {
        CancelDrag(sceneRuntime);
        return;
    }

    const bool useCenterPivot =
        sceneRuntime.GetToolPivotMode() == SceneToolPivotMode::Center;
    const bool localSpace =
        sceneRuntime.GetToolSpaceMode() == SceneToolSpaceMode::Local;
    const TransformGizmoSelectionState selection =
        BuildSelectionState(sceneRuntime, useCenterPivot);
    if (selection.primaryObject == nullptr) {
        CancelDrag(sceneRuntime);
        return;
    }

    const SceneToolMode toolMode = sceneRuntime.GetToolMode();
    if (toolMode == SceneToolMode::View) {
        CancelDrag(sceneRuntime);
        return;
    }

    Vector2 localPointer = ToLocalPoint(viewportRect, pointerScreen);
    const bool usingTransformTool = IsTransformToolMode(toolMode);
    const bool showingMoveGizmo = ShouldShowMoveGizmo(toolMode);
    const bool showingRotateGizmo = ShouldShowRotateGizmo(toolMode);
    const bool showingScaleGizmo = ShouldShowScaleGizmo(toolMode);
    SceneViewportTransformGizmoHandleBuildInputs inputs = {};
    if (showingMoveGizmo) {
        state.moveContext = BuildMoveContext(
            selection,
            overlay,
            viewportRect,
            localPointer,
            localSpace);
        if (state.moveGizmo.IsActive() &&
            state.moveContext.selectedObject != nullptr &&
            state.moveContext.selectedObject->GetID() !=
                state.moveGizmo.GetActiveEntityId()) {
            state.moveGizmo.CancelDrag(&state.undoBridge);
        }
    } else if (state.moveGizmo.IsActive()) {
        state.moveGizmo.CancelDrag(&state.undoBridge);
    }

    if (showingRotateGizmo) {
        state.rotateContext = BuildRotateContext(
            selection,
            overlay,
            viewportRect,
            localPointer,
            localSpace,
            useCenterPivot);
        if (state.rotateGizmo.IsActive() &&
            state.rotateContext.selectedObject != nullptr &&
            state.rotateContext.selectedObject->GetID() !=
                state.rotateGizmo.GetActiveEntityId()) {
            state.rotateGizmo.CancelDrag(&state.undoBridge);
        }
    } else if (state.rotateGizmo.IsActive()) {
        state.rotateGizmo.CancelDrag(&state.undoBridge);
    }

    if (showingScaleGizmo) {
        state.scaleContext = BuildScaleContext(
            selection,
            overlay,
            viewportRect,
            localPointer,
            localSpace);
        state.scaleContext.uniformOnly = usingTransformTool;
        if (state.scaleGizmo.IsActive() &&
            state.scaleContext.selectedObject != nullptr &&
            state.scaleContext.selectedObject->GetID() !=
                state.scaleGizmo.GetActiveEntityId()) {
            state.scaleGizmo.CancelDrag(&state.undoBridge);
        }
    } else if (state.scaleGizmo.IsActive()) {
        state.scaleGizmo.CancelDrag(&state.undoBridge);
    }

    const ActiveTransformGizmoKind activeKind = ResolveActiveGizmoKind(
        state.moveGizmo,
        state.rotateGizmo,
        state.scaleGizmo);

    if (showingMoveGizmo) {
        SceneViewportMoveGizmoContext updateContext = state.moveContext;
        updateContext.mousePosition = ResolveUpdatePointerPosition(
            state.moveContext.mousePosition,
            hoverEnabled,
            activeKind,
            ActiveTransformGizmoKind::Move);
        state.moveGizmo.Update(updateContext);
        inputs.moveGizmo = &state.moveGizmo.GetDrawData();
        inputs.moveEntityId = selection.primaryObject->GetID();
    }

    if (showingRotateGizmo) {
        SceneViewportRotateGizmoContext updateContext = state.rotateContext;
        updateContext.mousePosition = ResolveUpdatePointerPosition(
            state.rotateContext.mousePosition,
            hoverEnabled,
            activeKind,
            ActiveTransformGizmoKind::Rotate);
        state.rotateGizmo.Update(updateContext);
        inputs.rotateGizmo = &state.rotateGizmo.GetDrawData();
        inputs.rotateEntityId = selection.primaryObject->GetID();
    }

    if (showingScaleGizmo) {
        SceneViewportScaleGizmoContext updateContext = state.scaleContext;
        updateContext.mousePosition = ResolveUpdatePointerPosition(
            state.scaleContext.mousePosition,
            hoverEnabled,
            activeKind,
            ActiveTransformGizmoKind::Scale);
        state.scaleGizmo.Update(updateContext);
        inputs.scaleGizmo = &state.scaleGizmo.GetDrawData();
        inputs.scaleEntityId = selection.primaryObject->GetID();
    }

    const SceneViewportOverlayFrameData overlayFrame =
        BuildSceneViewportTransformGizmoOverlayFrameData(overlay, inputs);
    if (overlayFrame.screenTriangles.empty()) {
        return;
    }

    state.frame.visible = true;
    state.frame.triangles.reserve(overlayFrame.screenTriangles.size());
    for (const auto& triangle : overlayFrame.screenTriangles) {
        SceneViewportTransformGizmoTriangle triangleFrame = {};
        triangleFrame.a =
            ToScreenPoint(triangle.vertices[0].screenPosition, viewportRect);
        triangleFrame.b =
            ToScreenPoint(triangle.vertices[1].screenPosition, viewportRect);
        triangleFrame.c =
            ToScreenPoint(triangle.vertices[2].screenPosition, viewportRect);
        triangleFrame.color = ToUIColor(triangle.vertices[0].color);
        state.frame.triangles.push_back(std::move(triangleFrame));
    }
}

bool SceneViewportTransformGizmo::TryBeginDrag(EditorSceneRuntime& sceneRuntime) {
    State& state = *m_state;
    state.undoBridge.Bind(sceneRuntime);

    switch (sceneRuntime.GetToolMode()) {
    case SceneToolMode::Translate:
        return state.moveGizmo.TryBeginDrag(state.moveContext, state.undoBridge);
    case SceneToolMode::Rotate:
        return state.rotateGizmo.TryBeginDrag(state.rotateContext, state.undoBridge);
    case SceneToolMode::Scale:
        return state.scaleGizmo.TryBeginDrag(state.scaleContext, state.undoBridge);
    case SceneToolMode::Transform:
        if (state.scaleGizmo.EvaluateHit(state.scaleContext.mousePosition).HasHit()) {
            return state.scaleGizmo.TryBeginDrag(
                state.scaleContext,
                state.undoBridge);
        }
        if (state.moveGizmo.EvaluateHit(state.moveContext.mousePosition).HasHit()) {
            return state.moveGizmo.TryBeginDrag(
                state.moveContext,
                state.undoBridge);
        }
        if (state.rotateGizmo.EvaluateHit(state.rotateContext.mousePosition).HasHit()) {
            return state.rotateGizmo.TryBeginDrag(
                state.rotateContext,
                state.undoBridge);
        }
        return false;
    case SceneToolMode::View:
    default:
        return false;
    }
}

bool SceneViewportTransformGizmo::UpdateDrag(EditorSceneRuntime& sceneRuntime) {
    State& state = *m_state;
    state.undoBridge.Bind(sceneRuntime);

    switch (ResolveActiveGizmoKind(
        state.moveGizmo,
        state.rotateGizmo,
        state.scaleGizmo)) {
    case ActiveTransformGizmoKind::Move:
        state.moveGizmo.UpdateDrag(state.moveContext);
        return true;
    case ActiveTransformGizmoKind::Rotate:
        state.rotateGizmo.UpdateDrag(state.rotateContext);
        return true;
    case ActiveTransformGizmoKind::Scale:
        state.scaleGizmo.UpdateDrag(state.scaleContext);
        return true;
    case ActiveTransformGizmoKind::None:
    default:
        return false;
    }
}

bool SceneViewportTransformGizmo::EndDrag(EditorSceneRuntime& sceneRuntime) {
    State& state = *m_state;
    state.undoBridge.Bind(sceneRuntime);

    switch (ResolveActiveGizmoKind(
        state.moveGizmo,
        state.rotateGizmo,
        state.scaleGizmo)) {
    case ActiveTransformGizmoKind::Move:
        state.moveGizmo.EndDrag(state.undoBridge);
        return true;
    case ActiveTransformGizmoKind::Rotate:
        state.rotateGizmo.EndDrag(state.undoBridge);
        return true;
    case ActiveTransformGizmoKind::Scale:
        state.scaleGizmo.EndDrag(state.undoBridge);
        return true;
    case ActiveTransformGizmoKind::None:
    default:
        return false;
    }
}

void SceneViewportTransformGizmo::CancelDrag(EditorSceneRuntime& sceneRuntime) {
    State& state = *m_state;
    state.undoBridge.Bind(sceneRuntime);

    if (state.moveGizmo.IsActive()) {
        state.moveGizmo.CancelDrag(&state.undoBridge);
    }
    if (state.rotateGizmo.IsActive()) {
        state.rotateGizmo.CancelDrag(&state.undoBridge);
    }
    if (state.scaleGizmo.IsActive()) {
        state.scaleGizmo.CancelDrag(&state.undoBridge);
    }
}

void SceneViewportTransformGizmo::ResetVisualState() {
    if (m_state == nullptr) {
        return;
    }

    m_state->frame = {};
}

bool SceneViewportTransformGizmo::IsActive() const {
    if (m_state == nullptr) {
        return false;
    }

    return ResolveActiveGizmoKind(
               m_state->moveGizmo,
               m_state->rotateGizmo,
               m_state->scaleGizmo) != ActiveTransformGizmoKind::None;
}

bool SceneViewportTransformGizmo::IsHoveringHandle() const {
    if (m_state == nullptr) {
        return false;
    }

    return m_state->moveGizmo.IsHoveringHandle() ||
        m_state->rotateGizmo.IsHoveringHandle() ||
        m_state->scaleGizmo.IsHoveringHandle();
}

const SceneViewportTransformGizmoFrame& SceneViewportTransformGizmo::GetFrame() const {
    return m_state->frame;
}

void AppendSceneViewportTransformGizmo(
    ::XCEngine::UI::UIDrawList& drawList,
    const SceneViewportTransformGizmoFrame& frame) {
    if (!frame.visible || frame.triangles.empty()) {
        return;
    }

    drawList.PushClipRect(frame.clipRect);
    for (const SceneViewportTransformGizmoTriangle& triangle : frame.triangles) {
        drawList.AddFilledTriangle(
            triangle.a,
            triangle.b,
            triangle.c,
            triangle.color);
    }
    drawList.PopClipRect();
}

} // namespace XCEngine::UI::Editor::App