// This file is part of the 64k demo project.
// It implements BVH construction and traversal.

#include "3d/bvh.h"
#include <algorithm>

namespace {

struct ObjectInfo {
  int index;
  AABB aabb;
  vec3 centroid;
};

AABB get_world_aabb(const Object3D& obj) {
  BoundingVolume local = obj.get_local_bounds();
  mat4 model = obj.get_model_matrix();

  vec3 corners[8] = {
      {local.min.x, local.min.y, local.min.z},
      {local.max.x, local.min.y, local.min.z},
      {local.min.x, local.max.y, local.min.z},
      {local.max.x, local.max.y, local.min.z},
      {local.min.x, local.min.y, local.max.z},
      {local.max.x, local.min.y, local.max.z},
      {local.min.x, local.max.y, local.max.z},
      {local.max.x, local.max.y, local.max.z},
  };

  AABB world;
  for (int i = 0; i < 8; ++i) {
    vec4 p = model * vec4(corners[i].x, corners[i].y, corners[i].z, 1.0f);
    world.expand(p.xyz());
  }
  return world;
}

int build_recursive(std::vector<BVHNode>& nodes,
                    std::vector<ObjectInfo>& obj_info, int start, int end) {
  int node_idx = (int)nodes.size();
  nodes.emplace_back();

  AABB bounds;
  for (int i = start; i < end; ++i) {
    bounds.expand(obj_info[i].aabb);
  }

  int count = end - start;
  if (count == 1) {
    // Leaf node
    nodes[node_idx].min_x = bounds.min.x;
    nodes[node_idx].min_y = bounds.min.y;
    nodes[node_idx].min_z = bounds.min.z;
    nodes[node_idx].left_idx = -1;

    nodes[node_idx].max_x = bounds.max.x;
    nodes[node_idx].max_y = bounds.max.y;
    nodes[node_idx].max_z = bounds.max.z;
    nodes[node_idx].right_idx = obj_info[start].index;
  } else {
    // Internal node
    // Find axis with largest variance (or just largest extent of centroids)
    AABB centroid_bounds;
    for (int i = start; i < end; ++i) {
      centroid_bounds.expand(obj_info[i].centroid);
    }

    vec3 extent = centroid_bounds.max - centroid_bounds.min;
    int axis = 0;
    if (extent.y > extent.x)
      axis = 1;
    if (extent.z > (axis == 0 ? extent.x : extent.y))
      axis = 2;

    float split =
        (centroid_bounds.min[axis] + centroid_bounds.max[axis]) * 0.5f;

    // Partition
    int mid = start;
    for (int i = start; i < end; ++i) {
      if (obj_info[i].centroid[axis] < split) {
        std::swap(obj_info[i], obj_info[mid]);
        mid++;
      }
    }

    // Fallback if partition failed
    if (mid == start || mid == end) {
      mid = start + count / 2;
    }

    int left = build_recursive(nodes, obj_info, start, mid);
    int right = build_recursive(nodes, obj_info, mid, end);

    nodes[node_idx].min_x = bounds.min.x;
    nodes[node_idx].min_y = bounds.min.y;
    nodes[node_idx].min_z = bounds.min.z;
    nodes[node_idx].left_idx = left;

    nodes[node_idx].max_x = bounds.max.x;
    nodes[node_idx].max_y = bounds.max.y;
    nodes[node_idx].max_z = bounds.max.z;
    nodes[node_idx].right_idx = right;
  }

  return node_idx;
}

} // namespace

void BVHBuilder::build(BVH& out_bvh, const std::vector<Object3D>& objects) {
  out_bvh.nodes.clear();
  if (objects.empty())
    return;

  std::vector<ObjectInfo> obj_info;
  for (int i = 0; i < (int)objects.size(); ++i) {
    if (objects[i].type == ObjectType::SKYBOX)
      continue;
    AABB aabb = get_world_aabb(objects[i]);
    obj_info.push_back({i, aabb, aabb.center()});
  }

  if (obj_info.empty())
    return;

  out_bvh.nodes.reserve(obj_info.size() * 2);
  build_recursive(out_bvh.nodes, obj_info, 0, (int)obj_info.size());
}

void BVH::query(const AABB& box, std::vector<int>& out_indices) const {
  if (nodes.empty())
    return;

  std::vector<int> stack;
  stack.push_back(0);

  while (!stack.empty()) {
    int idx = stack.back();
    stack.pop_back();

    const BVHNode& node = nodes[idx];
    AABB node_aabb({node.min_x, node.min_y, node.min_z},
                   {node.max_x, node.max_y, node.max_z});

    if (node_aabb.intersects(box)) {
      if (node.left_idx < 0) {
        out_indices.push_back(node.right_idx);
      } else {
        stack.push_back(node.left_idx);
        stack.push_back(node.right_idx);
      }
    }
  }
}