#include "world.h"
#include "game.h"

size_t buildWorldBVHLeafs(BVHNode leafs[WORLD_ENTITY_MAX], const World* world)
{
  size_t leafsSize = 0;
  const Entity* entities = world->entities;
  bool grouped[WORLD_ENTITY_MAX];
  int ungroupedCount = WORLD_ENTITY_MAX;
  memset(grouped, 0, sizeof(grouped));

  while (ungroupedCount > 0)
  {
    BVHNode leaf;

    for (int leafIndex = 0; leafIndex < BVH_MAX; ++leafIndex)
    {
      int closest = -1;
      int closestGroupedIndex = 0;
      float closestDistance = world->size.x;

      // Find closest.
      for (int index = 0; index < WORLD_ENTITY_MAX; ++index)
      {
        if (grouped[index])
        {
          continue;
        }
        
        // First entity.
        if (leafIndex == 0)
        {
          closest = index;
          break;
        }

        float distance = 0.0;
        BoundingBox overlapBox;
        overlapBox.min = entities[index].position;
        overlapBox.max = overlapBox.min;

        for (int innerIndex = 0; innerIndex < leafIndex; ++innerIndex)
        {
          distance += Vector3Distance(
            entities[leaf.entities[innerIndex]].position,
            entities[index].position);
          overlapBox.min = Vector3Min(
            overlapBox.min,
            entities[leaf.entities[innerIndex]].box.min);
          overlapBox.min = Vector3Min(
            overlapBox.min,
            entities[index].box.min);
          overlapBox.max = Vector3Max(
            overlapBox.max,
            entities[leaf.entities[innerIndex]].box.max);
          overlapBox.max = Vector3Max(
            overlapBox.max,
            entities[index].box.max);
        }

        distance /= (float)leafIndex;

        bool overlaps = false;

        // Too big (will count it as a overlap).
        if (Vector3Distance(overlapBox.min, overlapBox.max) >= BVH_BOX_MAX)
        {
          overlaps = true;
        }

        // Check if overlap is already happening.
        for (int nodeIndex = 0; nodeIndex < leafsSize && !overlaps;
             ++nodeIndex)
        {
          if (CheckCollisionBoxes(overlapBox, leafs[nodeIndex].box))
          {
            overlaps = true;
            break;
          }
        }

        // Update distance.
        if (!overlaps && distance < closestDistance)
        {
          closestDistance = distance;
          closest = index;
        }
      }

      if (closest == -1)
      {
        leaf.entities[leafIndex] = -1;
      }
      else
      {
        leaf.entities[leafIndex] = closest;
        grouped[closest] = true;
        --ungroupedCount;
      }
    }

    // Get bounding box.
    leaf.box = world->entities[leaf.entities[0]].box;

    for (int index = 1; index < BVH_MAX; ++index)
    {
      if (leaf.entities[index] == -1)
      {
        continue;
      }

      leaf.box.min = Vector3Min(
        leaf.box.min,
        world->entities[leaf.entities[index]].box.min);
      leaf.box.max = Vector3Max(
        leaf.box.max,
        world->entities[leaf.entities[index]].box.max);
    }

    leaf.position = Vector3Scale(Vector3Add(leaf.box.min, leaf.box.max), 0.5);

    memset(leaf.branches, 0, BVH_MAX_BRANCH_COUNT * sizeof(BVHNode*));
    leafs[leafsSize] = leaf;
    ++leafsSize;
  }

#ifdef FT_DEBUG_MODE
  // Test if everything is grouped.
  for (int index = 0; index < WORLD_ENTITY_MAX; ++index)
  {
    if (!grouped[index])
    {
      printf("Ungrouped: %d\n", index);
    }
  }

  // Test for leaf collision.
  for (int outer = 0; outer < leafsSize; ++outer)
  {
    for (int inner = 0; inner < leafsSize; ++inner)
    {
      if (outer != inner &&
          CheckCollisionBoxes(leafs[outer].box, leafs[inner].box))
      {
        printf("Leaf collision: %d and %d\n", outer, inner);
      }
    }
  }

  printf("leaf count: %ld\n", leafsSize);
#endif

  return leafsSize;
}

BVHNode buildWorldBVHTree(BVHNode* nodes, size_t nodesSize,
                          int taken[BVH_MAX_BRANCH_COUNT],
                          BVHNode* level, size_t levelSize,
                          const World* world)
{
  BVHNode node;
  memset(&node, 0, sizeof(BVHNode));

  for (int index = 1; index < BVH_MAX_BRANCH_COUNT; ++index)
  {
    taken[index] = -1;
  }

  // Add first node to branch.
  node.branches[0] = (BVHNode*)FT_MALLOC(sizeof(BVHNode));

  if (node.branches[0] == NULL)
  {
    ALLOCATION_ERROR;
    return node;
  }

  memcpy(node.branches[0], &nodes[0], sizeof(BVHNode));
  taken[0] = 0;
  node.branchCount = 1;

  for (int branchIndex = 1; branchIndex < BVH_MAX_BRANCH_COUNT; ++branchIndex)
  {
    int closest = -1;
    float closestDistance = world->size.x;

    // Find closest node.
    for (int nodeIndex = 0; nodeIndex < nodesSize; ++nodeIndex)
    {
      bool alreadyUsed = false;
      
      for (int index = 0; index < node.branchCount; ++index)
      {
        if (taken[index] == nodeIndex)
        {
          alreadyUsed = true;
          break;
        }
      }

      if (alreadyUsed)
      {
        continue;
      }
      
      float distance = 0.0;
      BoundingBox box = node.branches[0]->box;

      // Get distance and bounding box.
      for (int index = 0; index < branchIndex; ++index)
      {
        distance += Vector3Distance(nodes[nodeIndex].position,
                                    node.branches[index]->position);
        box.min = Vector3Min(box.min, nodes[nodeIndex].box.min);
        box.min = Vector3Min(box.min, node.branches[index]->box.min);
        box.max = Vector3Max(box.max, nodes[nodeIndex].box.max);
        box.max = Vector3Max(box.max, node.branches[index]->box.max);
      }

      distance /= (float)branchIndex;

      // Check for overlap.
      bool overlaps = false;

      for (int index = 0; index < levelSize; ++index)
      {
        if (CheckCollisionBoxes(box, level[index].box))
        {
          overlaps = true;
          break;
        }
      }

      if (!overlaps && distance < closestDistance)
      {
        closest = nodeIndex;
        closestDistance = distance;
      }
    }

    if (closest == -1)
    {
      break;
    }

    node.branches[branchIndex] = (BVHNode*)FT_MALLOC(sizeof(BVHNode));

    if (node.branches[branchIndex] == NULL)
    {
      ALLOCATION_ERROR;
      break;
    }

    memcpy(node.branches[branchIndex], &nodes[closest], sizeof(BVHNode));
    taken[node.branchCount] = closest;
    ++node.branchCount;
  }

  // Create node bounding box.
  node.box = node.branches[0]->box;
  
  for (int index = 0; index < BVH_MAX_BRANCH_COUNT; ++index)
  {
    if (node.branches[index] != NULL)
    {
      node.box.min = Vector3Min(node.box.min, node.branches[index]->box.min);
      node.box.max = Vector3Max(node.box.max, node.branches[index]->box.max);
    }
  }

  node.position = Vector3Scale(Vector3Add(node.box.min, node.box.max), 0.5);

  return node;
}

void buildWorldBVH(World* world)
{
  Entity* entities = world->entities;

  // Get leafs.
  BVHNode nodes[WORLD_ENTITY_MAX];
  size_t nodesSize = buildWorldBVHLeafs(nodes, world);

  BVHNode level[WORLD_ENTITY_MAX];
  size_t levelSize = 0;
  int levelCount = 0;

  while (true)
  {
    while (nodesSize > 0)
    {
      int taken[BVH_MAX_BRANCH_COUNT];
      BVHNode node = buildWorldBVHTree(nodes, nodesSize, taken, level,
                                       levelSize, world);

      // Take out taken nodes.
      int nodeIndex = 0;
      int removedCount = 0;

      for (int nodeCount = 0; nodeCount < nodesSize; ++nodeCount)
      {
        bool isTaken = false;

        for (int takenIndex = 0; takenIndex < node.branchCount; ++takenIndex)
        {
          if (taken[takenIndex] == nodeCount)
          {
            isTaken = true;
            ++removedCount;
            break;
          }
        }

        if (!isTaken)
        {
          nodes[nodeIndex] = nodes[nodeCount];
          ++nodeIndex;
        }
      }

      level[levelSize] = node;
      ++levelSize;
      nodesSize -= removedCount;
    }

#ifdef FT_DEBUG_MODE
    // Check for overlap
    int overlapCount = 0;
    
    for (int outer = 0; outer < levelSize - 1; ++outer)
    {
      for (int inner = outer + 1; inner < levelSize; ++inner)
      {
        if (CheckCollisionBoxes(level[inner].box, level[outer].box))
        {
          ++overlapCount;
        }
      }
    }
    
    printf("BVH level: %d, Size: %ld, Overlap: %d\n", levelCount, levelSize,
           overlapCount);
    ++levelCount;
#endif

    // Reached root node.
    if (levelSize <= 1)
    {
      world->bvh = level[0];
      break;
    }

    nodesSize = levelSize;
    levelSize = 0;
    memcpy(nodes, level, nodesSize * sizeof(BVHNode));
  }
}

World createWorld(int seed)
{
  World world;
  world.size = WORLD_SIZE;

  // Heightmap image.
  int offsetX = FT_RANDOM16(seed);
  int offsetY = FT_RANDOM16(seed);
  Image image = GenImagePerlinNoise(WORLD_IMAGE_WIDTH, WORLD_IMAGE_HEIGHT,
                                    offsetX, offsetY, WORLD_IMAGE_SCALE);

  // Heightmap.
  Mesh mesh = GenMeshHeightmap(image, world.size);
  world.heightmap = LoadModelFromMesh(mesh);
  world.texture = LoadTextureFromImage(image);
  world.heightmap.materials[0].maps[MATERIAL_MAP_DIFFUSE].texture =
    world.texture;

  UnloadImage(image);

  // Entities.
  for (int index = 0; index < WORLD_ENTITY_MAX; ++index)
  {
    FT_RANDOM16(seed);
    
    Entity entity = createEntity(seed % ENTITY_COUNT, Vector3Zero());
    Vector3 position;
    position.x = FT_RANDOM16(seed) % (int)world.size.x;
    position.z = FT_RANDOM16(seed) % (int)world.size.z;
    position.y = getWorldHeightAtLocation(&world,
                                          position.x, position.z) + 1.0;
    setEntityPosition(&entity, position);
    
    world.entities[index] = entity;
  }

  double currentTime = GetTime();
  buildWorldBVH(&world);

#ifdef FT_DEBUG_MODE
  printf("BVH build time: %lf\n", GetTime() - currentTime);
#endif

  world.bvhDebugSelect = 0;

  return world;
}

void drawBVHDebug(BVHNode bvh, int level, int selected)
{
  Color colors[] = {RED, GREEN, BLUE, ORANGE, YELLOW, PINK};
  int colorSize = 6;

  if (level == selected)
  {
    DrawBoundingBox(bvh.box, colors[level % colorSize]);
    return;
  }

  for (int index = 0; index < bvh.branchCount; ++index)
  {
    drawBVHDebug(*bvh.branches[index], level + 1, selected);
  }
}

void updateWorld(World* world, Game* game)
{
  DrawModel(world->heightmap, Vector3Zero(), 1.0, WHITE);

  for (int index = 0; index < WORLD_ENTITY_MAX; ++index)
  {
    updateEntity(&world->entities[index], game);
  }

  // Draw BVH leafs.
#ifdef FT_DEBUG_MODE
  if (IsKeyPressed(KEY_RIGHT))
  {
    ++world->bvhDebugSelect;
  }
  if (IsKeyPressed(KEY_LEFT))
  {
    --world->bvhDebugSelect;
  }
  
  drawBVHDebug(world->bvh, 0, world->bvhDebugSelect);
#endif
}

void freeWorldBVH(BVHNode bvh)
{
  // Play it safe to prevent memory leaks.
  for (int index = 0; index < BVH_MAX_BRANCH_COUNT; ++index)
  {
    if (bvh.branches[index] != NULL)
    {
      freeWorldBVH(*bvh.branches[index]);
      FT_FREE(bvh.branches[index]);
    }
  }
}

void freeWorld(World world)
{
  UnloadTexture(world.texture);
  UnloadModel(world.heightmap);
  freeWorldBVH(world.bvh);
}

float getWorldHeightAtLocation(const World* world, float x, float y)
{
  float mapX = (float)world->texture.width / world->size.x * x;
  float mapY = (float)world->texture.height / world->size.z * y;
  RayCollision result;

  for (int yOffset = -1; yOffset < 2; ++yOffset)
  {
    for (int xOffset = -1; xOffset < 2; ++xOffset)
    {
      int pixelX = mapX + xOffset;
      int pixelY = mapY + yOffset;

      if (pixelX < 0 || pixelX >= world->texture.width ||
          pixelY < 0 || pixelY >= world->texture.height)
      {
        continue;
      }

      int verticeStart = (pixelY * (world->texture.width - 1) + pixelX) * 18;
      float* vertices = &world->heightmap.meshes[0].vertices[verticeStart];

      // Cast to triangles at pixel. Really hacky indeed.
      Ray ray = (Ray){
        .position = (Vector3){x, world->size.y * 2.0, y},
        .direction = (Vector3){0.0, -1.0, 0.0}
      };

      result = GetRayCollisionTriangle(
        ray,
        (Vector3){vertices[0], vertices[1], vertices[2]},
        (Vector3){vertices[3], vertices[4], vertices[5]},
        (Vector3){vertices[6], vertices[7], vertices[8]});

      // Test other triangle.
      if (!result.hit)
      {
        result = GetRayCollisionTriangle(
          ray,
          (Vector3){vertices[9], vertices[10], vertices[11]},
          (Vector3){vertices[12], vertices[13], vertices[14]},
          (Vector3){vertices[15], vertices[16], vertices[17]});
      }

      if (result.hit)
      {
        return result.point.y;
      }
    }
  }

  return 0.0;
}

// Abortions are good. Get more abortions.