#include "mocks/mock_robot_arm.h"

#include "rm_control/core/drag_teach_service.h"

#include <chrono>
#include <filesystem>
#include <iostream>
#include <string>
#include <thread>

#ifndef RM_TEST_OUTPUT_DIR
#define RM_TEST_OUTPUT_DIR "."
#endif

namespace {

struct TestContext {
    int failures = 0;

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

rm_control::core::HardwareExecutionGateInput AllowedGate() {
    rm_control::core::HardwareExecutionGateInput gate;
    gate.live_requested = true;
    gate.hardware_execution_enabled = true;
    gate.arm_connected = true;
    gate.hand_connected = true;
    gate.emergency_stop_verified = true;
    gate.safety_limits_loaded = true;
    gate.operator_confirmed = true;
    gate.trajectory_validated = true;
    gate.trajectory_sample_period_confirmed = true;
    gate.vendor_raw_units_per_degree_confirmed = true;
    gate.workspace_and_virtual_wall_verified = true;
    gate.configured_max_speed_percent = 20;
    gate.native_replay_without_speed_control_accepted = true;
    gate.native_replay_speed_within_config_verified = true;
    return gate;
}

rm_control::core::DragTeachServiceConfig ServiceConfig(const std::string &name) {
    rm_control::core::DragTeachServiceConfig config;
    config.storage_root =
        std::filesystem::path(RM_TEST_OUTPUT_DIR) / "drag_teach_service";
    config.session_name = name;
    config.sample_period = std::chrono::milliseconds(10);
    return config;
}

bool WaitForSampleCount(rm_control::core::DragTeachService *service,
                        std::size_t minimum_sample_count) {
    const auto deadline = std::chrono::steady_clock::now() +
                          std::chrono::milliseconds(500);
    while (std::chrono::steady_clock::now() < deadline) {
        if (service->recordedSampleCount() >= minimum_sample_count) {
            return true;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(2));
    }
    return false;
}

bool WaitForState(rm_control::core::DragTeachService *service,
                  rm_control::core::DragTeachServiceState expected) {
    const auto deadline = std::chrono::steady_clock::now() +
                          std::chrono::milliseconds(500);
    while (std::chrono::steady_clock::now() < deadline) {
        if (service->state() == expected) {
            return true;
        }
        std::this_thread::sleep_for(std::chrono::milliseconds(2));
    }
    return false;
}

void RecordAndSave(rm_control::core::DragTeachService *service,
                   const std::string &name,
                   TestContext *context) {
    const auto start = service->StartRecording(ServiceConfig(name), AllowedGate());
    context->Expect(start.ok, "StartRecording should succeed with complete gate");
    if (!start.ok) {
        return;
    }
    context->Expect(WaitForSampleCount(service, 3),
                    "recording should produce three samples before timeout");
    const auto stop = service->StopRecording();
    context->Expect(stop.ok, "StopRecording should validate stationary trajectory");
    if (!stop.ok) {
        return;
    }
    context->Expect(stop.value.recorded_sample_count >= 2,
                    "recording should contain at least two samples");

    const auto artifact = service->SaveRawTrajectory();
    context->Expect(artifact.ok, "SaveRawTrajectory should create manifest artifact");
    if (!artifact.ok) {
        return;
    }
    context->Expect(artifact.value.raw_file_present, "mock raw file should be present");
    context->Expect(artifact.value.raw_file_sha256.size() == 64,
                    "raw file SHA-256 should have 64 hex characters");
    context->Expect(std::filesystem::exists(artifact.value.paths.manifest),
                    "manifest should exist");
}

void TestStartRequiresHandConnection(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    rm_control::core::DragTeachService service(robot);
    auto gate = AllowedGate();
    gate.hand_connected = false;
    const auto result = service.StartRecording(ServiceConfig("missing_hand"), gate);
    context->Expect(!result.ok, "StartRecording must reject missing hand connection");
    context->Expect(robot.start_calls == 0, "rejected start must not call hardware");
}

void TestFullLifecycleAndContinueRecheck(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    rm_control::core::DragTeachService service(robot);
    RecordAndSave(&service, "full_lifecycle", context);
    if (service.state() != rm_control::core::DragTeachServiceState::Saved) {
        return;
    }

    const auto origin = service.MoveToTrajectoryOrigin(AllowedGate());
    const auto replay = service.ReplayLastDragTrajectory(AllowedGate());
    const auto pause = service.PauseReplay();
    context->Expect(origin.ok && replay.ok && pause.ok,
                    "origin, replay and pause should follow valid lifecycle");

    auto denied_gate = AllowedGate();
    denied_gate.live_requested = false;
    const auto denied_continue = service.ContinueReplay(denied_gate);
    context->Expect(!denied_continue.ok, "ContinueReplay must recheck live gate");
    context->Expect(robot.continue_calls == 0,
                    "denied ContinueReplay must not call hardware");

    const auto continued = service.ContinueReplay(AllowedGate());
    const auto stopped = service.StopReplay();
    context->Expect(continued.ok && stopped.ok,
                    "authorized continue and safety stop should succeed");
    context->Expect(robot.origin_calls == 1 && robot.replay_calls == 1 &&
                        robot.pause_calls == 1 && robot.continue_calls == 1 &&
                        robot.stop_replay_calls == 1,
                    "adapter replay methods should each be called once");
}

void TestWrongModelRejectedBeforeHardwareStart(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    robot.robot_info.is_rm75 = false;
    rm_control::core::DragTeachService service(robot);
    const auto result =
        service.StartRecording(ServiceConfig("wrong_model"), AllowedGate());
    context->Expect(!result.ok, "non-RM75 model must be rejected");
    context->Expect(robot.start_calls == 0, "wrong model must not start drag teach");
}

void TestJointLimitViolationRejectedBeforeHardwareStart(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    robot.arm_state.joints_deg[0] = 400.0F;
    rm_control::core::DragTeachService service(robot);
    const auto result =
        service.StartRecording(ServiceConfig("joint_limit"), AllowedGate());
    context->Expect(!result.ok, "current joint outside controller limits must be rejected");
    context->Expect(robot.start_calls == 0,
                    "joint limit violation must not start drag teach");
}

void TestSamplingFailureRetriesFailedStop(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    robot.fail_current_state_on_or_after_call = 2;
    robot.stop_recording_failures_remaining = 1;
    rm_control::core::DragTeachService service(robot);
    const auto start = service.StartRecording(
        ServiceConfig("sampling_failure"), AllowedGate());
    context->Expect(start.ok, "sampling failure test should initially start");
    if (!start.ok) {
        return;
    }
    context->Expect(WaitForState(&service, rm_control::core::DragTeachServiceState::Faulted),
                    "sampling read failure should enter Faulted");
    const auto stopped = service.StopRecording();
    context->Expect(!stopped.ok, "sampling failure should be reported to caller");
    context->Expect(robot.stop_recording_calls == 2,
                    "failed sampling-thread stop must be retried by StopRecording");
}

void TestUncertainStartAttemptsSafetyStop(TestContext *context) {
    rm_control::tests::MockRobotArm robot;
    robot.start_recording_failures_remaining = 1;
    rm_control::core::DragTeachService service(robot);
    const auto result = service.StartRecording(
        ServiceConfig("uncertain_start"), AllowedGate());
    context->Expect(!result.ok, "failed start must be reported");
    context->Expect(robot.start_calls == 1,
                    "failed start should call the adapter exactly once");
    context->Expect(robot.stop_recording_calls == 1,
                    "uncertain start must immediately attempt stopDragTeach");
}

} // namespace

int main() {
    TestContext context;
    TestStartRequiresHandConnection(&context);
    TestFullLifecycleAndContinueRecheck(&context);
    TestWrongModelRejectedBeforeHardwareStart(&context);
    TestJointLimitViolationRejectedBeforeHardwareStart(&context);
    TestSamplingFailureRetriesFailedStop(&context);
    TestUncertainStartAttemptsSafetyStop(&context);

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