#include "CollisionMesh.h"

#include <algorithm>
#include <cfloat>

static constexpr int LEAF_TRIANGLES = 4;

static glm::vec3 TriangleCentroid(const Triangle& tri)
{
    return (tri.a + tri.b + tri.c) / 3.0f;
}

// Bounds overlap plus the triangle's plane, which is what prunes huge floor triangles
static bool TriangleOverlapsBox(const Triangle& tri, const glm::vec3& boxMin, const glm::vec3& boxMax)
{
    const glm::vec3 triMin = glm::min(tri.a, glm::min(tri.b, tri.c));
    const glm::vec3 triMax = glm::max(tri.a, glm::max(tri.b, tri.c));
    if (glm::any(glm::lessThan(triMax, boxMin)) || glm::any(glm::greaterThan(triMin, boxMax)))
        return false;

    const glm::vec3 normal = glm::cross(tri.b - tri.a, tri.c - tri.a);
    const glm::vec3 center = (boxMin + boxMax) * 0.5f;
    const glm::vec3 halfSize = (boxMax - boxMin) * 0.5f;
    const float distance = glm::dot(normal, center - tri.a);
    const float radius = glm::dot(glm::abs(normal), halfSize);
    return glm::abs(distance) <= radius;
}

CollisionMesh::CollisionMesh(const std::vector<Triangle>& meshTriangles)
    : triangles(meshTriangles)
{
    if (!triangles.empty())
        Build(0, (int)triangles.size());
}

int CollisionMesh::Build(int start, int count)
{
    const int index = (int)nodes.size();
    nodes.push_back(Node());

    glm::vec3 boundsMin(FLT_MAX);
    glm::vec3 boundsMax(-FLT_MAX);
    glm::vec3 centroidMin(FLT_MAX);
    glm::vec3 centroidMax(-FLT_MAX);
    for (int i = start; i < start + count; i++) {
        const Triangle& tri = triangles[i];
        boundsMin = glm::min(boundsMin, glm::min(tri.a, glm::min(tri.b, tri.c)));
        boundsMax = glm::max(boundsMax, glm::max(tri.a, glm::max(tri.b, tri.c)));
        const glm::vec3 centroid = TriangleCentroid(tri);
        centroidMin = glm::min(centroidMin, centroid);
        centroidMax = glm::max(centroidMax, centroid);
    }
    nodes[index].boundsMin = boundsMin;
    nodes[index].boundsMax = boundsMax;

    if (count <= LEAF_TRIANGLES) {
        nodes[index].start = start;
        nodes[index].count = count;
        return index;
    }

    // Split at the median along the widest axis
    const glm::vec3 extent = centroidMax - centroidMin;
    const int axis = extent.x > extent.y ? (extent.x > extent.z ? 0 : 2) : (extent.y > extent.z ? 1 : 2);
    const int half = count / 2;
    std::nth_element(triangles.begin() + start, triangles.begin() + start + half, triangles.begin() + start + count,
        [axis](const Triangle& a, const Triangle& b) {
            return TriangleCentroid(a)[axis] < TriangleCentroid(b)[axis];
        });

    // Build grows nodes, so index stays valid but references into it wouldn't
    const int left = Build(start, half);
    const int right = Build(start + half, count - half);
    nodes[index].left = left;
    nodes[index].right = right;
    return index;
}

void CollisionMesh::Query(const Triangle& tri, std::vector<int>& hits) const
{
    if (nodes.empty())
        return;

    int stack[64];
    int stackSize = 0;
    stack[stackSize++] = 0;

    while (stackSize > 0) {
        const Node& node = nodes[stack[--stackSize]];
        if (!TriangleOverlapsBox(tri, node.boundsMin, node.boundsMax))
            continue;

        if (node.left == -1) {
            for (int i = node.start; i < node.start + node.count; i++)
                hits.push_back(i);
        }
        else {
            stack[stackSize++] = node.left;
            stack[stackSize++] = node.right;
        }
    }
}

const CollisionMesh& CollisionMesh::UnitBox()
{
    static const CollisionMesh box = []() {
        const glm::vec3 axes[3] = { {1.0f, 0.0f, 0.0f}, {0.0f, 1.0f, 0.0f}, {0.0f, 0.0f, 1.0f} };

        std::vector<Triangle> faces;
        for (int axis = 0; axis < 3; axis++) {
            for (float side : { 1.0f, -1.0f }) {
                const glm::vec3 normal = axes[axis] * side;
                const glm::vec3 u = axes[(axis + 1) % 3];
                const glm::vec3 v = glm::cross(normal, u);  // u x v = normal, so the winding faces outward

                const glm::vec3 center = normal * 0.5f;
                const glm::vec3 p0 = center - u * 0.5f - v * 0.5f;
                const glm::vec3 p1 = center + u * 0.5f - v * 0.5f;
                const glm::vec3 p2 = center + u * 0.5f + v * 0.5f;
                const glm::vec3 p3 = center - u * 0.5f + v * 0.5f;
                faces.push_back({ p0, p1, p2 });
                faces.push_back({ p0, p2, p3 });
            }
        }
        return CollisionMesh(faces);
    }();
    return box;
}