#pragma once

#include <vector>
#include "Intersection.hpp"

class VoxelGrid {
private:
    int width, height, depth, xySize;
    Vec3<float> cellSize;
    AABB aabb;
    std::vector<uint8_t> data; // 每个体素只占用1 bit

public:
    VoxelGrid(const AABB &aabb, const Vec3<float> &cellSize)
            : cellSize(cellSize) {
        // 根据 AABB 和 cellSize 计算宽、高、深
        width = static_cast<int>(std::lround((aabb.max.x - aabb.min.x) / cellSize.x));
        height = static_cast<int>(std::lround((aabb.max.y - aabb.min.y) / cellSize.y));
        depth = static_cast<int>(std::lround((aabb.max.z - aabb.min.z) / cellSize.z));

        xySize = width * height;
        int totalSize = (xySize * depth + 7) / 8; // 每8个体素占用一个字节
        data.resize(totalSize, 0);
        this->aabb = aabb; // 使用传入的 AABB
    }

    VoxelGrid() = default;

    // 根据索引获取单元格的 AABB
    AABB getCellAABB(int x, int y, int z) const {
        Vec3<float> cellMin(x * cellSize.x, y * cellSize.y, z * cellSize.z);
        Vec3<float> cellMax((x + 1) * cellSize.x, (y + 1) * cellSize.y, (z + 1) * cellSize.z);
        return AABB(cellMin, cellMax) + aabb.min;
    }

    // 根据索引写入
    void setVoxel(int x, int y, int z, bool value) {
        assert(x >= 0 && x < width && y >= 0 && y < height && z >= 0 && z < depth);
        int index = x + y * width + z * xySize;
        int byteIndex = index / 8;
        int bitIndex = index % 8;

        if (value) {
            data[byteIndex] |= (1 << bitIndex); // 设置为1
        } else {
            data[byteIndex] &= ~(1 << bitIndex); // 设置为0
        }
    }

    // 根据索引读取
    bool getVoxel(int x, int y, int z) const {
        assert(x >= 0 && x < width && y >= 0 && y < height && z >= 0 && z < depth);
        int index = x + y * width + z * xySize;
        int byteIndex = index / 8;
        int bitIndex = index % 8;

        return (data[byteIndex] >> bitIndex) & 1; // 返回位值
    }

    static VoxelGrid generateFromMesh(const Mesh &mesh, const Vec3<float> &cellSize) {
        auto sceneAABB = mesh.getAABB();
        auto voxelGrid = VoxelGrid(sceneAABB, cellSize);
        BVH bvh(mesh);
        int solvedVoxel = 0;
        for (int z = 0; z < voxelGrid.depth; ++z) {
            for (int y = 0; y < voxelGrid.height; ++y) {
                for (int x = 0; x < voxelGrid.width; ++x) {
                    auto start = voxelGrid.getCellAABB(x, y, z).center();
                    auto dir = Vec3f(1, 2, 3).norm();
                    if (bvh.intersectWithRay({start, dir}) % 2 == 1) {
                        voxelGrid.setVoxel(x, y, z, true);
                    } else {
                        voxelGrid.setVoxel(x, y, z, false);
                    }

                    solvedVoxel++;
                    if (solvedVoxel % int(1e5) == 0) {
                        std::cout << "Solved " << solvedVoxel << " voxels." << std::endl;
                    }
                }
            }
        }
        return voxelGrid;
    }
};