#include "rm_interface.h"

#include <algorithm>
#include <array>
#include <atomic>
#include <chrono>
#include <cmath>
#include <cstdarg>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <iomanip>
#include <iostream>
#include <string>
#include <thread>
#include <vector>
#include <csignal>

namespace {

constexpr const char *kRobotIp = "192.168.1.18";
constexpr int kRobotPort = 8080;
constexpr int kRobotApiTimeoutMs = 1500;

constexpr int kToolModbusPort = 1;
constexpr int kToolModbusBaudrate = 115200;
constexpr int kToolModbusTimeoutUnits = 10;
constexpr int kToolVoltage24V = 3;

constexpr int kRightHandDeviceId = 0x27;
constexpr int kO6IdentityRegisterAddress = 30;
constexpr int kO6JointPositionRegisterAddress = 0x01;
constexpr int kO6JointRegisterCount = 6;

constexpr int kFirstArmSwingJointOneBased = 2;
constexpr int kSecondArmSwingJointOneBased = 3;
constexpr float kDefaultAmplitudeDeg = 10.0F;
constexpr float kMaxAutomaticAmplitudeDeg = 12.0F;
constexpr int kDefaultSpeedPercent = 6;
constexpr int kMaxAutomaticSpeedPercent = 10;
constexpr int kDefaultCycleCount = 1;
constexpr int kMaxAutomaticCycleCount = 10;
constexpr int kSafetyCountdownSeconds = 5;
constexpr int kMoveTimeoutSeconds = 18;
constexpr int kPollIntervalMs = 100;
constexpr int kInterCyclePauseMs = 350;
constexpr float kFinalJointErrorDeg = 0.8F;

std::atomic_bool g_stop_requested{false};

struct ProgramOptions {
    bool execute = false;
    float amplitude_deg = kDefaultAmplitudeDeg;
    int speed_percent = kDefaultSpeedPercent;
    int cycles = kDefaultCycleCount;
    bool release_at_end = true;
};

struct ModbusLayout {
    int device_id = kRightHandDeviceId;
    int address = kO6JointPositionRegisterAddress;
    int register_count = kO6JointRegisterCount;
};

void signalHandler(int) {
    g_stop_requested.store(true);
}

void apiLogCallback(const char *message, va_list args) {
    if (message == nullptr) {
        return;
    }

    char buffer[1024] = {};
    std::vsnprintf(buffer, sizeof(buffer), message, args);
    std::cerr << "[rm_api] " << buffer << '\n';
}

void destroyApi() {
    const int destroy_result = rm_destroy();
    if (destroy_result != 0) {
        std::cerr << "rm_destroy failed, error code: " << destroy_result << '\n';
    }
}

void cleanup(rm_robot_handle *handle) {
    if (handle != nullptr) {
        const int delete_result = rm_delete_robot_arm(handle);
        if (delete_result != 0) {
            std::cerr << "rm_delete_robot_arm failed, error code: " << delete_result << '\n';
        }
    }
    destroyApi();
}

bool parseIntArgument(const char *value, int *out) {
    if (value == nullptr || out == nullptr) {
        return false;
    }
    char *end = nullptr;
    const long parsed = std::strtol(value, &end, 10);
    if (end == value || *end != '\0') {
        return false;
    }
    *out = static_cast<int>(parsed);
    return true;
}

bool parseFloatArgument(const char *value, float *out) {
    if (value == nullptr || out == nullptr) {
        return false;
    }
    char *end = nullptr;
    const float parsed = std::strtof(value, &end);
    if (end == value || *end != '\0') {
        return false;
    }
    *out = parsed;
    return true;
}

bool parseOptions(int argc, char **argv, ProgramOptions *options) {
    if (options == nullptr) {
        return false;
    }

    for (int i = 1; i < argc; ++i) {
        if (std::strcmp(argv[i], "--execute") == 0) {
            options->execute = true;
            continue;
        }
        if (std::strcmp(argv[i], "--hold-at-end") == 0) {
            options->release_at_end = false;
            continue;
        }
        if (std::strcmp(argv[i], "--amplitude-deg") == 0 && i + 1 < argc) {
            if (!parseFloatArgument(argv[++i], &options->amplitude_deg)) {
                std::cerr << "Invalid --amplitude-deg value.\n";
                return false;
            }
            continue;
        }
        if (std::strcmp(argv[i], "--speed-percent") == 0 && i + 1 < argc) {
            if (!parseIntArgument(argv[++i], &options->speed_percent)) {
                std::cerr << "Invalid --speed-percent value.\n";
                return false;
            }
            continue;
        }
        if (std::strcmp(argv[i], "--cycles") == 0 && i + 1 < argc) {
            if (!parseIntArgument(argv[++i], &options->cycles)) {
                std::cerr << "Invalid --cycles value.\n";
                return false;
            }
            continue;
        }

        std::cerr << "Unknown or incomplete argument: " << argv[i] << '\n';
        return false;
    }

    return true;
}

bool validateOptions(const ProgramOptions &options) {
    if (options.amplitude_deg <= 0.0F || options.amplitude_deg > kMaxAutomaticAmplitudeDeg) {
        std::cerr << "Refuse amplitude " << options.amplitude_deg
                  << " deg. Automatic limit is 0.." << kMaxAutomaticAmplitudeDeg << " deg.\n";
        return false;
    }
    if (options.speed_percent < 1 || options.speed_percent > kMaxAutomaticSpeedPercent) {
        std::cerr << "Refuse speed " << options.speed_percent
                  << "%. Automatic limit is 1.." << kMaxAutomaticSpeedPercent << "%.\n";
        return false;
    }
    if (options.cycles < 1 || options.cycles > kMaxAutomaticCycleCount) {
        std::cerr << "Refuse cycle count " << options.cycles
                  << ". Automatic limit is 1.." << kMaxAutomaticCycleCount << ".\n";
        return false;
    }
    return true;
}

bool connectRobot(rm_robot_handle **handle) {
    rm_set_log_call_back(apiLogCallback, 2);
    rm_set_timeout(kRobotApiTimeoutMs);

    const int init_result = rm_init(RM_TRIPLE_MODE_E);
    if (init_result != 0) {
        std::cerr << "rm_init failed, error code: " << init_result << '\n';
        return false;
    }

    std::cout << "RealMan API version: " << rm_api_version() << '\n';
    std::cout << "Connecting to " << kRobotIp << ':' << kRobotPort << "...\n";

    *handle = rm_create_robot_arm(kRobotIp, kRobotPort);
    if (*handle == nullptr || (*handle)->id < 0) {
        std::cerr << "Failed to connect RealMan arm at " << kRobotIp << ':' << kRobotPort << '\n';
        cleanup(*handle);
        *handle = nullptr;
        return false;
    }

    std::cout << "Connected, handle id: " << (*handle)->id << '\n';
    return true;
}

bool readArmState(rm_robot_handle *handle, rm_current_arm_state_t *state) {
    const int result = rm_get_current_arm_state(handle, state);
    if (result != 0) {
        std::cerr << "rm_get_current_arm_state failed, error code: " << result << '\n';
        return false;
    }
    return true;
}

bool allFinite(const float joints[ARM_DOF], int dof) {
    for (int i = 0; i < dof; ++i) {
        if (!std::isfinite(joints[i])) {
            std::cerr << "Joint J" << (i + 1) << " is not finite: " << joints[i] << '\n';
            return false;
        }
    }
    return true;
}

void copyJoints(const float source[ARM_DOF], std::array<float, ARM_DOF> *target) {
    for (int i = 0; i < ARM_DOF; ++i) {
        (*target)[static_cast<size_t>(i)] = source[i];
    }
}

void printJoints(const char *label, const float joints[ARM_DOF], int dof) {
    std::cout << label;
    for (int i = 0; i < dof; ++i) {
        std::cout << (i == 0 ? " [" : ", ")
                  << "J" << (i + 1) << "=" << std::fixed << std::setprecision(3) << joints[i];
    }
    std::cout << "] deg\n";
}

double maxAbsDiff(const float *left, const float *right, int count) {
    double max_diff = 0.0;
    for (int i = 0; i < count; ++i) {
        max_diff = std::max(max_diff, std::fabs(static_cast<double>(left[i] - right[i])));
    }
    return max_diff;
}

rm_peripheral_read_write_params_t makePeripheralParams(const ModbusLayout &layout) {
    rm_peripheral_read_write_params_t param = {};
    param.port = kToolModbusPort;
    param.address = layout.address;
    param.device = layout.device_id;
    param.num = layout.register_count;
    return param;
}

void printToolVoltage(rm_robot_handle *handle, const char *label) {
    int voltage_type = -1;
    const int result = rm_get_tool_voltage(handle, &voltage_type);
    if (result != 0) {
        std::cout << label << " rm_get_tool_voltage failed, error code: " << result << '\n';
        return;
    }
    std::cout << label << " tool voltage type=" << voltage_type
              << " (0=0V, 1=5V, 2=12V, 3=24V)\n";
}

bool configureToolPowerAndRs485(rm_robot_handle *handle) {
    printToolVoltage(handle, "Before");

    const int voltage_result = rm_set_tool_voltage(handle, kToolVoltage24V);
    if (voltage_result != 0) {
        std::cerr << "rm_set_tool_voltage(24V) failed, error code: " << voltage_result << '\n';
        return false;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(300));
    printToolVoltage(handle, "After");

    const int modbus_result = rm_set_modbus_mode(
        handle,
        kToolModbusPort,
        kToolModbusBaudrate,
        kToolModbusTimeoutUnits);
    if (modbus_result != 0) {
        std::cerr << "rm_set_modbus_mode(port=1, baud=115200) failed, error code: " << modbus_result << '\n';
        return false;
    }
    std::cout << "Tool RS485 configured as Modbus RTU master, baudrate="
              << kToolModbusBaudrate << ", timeout_units=" << kToolModbusTimeoutUnits << '\n';
    return true;
}

bool readO6IdentityRegisters(rm_robot_handle *handle, int device_id) {
    ModbusLayout layout = {device_id, kO6IdentityRegisterAddress, kO6JointRegisterCount};
    rm_peripheral_read_write_params_t param = makePeripheralParams(layout);

    int data[32] = {};
    const int result = rm_read_multiple_input_registers(handle, param, data);
    if (result != 0) {
        std::cerr << "Read O6 identity registers failed, device=0x" << std::hex << device_id
                  << std::dec << ", error code: " << result << '\n';
        return false;
    }

    std::cout << "O6 identity bytes, device=0x" << std::hex << device_id << std::dec << ':';
    for (int i = 0; i < kO6JointRegisterCount * 2; ++i) {
        std::cout << ' ' << data[i];
    }
    std::cout << " | ascii=";
    for (int i = 0; i < kO6JointRegisterCount * 2; ++i) {
        const int value = data[i] & 0xFF;
        std::cout << (value >= 32 && value <= 126 ? static_cast<char>(value) : '.');
    }
    std::cout << '\n';
    return true;
}

bool writeO6JointPose(rm_robot_handle *handle, const std::vector<int> &pose, const char *label) {
    if (pose.size() != kO6JointRegisterCount) {
        std::cerr << "Internal error: O6 pose must contain 6 joint values.\n";
        return false;
    }

    ModbusLayout layout = {};
    rm_peripheral_read_write_params_t param = makePeripheralParams(layout);
    int data[20] = {};
    for (int &value : data) {
        value = 255;
    }
    for (int i = 0; i < kO6JointRegisterCount; ++i) {
        data[i] = pose[static_cast<size_t>(i)];
    }

    const int result = rm_write_registers(handle, param, data);
    if (result != 0) {
        std::cerr << "Write O6 pose failed at " << label << ", error code: " << result << '\n';
        return false;
    }

    std::cout << "Wrote O6 " << label << " pose:";
    for (int value : pose) {
        std::cout << ' ' << value;
    }
    std::cout << '\n';
    return true;
}

bool startJointMove(rm_robot_handle *handle,
                    const std::array<float, ARM_DOF> &target,
                    int speed_percent,
                    const char *label,
                    int dof) {
    printJoints(label, target.data(), dof);
    const int result = rm_movej(handle, target.data(), speed_percent, 0, 0, 0);
    if (result != 0) {
        std::cerr << "rm_movej nonblocking failed at " << label << ", error code: " << result << '\n';
        rm_set_arm_stop(handle);
        return false;
    }
    return true;
}

bool waitUntilNearTarget(rm_robot_handle *handle,
                         const std::array<float, ARM_DOF> &target,
                         const char *label,
                         int dof) {
    const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(kMoveTimeoutSeconds);
    while (std::chrono::steady_clock::now() < deadline) {
        if (g_stop_requested.load()) {
            std::cerr << "Stop requested. Sending slow stop.\n";
            rm_set_arm_slow_stop(handle);
            return false;
        }

        rm_current_arm_state_t current = {};
        if (!readArmState(handle, &current)) {
            rm_set_arm_stop(handle);
            return false;
        }

        const double error_deg = maxAbsDiff(target.data(), current.joint, dof);
        if (error_deg <= kFinalJointErrorDeg) {
            std::cout << label << " reached, max_error_deg=" << std::fixed
                      << std::setprecision(3) << error_deg << '\n';
            return true;
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(kPollIntervalMs));
    }

    std::cerr << label << " timed out. Sending stop.\n";
    rm_set_arm_stop(handle);
    return false;
}

bool moveArmAndOptionallyWriteHand(rm_robot_handle *handle,
                                   const std::array<float, ARM_DOF> &target,
                                   const std::vector<int> &hand_pose_after_start,
                                   const char *arm_label,
                                   const char *hand_label,
                                   const ProgramOptions &options,
                                   int dof) {
    if (!startJointMove(handle, target, options.speed_percent, arm_label, dof)) {
        return false;
    }

    std::this_thread::sleep_for(std::chrono::milliseconds(350));
    if (!hand_pose_after_start.empty() && !writeO6JointPose(handle, hand_pose_after_start, hand_label)) {
        rm_set_arm_stop(handle);
        return false;
    }

    return waitUntilNearTarget(handle, target, arm_label, dof);
}

void countdownBeforeMotion() {
    std::cout << "\nSafety countdown before motion. Clear workspace and keep E-stop reachable.\n";
    for (int second = kSafetyCountdownSeconds; second > 0; --second) {
        if (g_stop_requested.load()) {
            break;
        }
        std::cout << "Moving in " << second << "...\n";
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

bool buildSwingTargets(const rm_current_arm_state_t &start_state,
                       const ProgramOptions &options,
                       int dof,
                       std::array<float, ARM_DOF> *positive,
                       std::array<float, ARM_DOF> *negative,
                       std::array<float, ARM_DOF> *start) {
    if (dof < kSecondArmSwingJointOneBased) {
        std::cerr << "Robot DOF is too small for J2/J3 swing: " << dof << '\n';
        return false;
    }
    if (!allFinite(start_state.joint, dof)) {
        return false;
    }

    copyJoints(start_state.joint, start);
    copyJoints(start_state.joint, positive);
    copyJoints(start_state.joint, negative);

    const size_t j2 = static_cast<size_t>(kFirstArmSwingJointOneBased - 1);
    const size_t j3 = static_cast<size_t>(kSecondArmSwingJointOneBased - 1);
    (*positive)[j2] += options.amplitude_deg;
    (*positive)[j3] -= options.amplitude_deg;
    (*negative)[j2] -= options.amplitude_deg;
    (*negative)[j3] += options.amplitude_deg;

    printJoints("Current joints:", start_state.joint, dof);
    printJoints("Planned positive swing:", positive->data(), dof);
    printJoints("Planned negative swing:", negative->data(), dof);
    printJoints("Planned return:", start->data(), dof);
    return true;
}

bool runOneCompleteCombinedCycle(rm_robot_handle *handle,
                                 const ProgramOptions &options,
                                 const std::array<float, ARM_DOF> &positive,
                                 const std::array<float, ARM_DOF> &negative,
                                 const std::array<float, ARM_DOF> &start,
                                 int cycle_index,
                                 int dof) {
    const std::vector<int> pre_grasp_pose = {220, 170, 230, 180, 180, 180};
    const std::vector<int> grasp_pose = {140, 100, 55, 55, 55, 55};
    const std::vector<int> hold_pose = {140, 100, 70, 70, 70, 70};
    const std::vector<int> safe_open_pose = {255, 255, 255, 255, 255, 255};

    std::cout << "Complete combined action cycle " << cycle_index << "/" << options.cycles << '\n';
    if (!writeO6JointPose(handle, pre_grasp_pose, "pre_grasp")) {
        return false;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(400));

    if (!moveArmAndOptionallyWriteHand(handle, positive, grasp_pose, "swing_positive", "grasp", options, dof)) {
        return false;
    }
    if (!moveArmAndOptionallyWriteHand(handle, negative, hold_pose, "swing_negative", "hold", options, dof)) {
        return false;
    }

    const bool is_last_cycle = cycle_index == options.cycles;
    const std::vector<int> final_pose = (!is_last_cycle || options.release_at_end) ? safe_open_pose : hold_pose;
    const char *final_hand_label = (!is_last_cycle || options.release_at_end) ? "release_safe_open" : "hold_at_end";
    if (!moveArmAndOptionallyWriteHand(handle,
                                       start,
                                       final_pose,
                                       "return_to_start",
                                       final_hand_label,
                                       options,
                                       dof)) {
        return false;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(kInterCyclePauseMs));
    return true;
}

bool runCombinedMotion(rm_robot_handle *handle,
                       const ProgramOptions &options,
                       const std::array<float, ARM_DOF> &positive,
                       const std::array<float, ARM_DOF> &negative,
                       const std::array<float, ARM_DOF> &start,
                       int dof) {
    const std::vector<int> safe_open_pose = {255, 255, 255, 255, 255, 255};

    countdownBeforeMotion();
    if (g_stop_requested.load()) {
        std::cerr << "Motion canceled before start.\n";
        return false;
    }

    if (!writeO6JointPose(handle, safe_open_pose, "safe_open")) {
        return false;
    }
    std::this_thread::sleep_for(std::chrono::milliseconds(500));

    for (int cycle = 1; cycle <= options.cycles; ++cycle) {
        if (!runOneCompleteCombinedCycle(handle, options, positive, negative, start, cycle, dof)) {
            return false;
        }
    }

    rm_current_arm_state_t final_state = {};
    if (readArmState(handle, &final_state)) {
        printJoints("Final joints:", final_state.joint, dof);
    }
    return true;
}

} // namespace

int main(int argc, char **argv) {
    std::cout.setf(std::ios::unitbuf);
    std::cerr.setf(std::ios::unitbuf);
    std::signal(SIGINT, signalHandler);
    std::signal(SIGTERM, signalHandler);

    ProgramOptions options = {};
    if (!parseOptions(argc, argv, &options) || !validateOptions(options)) {
        return 64;
    }

    rm_robot_handle *handle = nullptr;
    if (!connectRobot(&handle)) {
        return 1;
    }

    int exit_code = 0;
    rm_robot_info_t robot_info = {};
    int result = rm_get_robot_info(handle, &robot_info);
    if (result != 0) {
        std::cerr << "rm_get_robot_info failed, error code: " << result << '\n';
        exit_code = 2;
    }

    const int dof = static_cast<int>(robot_info.arm_dof);
    rm_current_arm_state_t start_state = {};
    std::array<float, ARM_DOF> positive = {};
    std::array<float, ARM_DOF> negative = {};
    std::array<float, ARM_DOF> start = {};

    if (exit_code == 0 && !readArmState(handle, &start_state)) {
        exit_code = 3;
    }
    if (exit_code == 0 && !buildSwingTargets(start_state, options, dof, &positive, &negative, &start)) {
        exit_code = 4;
    }

    if (exit_code == 0) {
        std::cout << "Options: execute=" << (options.execute ? "true" : "false")
                  << ", swing_joints=J" << kFirstArmSwingJointOneBased << "/J" << kSecondArmSwingJointOneBased
                  << ", amplitude_deg=" << options.amplitude_deg
                  << ", speed_percent=" << options.speed_percent
                  << ", cycles=" << options.cycles
                  << ", release_at_end=" << (options.release_at_end ? "true" : "false") << '\n';
    }

    if (exit_code == 0 && !options.execute) {
        std::cout << "Dry run only. Add --execute after confirming workspace clearance and E-stop readiness.\n";
    } else if (exit_code == 0) {
        if (!configureToolPowerAndRs485(handle)) {
            exit_code = 5;
        } else if (!readO6IdentityRegisters(handle, kRightHandDeviceId)) {
            exit_code = 6;
        } else if (!runCombinedMotion(handle, options, positive, negative, start, dof)) {
            exit_code = 7;
        }
    }

    if (exit_code == 0) {
        std::cout << "Combined arm swing and O6 grasp sequence finished.\n";
    }

    cleanup(handle);
    return exit_code;
}
