#include "SceneViewportOverlayProviders.h"

#include <XCEngine/Components/CameraComponent.h>
#include <XCEngine/Components/GameObject.h>
#include <XCEngine/Components/LightComponent.h>
#include <XCEngine/Components/TransformComponent.h>
#include <XCEngine/Core/Math/Rect.h>
#include <XCEngine/Scene/Scene.h>

#include "SceneViewportOverlayHandleBuilder.h"
#include "SceneViewportMath.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <utility>

namespace XCEngine {
namespace Editor {

namespace {

bool CanBuildOverlayForGameObject(const Components::GameObject* gameObject) {
    return gameObject != nullptr &&
        gameObject->GetTransform() != nullptr &&
        gameObject->IsActiveInHierarchy();
}

float ResolveCameraAspect(
    const Components::CameraComponent& camera,
    uint32_t viewportWidth,
    uint32_t viewportHeight) {
    const Math::Rect viewportRect = camera.GetViewportRect();
    const float resolvedWidth = static_cast<float>(viewportWidth) *
        (viewportRect.width > Math::EPSILON ? viewportRect.width : 1.0f);
    const float resolvedHeight = static_cast<float>(viewportHeight) *
        (viewportRect.height > Math::EPSILON ? viewportRect.height : 1.0f);
    return resolvedHeight > Math::EPSILON
        ? resolvedWidth / resolvedHeight
        : 1.0f;
}

float ComputeWorldUnitsPerPixel(
    const SceneViewportOverlayData& overlay,
    const Math::Vector3& worldPoint,
    uint32_t viewportHeight) {
    if (!overlay.valid || viewportHeight <= 1u) {
        return 0.0f;
    }

    const Math::Vector3 cameraForward = overlay.cameraForward.Normalized();
    const float depth = Math::Vector3::Dot(worldPoint - overlay.cameraPosition, cameraForward);
    if (depth <= Math::EPSILON) {
        return 0.0f;
    }

    return 2.0f * depth * std::tan(overlay.verticalFovDegrees * Math::DEG_TO_RAD * 0.5f) /
        static_cast<float>(viewportHeight);
}

void AppendWorldLine(
    SceneViewportOverlayFrameData& frameData,
    const Math::Vector3& startWorld,
    const Math::Vector3& endWorld,
    const Math::Color& color,
    float thicknessPixels,
    SceneViewportOverlayDepthMode depthMode) {
    SceneViewportOverlayLinePrimitive& line = frameData.worldLines.emplace_back();
    line.startWorld = startWorld;
    line.endWorld = endWorld;
    line.color = color;
    line.thicknessPixels = thicknessPixels;
    line.depthMode = depthMode;
}

void AppendWorldSprite(
    SceneViewportOverlayFrameData& frameData,
    const Math::Vector3& worldPosition,
    const Math::Vector2& sizePixels,
    const Math::Color& tintColor,
    float sortDepth,
    uint64_t entityId,
    SceneViewportOverlaySpriteTextureKind textureKind,
    SceneViewportOverlayDepthMode depthMode) {
    if (entityId == 0 || sizePixels.x <= Math::EPSILON || sizePixels.y <= Math::EPSILON) {
        return;
    }

    SceneViewportOverlaySpritePrimitive& sprite = frameData.worldSprites.emplace_back();
    sprite.worldPosition = worldPosition;
    sprite.sizePixels = sizePixels;
    sprite.tintColor = tintColor;
    sprite.sortDepth = sortDepth;
    sprite.entityId = entityId;
    sprite.textureKind = textureKind;
    sprite.depthMode = depthMode;
}

void AppendHandleRecord(
    SceneViewportOverlayFrameData& frameData,
    SceneViewportOverlayHandleKind kind,
    uint64_t handleId,
    uint64_t entityId,
    const Math::Vector3& worldPosition,
    const Math::Vector2& sizePixels,
    float sortDepth) {
    if (kind == SceneViewportOverlayHandleKind::None ||
        handleId == 0 ||
        entityId == 0 ||
        sizePixels.x <= Math::EPSILON ||
        sizePixels.y <= Math::EPSILON) {
        return;
    }

    SceneViewportOverlayHandleRecord& handleRecord = frameData.handleRecords.emplace_back();
    handleRecord.kind = kind;
    handleRecord.handleId = handleId;
    handleRecord.entityId = entityId;
    handleRecord.shape = SceneViewportOverlayHandleShape::WorldRect;
    handleRecord.priority = Detail::kSceneViewportHandlePrioritySceneIcon;
    handleRecord.worldPosition = worldPosition;
    handleRecord.sizePixels = sizePixels;
    handleRecord.sortDepth = sortDepth;
}

void AppendSceneIconOverlay(
    SceneViewportOverlayFrameData& frameData,
    const SceneViewportOverlayData& overlay,
    uint32_t viewportWidth,
    uint32_t viewportHeight,
    const Components::GameObject& gameObject,
    const Math::Vector2& sizePixels,
    SceneViewportOverlaySpriteTextureKind textureKind) {
    const Components::TransformComponent* transform = gameObject.GetTransform();
    if (transform == nullptr) {
        return;
    }

    const SceneViewportProjectedPoint projectedPoint = ProjectSceneViewportWorldPoint(
        overlay,
        static_cast<float>(viewportWidth),
        static_cast<float>(viewportHeight),
        transform->GetPosition());
    if (!projectedPoint.visible) {
        return;
    }

    AppendWorldSprite(
        frameData,
        transform->GetPosition(),
        sizePixels,
        Math::Color::White(),
        projectedPoint.ndcDepth,
        gameObject.GetID(),
        textureKind,
        SceneViewportOverlayDepthMode::AlwaysOnTop);
    AppendHandleRecord(
        frameData,
        SceneViewportOverlayHandleKind::SceneIcon,
        gameObject.GetID(),
        gameObject.GetID(),
        transform->GetPosition(),
        sizePixels,
        projectedPoint.ndcDepth);
}

void AppendCameraFrustumOverlay(
    SceneViewportOverlayFrameData& frameData,
    const Components::CameraComponent& camera,
    const Components::GameObject& gameObject,
    uint32_t viewportWidth,
    uint32_t viewportHeight) {
    const Components::TransformComponent* transform = gameObject.GetTransform();
    if (transform == nullptr) {
        return;
    }

    const Math::Vector3 position = transform->GetPosition();
    const Math::Vector3 forward = transform->GetForward().Normalized();
    const Math::Vector3 right = transform->GetRight().Normalized();
    const Math::Vector3 up = transform->GetUp().Normalized();
    if (forward.SqrMagnitude() <= Math::EPSILON ||
        right.SqrMagnitude() <= Math::EPSILON ||
        up.SqrMagnitude() <= Math::EPSILON) {
        return;
    }

    const float nearClip = (std::max)(camera.GetNearClipPlane(), 0.01f);
    const float farClip = (std::max)(camera.GetFarClipPlane(), nearClip + 0.01f);
    const float aspect = ResolveCameraAspect(camera, viewportWidth, viewportHeight);

    float nearHalfHeight = 0.0f;
    float nearHalfWidth = 0.0f;
    float farHalfHeight = 0.0f;
    float farHalfWidth = 0.0f;
    if (camera.GetProjectionType() == Components::CameraProjectionType::Perspective) {
        const float halfFovRadians =
            std::clamp(camera.GetFieldOfView(), 1.0f, 179.0f) * Math::DEG_TO_RAD * 0.5f;
        nearHalfHeight = std::tan(halfFovRadians) * nearClip;
        nearHalfWidth = nearHalfHeight * aspect;
        farHalfHeight = std::tan(halfFovRadians) * farClip;
        farHalfWidth = farHalfHeight * aspect;
    } else {
        const float halfHeight = (std::max)(camera.GetOrthographicSize(), 0.01f);
        const float halfWidth = halfHeight * aspect;
        nearHalfHeight = halfHeight;
        nearHalfWidth = halfWidth;
        farHalfHeight = halfHeight;
        farHalfWidth = halfWidth;
    }

    const Math::Vector3 nearCenter = position + forward * nearClip;
    const Math::Vector3 farCenter = position + forward * farClip;
    const std::array<Math::Vector3, 8> corners = {{
        nearCenter + up * nearHalfHeight - right * nearHalfWidth,
        nearCenter + up * nearHalfHeight + right * nearHalfWidth,
        nearCenter - up * nearHalfHeight + right * nearHalfWidth,
        nearCenter - up * nearHalfHeight - right * nearHalfWidth,
        farCenter + up * farHalfHeight - right * farHalfWidth,
        farCenter + up * farHalfHeight + right * farHalfWidth,
        farCenter - up * farHalfHeight + right * farHalfWidth,
        farCenter - up * farHalfHeight - right * farHalfWidth
    }};

    static constexpr std::array<std::pair<size_t, size_t>, 12> kFrustumEdges = {{
        { 0u, 1u }, { 1u, 2u }, { 2u, 3u }, { 3u, 0u },
        { 4u, 5u }, { 5u, 6u }, { 6u, 7u }, { 7u, 4u },
        { 0u, 4u }, { 1u, 5u }, { 2u, 6u }, { 3u, 7u }
    }};
    constexpr Math::Color kFrustumColor(1.0f, 1.0f, 1.0f, 1.0f);

    for (const auto& edge : kFrustumEdges) {
        AppendWorldLine(
            frameData,
            corners[edge.first],
            corners[edge.second],
            kFrustumColor,
            1.6f,
            SceneViewportOverlayDepthMode::AlwaysOnTop);
    }
}

void AppendDirectionalLightOverlay(
    SceneViewportOverlayFrameData& frameData,
    const Components::GameObject& gameObject,
    const SceneViewportOverlayData& overlay,
    uint32_t viewportHeight) {
    const Components::TransformComponent* transform = gameObject.GetTransform();
    if (transform == nullptr) {
        return;
    }

    const Math::Vector3 position = transform->GetPosition();
    const Math::Vector3 lightDirection = (transform->GetForward() * -1.0f).Normalized();
    const Math::Vector3 right = transform->GetRight().Normalized();
    const Math::Vector3 up = transform->GetUp().Normalized();
    if (lightDirection.SqrMagnitude() <= Math::EPSILON ||
        right.SqrMagnitude() <= Math::EPSILON ||
        up.SqrMagnitude() <= Math::EPSILON) {
        return;
    }

    const float worldUnitsPerPixel = ComputeWorldUnitsPerPixel(overlay, position, viewportHeight);
    if (worldUnitsPerPixel <= Math::EPSILON) {
        return;
    }

    constexpr Math::Color kDirectionalLightColor(1.0f, 0.92f, 0.24f, 1.0f);
    constexpr float kLineThickness = 1.8f;
    constexpr size_t kRingSegmentCount = 32u;
    constexpr std::array<float, 6> kRayAngles = {{
        0.0f,
        Math::PI / 3.0f,
        Math::PI * 2.0f / 3.0f,
        Math::PI,
        Math::PI * 4.0f / 3.0f,
        Math::PI * 5.0f / 3.0f
    }};

    const float ringRadius = worldUnitsPerPixel * 26.0f;
    const float ringOffset = worldUnitsPerPixel * 54.0f;
    const float innerRayRadius = ringRadius * 0.52f;
    const float rayLength = worldUnitsPerPixel * 96.0f;
    const Math::Vector3 ringCenter = position + lightDirection * ringOffset;

    for (size_t segmentIndex = 0; segmentIndex < kRingSegmentCount; ++segmentIndex) {
        const float angle0 =
            static_cast<float>(segmentIndex) / static_cast<float>(kRingSegmentCount) * Math::PI * 2.0f;
        const float angle1 =
            static_cast<float>(segmentIndex + 1u) / static_cast<float>(kRingSegmentCount) * Math::PI * 2.0f;
        const Math::Vector3 p0 =
            ringCenter + right * std::cos(angle0) * ringRadius + up * std::sin(angle0) * ringRadius;
        const Math::Vector3 p1 =
            ringCenter + right * std::cos(angle1) * ringRadius + up * std::sin(angle1) * ringRadius;
        AppendWorldLine(
            frameData,
            p0,
            p1,
            kDirectionalLightColor,
            kLineThickness,
            SceneViewportOverlayDepthMode::AlwaysOnTop);
    }

    AppendWorldLine(
        frameData,
        position,
        ringCenter,
        kDirectionalLightColor,
        kLineThickness,
        SceneViewportOverlayDepthMode::AlwaysOnTop);
    AppendWorldLine(
        frameData,
        ringCenter,
        ringCenter + lightDirection * rayLength,
        kDirectionalLightColor,
        kLineThickness,
        SceneViewportOverlayDepthMode::AlwaysOnTop);

    for (float angle : kRayAngles) {
        const Math::Vector3 rayStart =
            ringCenter + right * std::cos(angle) * innerRayRadius + up * std::sin(angle) * innerRayRadius;
        AppendWorldLine(
            frameData,
            rayStart,
            rayStart + lightDirection * rayLength,
            kDirectionalLightColor,
            kLineThickness,
            SceneViewportOverlayDepthMode::AlwaysOnTop);
    }
}

class SceneViewportCameraOverlayProvider final : public ISceneViewportOverlayProvider {
public:
    const char* GetName() const override {
        return "SceneViewportCameraOverlayProvider";
    }

    void AppendOverlay(
        const SceneViewportOverlayBuildContext& context,
        SceneViewportOverlayFrameData& frameData) const override {
        if (!context.IsValid()) {
            return;
        }

        constexpr Math::Vector2 kCameraIconSize(90.0f, 90.0f);

        for (Components::CameraComponent* camera : context.scene->FindObjectsOfType<Components::CameraComponent>()) {
            if (camera == nullptr || !camera->IsEnabled()) {
                continue;
            }

            Components::GameObject* gameObject = camera->GetGameObject();
            if (!CanBuildOverlayForGameObject(gameObject)) {
                continue;
            }

            AppendSceneIconOverlay(
                frameData,
                *context.overlay,
                context.viewportWidth,
                context.viewportHeight,
                *gameObject,
                kCameraIconSize,
                SceneViewportOverlaySpriteTextureKind::Camera);
        }

        for (uint64_t entityId : *context.selectedObjectIds) {
            if (entityId == 0) {
                continue;
            }

            Components::GameObject* gameObject = context.scene->FindByID(entityId);
            if (!CanBuildOverlayForGameObject(gameObject)) {
                continue;
            }

            Components::CameraComponent* camera = gameObject->GetComponent<Components::CameraComponent>();
            if (camera == nullptr || !camera->IsEnabled()) {
                continue;
            }

            AppendCameraFrustumOverlay(
                frameData,
                *camera,
                *gameObject,
                context.viewportWidth,
                context.viewportHeight);
        }
    }
};

class SceneViewportLightOverlayProvider final : public ISceneViewportOverlayProvider {
public:
    const char* GetName() const override {
        return "SceneViewportLightOverlayProvider";
    }

    void AppendOverlay(
        const SceneViewportOverlayBuildContext& context,
        SceneViewportOverlayFrameData& frameData) const override {
        if (!context.IsValid()) {
            return;
        }

        constexpr Math::Vector2 kLightIconSize(100.0f, 100.0f);

        for (Components::LightComponent* light : context.scene->FindObjectsOfType<Components::LightComponent>()) {
            if (light == nullptr || !light->IsEnabled()) {
                continue;
            }

            Components::GameObject* gameObject = light->GetGameObject();
            if (!CanBuildOverlayForGameObject(gameObject)) {
                continue;
            }

            AppendSceneIconOverlay(
                frameData,
                *context.overlay,
                context.viewportWidth,
                context.viewportHeight,
                *gameObject,
                kLightIconSize,
                SceneViewportOverlaySpriteTextureKind::Light);
        }

        for (uint64_t entityId : *context.selectedObjectIds) {
            if (entityId == 0) {
                continue;
            }

            Components::GameObject* gameObject = context.scene->FindByID(entityId);
            if (!CanBuildOverlayForGameObject(gameObject)) {
                continue;
            }

            Components::LightComponent* light = gameObject->GetComponent<Components::LightComponent>();
            if (light == nullptr ||
                !light->IsEnabled() ||
                light->GetLightType() != Components::LightType::Directional) {
                continue;
            }

            AppendDirectionalLightOverlay(
                frameData,
                *gameObject,
                *context.overlay,
                context.viewportHeight);
        }
    }
};

class SceneViewportTransformGizmoOverlayProvider final : public ISceneViewportOverlayProvider {
public:
    const char* GetName() const override {
        return "SceneViewportTransformGizmoOverlayProvider";
    }

    void AppendOverlay(
        const SceneViewportOverlayBuildContext& context,
        SceneViewportOverlayFrameData& frameData) const override {
        if (!context.IsValid() ||
            context.transformGizmoOverlayState == nullptr ||
            !context.transformGizmoOverlayState->HasAnyVisibleGizmo()) {
            return;
        }

        const SceneViewportTransformGizmoHandleBuildInputs inputs =
            BuildSceneViewportTransformGizmoHandleBuildInputs(*context.transformGizmoOverlayState);
        AppendTransformGizmoScreenTriangles(frameData, inputs);
        AppendTransformGizmoHandleRecords(frameData, inputs);
    }
};

} // namespace

void SceneViewportOverlayProviderRegistry::AddProvider(
    std::unique_ptr<ISceneViewportOverlayProvider> provider) {
    if (provider == nullptr) {
        return;
    }

    m_providers.emplace_back(std::move(provider));
}

size_t SceneViewportOverlayProviderRegistry::GetProviderCount() const {
    return m_providers.size();
}

const ISceneViewportOverlayProvider* SceneViewportOverlayProviderRegistry::GetProvider(size_t index) const {
    return index < m_providers.size() ? m_providers[index].get() : nullptr;
}

void SceneViewportOverlayProviderRegistry::AppendOverlay(
    const SceneViewportOverlayBuildContext& context,
    SceneViewportOverlayFrameData& frameData) const {
    if (!context.IsValid()) {
        return;
    }

    for (const auto& provider : m_providers) {
        if (provider != nullptr) {
            provider->AppendOverlay(context, frameData);
        }
    }
}

std::unique_ptr<ISceneViewportOverlayProvider> CreateSceneViewportCameraOverlayProvider() {
    return std::make_unique<SceneViewportCameraOverlayProvider>();
}

std::unique_ptr<ISceneViewportOverlayProvider> CreateSceneViewportLightOverlayProvider() {
    return std::make_unique<SceneViewportLightOverlayProvider>();
}

std::unique_ptr<ISceneViewportOverlayProvider> CreateSceneViewportTransformGizmoOverlayProvider() {
    return std::make_unique<SceneViewportTransformGizmoOverlayProvider>();
}

SceneViewportOverlayProviderRegistry BuildDefaultSceneViewportOverlayProviderRegistry() {
    SceneViewportOverlayProviderRegistry registry;
    registry.AddProvider(CreateSceneViewportCameraOverlayProvider());
    registry.AddProvider(CreateSceneViewportLightOverlayProvider());
    registry.AddProvider(CreateSceneViewportTransformGizmoOverlayProvider());
    return registry;
}

} // namespace Editor
} // namespace XCEngine