#include "realman_trajectory_common.hpp"

#include <thread>

namespace {

using realman_trajectory::RobotConnectionConfig;
using realman_trajectory::RobotSession;
using realman_trajectory::TrajectoryPoint;

constexpr int kDefaultDurationSeconds = 10;
constexpr int kDefaultSamplePeriodMs = 100;
constexpr int kMinimumSamplePeriodMs = 50;
constexpr int kMaximumDurationSeconds = 3600;

struct RecordConfig {
    std::string output_csv_path;
    int duration_seconds = kDefaultDurationSeconds;
    int sample_period_ms = kDefaultSamplePeriodMs;
    RobotConnectionConfig robot;
};

void printUsage(const char *program_name) {
    std::cout << "Usage: " << program_name
              << " [output_csv] [duration_seconds] [sample_period_ms] [robot_ip] [port] [timeout_ms]\n"
              << "\n"
              << "Examples:\n"
              << "  " << program_name << "\n"
              << "  " << program_name << " reports/realman_trajectories/demo.csv 15 100\n"
              << "\n"
              << "This program only records current joints and TCP pose. It does not send motion commands.\n";
}

bool parseRecordArguments(int argc, char **argv, RecordConfig *config) {
    if (config == nullptr) {
        return false;
    }
    if (argc > 1 && (std::string(argv[1]) == "-h" || std::string(argv[1]) == "--help")) {
        printUsage(argv[0]);
        std::exit(0);
    }
    if (argc > 1) {
        config->output_csv_path = argv[1];
    }
    if (argc > 2 && !realman_trajectory::parseIntArgument(argv[2], &config->duration_seconds)) {
        std::cerr << "Invalid duration_seconds: " << argv[2] << '\n';
        return false;
    }
    if (argc > 3 && !realman_trajectory::parseIntArgument(argv[3], &config->sample_period_ms)) {
        std::cerr << "Invalid sample_period_ms: " << argv[3] << '\n';
        return false;
    }
    if (argc > 4) {
        config->robot.robot_ip = argv[4];
    }
    if (argc > 5 && !realman_trajectory::parseIntArgument(argv[5], &config->robot.robot_port)) {
        std::cerr << "Invalid port: " << argv[5] << '\n';
        return false;
    }
    if (argc > 6 && !realman_trajectory::parseIntArgument(argv[6], &config->robot.timeout_ms)) {
        std::cerr << "Invalid timeout_ms: " << argv[6] << '\n';
        return false;
    }
    if (argc > 7) {
        std::cerr << "Too many arguments.\n";
        return false;
    }
    if (config->output_csv_path.empty()) {
        config->output_csv_path = realman_trajectory::defaultRecordedTrajectoryPath();
    }
    return true;
}

bool validateRecordConfig(const RecordConfig &config) {
    if (config.output_csv_path.empty()) {
        std::cerr << "output_csv must not be empty.\n";
        return false;
    }
    if (config.duration_seconds <= 0 || config.duration_seconds > kMaximumDurationSeconds) {
        std::cerr << "duration_seconds must be in range 1.." << kMaximumDurationSeconds << ".\n";
        return false;
    }
    if (config.sample_period_ms < kMinimumSamplePeriodMs) {
        std::cerr << "sample_period_ms must be at least " << kMinimumSamplePeriodMs << ".\n";
        return false;
    }
    if (config.robot.robot_port <= 0 || config.robot.robot_port > 65535) {
        std::cerr << "port must be in range 1..65535.\n";
        return false;
    }
    return config.robot.timeout_ms > 0;
}

TrajectoryPoint makeTrajectoryPoint(long long elapsed_ms, const rm_current_arm_state_t &state) {
    TrajectoryPoint point = {};
    point.elapsed_ms = elapsed_ms;
    point.joint_deg = realman_trajectory::copyJointArray(state.joint);
    point.tcp_pose = state.pose;
    return point;
}

bool recordTrajectoryCsv(rm_robot_handle *handle, const RecordConfig &config, int dof) {
    std::ofstream output(config.output_csv_path);
    if (!output) {
        std::cerr << "Failed to open output CSV: " << config.output_csv_path << '\n';
        return false;
    }

    output << realman_trajectory::trajectoryCsvHeader() << '\n';
    const auto start_time = std::chrono::steady_clock::now();
    const int total_samples = (config.duration_seconds * 1000) / config.sample_period_ms + 1;

    for (int sample_index = 0; sample_index < total_samples; ++sample_index) {
        const auto target_time = start_time + std::chrono::milliseconds(sample_index * config.sample_period_ms);
        std::this_thread::sleep_until(target_time);

        rm_current_arm_state_t state = {};
        if (!realman_trajectory::readCurrentArmState(handle, &state)) {
            return false;
        }

        const auto now = std::chrono::steady_clock::now();
        const long long elapsed_ms = std::chrono::duration_cast<std::chrono::milliseconds>(now - start_time).count();
        const TrajectoryPoint point = makeTrajectoryPoint(elapsed_ms, state);
        realman_trajectory::writeTrajectoryPointCsv(&output, sample_index, point);

        if (sample_index == 0 || sample_index == total_samples - 1) {
            std::cout << "Sample " << sample_index << " joints "
                      << realman_trajectory::formatJoints(point.joint_deg, dof) << '\n';
        }
    }

    std::cout << "Recorded " << total_samples << " samples to " << config.output_csv_path << '\n';
    return true;
}

} // namespace

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

    RecordConfig config;
    if (!parseRecordArguments(argc, argv, &config) || !validateRecordConfig(config)) {
        printUsage(argv[0]);
        return 64;
    }

    RobotSession session(config.robot);
    if (!session.connect()) {
        return 1;
    }

    int dof = 0;
    if (!realman_trajectory::readAndValidateRobotDof(session.handle(), &dof)) {
        return 2;
    }

    if (!recordTrajectoryCsv(session.handle(), config, dof)) {
        return 3;
    }

    return 0;
}
