#include <gtest/gtest.h>
#include "vde/brep/constraint_solver.h"
#include "vde/brep/modeling.h"
#include "vde/core/transform.h"
#include "vde/core/aabb.h"
#include <cmath>
#include <vector>

using namespace vde::brep;
using namespace vde::core;

// ════════════════════════════════════════════════════════════════
//  Helpers
// ════════════════════════════════════════════════════════════════

namespace {

/// World-space AABB of a node (transforming local model bounds).
AABB3D world_bounds(const AssemblyNode* node) {
    AABB3D bb_world;
    if (!node->model.has_value()) return bb_world;
    AABB3D bb_local = node->model->bounds();
    Point3D corners[8] = {
        bb_local.min(),
        Point3D(bb_local.max().x(), bb_local.min().y(), bb_local.min().z()),
        Point3D(bb_local.max().x(), bb_local.max().y(), bb_local.min().z()),
        Point3D(bb_local.min().x(), bb_local.max().y(), bb_local.min().z()),
        Point3D(bb_local.min().x(), bb_local.min().y(), bb_local.max().z()),
        Point3D(bb_local.max().x(), bb_local.min().y(), bb_local.max().z()),
        bb_local.max(),
        Point3D(bb_local.min().x(), bb_local.max().y(), bb_local.max().z()),
    };
    for (const auto& c : corners) {
        bb_world.expand(node->local_transform * c);
    }
    return bb_world;
}

} // anonymous namespace

// ════════════════════════════════════════════════════════════════
//  apply_coincident
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, CoincidentAlignsFacePlanes) {
    Assembly assy("test_coincident");
    auto* box_a = assy.root.add_part("box_a", make_box(2, 2, 2));
    auto* box_b = assy.root.add_part("box_b", make_box(2, 2, 2),
                                      translate(0, 0, 10));

    apply_coincident(*box_a, *box_b);

    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);

    // After coincident, box_b bottom should sit on box_a top (gap ~0)
    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 0.0, 1e-4);
}

TEST(ConstraintSolverTest, CoincidentPreservesLateralAlignment) {
    Assembly assy("test_coincident2");
    auto* box_a = assy.root.add_part("box_a", make_box(4, 4, 2));
    auto* box_b = assy.root.add_part("box_b", make_box(2, 2, 4),
                                      translate(0, 0, 8));

    apply_coincident(*box_a, *box_b);

    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);

    // Centers should align in X and Y
    EXPECT_NEAR(bb_b.center().x(), bb_a.center().x(), 1e-4);
    EXPECT_NEAR(bb_b.center().y(), bb_a.center().y(), 1e-4);
    EXPECT_NEAR(bb_b.min().z(), bb_a.max().z(), 1e-4);
}

TEST(ConstraintSolverTest, CoincidentNoModelReturnsIdentity) {
    Assembly assy("test_no_model");
    auto* node_a = assy.root.add_subassembly("empty_a");
    auto* node_b = assy.root.add_subassembly("empty_b");

    auto T = apply_coincident(*node_a, *node_b);
    EXPECT_TRUE(T.isApprox(Transform3D::Identity(), 1e-12));
}

// ════════════════════════════════════════════════════════════════
//  apply_concentric
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, ConcentricAlignsAxes) {
    Assembly assy("test_concentric");
    auto* cyl_a = assy.root.add_part("cyl_a", make_cylinder(1.0, 6.0));
    auto* cyl_b = assy.root.add_part("cyl_b", make_cylinder(0.5, 4.0),
                                      translate(5, 5, 0));

    apply_concentric(*cyl_a, *cyl_b);

    AABB3D bb_a = world_bounds(cyl_a);
    AABB3D bb_b = world_bounds(cyl_b);

    // Centers should align in X and Y after concentric constraint
    EXPECT_NEAR(bb_b.center().x(), bb_a.center().x(), 1e-4);
    EXPECT_NEAR(bb_b.center().y(), bb_a.center().y(), 1e-4);
}

TEST(ConstraintSolverTest, ConcentricPreservesZPosition) {
    Assembly assy("test_concentric_z");
    auto* cyl_a = assy.root.add_part("cyl_a", make_cylinder(1.0, 6.0));
    auto* cyl_b = assy.root.add_part("cyl_b", make_cylinder(0.5, 4.0),
                                      translate(5, 5, 3.0));

    double z_before = world_bounds(cyl_b).center().z();
    apply_concentric(*cyl_a, *cyl_b);
    double z_after = world_bounds(cyl_b).center().z();

    // Z should be preserved (concentric only affects radial directions)
    EXPECT_NEAR(z_after, z_before, 1e-4);
}

TEST(ConstraintSolverTest, ConcentricDifferentOrientationCylinders) {
    // TODO: This test requires rotated input shapes; for now we test
    // the axis-invariant case.
    Assembly assy("test_concentric_orient");
    auto* cyl_a = assy.root.add_part("cyl_a", make_cylinder(1.0, 6.0));
    auto* cyl_b = assy.root.add_part("cyl_b", make_cylinder(0.5, 4.0),
                                      translate(3, 0, 0));

    apply_concentric(*cyl_a, *cyl_b);

    AABB3D bb_a = world_bounds(cyl_a);
    AABB3D bb_b = world_bounds(cyl_b);

    EXPECT_NEAR(bb_b.center().x(), bb_a.center().x(), 1e-4);
    EXPECT_NEAR(bb_b.center().y(), bb_a.center().y(), 1e-4);
}

// ════════════════════════════════════════════════════════════════
//  apply_distance
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, DistanceProducesCorrectGap) {
    Assembly assy("test_distance");
    auto* box_a = assy.root.add_part("box_a", make_box(2, 2, 2));
    auto* box_b = assy.root.add_part("box_b", make_box(2, 2, 2),
                                      translate(0, 0, 10));

    apply_distance(*box_a, *box_b, 5.0);

    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);

    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 5.0, 1e-4);
}

TEST(ConstraintSolverTest, DistanceZeroIsCoincident) {
    Assembly assy("test_distance_zero");
    auto* box_a = assy.root.add_part("box_a", make_box(2, 2, 2));
    auto* box_b = assy.root.add_part("box_b", make_box(2, 2, 2),
                                      translate(0, 0, 8));

    apply_distance(*box_a, *box_b, 0.0);

    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);

    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 0.0, 1e-4);
}

TEST(ConstraintSolverTest, DistanceNegativePushesAway) {
    // Negative distance means overlap
    Assembly assy("test_distance_neg");
    auto* box_a = assy.root.add_part("box_a", make_box(2, 2, 2));
    auto* box_b = assy.root.add_part("box_b", make_box(2, 2, 2),
                                      translate(0, 0, 10));

    apply_distance(*box_a, *box_b, -1.0);

    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);

    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, -1.0, 1e-4);
}

// ════════════════════════════════════════════════════════════════
//  solve_constraints
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, SolveSingleConstraint) {
    Assembly assy("test_solve_single");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 2), translate(5, 3, 10));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Coincident}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-4));

    auto* box_a = assy.root.children[0].get();
    auto* box_b = assy.root.children[1].get();
    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);
    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 0.0, 1e-3);
}

TEST(ConstraintSolverTest, SolveMultipleConstraints) {
    Assembly assy("test_solve_multi");
    assy.root.add_part("box_a", make_box(4, 4, 2));
    assy.root.add_part("cyl_a", make_cylinder(1.0, 6.0), translate(5, 0, 0));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Coincident},   // faces coplanar
        {0, 1, Constraint3D::Concentric}    // axes aligned
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 100, 1e-3));

    auto* box_a = assy.root.children[0].get();
    auto* cyl_a = assy.root.children[1].get();
    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(cyl_a);

    // Centers aligned in X and Y (concentric)
    EXPECT_NEAR(bb_b.center().x(), bb_a.center().x(), 1e-3);
    EXPECT_NEAR(bb_b.center().y(), bb_a.center().y(), 1e-3);
}

TEST(ConstraintSolverTest, SolveWithDistance) {
    Assembly assy("test_solve_dist");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 2), translate(0, 0, 10));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Distance, 3.0}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-4));

    auto* box_a = assy.root.children[0].get();
    auto* box_b = assy.root.children[1].get();
    AABB3D bb_a = world_bounds(box_a);
    AABB3D bb_b = world_bounds(box_b);
    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 3.0, 1e-3);
}

TEST(ConstraintSolverTest, SolveEmptyConstraintsReturnsTrue) {
    Assembly assy("test_solve_empty");
    assy.root.add_part("box_a", make_box(2, 2, 2));

    std::vector<ConstraintEntry> constraints;
    EXPECT_TRUE(solve_constraints(assy, constraints));
}

TEST(ConstraintSolverTest, SolveConvergenceCheck) {
    Assembly assy("test_solve_conv");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 2), translate(10, 10, 10));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Coincident}
    };

    // Should converge well within 100 iterations
    bool converged = solve_constraints(assy, constraints, 100, 1e-6);
    EXPECT_TRUE(converged);
}

TEST(ConstraintSolverTest, SolveConflictingConstraintsHandledGracefully) {
    // Conflicting: box_b must be both coincident and at distance 10 from box_a
    Assembly assy("test_conflict");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 2), translate(0, 0, 10));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Coincident},
        {0, 1, Constraint3D::Distance, 10.0}
    };

    // Should NOT crash — just returns false (didn't converge)
    bool converged = solve_constraints(assy, constraints, 50, 1e-6);
    // Conflicting constraints won't converge; this is expected
    EXPECT_FALSE(converged);
}

// ════════════════════════════════════════════════════════════════
//  Constraint3D type coverage — Parallel, Perpendicular, Angle
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, ParallelConstraintAppliesTransform) {
    Assembly assy("test_parallel");
    assy.root.add_part("box_a", make_box(2, 2, 8));  // tall in Z
    assy.root.add_part("box_b", make_box(2, 2, 8),   // tall in Z
                        translate(0, 5, 0));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Parallel}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-6));
    // Parallel constraint should not crash; it applies rotation to align axes
    SUCCEED();
}

TEST(ConstraintSolverTest, PerpendicularConstraintAppliesTransform) {
    Assembly assy("test_perp");
    assy.root.add_part("box_a", make_box(2, 2, 8));  // tall in Z
    assy.root.add_part("box_b", make_box(2, 2, 8),
                        translate(0, 5, 0));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Perpendicular}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-6));

    // Verify transform is non-identity
    auto* box_b = assy.root.children[1].get();
    EXPECT_FALSE(box_b->local_transform.isApprox(Transform3D::Identity(), 1e-3));
}

TEST(ConstraintSolverTest, AngleConstraintAppliesRotation) {
    Assembly assy("test_angle");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 8));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Angle, M_PI / 2.0}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-6));

    auto* box_b = assy.root.children[1].get();
    EXPECT_FALSE(box_b->local_transform.isApprox(Transform3D::Identity(), 1e-6));
}

TEST(ConstraintSolverTest, TangentConstraintApplies) {
    Assembly assy("test_tangent");
    assy.root.add_part("box_a", make_box(2, 2, 2));
    assy.root.add_part("box_b", make_box(2, 2, 2),
                        translate(0, 0, 10));

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Tangent}
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 50, 1e-4));

    AABB3D bb_a = world_bounds(assy.root.children[0].get());
    AABB3D bb_b = world_bounds(assy.root.children[1].get());
    double gap = bb_b.min().z() - bb_a.max().z();
    EXPECT_NEAR(gap, 0.0, 1e-3);
}

// ════════════════════════════════════════════════════════════════
//  Three-node assembly scenario
// ════════════════════════════════════════════════════════════════

TEST(ConstraintSolverTest, ThreeNodeSolve) {
    Assembly assy("test_three_node");
    assy.root.add_part("base", make_box(4, 4, 1));             // node 0
    assy.root.add_part("mid", make_box(2, 2, 2),
                        translate(1, 1, 10));                   // node 1
    assy.root.add_part("top", make_box(1, 1, 3),
                        translate(1, 1, 15));                   // node 2

    std::vector<ConstraintEntry> constraints = {
        {0, 1, Constraint3D::Coincident},   // mid sits on base
        {1, 2, Constraint3D::Coincident},   // top sits on mid
        {0, 1, Constraint3D::Concentric},   // centers aligned
        {1, 2, Constraint3D::Concentric}    // centers aligned
    };

    EXPECT_TRUE(solve_constraints(assy, constraints, 200, 1e-3));

    // All three should now be stacked and centered
    AABB3D bb0 = world_bounds(assy.root.children[0].get());
    AABB3D bb1 = world_bounds(assy.root.children[1].get());
    AABB3D bb2 = world_bounds(assy.root.children[2].get());

    // Centers aligned in X and Y
    EXPECT_NEAR(bb1.center().x(), bb0.center().x(), 1e-3);
    EXPECT_NEAR(bb1.center().y(), bb0.center().y(), 1e-3);
    EXPECT_NEAR(bb2.center().x(), bb0.center().x(), 1e-3);
    EXPECT_NEAR(bb2.center().y(), bb0.center().y(), 1e-3);

    // Stacked: bb1 sits on bb0, bb2 sits on bb1
    double gap_01 = bb1.min().z() - bb0.max().z();
    double gap_12 = bb2.min().z() - bb1.max().z();
    EXPECT_NEAR(gap_01, 0.0, 1e-3);
    EXPECT_NEAR(gap_12, 0.0, 1e-3);
}