// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// bvh.hpp — Bounding Volume Hierarchy (SAH ベース)
//
// 三角形メッシュの高速レイ交差判定。
// SAH (Surface Area Heuristic) で構築し、再帰的にトラバース。

#ifndef CALX_IMAGE_RENDER_BVH_HPP
#define CALX_IMAGE_RENDER_BVH_HPP

#include <image/geometry/primitives.hpp>
#include <image/types/RGB.hpp>
#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <vector>

namespace calx::image {

/// 頂点カラー付き三角形
template<typename T>
struct ColoredTriangle {
    Point3D<T> v0, v1, v2;
    RGB<T> c0, c1, c2;      // 頂点カラー
    Point3D<T> n0, n1, n2;  // 頂点法線 (スムーズシェーディング)
    T opacity = T(1);        // 不透明度 (0=完全透明, 1=不透明)
};

/// BVH ノード
template<typename T>
struct BVHNode {
    AABB<T> bounds;
    int left = -1, right = -1;   // 子ノードインデックス (-1 = リーフ)
    int triStart = 0, triCount = 0; // リーフの三角形範囲
};

/// BVH (三角形メッシュ用)
template<typename T>
class BVH {
public:
    struct HitInfo {
        T t;
        Point3D<T> normal;
        RGB<T> color;       // 補間された頂点カラー
        T opacity = T(1);   // 補間された不透明度
    };

    /// 三角形リストから BVH を構築
    void build(std::vector<ColoredTriangle<T>> triangles) {
        tris_ = std::move(triangles);
        nodes_.clear();
        if (tris_.empty()) return;

        // 三角形インデックス
        indices_.resize(tris_.size());
        std::iota(indices_.begin(), indices_.end(), 0);

        // 各三角形の AABB と重心を事前計算
        centroids_.resize(tris_.size());
        triBounds_.resize(tris_.size());
        for (std::size_t i = 0; i < tris_.size(); ++i) {
            const auto& tri = tris_[i];
            triBounds_[i] = triAABB(tri);
            centroids_[i] = {
                (tri.v0.x + tri.v1.x + tri.v2.x) / T(3),
                (tri.v0.y + tri.v1.y + tri.v2.y) / T(3),
                (tri.v0.z + tri.v1.z + tri.v2.z) / T(3)
            };
        }

        // 再帰的に構築
        nodes_.reserve(tris_.size() * 2);
        buildNode(0, static_cast<int>(tris_.size()), 0);
    }

    /// レイとの最近交差判定
    bool intersect(const Point3D<T>& origin, const Point3D<T>& dir,
                   T tMin, T tMax, HitInfo& hit) const {
        if (nodes_.empty()) return false;
        return traverseNode(0, origin, dir, tMin, tMax, hit);
    }

    /// レイ上の全交差点を列挙 (半透明合成用)
    /// 結果は t の昇順にソートされる
    void intersectAll(const Point3D<T>& origin, const Point3D<T>& dir,
                      T tMin, T tMax, std::vector<HitInfo>& hits) const {
        if (nodes_.empty()) return;
        Point3D<T> invDir{T(1) / dir.x, T(1) / dir.y, T(1) / dir.z};
        traverseAllNode(0, origin, dir, invDir, tMin, tMax, hits);
        std::sort(hits.begin(), hits.end(),
                  [](const HitInfo& a, const HitInfo& b) { return a.t < b.t; });
    }

    std::size_t triangleCount() const { return tris_.size(); }

private:
    std::vector<ColoredTriangle<T>> tris_;
    std::vector<int> indices_;
    std::vector<Point3D<T>> centroids_;
    std::vector<AABB<T>> triBounds_;
    std::vector<BVHNode<T>> nodes_;

    static AABB<T> triAABB(const ColoredTriangle<T>& tri) {
        return {
            {std::min({tri.v0.x, tri.v1.x, tri.v2.x}),
             std::min({tri.v0.y, tri.v1.y, tri.v2.y}),
             std::min({tri.v0.z, tri.v1.z, tri.v2.z})},
            {std::max({tri.v0.x, tri.v1.x, tri.v2.x}),
             std::max({tri.v0.y, tri.v1.y, tri.v2.y}),
             std::max({tri.v0.z, tri.v1.z, tri.v2.z})}
        };
    }

    static AABB<T> mergeAABB(const AABB<T>& a, const AABB<T>& b) {
        return {
            {std::min(a.min.x, b.min.x), std::min(a.min.y, b.min.y), std::min(a.min.z, b.min.z)},
            {std::max(a.max.x, b.max.x), std::max(a.max.y, b.max.y), std::max(a.max.z, b.max.z)}
        };
    }

    static T surfaceArea(const AABB<T>& b) {
        T dx = b.max.x - b.min.x;
        T dy = b.max.y - b.min.y;
        T dz = b.max.z - b.min.z;
        return T(2) * (dx * dy + dy * dz + dz * dx);
    }

    int buildNode(int start, int end, int depth) {
        int idx = static_cast<int>(nodes_.size());
        nodes_.push_back({});
        auto& node = nodes_[idx];

        // バウンディングボックス計算
        node.bounds = triBounds_[indices_[start]];
        for (int i = start + 1; i < end; ++i)
            node.bounds = mergeAABB(node.bounds, triBounds_[indices_[i]]);

        int count = end - start;

        // リーフ条件
        if (count <= 4 || depth >= 32) {
            node.triStart = start;
            node.triCount = count;
            return idx;
        }

        // SAH で最適分割を探索
        T bestCost = std::numeric_limits<T>::max();
        int bestAxis = -1, bestSplit = -1;
        T parentArea = surfaceArea(node.bounds);

        for (int axis = 0; axis < 3; ++axis) {
            // 重心でソート
            std::sort(indices_.begin() + start, indices_.begin() + end,
                [&](int a, int b) {
                    auto ca = (axis == 0) ? centroids_[a].x :
                              (axis == 1) ? centroids_[a].y : centroids_[a].z;
                    auto cb = (axis == 0) ? centroids_[b].x :
                              (axis == 1) ? centroids_[b].y : centroids_[b].z;
                    return ca < cb;
                });

            // sweep: 左から累積 AABB
            constexpr int MAX_BINS = 32;
            int nBins = std::min(count - 1, MAX_BINS);

            for (int b = 0; b < nBins; ++b) {
                int split = start + (count * (b + 1)) / (nBins + 1);
                if (split <= start || split >= end) continue;

                AABB<T> leftBox = triBounds_[indices_[start]];
                for (int i = start + 1; i < split; ++i)
                    leftBox = mergeAABB(leftBox, triBounds_[indices_[i]]);

                AABB<T> rightBox = triBounds_[indices_[split]];
                for (int i = split + 1; i < end; ++i)
                    rightBox = mergeAABB(rightBox, triBounds_[indices_[i]]);

                T cost = T(split - start) * surfaceArea(leftBox) +
                         T(end - split) * surfaceArea(rightBox);
                cost /= parentArea;

                if (cost < bestCost) {
                    bestCost = cost;
                    bestAxis = axis;
                    bestSplit = split;
                }
            }
        }

        if (bestAxis < 0 || bestSplit <= start || bestSplit >= end) {
            node.triStart = start;
            node.triCount = count;
            return idx;
        }

        // 最適軸で再ソート
        std::sort(indices_.begin() + start, indices_.begin() + end,
            [&](int a, int b) {
                auto ca = (bestAxis == 0) ? centroids_[a].x :
                          (bestAxis == 1) ? centroids_[a].y : centroids_[a].z;
                auto cb = (bestAxis == 0) ? centroids_[b].x :
                          (bestAxis == 1) ? centroids_[b].y : centroids_[b].z;
                return ca < cb;
            });

        // nodes_ の参照は子ノード構築で無効になるのでインデックスで操作
        int leftChild = buildNode(start, bestSplit, depth + 1);
        int rightChild = buildNode(bestSplit, end, depth + 1);
        nodes_[idx].left = leftChild;
        nodes_[idx].right = rightChild;
        nodes_[idx].triCount = 0;

        return idx;
    }

    /// AABB-Ray 交差 (slab method, 高速版)
    static bool intersectAABB(const AABB<T>& box,
                              const Point3D<T>& origin,
                              const Point3D<T>& invDir,
                              T tMin, T tMax) {
        T tx1 = (box.min.x - origin.x) * invDir.x;
        T tx2 = (box.max.x - origin.x) * invDir.x;
        tMin = std::max(tMin, std::min(tx1, tx2));
        tMax = std::min(tMax, std::max(tx1, tx2));

        T ty1 = (box.min.y - origin.y) * invDir.y;
        T ty2 = (box.max.y - origin.y) * invDir.y;
        tMin = std::max(tMin, std::min(ty1, ty2));
        tMax = std::min(tMax, std::max(ty1, ty2));

        T tz1 = (box.min.z - origin.z) * invDir.z;
        T tz2 = (box.max.z - origin.z) * invDir.z;
        tMin = std::max(tMin, std::min(tz1, tz2));
        tMax = std::min(tMax, std::max(tz1, tz2));

        return tMin <= tMax;
    }

    /// Möller–Trumbore + 重心座標
    bool intersectTri(int triIdx, const Point3D<T>& origin,
                      const Point3D<T>& dir, T tMin, T tMax,
                      T& tHit, T& u, T& v) const {
        const auto& tri = tris_[triIdx];
        auto e1 = tri.v1 - tri.v0;
        auto e2 = tri.v2 - tri.v0;
        auto h = cross(dir, e2);
        T a = dot(e1, h);
        if (std::abs(a) < T(1e-10)) return false;

        T f = T(1) / a;
        auto s = origin - tri.v0;
        u = f * dot(s, h);
        if (u < T(0) || u > T(1)) return false;

        auto q = cross(s, e1);
        v = f * dot(dir, q);
        if (v < T(0) || u + v > T(1)) return false;

        T t = f * dot(e2, q);
        if (t < tMin || t > tMax) return false;
        tHit = t;
        return true;
    }

    bool traverseNode(int nodeIdx, const Point3D<T>& origin,
                      const Point3D<T>& dir, T tMin, T& tMax,
                      HitInfo& hit) const {
        const auto& node = nodes_[nodeIdx];

        Point3D<T> invDir{T(1) / dir.x, T(1) / dir.y, T(1) / dir.z};
        if (!intersectAABB(node.bounds, origin, invDir, tMin, tMax))
            return false;

        if (node.triCount > 0) {
            // リーフ: 三角形と交差判定
            bool found = false;
            for (int i = node.triStart; i < node.triStart + node.triCount; ++i) {
                T t, u, v;
                if (intersectTri(indices_[i], origin, dir, tMin, tMax, t, u, v)) {
                    tMax = t;
                    const auto& tri = tris_[indices_[i]];
                    T w = T(1) - u - v;

                    hit.t = t;
                    // スムーズ法線 (頂点法線の補間)
                    hit.normal = {
                        w * tri.n0.x + u * tri.n1.x + v * tri.n2.x,
                        w * tri.n0.y + u * tri.n1.y + v * tri.n2.y,
                        w * tri.n0.z + u * tri.n1.z + v * tri.n2.z
                    };
                    T nlen = hit.normal.norm();
                    if (nlen > T(1e-10))
                        hit.normal = hit.normal * (T(1) / nlen);

                    // 頂点カラーの補間
                    hit.color = {
                        w * tri.c0.r + u * tri.c1.r + v * tri.c2.r,
                        w * tri.c0.g + u * tri.c1.g + v * tri.c2.g,
                        w * tri.c0.b + u * tri.c1.b + v * tri.c2.b
                    };
                    hit.opacity = tri.opacity;
                    found = true;
                }
            }
            return found;
        }

        // 内部ノード: 両方の子をトラバース
        bool hitLeft = traverseNode(node.left, origin, dir, tMin, tMax, hit);
        bool hitRight = traverseNode(node.right, origin, dir, tMin, tMax, hit);
        return hitLeft || hitRight;
    }

    /// 全交差点列挙用トラバーサル (tMax を縮めない)
    void traverseAllNode(int nodeIdx, const Point3D<T>& origin,
                         const Point3D<T>& dir, const Point3D<T>& invDir,
                         T tMin, T tMax,
                         std::vector<HitInfo>& hits) const {
        const auto& node = nodes_[nodeIdx];

        if (!intersectAABB(node.bounds, origin, invDir, tMin, tMax))
            return;

        if (node.triCount > 0) {
            for (int i = node.triStart; i < node.triStart + node.triCount; ++i) {
                T t, u, v;
                if (intersectTri(indices_[i], origin, dir, tMin, tMax, t, u, v)) {
                    const auto& tri = tris_[indices_[i]];
                    T w = T(1) - u - v;
                    HitInfo h;
                    h.t = t;
                    h.normal = {
                        w * tri.n0.x + u * tri.n1.x + v * tri.n2.x,
                        w * tri.n0.y + u * tri.n1.y + v * tri.n2.y,
                        w * tri.n0.z + u * tri.n1.z + v * tri.n2.z
                    };
                    T nlen = h.normal.norm();
                    if (nlen > T(1e-10)) h.normal = h.normal * (T(1) / nlen);
                    h.color = {
                        w * tri.c0.r + u * tri.c1.r + v * tri.c2.r,
                        w * tri.c0.g + u * tri.c1.g + v * tri.c2.g,
                        w * tri.c0.b + u * tri.c1.b + v * tri.c2.b
                    };
                    h.opacity = tri.opacity;
                    hits.push_back(h);
                }
            }
            return;
        }

        traverseAllNode(node.left, origin, dir, invDir, tMin, tMax, hits);
        traverseAllNode(node.right, origin, dir, invDir, tMin, tMax, hits);
    }
};

} // namespace calx::image

#endif // CALX_IMAGE_RENDER_BVH_HPP
