// SPDX-License-Identifier: LGPL-2.1-or-later
/****************************************************************************
 *                                                                          *
 *   Copyright (c) 2023 Ondsel <development@ondsel.com>                     *
 *                                                                          *
 *   This file is part of FreeCAD.                                          *
 *                                                                          *
 *   FreeCAD is free software: you can redistribute it and/or modify it     *
 *   under the terms of the GNU Lesser General Public License as            *
 *   published by the Free Software Foundation, either version 2.1 of the   *
 *   License, or (at your option) any later version.                        *
 *                                                                          *
 *   FreeCAD is distributed in the hope that it will be useful, but         *
 *   WITHOUT ANY WARRANTY; without even the implied warranty of             *
 *   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU       *
 *   Lesser General Public License for more details.                        *
 *                                                                          *
 *   You should have received a copy of the GNU Lesser General Public       *
 *   License along with FreeCAD. If not, see                                *
 *   <https://www.gnu.org/licenses/>.                                       *
 *                                                                          *
 ***************************************************************************/

#include <boost/core/ignore_unused.hpp>
#include <cmath>
#include <numbers>
#include <vector>
#include <unordered_map>


#include <App/Application.h>
#include <App/Datums.h>
#include <App/Document.h>
#include <App/DocumentObjectGroup.h>
#include <App/FeaturePythonPyImp.h>
#include <App/Link.h>
#include <App/PropertyPythonObject.h>
#include <Base/Console.h>
#include <Base/Placement.h>
#include <Base/Rotation.h>
#include <Base/Tools.h>
#include <Base/Interpreter.h>

#include <Mod/Part/App/TopoShape.h>
#include <Mod/Part/App/AttachExtension.h>

#include <Mod/Assembly/Solver/IKCSolver.h>
#include <Mod/Assembly/Solver/SolverRegistry.h>

#include "AssemblyLink.h"
#include "AssemblyObject.h"
#include "AssemblyObjectPy.h"
#include "AssemblyUtils.h"
#include "JointGroup.h"
#include "ViewGroup.h"

FC_LOG_LEVEL_INIT("Assembly", true, true, true)

using namespace Assembly;

namespace PartApp = Part;


// ── Transform conversion helpers ───────────────────────────────────

static KCSolve::Transform placementToTransform(const Base::Placement& plc)
{
    KCSolve::Transform tf;
    Base::Vector3d pos = plc.getPosition();
    tf.position = {pos.x, pos.y, pos.z};

    // Base::Rotation(q0,q1,q2,q3) = (x,y,z,w)
    // KCSolve::Transform quaternion = (w,x,y,z)
    double q0, q1, q2, q3;
    plc.getRotation().getValue(q0, q1, q2, q3);
    tf.quaternion = {q3, q0, q1, q2};

    return tf;
}

static Base::Placement transformToPlacement(const KCSolve::Transform& tf)
{
    Base::Vector3d pos(tf.position[0], tf.position[1], tf.position[2]);
    // KCSolve (w,x,y,z) → Base::Rotation(x,y,z,w)
    Base::Rotation rot(tf.quaternion[1], tf.quaternion[2], tf.quaternion[3], tf.quaternion[0]);
    return Base::Placement(pos, rot);
}


// ================================ Assembly Object ============================

PROPERTY_SOURCE(Assembly::AssemblyObject, App::Part)

AssemblyObject::AssemblyObject()
    : bundleFixed(false)
    , lastDoF(0)
    , lastHasConflict(false)
    , lastHasRedundancies(false)
    , lastHasPartialRedundancies(false)
    , lastHasMalformedConstraints(false)
    , lastSolverStatus(0)
{
    lastDoF = numberOfComponents() * 6;
    signalSolverUpdate();
}

AssemblyObject::~AssemblyObject() = default;

PyObject* AssemblyObject::getPyObject()
{
    if (PythonObject.is(Py::_None())) {
        // ref counter is set to 1
        PythonObject = Py::Object(new AssemblyObjectPy(this), true);
    }
    return Py::new_reference_to(PythonObject);
}

App::DocumentObjectExecReturn* AssemblyObject::execute()
{
    App::DocumentObjectExecReturn* ret = App::Part::execute();

    ParameterGrp::handle hGrp = App::GetApplication().GetParameterGroupByPath(
        "User parameter:BaseApp/Preferences/Mod/Assembly"
    );
    if (hGrp->GetBool("SolveOnRecompute", true)) {
        solve(false, false);  // No need to update jcs since recompute updated them.
    }
    return ret;
}

void AssemblyObject::onChanged(const App::Property* prop)
{
    if (App::GetApplication().isRestoring()) {
        App::Part::onChanged(prop);
        return;
    }

    if (prop == &Group) {
        updateSolveStatus();
    }
    App::Part::onChanged(prop);
}


// ── Solver integration ─────────────────────────────────────────────

KCSolve::IKCSolver* AssemblyObject::getOrCreateSolver()
{
    if (!solver_) {
        solver_ = KCSolve::SolverRegistry::instance().get("ondsel");
    }
    return solver_.get();
}

int AssemblyObject::solve(bool enableRedo, bool updateJCS)
{
    ensureIdentityPlacements();

    auto* solver = getOrCreateSolver();
    if (!solver) {
        FC_ERR("No solver available");
        lastSolverStatus = -1;
        return -1;
    }

    partIdToObjs_.clear();
    objToPartId_.clear();

    auto groundedObjs = getGroundedParts();
    if (groundedObjs.empty()) {
        return -6;
    }

    std::vector<App::DocumentObject*> joints = getJoints(updateJCS);
    removeUnconnectedJoints(joints, groundedObjs);

    KCSolve::SolveContext ctx = buildSolveContext(joints);

    // Always save placements to enable orientation flip detection
    savePlacementsForUndo();

    try {
        lastResult_ = solver->solve(ctx);
        lastSolverStatus = static_cast<int>(lastResult_.status);
    }
    catch (const std::exception& e) {
        FC_ERR("Solve failed: " << e.what());
        lastSolverStatus = -1;
        updateSolveStatus();
        return -1;
    }
    catch (...) {
        FC_ERR("Solve failed: unhandled exception");
        lastSolverStatus = -1;
        updateSolveStatus();
        return -1;
    }

    if (lastResult_.status == KCSolve::SolveStatus::Failed) {
        updateSolveStatus();
        return -1;
    }

    // Validate that the solve didn't cause any parts to flip orientation
    if (!validateNewPlacements()) {
        // Restore previous placements - the solve found an invalid configuration
        undoSolve();
        lastSolverStatus = -2;
        updateSolveStatus();
        return -2;
    }

    setNewPlacements();

    // Clear undo history if the caller didn't want redo capability
    if (!enableRedo) {
        clearUndo();
    }

    redrawJointPlacements(joints);

    updateSolveStatus();

    return 0;
}

void AssemblyObject::updateSolveStatus()
{
    lastRedundantJoints.clear();
    lastHasRedundancies = false;
    //+1 because there's a grounded joint to origin
    lastDoF = (1 + numberOfComponents()) * 6;

    if (!solver_ || lastResult_.placements.empty()) {
        solve();
    }

    if (!solver_) {
        return;
    }

    // Use DOF from the solver result if available
    if (lastResult_.dof >= 0) {
        lastDoF = lastResult_.dof;
    }

    // Process diagnostics from the solver result
    for (const auto& diag : lastResult_.diagnostics) {
        // Filter to only actual FreeCAD joint objects (not grounding joints)
        App::DocumentObject* docObj = getDocument()->getObject(diag.constraint_id.c_str());
        if (!docObj || !docObj->getPropertyByName("Reference1")) {
            continue;
        }

        std::string objName = docObj->getNameInDocument();

        switch (diag.kind) {
            case KCSolve::ConstraintDiagnostic::Kind::Redundant:
                if (std::find(lastRedundantJoints.begin(), lastRedundantJoints.end(), objName)
                    == lastRedundantJoints.end()) {
                    lastRedundantJoints.push_back(objName);
                }
                break;
            case KCSolve::ConstraintDiagnostic::Kind::Conflicting:
                if (std::find(lastConflictingJoints.begin(), lastConflictingJoints.end(), objName)
                    == lastConflictingJoints.end()) {
                    lastConflictingJoints.push_back(objName);
                }
                break;
            case KCSolve::ConstraintDiagnostic::Kind::PartiallyRedundant:
                if (std::find(lastPartialRedundantJoints.begin(), lastPartialRedundantJoints.end(), objName)
                    == lastPartialRedundantJoints.end()) {
                    lastPartialRedundantJoints.push_back(objName);
                }
                break;
            case KCSolve::ConstraintDiagnostic::Kind::Malformed:
                if (std::find(lastMalformedJoints.begin(), lastMalformedJoints.end(), objName)
                    == lastMalformedJoints.end()) {
                    lastMalformedJoints.push_back(objName);
                }
                break;
        }
    }

    lastHasRedundancies = !lastRedundantJoints.empty();
    lastHasConflict = !lastConflictingJoints.empty();
    lastHasPartialRedundancies = !lastPartialRedundantJoints.empty();
    lastHasMalformedConstraints = !lastMalformedJoints.empty();

    signalSolverUpdate();
}

int AssemblyObject::generateSimulation(App::DocumentObject* sim)
{
    auto* solver = getOrCreateSolver();
    if (!solver) {
        return -1;
    }

    partIdToObjs_.clear();
    objToPartId_.clear();

    auto groundedObjs = getGroundedParts();
    if (groundedObjs.empty()) {
        return -6;
    }

    std::vector<App::DocumentObject*> joints = getJoints();
    removeUnconnectedJoints(joints, groundedObjs);

    KCSolve::SolveContext ctx = buildSolveContext(joints, true, sim);

    try {
        lastResult_ = solver->run_kinematic(ctx);
    }
    catch (...) {
        Base::Console().error("Generation of simulation failed\n");
        return -1;
    }

    if (lastResult_.status == KCSolve::SolveStatus::Failed) {
        return -1;
    }

    return 0;
}

std::vector<App::DocumentObject*> AssemblyObject::getMotionsFromSimulation(App::DocumentObject* sim)
{
    if (!sim) {
        return {};
    }

    auto* prop = dynamic_cast<App::PropertyLinkList*>(sim->getPropertyByName("Group"));
    if (!prop) {
        return {};
    }

    return prop->getValue();
}

int Assembly::AssemblyObject::updateForFrame(size_t index, bool updateJCS)
{
    if (!solver_) {
        return -1;
    }

    auto nfrms = solver_->num_frames();
    if (index >= nfrms) {
        return -1;
    }

    lastResult_ = solver_->update_for_frame(index);
    setNewPlacements();
    auto jointDocs = getJoints(updateJCS);
    redrawJointPlacements(jointDocs);
    return 0;
}

size_t Assembly::AssemblyObject::numberOfFrames()
{
    return solver_ ? solver_->num_frames() : 0;
}

void AssemblyObject::preDrag(std::vector<App::DocumentObject*> dragParts)
{
    bundleFixed = true;

    auto* solver = getOrCreateSolver();
    if (!solver) {
        bundleFixed = false;
        return;
    }

    partIdToObjs_.clear();
    objToPartId_.clear();

    auto groundedObjs = getGroundedParts();
    if (groundedObjs.empty()) {
        bundleFixed = false;
        return;
    }

    std::vector<App::DocumentObject*> joints = getJoints();
    removeUnconnectedJoints(joints, groundedObjs);

    KCSolve::SolveContext ctx = buildSolveContext(joints);
    bundleFixed = false;

    draggedParts.clear();
    for (auto part : dragParts) {
        // make sure no duplicate
        if (std::ranges::find(draggedParts, part) != draggedParts.end()) {
            continue;
        }

        // Free-floating parts should not be added since they are ignored by the solver!
        if (!isPartConnected(part)) {
            continue;
        }

        // Some objects have been bundled, we don't want to add these to dragged parts
        auto it = objToPartId_.find(part);
        if (it == objToPartId_.end()) {
            continue;
        }

        // Check if this is a bundled (non-primary) object
        const auto& mappings = partIdToObjs_[it->second];
        bool isBundled = false;
        for (const auto& m : mappings) {
            if (m.obj == part && !m.offset.isIdentity()) {
                isBundled = true;
                break;
            }
        }
        if (isBundled) {
            continue;
        }

        draggedParts.push_back(part);
    }

    // Build drag part IDs for the solver
    std::vector<std::string> dragPartIds;
    for (auto* part : draggedParts) {
        auto idIt = objToPartId_.find(part);
        if (idIt != objToPartId_.end()) {
            dragPartIds.push_back(idIt->second);
        }
    }

    savePlacementsForUndo();

    try {
        lastResult_ = solver->pre_drag(ctx, dragPartIds);
        setNewPlacements();
    }
    catch (...) {
        // If pre_drag fails, we still need to be in a valid state
    }
}

void AssemblyObject::doDragStep()
{
    try {
        std::vector<KCSolve::SolveResult::PartResult> dragPlacements;

        for (auto& part : draggedParts) {
            if (!part) {
                continue;
            }

            auto idIt = objToPartId_.find(part);
            if (idIt == objToPartId_.end()) {
                continue;
            }

            Base::Placement plc = getPlacementFromProp(part, "Placement");
            dragPlacements.push_back({idIt->second, placementToTransform(plc)});
        }

        lastResult_ = solver_->drag_step(dragPlacements);

        if (validateNewPlacements()) {
            setNewPlacements();

            auto joints = getJoints(false);
            for (auto* joint : joints) {
                if (joint->Visibility.getValue()) {
                    // redraw only the moving joint as its quite slow as its python code.
                    redrawJointPlacement(joint);
                }
            }
        }
    }
    catch (...) {
        // We do nothing if a solve step fails.
    }
}

bool AssemblyObject::validateNewPlacements()
{
    // First we check if a grounded object has moved. It can happen that they flip.
    auto groundedParts = getGroundedParts();
    for (auto* obj : groundedParts) {
        auto* propPlacement = dynamic_cast<App::PropertyPlacement*>(
            obj->getPropertyByName("Placement")
        );
        if (propPlacement) {
            Base::Placement oldPlc = propPlacement->getValue();

            auto idIt = objToPartId_.find(obj);
            if (idIt == objToPartId_.end()) {
                continue;
            }

            // Find the new placement from lastResult_
            for (const auto& pr : lastResult_.placements) {
                if (pr.id != idIt->second) {
                    continue;
                }

                Base::Placement newPlacement = transformToPlacement(pr.placement);

                // Apply bundle offset if present
                const auto& mappings = partIdToObjs_[pr.id];
                for (const auto& m : mappings) {
                    if (m.obj == obj && !m.offset.isIdentity()) {
                        newPlacement = newPlacement * m.offset;
                        break;
                    }
                }

                if (!oldPlc.isSame(newPlacement, Precision::Confusion())) {
                    Base::Console().warning(
                        "Assembly : Ignoring bad solve, a grounded object (%s) moved.\n",
                        obj->getFullLabel()
                    );
                    return false;
                }
                break;
            }
        }
    }

    // Check if any part has flipped orientation (rotation > 90 degrees from original)
    for (const auto& savedPair : previousPositions) {
        App::DocumentObject* obj = savedPair.first;
        if (!obj) {
            continue;
        }

        auto idIt = objToPartId_.find(obj);
        if (idIt == objToPartId_.end()) {
            continue;
        }

        // Find the new placement from lastResult_
        for (const auto& pr : lastResult_.placements) {
            if (pr.id != idIt->second) {
                continue;
            }

            Base::Placement newPlacement = transformToPlacement(pr.placement);

            // Apply bundle offset if present
            const auto& mappings = partIdToObjs_[pr.id];
            for (const auto& m : mappings) {
                if (m.obj == obj && !m.offset.isIdentity()) {
                    newPlacement = newPlacement * m.offset;
                    break;
                }
            }

            const Base::Placement& oldPlc = savedPair.second;

            // Calculate the rotation difference between old and new orientations
            Base::Rotation oldRot = oldPlc.getRotation();
            Base::Rotation newRot = newPlacement.getRotation();

            // Get the relative rotation: how much did the part rotate?
            Base::Rotation relativeRot = newRot * oldRot.inverse();

            // Get the angle of this rotation
            Base::Vector3d axis;
            double angle;
            relativeRot.getRawValue(axis, angle);

            // If the part rotated more than 90 degrees, consider it a flip
            // Use 91 degrees to allow for small numerical errors
            constexpr double maxAngle = 91.0 * M_PI / 180.0;
            if (std::abs(angle) > maxAngle) {
                Base::Console().warning(
                    "Assembly : Ignoring bad solve, part (%s) flipped orientation (%.1f degrees).\n",
                    obj->getFullLabel(),
                    std::abs(angle) * 180.0 / M_PI
                );
                return false;
            }
            break;
        }
    }

    return true;
}

void AssemblyObject::postDrag()
{
    if (solver_) {
        solver_->post_drag();
    }
    purgeTouched();
}

void AssemblyObject::savePlacementsForUndo()
{
    previousPositions.clear();

    for (const auto& [partId, mappings] : partIdToObjs_) {
        for (const auto& mapping : mappings) {
            App::DocumentObject* obj = mapping.obj;
            if (!obj) {
                continue;
            }

            auto* propPlc = dynamic_cast<App::PropertyPlacement*>(obj->getPropertyByName("Placement"));
            if (!propPlc) {
                continue;
            }

            previousPositions.push_back({obj, propPlc->getValue()});
        }
    }
}

void AssemblyObject::undoSolve()
{
    if (previousPositions.size() == 0) {
        return;
    }

    for (auto& pair : previousPositions) {
        App::DocumentObject* obj = pair.first;
        if (!obj) {
            continue;
        }

        // Check if the object has a "Placement" property
        auto* propPlacement = dynamic_cast<App::PropertyPlacement*>(
            obj->getPropertyByName("Placement")
        );
        if (!propPlacement) {
            continue;
        }

        propPlacement->setValue(pair.second);
    }
    previousPositions.clear();

    // update joint placements:
    getJoints(/*updateJCS*/ true, /*delBadJoints*/ false);
}

void AssemblyObject::clearUndo()
{
    previousPositions.clear();
}

void AssemblyObject::exportAsASMT(std::string fileName)
{
    auto* solver = getOrCreateSolver();
    if (!solver) {
        return;
    }

    partIdToObjs_.clear();
    objToPartId_.clear();

    auto groundedObjs = getGroundedParts();
    std::vector<App::DocumentObject*> joints = getJoints();
    removeUnconnectedJoints(joints, groundedObjs);

    KCSolve::SolveContext ctx = buildSolveContext(joints);

    try {
        solver->solve(ctx);
    }
    catch (...) {
        // Build anyway for export
    }

    solver->export_native(fileName);
}

void AssemblyObject::setNewPlacements()
{
    for (const auto& pr : lastResult_.placements) {
        auto it = partIdToObjs_.find(pr.id);
        if (it == partIdToObjs_.end()) {
            continue;
        }

        Base::Placement basePlc = transformToPlacement(pr.placement);

        for (const auto& mapping : it->second) {
            App::DocumentObject* obj = mapping.obj;
            if (!obj) {
                continue;
            }

            auto* propPlacement = dynamic_cast<App::PropertyPlacement*>(
                obj->getPropertyByName("Placement")
            );
            if (!propPlacement) {
                continue;
            }

            Base::Placement newPlacement = basePlc;
            if (!mapping.offset.isIdentity()) {
                newPlacement = basePlc * mapping.offset;
            }
            if (!propPlacement->getValue().isSame(newPlacement)) {
                propPlacement->setValue(newPlacement);
                obj->purgeTouched();
            }
        }
    }
}

void AssemblyObject::redrawJointPlacements(std::vector<App::DocumentObject*> joints)
{
    // Notify the joint objects that the transform of the coin object changed.
    for (auto* joint : joints) {
        if (!joint) {
            continue;
        }
        redrawJointPlacement(joint);
    }
}

void AssemblyObject::redrawJointPlacement(App::DocumentObject* joint)
{
    if (!joint) {
        return;
    }

    Base::PyGILStateLocker lock;

    App::PropertyPythonObject* proxy = joint
        ? dynamic_cast<App::PropertyPythonObject*>(joint->getPropertyByName("Proxy"))
        : nullptr;

    if (!proxy) {
        return;
    }

    Py::Object jointPy = proxy->getValue();

    if (!jointPy.hasAttr("redrawJointPlacements")) {
        return;
    }

    Py::Object attr = jointPy.getAttr("redrawJointPlacements");
    if (attr.ptr() && attr.isCallable()) {
        Py::Tuple args(1);
        args.setItem(0, Py::asObject(joint->getPyObject()));
        Py::Callable(attr).apply(args);
    }
}


// ── SolveContext building ──────────────────────────────────────────

std::string AssemblyObject::registerPart(App::DocumentObject* obj)
{
    // Check if already registered
    auto it = objToPartId_.find(obj);
    if (it != objToPartId_.end()) {
        return it->second;
    }

    std::string partId = obj->getFullName();
    Base::Placement plc = getPlacementFromProp(obj, "Placement");

    objToPartId_[obj] = partId;
    partIdToObjs_[partId].push_back({obj, Base::Placement()});

    // When bundling fixed joints, recursively discover connected parts
    if (bundleFixed) {
        auto addConnectedFixedParts = [&](App::DocumentObject* currentPart, auto& self) -> void {
            std::vector<App::DocumentObject*> joints = getJointsOfPart(currentPart);
            for (auto* joint : joints) {
                JointType jointType = getJointType(joint);
                if (jointType == JointType::Fixed) {
                    App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
                    App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
                    App::DocumentObject* partToAdd = currentPart == part1 ? part2 : part1;

                    if (objToPartId_.find(partToAdd) != objToPartId_.end()) {
                        continue;  // already registered
                    }

                    Base::Placement plci = getPlacementFromProp(partToAdd, "Placement");
                    Base::Placement offset = plc.inverse() * plci;

                    objToPartId_[partToAdd] = partId;
                    partIdToObjs_[partId].push_back({partToAdd, offset});

                    self(partToAdd, self);
                }
            }
        };

        addConnectedFixedParts(obj, addConnectedFixedParts);
    }

    return partId;
}

KCSolve::SolveContext AssemblyObject::buildSolveContext(
    const std::vector<App::DocumentObject*>& joints,
    bool forSimulation,
    App::DocumentObject* sim
)
{
    KCSolve::SolveContext ctx;
    ctx.bundle_fixed = bundleFixed;

    // ── Parts: register grounded parts ─────────────────────────────

    auto groundedObjs = getGroundedParts();
    for (auto* obj : groundedObjs) {
        if (!obj) {
            continue;
        }

        std::string partId = registerPart(obj);
        Base::Placement plc = getPlacementFromProp(obj, "Placement");

        KCSolve::Part part;
        part.id = partId;
        part.placement = placementToTransform(plc);
        part.mass = getObjMass(obj);
        part.grounded = true;
        ctx.parts.push_back(std::move(part));
    }

    // ── Constraints: process each joint ────────────────────────────

    // Collect motions for simulation
    std::vector<App::DocumentObject*> motionObjs;
    if (forSimulation && sim) {
        motionObjs = getMotionsFromSimulation(sim);
    }

    for (auto* joint : joints) {
        if (!joint) {
            continue;
        }

        JointType jointType = getJointType(joint);

        // When bundling fixed joints, skip Fixed type (parts are already bundled)
        if (bundleFixed && jointType == JointType::Fixed) {
            continue;
        }

        // Determine BaseJointKind and params
        KCSolve::BaseJointKind kind;
        std::vector<double> params;

        switch (jointType) {
            case JointType::Fixed:
                kind = KCSolve::BaseJointKind::Fixed;
                break;
            case JointType::Revolute:
                kind = KCSolve::BaseJointKind::Revolute;
                break;
            case JointType::Cylindrical:
                kind = KCSolve::BaseJointKind::Cylindrical;
                break;
            case JointType::Slider:
                kind = KCSolve::BaseJointKind::Slider;
                break;
            case JointType::Ball:
                kind = KCSolve::BaseJointKind::Ball;
                break;
            case JointType::Parallel:
                kind = KCSolve::BaseJointKind::Parallel;
                break;
            case JointType::Perpendicular:
                kind = KCSolve::BaseJointKind::Perpendicular;
                break;
            case JointType::Angle: {
                double angle = fabs(Base::toRadians(getJointAngle(joint)));
                if (fmod(angle, 2 * std::numbers::pi) < Precision::Confusion()) {
                    kind = KCSolve::BaseJointKind::Parallel;
                }
                else {
                    kind = KCSolve::BaseJointKind::Angle;
                    params.push_back(angle);
                }
                break;
            }
            case JointType::RackPinion: {
                kind = KCSolve::BaseJointKind::RackPinion;
                params.push_back(getJointDistance(joint));
                break;
            }
            case JointType::Screw: {
                int slidingIndex = slidingPartIndex(joint);
                if (slidingIndex == 0) {
                    continue;  // invalid — needs a slider
                }
                if (slidingIndex != 1) {
                    swapJCS(joint);
                }
                kind = KCSolve::BaseJointKind::Screw;
                params.push_back(getJointDistance(joint));
                break;
            }
            case JointType::Gears: {
                kind = KCSolve::BaseJointKind::Gear;
                params.push_back(getJointDistance(joint));
                params.push_back(getJointDistance2(joint));
                break;
            }
            case JointType::Belt: {
                kind = KCSolve::BaseJointKind::Gear;
                params.push_back(getJointDistance(joint));
                params.push_back(-getJointDistance2(joint));
                break;
            }
            case JointType::Distance: {
                // Decompose based on geometry classification
                DistanceType distType = getDistanceType(joint);
                std::string elt1 = getElementFromProp(joint, "Reference1");
                std::string elt2 = getElementFromProp(joint, "Reference2");
                auto* obj1 = getLinkedObjFromRef(joint, "Reference1");
                auto* obj2 = getLinkedObjFromRef(joint, "Reference2");
                double distance = getJointDistance(joint);

                switch (distType) {
                    case DistanceType::PointPoint:
                        if (distance < Precision::Confusion()) {
                            kind = KCSolve::BaseJointKind::Coincident;
                        }
                        else {
                            kind = KCSolve::BaseJointKind::DistancePointPoint;
                            params.push_back(distance);
                        }
                        break;

                    case DistanceType::LineLine:
                        kind = KCSolve::BaseJointKind::Concentric;
                        params.push_back(distance);
                        break;
                    case DistanceType::LineCircle:
                        kind = KCSolve::BaseJointKind::Concentric;
                        params.push_back(distance + getEdgeRadius(obj2, elt2));
                        break;
                    case DistanceType::CircleCircle:
                        kind = KCSolve::BaseJointKind::Concentric;
                        params.push_back(distance + getEdgeRadius(obj1, elt1) + getEdgeRadius(obj2, elt2));
                        break;

                    case DistanceType::PlanePlane:
                        kind = KCSolve::BaseJointKind::Planar;
                        params.push_back(distance);
                        break;
                    case DistanceType::PlaneCylinder:
                        kind = KCSolve::BaseJointKind::LineInPlane;
                        params.push_back(distance + getFaceRadius(obj2, elt2));
                        break;
                    case DistanceType::PlaneSphere:
                        kind = KCSolve::BaseJointKind::PointInPlane;
                        params.push_back(distance + getFaceRadius(obj2, elt2));
                        break;
                    case DistanceType::PlaneTorus:
                        kind = KCSolve::BaseJointKind::Planar;
                        params.push_back(distance);
                        break;

                    case DistanceType::CylinderCylinder:
                        kind = KCSolve::BaseJointKind::Concentric;
                        params.push_back(distance + getFaceRadius(obj1, elt1) + getFaceRadius(obj2, elt2));
                        break;
                    case DistanceType::CylinderSphere:
                        kind = KCSolve::BaseJointKind::DistanceCylSph;
                        params.push_back(distance + getFaceRadius(obj1, elt1) + getFaceRadius(obj2, elt2));
                        break;
                    case DistanceType::CylinderTorus:
                        kind = KCSolve::BaseJointKind::Concentric;
                        params.push_back(distance + getFaceRadius(obj1, elt1) + getFaceRadius(obj2, elt2));
                        break;

                    case DistanceType::TorusTorus:
                        kind = KCSolve::BaseJointKind::Planar;
                        params.push_back(distance);
                        break;
                    case DistanceType::TorusSphere:
                        kind = KCSolve::BaseJointKind::DistanceCylSph;
                        params.push_back(distance + getFaceRadius(obj1, elt1) + getFaceRadius(obj2, elt2));
                        break;
                    case DistanceType::SphereSphere:
                        kind = KCSolve::BaseJointKind::DistancePointPoint;
                        params.push_back(distance + getFaceRadius(obj1, elt1) + getFaceRadius(obj2, elt2));
                        break;

                    case DistanceType::PointPlane:
                        kind = KCSolve::BaseJointKind::PointInPlane;
                        params.push_back(distance);
                        break;
                    case DistanceType::PointCylinder:
                        kind = KCSolve::BaseJointKind::DistanceCylSph;
                        params.push_back(distance + getFaceRadius(obj1, elt1));
                        break;
                    case DistanceType::PointSphere:
                        kind = KCSolve::BaseJointKind::DistancePointPoint;
                        params.push_back(distance + getFaceRadius(obj1, elt1));
                        break;

                    case DistanceType::LinePlane:
                        kind = KCSolve::BaseJointKind::LineInPlane;
                        params.push_back(distance);
                        break;

                    case DistanceType::PointLine:
                        kind = KCSolve::BaseJointKind::DistanceCylSph;
                        params.push_back(distance);
                        break;
                    case DistanceType::PointCurve:
                        kind = KCSolve::BaseJointKind::PointInPlane;
                        params.push_back(distance);
                        break;

                    default:
                        kind = KCSolve::BaseJointKind::Planar;
                        params.push_back(distance);
                        break;
                }
                break;
            }
            default:
                continue;
        }

        // Validate the joint (skip self-referential bundled parts)
        if (!isJointValid(joint)) {
            continue;
        }

        // Compute marker transforms
        std::string partIdI, partIdJ;
        KCSolve::Transform markerI, markerJ;

        if (jointType == JointType::RackPinion) {
            auto rp = computeRackPinionMarkers(joint);
            partIdI = rp.partIdI;
            markerI = rp.markerI;
            partIdJ = rp.partIdJ;
            markerJ = rp.markerJ;

            if (partIdI.empty() || partIdJ.empty()) {
                continue;
            }
        }
        else {
            // Resolve part IDs from joint references
            App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
            App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
            if (!part1 || !part2) {
                continue;
            }

            // Ensure both parts are registered
            partIdI = registerPart(part1);
            partIdJ = registerPart(part2);

            markerI = computeMarkerTransform(joint, "Reference1", "Placement1");
            markerJ = computeMarkerTransform(joint, "Reference2", "Placement2");
        }

        // Build the constraint
        KCSolve::Constraint c;
        c.id = joint->getNameInDocument();
        c.part_i = partIdI;
        c.marker_i = markerI;
        c.part_j = partIdJ;
        c.marker_j = markerJ;
        c.type = kind;
        c.params = std::move(params);

        // Add limits (only if not in simulation mode — motions may clash)
        if (motionObjs.empty()) {
            if (jointType == JointType::Slider || jointType == JointType::Cylindrical) {
                auto* pLenMin = dynamic_cast<App::PropertyFloat*>(joint->getPropertyByName("LengthMin"));
                auto* pLenMax = dynamic_cast<App::PropertyFloat*>(joint->getPropertyByName("LengthMax"));
                auto* pMinEnabled = dynamic_cast<App::PropertyBool*>(
                    joint->getPropertyByName("EnableLengthMin")
                );
                auto* pMaxEnabled = dynamic_cast<App::PropertyBool*>(
                    joint->getPropertyByName("EnableLengthMax")
                );

                if (pLenMin && pLenMax && pMinEnabled && pMaxEnabled) {
                    bool minEnabled = pMinEnabled->getValue();
                    bool maxEnabled = pMaxEnabled->getValue();
                    double minLength = pLenMin->getValue();
                    double maxLength = pLenMax->getValue();

                    if ((minLength > maxLength) && minEnabled && maxEnabled) {
                        pLenMin->setValue(maxLength);
                        pLenMax->setValue(minLength);
                        minLength = maxLength;
                        maxLength = pLenMax->getValue();

                        pMinEnabled->setValue(maxEnabled);
                        pMaxEnabled->setValue(minEnabled);
                        minEnabled = maxEnabled;
                        maxEnabled = pMaxEnabled->getValue();
                    }

                    if (minEnabled) {
                        c.limits.push_back({
                            KCSolve::Constraint::Limit::Kind::TranslationMin,
                            minLength,
                            1.0e-9
                        });
                    }
                    if (maxEnabled) {
                        c.limits.push_back({
                            KCSolve::Constraint::Limit::Kind::TranslationMax,
                            maxLength,
                            1.0e-9
                        });
                    }
                }
            }
            if (jointType == JointType::Revolute || jointType == JointType::Cylindrical) {
                auto* pRotMin = dynamic_cast<App::PropertyFloat*>(joint->getPropertyByName("AngleMin"));
                auto* pRotMax = dynamic_cast<App::PropertyFloat*>(joint->getPropertyByName("AngleMax"));
                auto* pMinEnabled = dynamic_cast<App::PropertyBool*>(
                    joint->getPropertyByName("EnableAngleMin")
                );
                auto* pMaxEnabled = dynamic_cast<App::PropertyBool*>(
                    joint->getPropertyByName("EnableAngleMax")
                );

                if (pRotMin && pRotMax && pMinEnabled && pMaxEnabled) {
                    bool minEnabled = pMinEnabled->getValue();
                    bool maxEnabled = pMaxEnabled->getValue();
                    double minAngle = pRotMin->getValue();
                    double maxAngle = pRotMax->getValue();

                    if ((minAngle > maxAngle) && minEnabled && maxEnabled) {
                        pRotMin->setValue(maxAngle);
                        pRotMax->setValue(minAngle);
                        minAngle = maxAngle;
                        maxAngle = pRotMax->getValue();

                        pMinEnabled->setValue(maxEnabled);
                        pMaxEnabled->setValue(minEnabled);
                        minEnabled = maxEnabled;
                        maxEnabled = pMaxEnabled->getValue();
                    }

                    if (minEnabled) {
                        c.limits.push_back({
                            KCSolve::Constraint::Limit::Kind::RotationMin,
                            minAngle,
                            1.0e-9
                        });
                    }
                    if (maxEnabled) {
                        c.limits.push_back({
                            KCSolve::Constraint::Limit::Kind::RotationMax,
                            maxAngle,
                            1.0e-9
                        });
                    }
                }
            }
        }

        ctx.constraints.push_back(std::move(c));

        // Add motions for simulation
        if (forSimulation) {
            std::vector<App::DocumentObject*> done;
            for (auto* motion : motionObjs) {
                if (std::ranges::find(done, motion) != done.end()) {
                    continue;
                }

                auto* pJoint = dynamic_cast<App::PropertyXLinkSub*>(motion->getPropertyByName("Joint"));
                if (!pJoint) {
                    continue;
                }
                App::DocumentObject* motionJoint = pJoint->getValue();
                if (joint != motionJoint) {
                    continue;
                }

                auto* pType = dynamic_cast<App::PropertyEnumeration*>(motion->getPropertyByName("MotionType"));
                auto* pFormula = dynamic_cast<App::PropertyString*>(motion->getPropertyByName("Formula"));
                if (!pType || !pFormula) {
                    continue;
                }
                std::string formula = pFormula->getValue();
                if (formula.empty()) {
                    continue;
                }
                std::string motionType = pType->getValueAsString();

                // Check for paired motion (cylindrical joints can have both angular + linear)
                for (auto* motion2 : motionObjs) {
                    auto* pJoint2 = dynamic_cast<App::PropertyXLinkSub*>(motion2->getPropertyByName("Joint"));
                    if (!pJoint2) {
                        continue;
                    }
                    App::DocumentObject* motionJoint2 = pJoint2->getValue();
                    if (joint != motionJoint2 || motion2 == motion) {
                        continue;
                    }

                    auto* pType2 = dynamic_cast<App::PropertyEnumeration*>(
                        motion2->getPropertyByName("MotionType")
                    );
                    auto* pFormula2 = dynamic_cast<App::PropertyString*>(motion2->getPropertyByName("Formula"));
                    if (!pType2 || !pFormula2) {
                        continue;
                    }
                    std::string formula2 = pFormula2->getValue();
                    if (formula2.empty()) {
                        continue;
                    }
                    std::string motionType2 = pType2->getValueAsString();
                    if (motionType2 == motionType) {
                        continue;
                    }

                    // Two different motion types on same joint → General motion
                    KCSolve::MotionDef md;
                    md.kind = KCSolve::MotionDef::Kind::General;
                    md.joint_id = joint->getNameInDocument();
                    md.marker_i = "";  // Adapter resolves from joint_id
                    md.marker_j = "";
                    md.rotation_expr = motionType == "Angular" ? formula : formula2;
                    md.translation_expr = motionType == "Angular" ? formula2 : formula;
                    ctx.motions.push_back(std::move(md));

                    done.push_back(motion2);
                }

                // Single motion
                KCSolve::MotionDef md;
                md.joint_id = joint->getNameInDocument();
                md.marker_i = "";
                md.marker_j = "";
                if (motionType == "Angular") {
                    md.kind = KCSolve::MotionDef::Kind::Rotational;
                    md.rotation_expr = formula;
                }
                else {
                    md.kind = KCSolve::MotionDef::Kind::Translational;
                    md.translation_expr = formula;
                }
                ctx.motions.push_back(std::move(md));
            }
        }
    }

    // ── Parts: ensure all referenced parts are in the context ──────

    // Some parts may have been registered during constraint processing
    // but not yet added to ctx.parts
    for (const auto& [partId, mappings] : partIdToObjs_) {
        bool alreadyInCtx = false;
        for (const auto& p : ctx.parts) {
            if (p.id == partId) {
                alreadyInCtx = true;
                break;
            }
        }
        if (!alreadyInCtx) {
            App::DocumentObject* primaryObj = mappings[0].obj;
            Base::Placement plc = getPlacementFromProp(primaryObj, "Placement");

            KCSolve::Part part;
            part.id = partId;
            part.placement = placementToTransform(plc);
            part.mass = getObjMass(primaryObj);
            part.grounded = false;
            ctx.parts.push_back(std::move(part));
        }
    }

    // ── Simulation parameters ──────────────────────────────────────

    if (forSimulation && sim) {
        auto valueOf = [](App::DocumentObject* docObj, const char* propName) {
            auto* prop = dynamic_cast<App::PropertyFloat*>(docObj->getPropertyByName(propName));
            if (!prop) {
                return 0.0;
            }
            return prop->getValue();
        };

        KCSolve::SimulationParams sp;
        sp.t_start = valueOf(sim, "aTimeStart");
        sp.t_end = valueOf(sim, "bTimeEnd");
        sp.h_out = valueOf(sim, "cTimeStepOutput");
        sp.h_min = 1.0e-9;
        sp.h_max = 1.0;
        sp.error_tol = valueOf(sim, "fGlobalErrorTolerance");
        ctx.simulation = sp;
    }

    return ctx;
}


// ── Marker transform computation ───────────────────────────────────

KCSolve::Transform AssemblyObject::computeMarkerTransform(
    App::DocumentObject* joint,
    const char* propRefName,
    const char* propPlcName
)
{
    App::DocumentObject* part = getMovingPartFromRef(joint, propRefName);
    App::DocumentObject* obj = getObjFromJointRef(joint, propRefName);

    if (!part || !obj) {
        Base::Console()
            .warning("The property %s of Joint %s is bad.\n", propRefName, joint->getFullName());
        return KCSolve::Transform::identity();
    }

    Base::Placement plc = getPlacementFromProp(joint, propPlcName);
    // Now we have plc which is the JCS placement, but its relative to the Object, not to the
    // containing Part.

    if (obj->getNameInDocument() != part->getNameInDocument()) {
        auto* ref = dynamic_cast<App::PropertyXLinkSub*>(joint->getPropertyByName(propRefName));
        if (!ref) {
            return KCSolve::Transform::identity();
        }

        Base::Placement obj_global_plc = getGlobalPlacement(obj, ref);
        plc = obj_global_plc * plc;

        Base::Placement part_global_plc = getGlobalPlacement(part, ref);
        plc = part_global_plc.inverse() * plc;
    }

    // Apply bundle offset if present
    auto idIt = objToPartId_.find(part);
    if (idIt != objToPartId_.end()) {
        const auto& mappings = partIdToObjs_[idIt->second];
        for (const auto& m : mappings) {
            if (m.obj == part && !m.offset.isIdentity()) {
                plc = m.offset * plc;
                break;
            }
        }
    }

    return placementToTransform(plc);
}

AssemblyObject::RackPinionResult AssemblyObject::computeRackPinionMarkers(App::DocumentObject* joint)
{
    RackPinionResult result;

    // ASMT rack pinion joint must get the rack as I and pinion as J.
    int slidingIndex = slidingPartIndex(joint);
    if (slidingIndex == 0) {
        return result;
    }

    if (slidingIndex != 1) {
        swapJCS(joint);  // make sure that rack is first.
    }

    App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
    App::DocumentObject* obj1 = getObjFromJointRef(joint, "Reference1");
    Base::Placement plc1 = getPlacementFromProp(joint, "Placement1");

    App::DocumentObject* obj2 = getObjFromJointRef(joint, "Reference2");
    Base::Placement plc2 = getPlacementFromProp(joint, "Placement2");

    if (!part1 || !obj1) {
        Base::Console().warning("Reference1 of Joint %s is bad.\n", joint->getFullName());
        return result;
    }

    // Ensure parts are registered
    result.partIdI = registerPart(part1);

    // For the pinion — use standard marker computation
    result.markerJ = computeMarkerTransform(joint, "Reference2", "Placement2");

    App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
    if (part2) {
        result.partIdJ = registerPart(part2);
    }

    // For the rack — need to adjust placement so X axis aligns with slider axis
    auto* ref1 = dynamic_cast<App::PropertyXLinkSub*>(joint->getPropertyByName("Reference1"));
    auto* ref2 = dynamic_cast<App::PropertyXLinkSub*>(joint->getPropertyByName("Reference2"));
    if (!ref1 || !ref2) {
        return result;
    }

    // Make the pinion plc relative to the rack placement
    Base::Placement pinion_global_plc = getGlobalPlacement(obj2, ref2);
    plc2 = pinion_global_plc * plc2;
    Base::Placement rack_global_plc = getGlobalPlacement(obj1, ref1);
    plc2 = rack_global_plc.inverse() * plc2;

    // The rot of the rack placement should be the same as the pinion, but with X axis along the
    // slider axis.
    Base::Rotation rot = plc2.getRotation();
    Base::Vector3d currentZAxis = rot.multVec(Base::Vector3d(0, 0, 1));
    Base::Vector3d currentXAxis = rot.multVec(Base::Vector3d(1, 0, 0));
    Base::Vector3d targetXAxis = plc1.getRotation().multVec(Base::Vector3d(0, 0, 1));

    double yawAdjustment = currentXAxis.GetAngle(targetXAxis);

    Base::Vector3d crossProd = currentXAxis.Cross(targetXAxis);
    if (currentZAxis * crossProd < 0) {
        yawAdjustment = -yawAdjustment;
    }

    Base::Rotation yawRotation(currentZAxis, yawAdjustment);
    Base::Rotation adjustedRotation = rot * yawRotation;
    plc1.setRotation(adjustedRotation);

    // Transform to part-relative coordinates (same as handleOneSideOfJoint end logic)
    if (obj1->getNameInDocument() != part1->getNameInDocument()) {
        plc1 = rack_global_plc * plc1;
        Base::Placement part_global_plc = getGlobalPlacement(part1, ref1);
        plc1 = part_global_plc.inverse() * plc1;
    }

    // Apply bundle offset if present
    auto idIt = objToPartId_.find(part1);
    if (idIt != objToPartId_.end()) {
        const auto& mappings = partIdToObjs_[idIt->second];
        for (const auto& m : mappings) {
            if (m.obj == part1 && !m.offset.isIdentity()) {
                plc1 = m.offset * plc1;
                break;
            }
        }
    }

    result.markerI = placementToTransform(plc1);

    return result;
}

bool AssemblyObject::isJointValid(App::DocumentObject* joint)
{
    // When dragging, parts connected by fixed joints are bundled.
    // A joint that references two parts in the same bundle is self-referential and must be skipped.
    App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
    App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
    if (!part1 || !part2) {
        return false;
    }

    auto it1 = objToPartId_.find(part1);
    auto it2 = objToPartId_.find(part2);
    if (it1 != objToPartId_.end() && it2 != objToPartId_.end() && it1->second == it2->second) {
        Base::Console().warning(
            "Assembly: Ignoring joint (%s) because its parts are connected by a fixed "
            "joint bundle. This joint is a conflicting or redundant constraint.\n",
            joint->getFullLabel()
        );
        return false;
    }
    return true;
}


// ── Joint / Part graph helpers (unchanged) ─────────────────────────

App::DocumentObject* AssemblyObject::getJointOfPartConnectingToGround(
    App::DocumentObject* part,
    std::string& name,
    const std::vector<App::DocumentObject*>& excludeJoints
)
{
    if (!part) {
        return nullptr;
    }

    std::vector<App::DocumentObject*> joints = getJointsOfPart(part);

    for (auto joint : joints) {
        if (!joint) {
            continue;
        }

        if (std::ranges::find(excludeJoints, joint) != excludeJoints.end()) {
            continue;
        }

        App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
        App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
        if (!part1 || !part2) {
            continue;
        }

        if (part == part1 && isJointConnectingPartToGround(joint, "Reference1")) {
            name = "Reference1";
            return joint;
        }
        if (part == part2 && isJointConnectingPartToGround(joint, "Reference2")) {
            name = "Reference2";
            return joint;
        }
    }
    return nullptr;
}

template<typename T>
T* AssemblyObject::getGroup()
{
    App::Document* doc = getDocument();

    std::vector<DocumentObject*> groups = doc->getObjectsOfType(T::getClassTypeId());
    if (groups.empty()) {
        return nullptr;
    }
    for (auto group : groups) {
        if (hasObject(group)) {
            return freecad_cast<T*>(group);
        }
    }
    return nullptr;
}

JointGroup* AssemblyObject::getJointGroup() const
{
    return Assembly::getJointGroup(this);
}

ViewGroup* AssemblyObject::getExplodedViewGroup() const
{
    App::Document* doc = getDocument();

    std::vector<DocumentObject*> viewGroups = doc->getObjectsOfType(ViewGroup::getClassTypeId());
    if (viewGroups.empty()) {
        return nullptr;
    }
    for (auto viewGroup : viewGroups) {
        if (hasObject(viewGroup)) {
            return freecad_cast<ViewGroup*>(viewGroup);
        }
    }
    return nullptr;
}

std::vector<App::DocumentObject*> AssemblyObject::getJoints(bool updateJCS, bool delBadJoints, bool subJoints)
{
    std::vector<App::DocumentObject*> joints = {};

    JointGroup* jointGroup = getJointGroup();
    if (!jointGroup) {
        return {};
    }

    Base::PyGILStateLocker lock;
    for (auto joint : jointGroup->getObjects()) {
        if (!joint) {
            continue;
        }

        auto* prop = dynamic_cast<App::PropertyBool*>(joint->getPropertyByName("Suppressed"));
        if (joint->isError() || !prop || prop->getValue()) {
            // Filter grounded joints and deactivated joints.
            continue;
        }

        auto* part1 = getMovingPartFromRef(joint, "Reference1");
        auto* part2 = getMovingPartFromRef(joint, "Reference2");
        if (!part1 || !part2 || part1->getFullName() == part2->getFullName()) {
            // Remove incomplete joints. Left-over when the user deletes a part.
            // Remove incoherent joints (self-pointing joints)
            if (delBadJoints) {
                getDocument()->removeObject(joint->getNameInDocument());
            }
            continue;
        }

        auto proxy = dynamic_cast<App::PropertyPythonObject*>(joint->getPropertyByName("Proxy"));
        if (proxy) {
            if (proxy->getValue().hasAttr("setJointConnectors")) {
                joints.push_back(joint);
            }
        }
    }

    // add sub assemblies joints.
    if (subJoints) {
        for (auto& assembly : getSubAssemblies()) {
            auto subJoints = assembly->getJoints();
            joints.insert(joints.end(), subJoints.begin(), subJoints.end());
        }
    }

    return joints;
}

std::vector<App::DocumentObject*> AssemblyObject::getGroundedJoints()
{
    std::vector<App::DocumentObject*> joints = {};

    JointGroup* jointGroup = getJointGroup();
    if (!jointGroup) {
        return {};
    }

    Base::PyGILStateLocker lock;
    for (auto obj : jointGroup->getObjects()) {
        if (!obj) {
            continue;
        }

        auto* propObj = dynamic_cast<App::PropertyLink*>(obj->getPropertyByName("ObjectToGround"));

        if (propObj) {
            joints.push_back(obj);
        }
    }

    return joints;
}

std::vector<App::DocumentObject*> AssemblyObject::getJointsOfObj(App::DocumentObject* obj)
{
    if (!obj) {
        return {};
    }

    std::vector<App::DocumentObject*> joints = getJoints(false);
    std::vector<App::DocumentObject*> jointsOf;

    for (auto joint : joints) {
        App::DocumentObject* obj1 = getObjFromJointRef(joint, "Reference1");
        App::DocumentObject* obj2 = getObjFromJointRef(joint, "Reference2");
        if (obj == obj1 || obj == obj2) {
            jointsOf.push_back(joint);
        }
    }

    return jointsOf;
}

std::vector<App::DocumentObject*> AssemblyObject::getJointsOfPart(App::DocumentObject* part)
{
    if (!part) {
        return {};
    }

    std::vector<App::DocumentObject*> joints = getJoints(false);
    std::vector<App::DocumentObject*> jointsOf;

    for (auto joint : joints) {
        App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
        App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
        if (part == part1 || part == part2) {
            jointsOf.push_back(joint);
        }
    }
    return jointsOf;
}

std::unordered_set<App::DocumentObject*> AssemblyObject::getGroundedParts()
{
    std::vector<App::DocumentObject*> groundedJoints = getGroundedJoints();

    std::unordered_set<App::DocumentObject*> groundedSet;
    for (auto gJoint : groundedJoints) {
        if (!gJoint) {
            continue;
        }

        auto* propObj = dynamic_cast<App::PropertyLink*>(gJoint->getPropertyByName("ObjectToGround"));

        if (propObj) {
            App::DocumentObject* objToGround = propObj->getValue();
            if (objToGround) {
                if (auto* asmLink = dynamic_cast<AssemblyLink*>(objToGround)) {
                    if (!asmLink->isRigid()) {
                        continue;
                    }
                }
                groundedSet.insert(objToGround);
            }
        }
    }

    // We also need to add all the root-level datums objects that are not attached.
    std::vector<App::DocumentObject*> objs = Group.getValues();
    for (auto* obj : objs) {
        if (obj->isDerivedFrom<App::LocalCoordinateSystem>()
            || obj->isDerivedFrom<App::DatumElement>()) {
            auto* pcAttach = obj->getExtensionByType<PartApp::AttachExtension>();
            if (pcAttach) {
                // If it's a Part datums, we check if it's attached. If yes then we ignore it.
                std::string mode = pcAttach->MapMode.getValueAsString();
                if (mode != "Deactivated") {
                    continue;
                }
            }
            groundedSet.insert(obj);
        }
    }

    // Origin is not in Group so we add it separately
    groundedSet.insert(Origin.getValue());

    return groundedSet;
}

bool AssemblyObject::isJointConnectingPartToGround(App::DocumentObject* joint, const char* propname)
{
    if (!joint || !isJointTypeConnecting(joint)) {
        return false;
    }

    App::DocumentObject* part = getMovingPartFromRef(joint, propname);
    if (!part) {
        return false;
    }

    // Check if the part is grounded.
    bool isGrounded = isPartGrounded(part);
    if (isGrounded) {
        return false;
    }

    // Check if the part is disconnected even with the joint
    bool isConnected = isPartConnected(part);
    if (!isConnected) {
        return false;
    }

    // to know if a joint is connecting to ground we disable all the other joints
    std::vector<App::DocumentObject*> jointsOfPart = getJointsOfPart(part);
    std::vector<bool> activatedStates;

    for (auto jointi : jointsOfPart) {
        if (jointi->getFullName() == joint->getFullName()) {
            continue;
        }

        activatedStates.push_back(getJointActivated(jointi));
        setJointActivated(jointi, false);
    }

    isConnected = isPartConnected(part);

    // restore activation states
    for (auto jointi : jointsOfPart) {
        if (jointi->getFullName() == joint->getFullName() || activatedStates.empty()) {
            continue;
        }

        setJointActivated(jointi, activatedStates[0]);
        activatedStates.erase(activatedStates.begin());
    }

    return isConnected;
}

bool AssemblyObject::isJointTypeConnecting(App::DocumentObject* joint)
{
    if (!joint) {
        return false;
    }

    JointType jointType = getJointType(joint);
    return jointType != JointType::RackPinion && jointType != JointType::Screw
        && jointType != JointType::Gears && jointType != JointType::Belt;
}


bool AssemblyObject::isObjInSetOfObjRefs(App::DocumentObject* obj, const std::vector<ObjRef>& set)
{
    if (!obj) {
        return false;
    }

    for (const auto& pair : set) {
        if (pair.obj == obj) {
            return true;
        }
    }
    return false;
}

void AssemblyObject::removeUnconnectedJoints(
    std::vector<App::DocumentObject*>& joints,
    std::unordered_set<App::DocumentObject*> groundedObjs
)
{
    std::vector<ObjRef> connectedParts;

    // Initialize connectedParts with groundedObjs
    for (auto* groundedObj : groundedObjs) {
        connectedParts.push_back({groundedObj, nullptr});
    }

    // Perform a traversal from each grounded object
    for (auto* groundedObj : groundedObjs) {
        traverseAndMarkConnectedParts(groundedObj, connectedParts, joints);
    }

    // Filter out unconnected joints
    joints.erase(
        std::remove_if(
            joints.begin(),
            joints.end(),
            [&](App::DocumentObject* joint) {
                App::DocumentObject* obj1 = getMovingPartFromRef(joint, "Reference1");
                App::DocumentObject* obj2 = getMovingPartFromRef(joint, "Reference2");
                return (
                    !isObjInSetOfObjRefs(obj1, connectedParts)
                    || !isObjInSetOfObjRefs(obj2, connectedParts)
                );
            }
        ),
        joints.end()
    );
}

void AssemblyObject::traverseAndMarkConnectedParts(
    App::DocumentObject* currentObj,
    std::vector<ObjRef>& connectedParts,
    const std::vector<App::DocumentObject*>& joints
)
{
    // getConnectedParts returns the objs connected to the currentObj by any joint
    auto connectedObjs = getConnectedParts(currentObj, joints);
    for (auto& nextObjRef : connectedObjs) {
        if (!isObjInSetOfObjRefs(nextObjRef.obj, connectedParts)) {
            // Create a new ObjRef with the nextObj and a nullptr for PropertyXLinkSub*
            connectedParts.push_back(nextObjRef);
            traverseAndMarkConnectedParts(nextObjRef.obj, connectedParts, joints);
        }
    }
}

std::vector<ObjRef> AssemblyObject::getConnectedParts(
    App::DocumentObject* part,
    const std::vector<App::DocumentObject*>& joints
)
{
    if (!part) {
        return {};
    }

    std::vector<ObjRef> connectedParts;

    for (auto joint : joints) {
        if (!isJointTypeConnecting(joint)) {
            continue;
        }

        App::DocumentObject* obj1 = getMovingPartFromRef(joint, "Reference1");
        App::DocumentObject* obj2 = getMovingPartFromRef(joint, "Reference2");

        if (!obj1 || !obj2) {
            continue;
        }

        if (obj1 == part) {
            auto* ref = dynamic_cast<App::PropertyXLinkSub*>(joint->getPropertyByName("Reference2"));
            if (!ref) {
                continue;
            }
            connectedParts.push_back({obj2, ref});
        }
        else if (obj2 == part) {
            auto* ref = dynamic_cast<App::PropertyXLinkSub*>(joint->getPropertyByName("Reference1"));
            if (!ref) {
                continue;
            }
            connectedParts.push_back({obj1, ref});
        }
    }
    return connectedParts;
}

bool AssemblyObject::isPartGrounded(App::DocumentObject* obj)
{
    if (!obj) {
        return false;
    }

    auto groundedObjs = getGroundedParts();

    for (auto* groundedObj : groundedObjs) {
        if (groundedObj->getFullName() == obj->getFullName()) {
            return true;
        }
    }

    return false;
}

bool AssemblyObject::isPartConnected(App::DocumentObject* obj)
{
    if (!obj) {
        return false;
    }

    auto groundedObjs = getGroundedParts();
    std::vector<App::DocumentObject*> joints = getJoints(false);

    std::vector<ObjRef> connectedParts;

    // Initialize connectedParts with groundedObjs
    for (auto* groundedObj : groundedObjs) {
        connectedParts.push_back({groundedObj, nullptr});
    }

    // Perform a traversal from each grounded object
    for (auto* groundedObj : groundedObjs) {
        traverseAndMarkConnectedParts(groundedObj, connectedParts, joints);
    }

    for (auto& objRef : connectedParts) {
        if (obj == objRef.obj) {
            return true;
        }
    }

    return false;
}

int AssemblyObject::slidingPartIndex(App::DocumentObject* joint)
{
    App::DocumentObject* part1 = getMovingPartFromRef(joint, "Reference1");
    App::DocumentObject* obj1 = getObjFromJointRef(joint, "Reference1");
    boost::ignore_unused(obj1);
    Base::Placement plc1 = getPlacementFromProp(joint, "Placement1");

    App::DocumentObject* part2 = getMovingPartFromRef(joint, "Reference2");
    App::DocumentObject* obj2 = getObjFromJointRef(joint, "Reference2");
    boost::ignore_unused(obj2);
    Base::Placement plc2 = getPlacementFromProp(joint, "Placement2");

    int slidingFound = 0;
    for (auto* jt : getJoints(false, false)) {
        if (getJointType(jt) == JointType::Slider) {
            App::DocumentObject* jpart1 = getMovingPartFromRef(jt, "Reference1");
            App::DocumentObject* jpart2 = getMovingPartFromRef(jt, "Reference2");
            int found = 0;
            Base::Placement plcjt, plci;
            if (jpart1 == part1 || jpart1 == part2) {
                found = (jpart1 == part1) ? 1 : 2;
                plci = (jpart1 == part1) ? plc1 : plc2;
                plcjt = getPlacementFromProp(jt, "Placement1");
            }
            else if (jpart2 == part1 || jpart2 == part2) {
                found = (jpart2 == part1) ? 1 : 2;
                plci = (jpart2 == part1) ? plc1 : plc2;
                plcjt = getPlacementFromProp(jt, "Placement2");
            }

            if (found != 0) {
                double y1, p1, r1, y2, p2, r2;
                plcjt.getRotation().getYawPitchRoll(y1, p1, r1);
                plci.getRotation().getYawPitchRoll(y2, p2, r2);
                if (fabs(p1 - p2) < Precision::Confusion() && fabs(r1 - r2) < Precision::Confusion()) {
                    slidingFound = found;
                }
            }
        }
    }
    return slidingFound;
}

std::vector<ObjRef> AssemblyObject::getDownstreamParts(
    App::DocumentObject* part,
    App::DocumentObject* joint
)
{
    if (!part) {
        return {};
    }

    // First we deactivate the joint
    bool state = false;
    if (joint) {
        state = getJointActivated(joint);
        setJointActivated(joint, false);
    }

    std::vector<App::DocumentObject*> joints = getJoints(false);

    std::vector<ObjRef> connectedParts = {{part, nullptr}};
    traverseAndMarkConnectedParts(part, connectedParts, joints);

    std::vector<ObjRef> downstreamParts;
    for (auto& parti : connectedParts) {
        if (!isPartConnected(parti.obj) && (parti.obj != part)) {
            downstreamParts.push_back(parti);
        }
    }

    if (joint) {
        setJointActivated(joint, state);
    }

    return downstreamParts;
}

App::DocumentObject* AssemblyObject::getUpstreamMovingPart(
    App::DocumentObject* part,
    App::DocumentObject*& joint,
    std::string& name,
    std::vector<App::DocumentObject*> excludeJoints
)
{
    if (!part || isPartGrounded(part)) {
        return nullptr;
    }

    excludeJoints.push_back(joint);

    joint = getJointOfPartConnectingToGround(part, name, excludeJoints);
    JointType jointType = getJointType(joint);
    if (jointType != JointType::Fixed) {
        return part;
    }

    part = getMovingPartFromRef(joint, name == "Reference1" ? "Reference2" : "Reference1");

    return getUpstreamMovingPart(part, joint, name);
}

double AssemblyObject::getObjMass(App::DocumentObject* obj)
{
    if (!obj) {
        return 0.0;
    }

    for (auto& pair : objMasses) {
        if (pair.first == obj) {
            return pair.second;
        }
    }
    return 1.0;
}

void AssemblyObject::setObjMasses(std::vector<std::pair<App::DocumentObject*, double>> objectMasses)
{
    objMasses = objectMasses;
}

std::vector<AssemblyLink*> AssemblyObject::getSubAssemblies()
{
    std::vector<AssemblyLink*> subAssemblies = {};

    App::Document* doc = getDocument();

    std::vector<DocumentObject*> assemblies = doc->getObjectsOfType(
        Assembly::AssemblyLink::getClassTypeId()
    );
    for (auto assembly : assemblies) {
        if (hasObject(assembly)) {
            subAssemblies.push_back(freecad_cast<AssemblyLink*>(assembly));
        }
    }

    return subAssemblies;
}

void AssemblyObject::ensureIdentityPlacements()
{
    std::vector<App::DocumentObject*> group = Group.getValues();
    for (auto* obj : group) {
        // When used in assembly, link groups must have identity placements.
        if (obj->isLinkGroup()) {
            auto* link = dynamic_cast<App::Link*>(obj);
            auto* pPlc = dynamic_cast<App::PropertyPlacement*>(obj->getPropertyByName("Placement"));
            if (!pPlc || !link) {
                continue;
            }

            Base::Placement plc = pPlc->getValue();
            if (plc.isIdentity()) {
                continue;
            }

            pPlc->setValue(Base::Placement());
            obj->purgeTouched();

            // To keep the LinkElement positions, we apply plc to their placements
            std::vector<App::DocumentObject*> elts = link->ElementList.getValues();
            for (auto* elt : elts) {
                pPlc = dynamic_cast<App::PropertyPlacement*>(elt->getPropertyByName("Placement"));
                pPlc->setValue(plc * pPlc->getValue());
                elt->purgeTouched();
            }
        }
    }
}

int AssemblyObject::numberOfComponents() const
{
    return getAssemblyComponents(this).size();
}

bool AssemblyObject::isEmpty() const
{
    return numberOfComponents() == 0;
}