#include "rm_control/core/drag_teach_recorder.h"

#include <chrono>
#include <iostream>
#include <thread>
#include <vector>

namespace rm_control::core {
namespace {

constexpr int kMinimumSamplePeriodMs = 50;
constexpr int kMaximumDurationSeconds = 3600;
constexpr int kSafetyCountdownSeconds = 5;

common::TrajectorySample makeSample(long long elapsed_ms, const common::ArmState &state) {
    common::TrajectorySample sample;
    sample.elapsed_ms = elapsed_ms;
    sample.joints_deg = state.joints_deg;
    sample.tcp_pose = state.tcp_pose;
    return sample;
}

void printCountdown(std::ostream &out) {
    out << "Prepare for six-dof force drag teach. Clear workspace and keep E-stop reachable.\n";
    for (int second = kSafetyCountdownSeconds; second > 0; --second) {
        out << "Drag teach starts in " << second << "...\n";
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

} // namespace

DragTeachRecorder::DragTeachRecorder(interfaces::IRobotArm &robot)
    : robot_(robot) {}

common::Result<DragTeachRecordSummary> DragTeachRecorder::record(const DragTeachRecordConfig &config,
                                                                 std::ostream &out) {
    auto validation = validateConfig(config);
    if (!validation.ok) {
        return common::Result<DragTeachRecordSummary>::failure(validation.error_message);
    }

    auto files = modules::trajectory::makeTrajectoryFileSet(config.storage_root, config.session_name);
    if (!files.ok) {
        return common::Result<DragTeachRecordSummary>::failure(files.error_message);
    }

    printCountdown(out);

    auto start_result = robot_.startDragTeach(config.drag_options);
    if (!start_result.ok) {
        return common::Result<DragTeachRecordSummary>::failure(start_result.error_message);
    }

    std::vector<common::TrajectorySample> samples;
    const auto start_time = std::chrono::steady_clock::now();
    const int total_samples = (config.duration_seconds * 1000) / config.sample_period_ms + 1;
    common::Result<void> sample_result = common::Result<void>::success();

    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);

        auto state = robot_.getCurrentArmState();
        if (!state.ok) {
            sample_result = common::Result<void>::failure(state.error_message);
            break;
        }

        const auto now = std::chrono::steady_clock::now();
        const long long elapsed_ms =
            std::chrono::duration_cast<std::chrono::milliseconds>(now - start_time).count();
        samples.push_back(makeSample(elapsed_ms, state.value));
        if (sample_index == 0 || sample_index == total_samples - 1) {
            out << "sample " << sample_index << "/" << (total_samples - 1)
                << " elapsed_ms=" << elapsed_ms << '\n';
        }
    }

    auto stop_result = robot_.stopDragTeach();
    if (!sample_result.ok) {
        return common::Result<DragTeachRecordSummary>::failure(sample_result.error_message);
    }
    if (!stop_result.ok) {
        return common::Result<DragTeachRecordSummary>::failure(stop_result.error_message);
    }

    auto write_result = modules::trajectory::writeTrajectoryCsv(files.value.csv_path, samples);
    if (!write_result.ok) {
        return common::Result<DragTeachRecordSummary>::failure(write_result.error_message);
    }

    DragTeachRecordSummary summary;
    summary.files = files.value;
    summary.sample_count = static_cast<int>(samples.size());

    if (config.save_vendor_trajectory) {
        auto vendor_result = robot_.saveDragTeachTrajectory(files.value.vendor_path);
        if (vendor_result.ok) {
            summary.vendor_trajectory_saved = true;
            summary.vendor_save_message = "saved " + std::to_string(vendor_result.value.point_count) +
                                          " vendor points";
        } else {
            summary.vendor_save_message = vendor_result.error_message;
        }
    } else {
        summary.vendor_save_message = "vendor trajectory save disabled";
    }

    out << "CSV saved: " << summary.files.csv_path << '\n';
    out << "Vendor trajectory: " << summary.vendor_save_message << '\n';
    return common::Result<DragTeachRecordSummary>::success(summary);
}

common::Result<void> DragTeachRecorder::validateConfig(const DragTeachRecordConfig &config) const {
    if (config.storage_root.empty()) {
        return common::Result<void>::failure("storage_root must not be empty");
    }
    if (config.session_name.empty()) {
        return common::Result<void>::failure("session_name must not be empty");
    }
    if (config.duration_seconds <= 0 || config.duration_seconds > kMaximumDurationSeconds) {
        return common::Result<void>::failure("duration_seconds must be in range 1.." +
                                            std::to_string(kMaximumDurationSeconds));
    }
    if (config.sample_period_ms < kMinimumSamplePeriodMs) {
        return common::Result<void>::failure("sample_period_ms must be at least " +
                                            std::to_string(kMinimumSamplePeriodMs));
    }
    return common::Result<void>::success();
}

} // namespace rm_control::core
