#include "Rendering/SceneRenderRequestPlanner.h"

#include "Components/GameObject.h"
#include "Components/LightComponent.h"
#include "Components/MeshFilterComponent.h"
#include "Components/MeshRendererComponent.h"
#include "Core/Math/Matrix4.h"
#include "Core/Math/Quaternion.h"
#include "Rendering/RenderSceneUtility.h"
#include "Rendering/SceneRenderRequestUtils.h"

#include "Scene/Scene.h"

#include <array>
#include <algorithm>
#include <limits>

namespace XCEngine {
namespace Rendering {

namespace {

constexpr uint32_t kDirectionalShadowMapDimension = 1024;
constexpr float kMinShadowFocusDistance = 5.0f;
constexpr float kMaxShadowFocusDistance = 64.0f;
constexpr float kPerspectiveShadowFocusFactor = 1.0f;
constexpr float kOrthographicShadowFocusFactor = 2.0f;
constexpr float kMinShadowDepthRange = 20.0f;
constexpr float kShadowBoundsPadding = 1.0f;
constexpr float kMinShadowDepthPadding = 2.0f;

bool ShouldPlanDirectionalShadowForCamera(
    const Components::CameraComponent& camera,
    size_t renderedBaseCameraCount,
    size_t renderedRequestCount) {
    if (camera.GetStackType() == Components::CameraStackType::Overlay) {
        return renderedBaseCameraCount == 0u &&
               renderedRequestCount == 0u;
    }

    return true;
}

std::array<Math::Vector3, 8> BuildBoundsCorners(const Math::Bounds& bounds) {
    const Math::Vector3 min = bounds.GetMin();
    const Math::Vector3 max = bounds.GetMax();
    return {
        min,
        Math::Vector3(max.x, min.y, min.z),
        Math::Vector3(min.x, max.y, min.z),
        Math::Vector3(min.x, min.y, max.z),
        max,
        Math::Vector3(min.x, max.y, max.z),
        Math::Vector3(max.x, min.y, max.z),
        Math::Vector3(max.x, max.y, min.z)
    };
}

Math::Bounds TransformBoundsToWorldSpace(
    const Math::Bounds& localBounds,
    const Math::Matrix4x4& localToWorld,
    std::array<Math::Vector3, 8>* outWorldCorners = nullptr) {
    const std::array<Math::Vector3, 8> localCorners = BuildBoundsCorners(localBounds);
    std::array<Math::Vector3, 8> worldCorners = {};
    for (size_t index = 0; index < localCorners.size(); ++index) {
        worldCorners[index] = localToWorld.MultiplyPoint(localCorners[index]);
    }

    Math::Bounds worldBounds(worldCorners[0], Math::Vector3::Zero());
    for (size_t index = 1; index < worldCorners.size(); ++index) {
        worldBounds.Encapsulate(worldCorners[index]);
    }

    if (outWorldCorners != nullptr) {
        *outWorldCorners = worldCorners;
    }

    return worldBounds;
}

void ExpandLightSpaceBounds(
    const Math::Matrix4x4& view,
    const std::array<Math::Vector3, 8>& worldCorners,
    float& minX,
    float& maxX,
    float& minY,
    float& maxY,
    float& minZ,
    float& maxZ) {
    for (const Math::Vector3& corner : worldCorners) {
        const Math::Vector3 cornerLS = view.MultiplyPoint(corner);
        minX = std::min(minX, cornerLS.x);
        maxX = std::max(maxX, cornerLS.x);
        minY = std::min(minY, cornerLS.y);
        maxY = std::max(maxY, cornerLS.y);
        minZ = std::min(minZ, cornerLS.z);
        maxZ = std::max(maxZ, cornerLS.z);
    }
}

DirectionalShadowRenderPlan BuildDirectionalShadowRenderPlan(
    const Components::Scene& scene,
    const Components::CameraComponent& camera,
    const Components::LightComponent& light,
    float viewportAspect) {
    DirectionalShadowRenderPlan plan = {};
    if (!light.GetCastsShadows()) {
        return plan;
    }

    Math::Vector3 lightDirection = light.transform().GetForward() * -1.0f;
    if (lightDirection.SqrMagnitude() <= Math::EPSILON) {
        lightDirection = Math::Vector3::Back();
    } else {
        lightDirection = lightDirection.Normalized();
    }

    const Math::Vector3 viewForward = camera.transform().GetForward().SqrMagnitude() <= Math::EPSILON
        ? Math::Vector3::Forward()
        : camera.transform().GetForward().Normalized();
    const Math::Vector3 viewRight = camera.transform().GetRight().SqrMagnitude() <= Math::EPSILON
        ? Math::Vector3::Right()
        : camera.transform().GetRight().Normalized();
    const Math::Vector3 viewUp = camera.transform().GetUp().SqrMagnitude() <= Math::EPSILON
        ? Math::Vector3::Up()
        : camera.transform().GetUp().Normalized();
    const Math::Vector3 cameraPosition = camera.transform().GetPosition();
    const float aspect = viewportAspect > Math::EPSILON ? viewportAspect : 1.0f;

    const float shadowDistance = camera.GetProjectionType() == Components::CameraProjectionType::Perspective
        ? std::clamp(
            camera.GetFarClipPlane() * kPerspectiveShadowFocusFactor,
            kMinShadowFocusDistance,
            kMaxShadowFocusDistance)
        : std::clamp(
            camera.GetOrthographicSize() * kOrthographicShadowFocusFactor,
            kMinShadowFocusDistance,
            kMaxShadowFocusDistance);
    const float sliceNear = std::max(camera.GetNearClipPlane(), 0.1f);
    const float sliceFar = std::max(sliceNear + 0.1f, shadowDistance);

    std::array<Math::Vector3, 8> frustumCorners = {};
    if (camera.GetProjectionType() == Components::CameraProjectionType::Perspective) {
        const float tanHalfFov = std::tan(camera.GetFieldOfView() * Math::DEG_TO_RAD * 0.5f);
        const float nearHalfHeight = tanHalfFov * sliceNear;
        const float nearHalfWidth = nearHalfHeight * aspect;
        const float farHalfHeight = tanHalfFov * sliceFar;
        const float farHalfWidth = farHalfHeight * aspect;

        const Math::Vector3 nearCenter = cameraPosition + viewForward * sliceNear;
        const Math::Vector3 farCenter = cameraPosition + viewForward * sliceFar;

        frustumCorners[0] = nearCenter - viewRight * nearHalfWidth - viewUp * nearHalfHeight;
        frustumCorners[1] = nearCenter + viewRight * nearHalfWidth - viewUp * nearHalfHeight;
        frustumCorners[2] = nearCenter - viewRight * nearHalfWidth + viewUp * nearHalfHeight;
        frustumCorners[3] = nearCenter + viewRight * nearHalfWidth + viewUp * nearHalfHeight;
        frustumCorners[4] = farCenter - viewRight * farHalfWidth - viewUp * farHalfHeight;
        frustumCorners[5] = farCenter + viewRight * farHalfWidth - viewUp * farHalfHeight;
        frustumCorners[6] = farCenter - viewRight * farHalfWidth + viewUp * farHalfHeight;
        frustumCorners[7] = farCenter + viewRight * farHalfWidth + viewUp * farHalfHeight;
    } else {
        const float orthoHalfHeight = camera.GetOrthographicSize();
        const float orthoHalfWidth = orthoHalfHeight * aspect;
        const Math::Vector3 nearCenter = cameraPosition + viewForward * sliceNear;
        const Math::Vector3 farCenter = cameraPosition + viewForward * sliceFar;

        frustumCorners[0] = nearCenter - viewRight * orthoHalfWidth - viewUp * orthoHalfHeight;
        frustumCorners[1] = nearCenter + viewRight * orthoHalfWidth - viewUp * orthoHalfHeight;
        frustumCorners[2] = nearCenter - viewRight * orthoHalfWidth + viewUp * orthoHalfHeight;
        frustumCorners[3] = nearCenter + viewRight * orthoHalfWidth + viewUp * orthoHalfHeight;
        frustumCorners[4] = farCenter - viewRight * orthoHalfWidth - viewUp * orthoHalfHeight;
        frustumCorners[5] = farCenter + viewRight * orthoHalfWidth - viewUp * orthoHalfHeight;
        frustumCorners[6] = farCenter - viewRight * orthoHalfWidth + viewUp * orthoHalfHeight;
        frustumCorners[7] = farCenter + viewRight * orthoHalfWidth + viewUp * orthoHalfHeight;
    }

    Math::Vector3 focusPoint = Math::Vector3::Zero();
    for (const Math::Vector3& corner : frustumCorners) {
        focusPoint += corner;
    }
    focusPoint /= static_cast<float>(frustumCorners.size());

    Math::Bounds frustumWorldBounds(frustumCorners[0], Math::Vector3::Zero());
    for (size_t index = 1; index < frustumCorners.size(); ++index) {
        frustumWorldBounds.Encapsulate(frustumCorners[index]);
    }

    const float shadowViewDistance = std::max(sliceFar, kMinShadowDepthRange * 0.5f);
    const Math::Vector3 shadowWorldPosition =
        focusPoint + lightDirection * shadowViewDistance;

    Math::Vector3 shadowUp = Math::Vector3::Up();
    if (std::abs(Math::Vector3::Dot(lightDirection, shadowUp)) > 0.98f) {
        shadowUp = Math::Vector3::Forward();
    }

    const Math::Vector3 shadowForward = (focusPoint - shadowWorldPosition).SqrMagnitude() <= Math::EPSILON
        ? (lightDirection * -1.0f)
        : (focusPoint - shadowWorldPosition).Normalized();
    const Math::Quaternion shadowRotation = Math::Quaternion::LookRotation(shadowForward, shadowUp);
    const Math::Matrix4x4 view = Math::Matrix4x4::TRS(
        shadowWorldPosition,
        shadowRotation,
        Math::Vector3::One()).Inverse();

    float minX = std::numeric_limits<float>::max();
    float maxX = std::numeric_limits<float>::lowest();
    float minY = std::numeric_limits<float>::max();
    float maxY = std::numeric_limits<float>::lowest();
    float minZ = std::numeric_limits<float>::max();
    float maxZ = std::numeric_limits<float>::lowest();
    ExpandLightSpaceBounds(view, frustumCorners, minX, maxX, minY, maxY, minZ, maxZ);

    const uint32_t cullingMask = camera.GetCullingMask();
    const std::vector<Components::MeshFilterComponent*> meshFilters =
        scene.FindObjectsOfType<Components::MeshFilterComponent>();
    for (Components::MeshFilterComponent* meshFilter : meshFilters) {
        if (meshFilter == nullptr ||
            !meshFilter->IsEnabled() ||
            meshFilter->GetGameObject() == nullptr) {
            continue;
        }

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

        const uint32_t gameObjectLayerMask = 1u << gameObject->GetLayer();
        if ((cullingMask & gameObjectLayerMask) == 0u) {
            continue;
        }

        Components::MeshRendererComponent* meshRenderer =
            gameObject->GetComponent<Components::MeshRendererComponent>();
        if (meshRenderer == nullptr ||
            !meshRenderer->IsEnabled() ||
            (!meshRenderer->GetCastShadows() && !meshRenderer->GetReceiveShadows())) {
            continue;
        }

        Resources::Mesh* mesh = meshFilter->GetMesh();
        if (mesh == nullptr || !mesh->IsValid()) {
            continue;
        }

        std::array<Math::Vector3, 8> worldCorners = {};
        const Math::Bounds worldBounds = TransformBoundsToWorldSpace(
            mesh->GetBounds(),
            gameObject->GetTransform()->GetLocalToWorldMatrix(),
            &worldCorners);
        if (!frustumWorldBounds.Intersects(worldBounds)) {
            continue;
        }

        ExpandLightSpaceBounds(view, worldCorners, minX, maxX, minY, maxY, minZ, maxZ);
    }

    minX -= kShadowBoundsPadding;
    maxX += kShadowBoundsPadding;
    minY -= kShadowBoundsPadding;
    maxY += kShadowBoundsPadding;
    minZ -= kMinShadowDepthPadding;
    maxZ += kMinShadowDepthPadding;

    const float shadowHalfExtent = std::max(
        std::max(std::abs(minX), std::abs(maxX)),
        std::max(std::abs(minY), std::abs(maxY)));
    const float shadowDepthRange = std::max(maxZ - minZ, kMinShadowDepthRange);
    const Math::Matrix4x4 projection = Math::Matrix4x4::Orthographic(
        minX,
        maxX,
        minY,
        maxY,
        minZ,
        maxZ);

    plan.enabled = true;
    plan.lightDirection = lightDirection;
    plan.focusPoint = focusPoint;
    plan.orthographicHalfExtent = shadowHalfExtent;
    plan.nearClipPlane = minZ;
    plan.farClipPlane = maxZ;
    plan.mapWidth = kDirectionalShadowMapDimension;
    plan.mapHeight = kDirectionalShadowMapDimension;
    plan.cameraData.view = view.Transpose();
    plan.cameraData.projection = projection.Transpose();
    plan.cameraData.viewProjection = (projection * view).Transpose();
    plan.cameraData.worldPosition = shadowWorldPosition;
    plan.cameraData.clearColor = Math::Color::Black();
    plan.cameraData.clearFlags = RenderClearFlags::Depth;
    plan.cameraData.viewportWidth = plan.mapWidth;
    plan.cameraData.viewportHeight = plan.mapHeight;
    return plan;
}

} // namespace

std::vector<Components::CameraComponent*> SceneRenderRequestPlanner::CollectCameras(
    const Components::Scene& scene,
    Components::CameraComponent* overrideCamera) const {
    std::vector<Components::CameraComponent*> cameras;

    if (SceneRenderRequestUtils::IsUsableCamera(overrideCamera)) {
        cameras.push_back(overrideCamera);
        return cameras;
    }

    cameras = scene.FindObjectsOfType<Components::CameraComponent>();
    cameras.erase(
        std::remove_if(
            cameras.begin(),
            cameras.end(),
            [](const Components::CameraComponent* camera) {
                return !SceneRenderRequestUtils::IsUsableCamera(camera);
            }),
        cameras.end());

    SceneRenderRequestUtils::SortSceneCamerasForRendering(cameras);
    return cameras;
}

std::vector<CameraRenderRequest> SceneRenderRequestPlanner::BuildRequests(
    const Components::Scene& scene,
    Components::CameraComponent* overrideCamera,
    const RenderContext& context,
    const RenderSurface& surface) const {
    std::vector<CameraRenderRequest> requests;
    const std::vector<Components::CameraComponent*> cameras =
        CollectCameras(scene, overrideCamera);

    size_t renderedBaseCameraCount = 0;
    for (Components::CameraComponent* camera : cameras) {
        CameraRenderRequest request;
        if (!SceneRenderRequestUtils::BuildCameraRenderRequest(
                scene,
                *camera,
                context,
                surface,
                renderedBaseCameraCount,
                requests.size(),
                request)) {
            continue;
        }

        if (ShouldPlanDirectionalShadowForCamera(
                *camera,
                renderedBaseCameraCount,
                requests.size())) {
            if (Components::LightComponent* mainDirectionalLight =
                    FindMainDirectionalLight(scene, camera->GetCullingMask());
                mainDirectionalLight != nullptr &&
                mainDirectionalLight->GetCastsShadows()) {
                const float viewportAspect = surface.GetRenderAreaHeight() > 0
                    ? static_cast<float>(surface.GetRenderAreaWidth()) /
                        static_cast<float>(surface.GetRenderAreaHeight())
                    : 1.0f;
                request.directionalShadow =
                    BuildDirectionalShadowRenderPlan(scene, *camera, *mainDirectionalLight, viewportAspect);
                if (request.directionalShadow.IsValid()) {
                    request.shadowCaster.clearFlags = RenderClearFlags::Depth;
                    request.shadowCaster.hasCameraDataOverride = true;
                    request.shadowCaster.cameraDataOverride = request.directionalShadow.cameraData;
                }
            }
        }

        requests.push_back(request);
        if (camera->GetStackType() == Components::CameraStackType::Base) {
            ++renderedBaseCameraCount;
        }
    }

    return requests;
}

} // namespace Rendering
} // namespace XCEngine