#include <data/data_center.hpp>
#include <data/data_type.hpp>
#include <base/subface.hpp>
#include <base/primitive.hpp>

dynamic_bitset_mp<> primitive::judge_sign_by_subface_sign(const stl_vector_mp<dynamic_bitset_mp<>>& subface_signs) const
{
    // NOTE: This overload assumes subface_signs[i] is already in the same order as
    // get_subfaces()[i] (i.e. local init order, not baked order). Only use when
    // signs have NOT been reordered by bake_blobtree.
    auto subfaces = get_subfaces();

    // Apply mark for i=0 as well (was incorrectly skipped before)
    dynamic_bitset_mp<> res = subfaces[0].get_mark() ? ~subface_signs[0] : subface_signs[0];
    for (size_t i = 1; i < subfaces.size(); ++i) {
        if (!subfaces[i].get_mark())
            res &= subface_signs[i];
        else
            res &= ~subface_signs[i];
    }

    return res;
}

dynamic_bitset_mp<> primitive::judge_sign_by_subface_sign(const stl_vector_mp<dynamic_bitset_mp<>>& subface_signs,
                                                          const stl_vector_mp<uint32_t>&            subface_indices,
                                                          const stl_vector_mp<bool>&                invert_flags) const
{
    // subface_indices[i] and invert_flags[i] are both in BAKED order (as stored in
    // baked_blobtree_t::subfaces_of_primitives / invert_flags_of_primitives).
    // invert_flags[i] is the mark that corresponds to baked subface subface_indices[i],
    // populated by bake_blobtree via pointer-based lookup so the order is always correct
    // regardless of how baking reorders subfaces relative to initialize() order.
    assert(subface_indices.size() == invert_flags.size());

    dynamic_bitset_mp<> res = invert_flags[0] ? ~subface_signs[subface_indices[0]] : subface_signs[subface_indices[0]];
    for (size_t i = 1; i < subface_indices.size(); ++i) {
        if (!invert_flags[i])
            res &= subface_signs[subface_indices[i]];
        else
            res &= ~subface_signs[subface_indices[i]];
    }

    return res;
}

// void primitive::initialize(primitive_data_center_t &data_center)
// {
//     this->data_center = make_pointer_wrapper(data_center);

//     auto subface_init_model_matrices = get_subface_init_model_matrices();
//     initialize(subface_init_model_matrices);
// }

void primitive::initialize(const stl_vector_mp<internal::paired_model_matrix_ptr_t>& matrices,
                           const stl_vector_mp<bool>&                                reversed_flags,
                           const stl_vector_mp<const void*>&                         geometry,
                           const aabb_t&                                             aabb)
{
    auto                       subfaces      = get_subfaces();
    auto                       subface_types = get_subface_types();
    // 阶段2: 创建 surface（geometry 所有权转移到 surface）
    stl_vector_mp<const void*> raw_geometry  = geometry.empty() ? get_subface_geometries() : geometry;

    stl_vector_mp<const void*> canonical_geometry;
    canonical_geometry.reserve(raw_geometry.size());

    for (size_t i = 0; i < subface_types.size(); ++i) {
        const void* canonical_ptr = nullptr;

        if (raw_geometry[i] != nullptr) {
            switch (subface_types[i]) {
                case surface_type::plane: {
                    canonical_ptr = data_center->acquire_plane_geometry(raw_geometry[i]);
                    std::cout << "[PRIMITIVE_INIT] Acquired plane geometry for subface " << i << "\n";
                    std::cout << "[PRIMITIVE_INIT] Plane geometry descriptor: local_aabb_min=("
                              << static_cast<const plane_descriptor_t*>(canonical_ptr)->local_aabb_min.x << ", "
                              << static_cast<const plane_descriptor_t*>(canonical_ptr)->local_aabb_min.y << ")\n";
                    break;
                }
                case surface_type::cylinder: {
                    // 将原始 geometry 注册到 data_center，获取规范化指针
                    canonical_ptr = data_center->acquire_cylinder_geometry(raw_geometry[i]);
                    break;
                }
                case surface_type::sphere: {
                    // TODO: 添加 sphere 的处理
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    canonical_ptr = data_center->acquire_extrude_polyline_geometry(raw_geometry[i]);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    canonical_ptr = data_center->acquire_extrude_helixline_geometry(raw_geometry[i]);
                    break;
                }
                default:                  canonical_ptr = nullptr; break;
            }
        }
        canonical_geometry.push_back(canonical_ptr);
    }
    for (size_t i = 0; i < subfaces.size(); ++i) {
        data_center->require_surface(subface_types[i], *matrices[i], canonical_geometry[i], subfaces[i]);
        subfaces[i].set_mark(static_cast<size_t>(reversed_flags[i]));
    }
    m_aabb = aabb;
}

void primitive::destroy()
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();
    assert(subfaces.front().get_ptr() != nullptr);

    // for (size_t i = 0; i < subfaces.size(); ++i) data_center->release_surface(subface_types[i], subfaces[i]);
    for (size_t i = 0; i < subfaces.size(); ++i) {
        auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[i]);
        const void* geometry_ptr  = get_geom_func(make_pointer_wrapper(subfaces[i].get_ptr()));

        data_center->release_surface(subface_types[i], subfaces[i]);
        subfaces[i].set_ptr(nullptr);

        // 释放descriptor
        if (geometry_ptr) {
            switch (subface_types[i]) {
                case surface_type::plane: {
                    data_center->release_plane_geometry(geometry_ptr);
                    break;
                }
                case surface_type::cylinder: {
                    data_center->release_cylinder_geometry(geometry_ptr);
                    break;
                }
                case surface_type::sphere: {
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    data_center->release_extrude_polyline_geometry(geometry_ptr);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    data_center->release_extrude_helixline_geometry(geometry_ptr);
                    break;
                }
                default: break;
            }
        }
    }
}

void primitive::apply_transform(internal::transform_type type, Eigen::Vector4d param)
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    // get a copy of old model matrices
    stl_vector_mp<internal::paired_model_matrix> new_model_matrices{};
    new_model_matrices.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        new_model_matrices.emplace_back(static_cast<internal::paired_model_matrix>(*subfaces[i]));
    }

    stl_vector_mp<const void*> geometry_ptrs;
    geometry_ptrs.reserve(subfaces.size());

    for (size_t i = 0; i < subfaces.size(); ++i) {
        if (subface_types[i] == surface_type::plane) {
            geometry_ptrs.push_back(nullptr);
        } else {
            auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[i]);
            const void* geom_ptr      = get_geom_func(make_pointer_wrapper(subfaces[i].get_ptr()));
            geometry_ptrs.push_back(geom_ptr);
        }
    }

    // make affine transform matrix from type and param
    // and apply transform to model matrices by the way, since this is faster than applying affine matrix
    Eigen::Affine3d affine = Eigen::Affine3d::Identity();
    switch (type) {
        case internal::transform_type::scale: {
            affine.linear() = param.head<3>().asDiagonal();
            Eigen::Affine3d affine_inv = affine.inverse();
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = affine * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * affine_inv;
            }
        } break;
        case internal::transform_type::rotation: {
            affine.linear() = Eigen::Quaterniond(param).toRotationMatrix();
            Eigen::Affine3d affine_inv = affine.inverse();
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = affine * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * affine_inv;
            }
        } break;
        case internal::transform_type::translation: {
            Eigen::Translation3d translation(param.head<3>());
            affine = Eigen::Affine3d(translation);
            for (auto& model_matrix : new_model_matrices) {
                model_matrix.local_to_world = translation * model_matrix.local_to_world;
                model_matrix.world_to_local = model_matrix.world_to_local * translation.inverse();
            }
        } break;
        default: throw std::invalid_argument("Invalid transform type");
    }
    // apply affine transform to aabb
    // m_aabb = m_aabb.transformed(affine);

    stl_vector_mp<size_t> saved_marks;
    saved_marks.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) { saved_marks.push_back(subfaces[i].get_mark()); }

    for (size_t i = 0; i < subfaces.size(); ++i) {
        data_center->require_surface(subface_types[i], new_model_matrices[i], geometry_ptrs[i], subfaces[i]);
        // Restore the mark that was set during initialize() via reversed_flags.
        subfaces[i].set_mark(saved_marks[i]);
    }
    recompute_aabb();
}

// 后续可能需要删除这个函数
aabb_t compute_local_aabb_from_descriptor(surface_type type, const void* descriptor_ptr)
{
    switch (type) {
        case surface_type::cylinder: {
            auto*  desc = static_cast<const cylinder_descriptor_t*>(descriptor_ptr);
            double r    = desc->radius;
                double h    = std::sqrt(desc->offset.x * desc->offset.x +
                                        desc->offset.y * desc->offset.y +
                                        desc->offset.z * desc->offset.z);
            return aabb_t(Eigen::Vector3d(-r, -r, 0), Eigen::Vector3d(r, r, h));
        }
        case surface_type::sphere: {
            // return aabb_t();
            return aabb_t(Eigen::Vector3d(-1, -1, -1), Eigen::Vector3d(2, 2, 2));
        }
        case surface_type::plane:                  return aabb_t();
        case surface_type::extrude_polyline_side:
        case surface_type::extrude_helixline_side:
        default:                                   std::cerr << "[WARNING] Unknown surface type: " << static_cast<int>(type) << "\n"; return aabb_t();
    }
}

// 重新计算 Primitive 的 AABB
void primitive::recompute_aabb()
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    // 重置为空 AABB
    m_aabb.setEmpty();
    bool has_valid_aabb = false;

    for (size_t i = 0; i < subfaces.size(); ++i) {
        aabb_t local_aabb = aabb_t();
        if (subface_types[i] == surface_type::extrude_polyline_side) {
            auto* face = static_cast<internal::extrude_polyline_side_face_t*>(subfaces[i].get_ptr());
            local_aabb.extend(face->get_aabb()); // 通过访问器读取
        } else if (subface_types[i] == surface_type::extrude_helixline_side) {
            auto* face = static_cast<internal::extrude_helixline_side_face_t*>(subfaces[i].get_ptr());
            local_aabb.extend(face->get_aabb());
        } else {
            const void* descriptor_ptr = nullptr;
            if (subface_types[i] != surface_type::plane && subfaces[i].get_ptr() != nullptr) {
                auto get_desc_func = internal::get_geometry_accessor_ptr(subface_types[i]);
                descriptor_ptr     = get_desc_func(make_pointer_wrapper(subfaces[i].get_ptr()));
            }
            local_aabb = compute_local_aabb_from_descriptor(subface_types[i], descriptor_ptr);

            if (local_aabb.isEmpty()) { continue; }
            if (!local_aabb.min().allFinite() || !local_aabb.max().allFinite()) {
                std::cerr << "[RECOMPUTE-AABB-WARNING] Invalid local AABB - skipping\n";
                continue;
            }
        }
        // 变换到世界坐标系
        aabb_t world_aabb = local_aabb;
        world_aabb.transform(subfaces[i]->local_to_world);
        if (!world_aabb.min().allFinite() || !world_aabb.max().allFinite()) {
            std::cerr << "[RECOMPUTE-AABB-WARNING] Invalid world AABB after transform\n";
            continue;
        }
        // 合并到总 AABB
        if (!has_valid_aabb) {
            m_aabb         = world_aabb;
            has_valid_aabb = true;
        } else {
            m_aabb.extend(world_aabb);
        }
    }
    if (!has_valid_aabb) {
        std::cerr << "[RECOMPUTE-AABB-WARNING] No valid AABB computed for primitive\n";
        {
            m_aabb = internal::k_aabb_unit;
        }
    }
}

void primitive::update_geometry(const void* new_geometry_ptr, size_t subface_index)
{
    auto subfaces      = get_subfaces();
    auto subface_types = get_subface_types();

    if (subface_types[subface_index] == surface_type::plane) {
        throw std::invalid_argument("Cannot update descriptor for plane surface");
    }

    auto        get_geom_func = internal::get_geometry_accessor_ptr(subface_types[subface_index]);
    const void* current_geom  = get_geom_func(make_pointer_wrapper(subfaces[subface_index].get_ptr()));
    if (current_geom == new_geometry_ptr) { return; }
    stl_vector_mp<internal::paired_model_matrix> current_matrices;
    current_matrices.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        current_matrices.emplace_back(static_cast<internal::paired_model_matrix>(*subfaces[i]));
    }

    stl_vector_mp<const void*> new_geometry;
    new_geometry.reserve(subfaces.size());
    for (size_t i = 0; i < subfaces.size(); ++i) {
        if (i == subface_index) {
            //  规范化新 descriptor
            const void* canonical_ptr = nullptr;

            switch (subface_types[i]) {
                case surface_type::plane: {
                    canonical_ptr = data_center->acquire_plane_geometry(new_geometry_ptr);
                }
                case surface_type::cylinder: {
                    // auto* desc    = static_cast<const cylinder_descriptor_t*>(new_descriptor_ptr);
                    canonical_ptr = data_center->acquire_cylinder_geometry(new_geometry_ptr);
                    break;
                }
                case surface_type::extrude_polyline_side: {
                    // const auto* raw = static_cast<const extrude_polyline_descriptor_t*>(new_descriptor_ptr);
                    canonical_ptr = data_center->acquire_extrude_polyline_geometry(new_geometry_ptr);
                    break;
                }
                case surface_type::extrude_helixline_side: {
                    // const auto* raw = static_cast<const extrude_helixline_descriptor_t*>(new_descriptor_ptr);
                    canonical_ptr = data_center->acquire_extrude_helixline_geometry(new_geometry_ptr);
                    break;
                }
                default: throw std::invalid_argument("Unsupported surface type for descriptor update");
            }
            new_geometry.push_back(canonical_ptr);
        } else {
            // 保持原有 descriptor
            if (subface_types[i] == surface_type::plane) {
                new_geometry.push_back(nullptr);
            } else {
                auto        get_desc_func     = internal::get_geometry_accessor_ptr(subface_types[i]);
                const void* existing_desc_ptr = get_desc_func(make_pointer_wrapper(subfaces[i].get_ptr()));
                new_geometry.push_back(existing_desc_ptr);
            }
        }
    }
    // 先释放 surface
    data_center->release_surface(subface_types[subface_index], subfaces[subface_index]);
    subfaces[subface_index].set_ptr(nullptr); // 清零，避免 require_surface 内双重释放

    if (current_geom && current_geom != new_geometry[subface_index]) {
        switch (subface_types[subface_index]) {
            case surface_type::plane: {
                data_center->release_plane_geometry(current_geom);
                break;
            }
            case surface_type::cylinder: {
                data_center->release_cylinder_geometry(current_geom);
                break;
            }
            case surface_type::extrude_polyline_side: {
                data_center->release_extrude_polyline_geometry(current_geom);
                break;
            }
            case surface_type::extrude_helixline_side: {
                data_center->release_extrude_helixline_geometry(current_geom);
                break;
            }
            default: break;
        }
    }
    const size_t saved_mark = subfaces[subface_index].get_mark();
    data_center->require_surface(subface_types[subface_index],
                                 current_matrices[subface_index],
                                 new_geometry[subface_index],
                                 subfaces[subface_index]);
    subfaces[subface_index].set_mark(saved_mark); 
    // 重新计算 AABB
    recompute_aabb();
}