#include "rm_control/modules/trajectory/drag_trajectory_validator.h"

#include <cstddef>
#include <iostream>
#include <limits>
#include <string>
#include <vector>

namespace {

struct TestContext {
    int failures = 0;

    void Expect(bool condition, const std::string &message) {
        if (!condition) {
            ++failures;
            std::cerr << "[FAIL] " << message << '\n';
        }
    }
};

rm_control::common::TrajectorySample MakeSample(std::size_t sample_index,
                                                long long elapsed_ms,
                                                float joint_1_deg) {
    rm_control::common::TrajectorySample sample;
    sample.sample_index = sample_index;
    sample.elapsed_ms = elapsed_ms;
    sample.joints_deg = {joint_1_deg, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F, 0.0F};
    return sample;
}

rm_control::modules::trajectory::DragTrajectoryValidationLimits MakePermissiveLimits() {
    rm_control::modules::trajectory::DragTrajectoryValidationLimits limits;
    limits.maximum_joint_jump_deg = 100.0;
    limits.maximum_joint_velocity_deg_per_second = 1000.0;
    limits.maximum_joint_acceleration_deg_per_second_squared = 10000.0;
    return limits;
}

void TestReportsPeakJointMotionForValidTrajectory(TestContext *context) {
    const std::vector<rm_control::common::TrajectorySample> samples = {
        MakeSample(0, 0, 0.0F),
        MakeSample(1, 1000, 1.0F),
        MakeSample(2, 2000, 3.0F),
    };

    const auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        samples, MakePermissiveLimits());
    context->Expect(result.ok, "valid seven-joint trajectory should pass");
    if (!result.ok) {
        return;
    }
    context->Expect(result.value.sample_count == 3, "report sample count should match input");
    context->Expect(result.value.duration_ms == 2000,
                    "report duration should use elapsed milliseconds");
    context->Expect(result.value.observed_maximum_joint_jump_deg == 2.0,
                    "report should contain two-degree maximum jump");
    context->Expect(result.value.observed_maximum_joint_velocity_deg_per_second == 2.0,
                    "report should contain two-degree-per-second maximum velocity");
    context->Expect(
        result.value.observed_maximum_joint_acceleration_deg_per_second_squared == 1.0,
        "report should calculate acceleration between interval midpoints");
}

void TestRejectsEmptyAndWrongJointCountTrajectories(TestContext *context) {
    const std::vector<rm_control::common::TrajectorySample> empty_samples;
    auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        empty_samples, MakePermissiveLimits());
    context->Expect(!result.ok, "empty trajectory must be rejected");

    auto sample = MakeSample(0, 0, 0.0F);
    sample.joints_deg.pop_back();
    result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {sample}, MakePermissiveLimits());
    context->Expect(!result.ok, "RM75 trajectory point with six joints must be rejected");
}

void TestRejectsNonMonotonicTimestamps(TestContext *context) {
    const std::vector<rm_control::common::TrajectorySample> samples = {
        MakeSample(0, 100, 0.0F),
        MakeSample(1, 100, 1.0F),
    };
    const auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        samples, MakePermissiveLimits());
    context->Expect(!result.ok,
                    "equal timestamps must be rejected to prevent division by zero");
}

void TestRejectsNonFiniteJointAndTcpValues(TestContext *context) {
    auto non_finite_joint = MakeSample(0, 0, 0.0F);
    non_finite_joint.joints_deg[2] = std::numeric_limits<float>::quiet_NaN();
    auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {non_finite_joint, MakeSample(1, 1000, 0.0F)}, MakePermissiveLimits());
    context->Expect(!result.ok, "non-finite joint value must be rejected");

    auto non_finite_pose = MakeSample(0, 0, 0.0F);
    non_finite_pose.tcp_pose.z = std::numeric_limits<float>::infinity();
    result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {non_finite_pose, MakeSample(1, 1000, 0.0F)}, MakePermissiveLimits());
    context->Expect(!result.ok, "non-finite TCP pose value must be rejected");
}

void TestRejectsJumpVelocityAndAccelerationLimitViolations(TestContext *context) {
    auto limits = MakePermissiveLimits();
    limits.maximum_joint_jump_deg = 1.0;
    auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {MakeSample(0, 0, 0.0F), MakeSample(1, 1000, 2.0F)}, limits);
    context->Expect(!result.ok, "joint jump above configured maximum must be rejected");

    limits = MakePermissiveLimits();
    limits.maximum_joint_velocity_deg_per_second = 1.0;
    result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {MakeSample(0, 0, 0.0F), MakeSample(1, 1000, 2.0F)}, limits);
    context->Expect(!result.ok, "joint velocity above configured maximum must be rejected");

    limits = MakePermissiveLimits();
    limits.maximum_joint_acceleration_deg_per_second_squared = 1.0;
    result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {MakeSample(0, 0, 0.0F), MakeSample(1, 1000, 0.0F),
         MakeSample(2, 2000, 2.0F)},
        limits);
    context->Expect(!result.ok,
                    "joint acceleration above configured maximum must be rejected");
}

void TestRejectsInvalidLimits(TestContext *context) {
    auto limits = MakePermissiveLimits();
    limits.maximum_joint_jump_deg = 0.0;
    const auto result = rm_control::modules::trajectory::ValidateDragTrajectory(
        {MakeSample(0, 0, 0.0F)}, limits);
    context->Expect(!result.ok, "zero safety limit must be rejected");
}

} // namespace

int main() {
    TestContext context;
    TestReportsPeakJointMotionForValidTrajectory(&context);
    TestRejectsEmptyAndWrongJointCountTrajectories(&context);
    TestRejectsNonMonotonicTimestamps(&context);
    TestRejectsNonFiniteJointAndTcpValues(&context);
    TestRejectsJumpVelocityAndAccelerationLimitViolations(&context);
    TestRejectsInvalidLimits(&context);

    if (context.failures == 0) {
        std::cout << "[PASS] rm_drag_trajectory_validator_tests\n";
        return 0;
    }
    std::cerr << context.failures << " test(s) failed\n";
    return 1;
}
