#include "fgc/CalibrationRoutine.h"

#include "fgc/Calibration.h"
#include "fgc/DumpParser.h"
#include "fgc/IImuSource.h"
#include "fgc/IMotorController.h"
#include "fgc/Logger.h"
#include "fgc/Paths.h"

#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdlib>
#include <sstream>
#include <utility>
#include <vector>

namespace fgc {

namespace {
using clock = std::chrono::steady_clock;

void msleep(int ms) { std::this_thread::sleep_for(std::chrono::milliseconds(ms)); }
}  // namespace

CalibrationRoutine::CalibrationRoutine(IMotorController& motor, IImuSource& imu, Geometry initial,
                                       CalibParams params)
    : motor_(motor), imu_(imu), initial_(std::move(initial)), params_(params) {}

CalibrationRoutine::~CalibrationRoutine() {
    cancel();
    if (thread_.joinable()) thread_.join();
}

bool CalibrationRoutine::start() {
    if (running_.exchange(true)) {
        LOG_WARN << "calibration already running";
        return false;
    }
    cancel_ = false;
    if (thread_.joinable()) thread_.join();
    thread_ = std::thread([this] { run(); });
    return true;
}

void CalibrationRoutine::cancel() { cancel_ = true; }

std::optional<Geometry> CalibrationRoutine::takeResult() {
    std::lock_guard<std::mutex> lk(result_mutex_);
    std::optional<Geometry> r;
    r.swap(result_);
    return r;
}

CalibProgress CalibrationRoutine::progress() const {
    std::lock_guard<std::mutex> lk(progress_mutex_);
    return progress_;
}

CalibReport CalibrationRoutine::report() const {
    std::lock_guard<std::mutex> lk(result_mutex_);
    return report_;
}

void CalibrationRoutine::setProgress(char axis, int step, int total, const char* phase) {
    std::lock_guard<std::mutex> lk(progress_mutex_);
    progress_.running = true;
    progress_.axis    = axis;
    progress_.step    = step;
    progress_.total   = total;
    progress_.phase   = phase;
}

namespace {
long long nowMs() {
    return std::chrono::duration_cast<std::chrono::milliseconds>(
               std::chrono::system_clock::now().time_since_epoch())
        .count();
}
}  // namespace

namespace {

// Wait until the axis is at `target` (standstill + within tol), or timeout/cancel.
bool waitSettle(IMotorController& motor, char axis, long target, long tol, int timeout_ms,
                const std::atomic<bool>& cancel) {
    const auto deadline = clock::now() + std::chrono::milliseconds(timeout_ms);
    while (clock::now() < deadline) {
        if (cancel) return false;
        MotorTelemetry t = motor.telemetry();
        const AxisTelemetry& a = (axis == 'Y') ? t.yaw : t.pitch;
        if (a.standstill && std::labs(a.xenc - target) <= tol) return true;
        msleep(100);
    }
    return false;
}

// Wait until both present axes report READY (homing complete), or timeout/cancel.
bool waitHomed(IMotorController& motor, int timeout_ms, const std::atomic<bool>& cancel) {
    const auto deadline = clock::now() + std::chrono::milliseconds(timeout_ms);
    while (clock::now() < deadline) {
        if (cancel) return false;
        MotorTelemetry t = motor.telemetry();
        if (t.yaw.ready() && (!t.pitch_present || t.pitch.ready())) return true;
        msleep(200);
    }
    return false;
}

// Dwell `dwell_ms` sampling the IMU; return the (circular for yaw) mean of the
// requested Euler angle in degrees.
double dwellAndMeasure(IImuSource& imu, char axis, int dwell_ms, int sample_ms,
                       const std::atomic<bool>& cancel) {
    std::vector<double> vals;
    const auto end = clock::now() + std::chrono::milliseconds(dwell_ms);
    while (clock::now() < end) {
        if (cancel) break;
        if (auto s = imu.sample())
            vals.push_back(axis == 'Y' ? s->yaw_deg : s->pitch_deg);
        msleep(sample_ms);
    }
    if (vals.empty()) return 0.0;
    if (axis == 'Y') {
        double m = circularMeanDeg(vals);   // robust to wrap within the dwell
        if (m < 0.0) m += 360.0;            // report a 0..360 heading (matches IMU)
        return m;
    }
    double m = 0;
    for (double v : vals) m += v;
    return m / vals.size();
}

}  // namespace

void CalibrationRoutine::run() {
    struct Done {
        CalibrationRoutine* self;
        ~Done() {
            self->running_ = false;
            std::lock_guard<std::mutex> lk(self->progress_mutex_);
            self->progress_.running = false;
        }
    } done{this};

    const int N = params_.positions;
    LOG_INFO << "=== gimbal calibration starting ===";
    setProgress('?', 0, N, "starting");
    CalibReport report;

    // 1. Preconditions.
    if (!imu_.connected()) {
        LOG_WARN << "calibration aborted: IMU not connected";
        return;
    }

    // 2. Home first if the gimbal is not already READY.
    MotorTelemetry t0 = motor_.telemetry();
    bool homed = t0.yaw.ready() && (!t0.pitch_present || t0.pitch.ready());
    if (!homed) {
        LOG_INFO << "calibration: gimbal not homed; homing first...";
        setProgress('?', 0, N, "homing");
        motor_.sendCommand("HOME");
        if (!waitHomed(motor_, params_.home_timeout_ms, cancel_)) {
            if (cancel_) { LOG_WARN << "calibration cancelled"; return; }
            LOG_WARN << "calibration aborted: homing did not reach READY";
            return;
        }
        t0 = motor_.telemetry();
    }
    const bool pitch_present = t0.pitch_present;

    // Fetch soft limits from a fresh dump.
    LOG_INFO << "calibration: requesting soft limits (dump)...";
    motor_.sendCommand("DUMP");
    DumpData dump;
    for (int i = 0; i < 50 && !cancel_; ++i) {  // up to ~5 s
        dump = parseDump(motor_.lastDump());
        if (dump.valid && !dump.axes.empty()) break;
        msleep(100);
    }
    if (!dump.valid) {
        LOG_WARN << "calibration aborted: could not read soft limits";
        return;
    }

    Geometry result = initial_;
    std::ostringstream log;
    log << "gimbal calibration\n==================\n";

    // Per-axis sweep span [a..b] (inset off the hard endstops) + a settle tolerance.
    auto axisSpan = [&](char axis, long& a, long& b, long& tol) -> bool {
        const DumpAxis* da = nullptr;
        for (const auto& ax : dump.axes) if (ax.axis == axis) da = &ax;
        if (!da) return false;
        const long lo = da->lim_neg, hi = da->lim_pos;
        const long travel = std::labs(hi - lo);
        const long inset = static_cast<long>(params_.inset_frac * travel);
        a = lo + inset;
        b = hi - inset;
        tol = std::max<long>(500, travel / 200);
        return true;
    };

    // Sweep one axis across [a..b] in N steps, dwelling to measure the IMU angle,
    // then least-squares fit counts vs degrees and record the result. Returns
    // false on cancel; on a degenerate fit it records a failed axis and returns true.
    auto sweepAxis = [&](char axis, AxisMap& map, long a, long b, long tol) -> bool {
        LOG_INFO << "calibrating " << (axis == 'Y' ? "YAW" : "PITCH")
                 << " over [" << a << ".." << b << "] counts in " << N << " steps";
        log << "\n[" << axis << "] sweep=[" << a << ".." << b << "] tol=" << tol
            << "\n target_counts, imu_deg\n";

        std::vector<std::pair<double, double>> pts;  // (deg, counts)
        double prev_deg = 0.0;
        bool   have_prev = false;
        for (int i = 0; i < N; ++i) {
            if (cancel_) { LOG_WARN << "calibration cancelled"; return false; }
            long target = a + (b - a) * i / (N - 1);
            setProgress(axis, i + 1, N, "moving");
            motor_.sendCommand(std::string("MOVE ") + axis + " " + std::to_string(target));
            if (!waitSettle(motor_, axis, target, tol, params_.settle_timeout_ms, cancel_)) {
                if (cancel_) return false;
                LOG_WARN << "  position " << (i + 1) << "/" << N
                         << " did not settle near " << target << " (continuing)";
            }
            setProgress(axis, i + 1, N, "dwelling");
            double deg = dwellAndMeasure(imu_, axis, params_.dwell_ms, params_.sample_ms, cancel_);
            // Unwrap against the previous sample so a sweep crossing 0/360 stays a
            // continuous line for the fit (no ±360 jump).
            if (have_prev) deg = unwrapNear(prev_deg, deg);
            prev_deg = deg;
            have_prev = true;
            pts.emplace_back(deg, static_cast<double>(target));
            LOG_INFO << "  " << (i + 1) << "/" << N << "  counts=" << target
                     << "  imu=" << deg << " deg";
            log << "  " << target << ", " << deg << "\n";
        }

        setProgress(axis, N, N, "fitting");
        LinearFit fit = linearFit(pts);
        if (!fit.ok) {
            LOG_WARN << "  fit failed for axis " << axis << " (degenerate data)";
            log << "  FIT FAILED\n";
            report.axes.push_back({axis, false, 0, 0, 0, static_cast<int>(pts.size())});
            return true;
        }
        const double old_cpd = map.counts_per_deg;
        const long   old_zc  = map.zero_count;
        map.counts_per_deg = fit.slope;
        map.zero_count     = static_cast<long>(std::lround(fit.intercept));
        LOG_INFO << "  fit: counts_per_deg " << old_cpd << " -> " << fit.slope
                 << ", zero_count " << old_zc << " -> " << map.zero_count
                 << "  (R^2=" << fit.r2 << ")";
        log << "  fit: counts_per_deg=" << fit.slope << " zero_count=" << map.zero_count
            << " R2=" << fit.r2 << " (was cpd=" << old_cpd << " zc=" << old_zc << ")\n";
        report.axes.push_back({axis, true, fit.slope, map.zero_count, fit.r2, fit.n});
        return true;
    };

    // Resolve the yaw sweep span up front — we move to its first position before
    // the pitch sweep so the whole pitch calibration happens at that fixed yaw.
    long ya = 0, yb = 0, ytol = 0;
    if (!axisSpan('Y', ya, yb, ytol)) {
        LOG_WARN << "calibration aborted: no yaw soft limits in dump";
        return;
    }

    // 3. Move yaw to its first sweep position and hold it there.
    LOG_INFO << "calibration: moving yaw to first position (counts=" << ya << ")";
    setProgress('Y', 1, N, "moving");
    motor_.sendCommand(std::string("MOVE Y ") + std::to_string(ya));
    waitSettle(motor_, 'Y', ya, ytol, params_.settle_timeout_ms, cancel_);
    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }

    // 4. Calibrate PITCH at this yaw position, then 5. move pitch to 0 deg.
    if (pitch_present) {
        long pa = 0, pb = 0, ptol = 0;
        if (axisSpan('P', pa, pb, ptol)) {
            if (!sweepAxis('P', result.pitch, pa, pb, ptol)) { return; }  // cancelled
            const long pzero = result.pitch.toCounts(0.0);
            LOG_INFO << "calibration: moving pitch to 0 deg (counts=" << pzero << ")";
            setProgress('P', N, N, "pitch->0");
            motor_.sendCommand(std::string("MOVE P ") + std::to_string(pzero));
            waitSettle(motor_, 'P', pzero, ptol, params_.settle_timeout_ms, cancel_);
        } else {
            LOG_WARN << "calibration: no pitch soft limits in dump; skipping pitch";
        }
    }
    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }

    // 6. Switch the IMU to a no-magnetometer XKF profile so its heading stops
    //    chasing the (stepper-distorted) magnetic field. Selected from the device's
    //    own available list; persists on the device. Best-effort — if none is
    //    available we log and continue with the drift correction below.
    if (auto cfg = imu_.config()) {
        int nomag = pickNoMagProfile(cfg->available_profiles);
        if (nomag >= 0 && (!cfg->has_scenario || cfg->scenario_type != nomag)) {
            LOG_INFO << "calibration: switching IMU to no-mag XKF profile (type " << nomag << ")";
            setProgress('Y', 0, N, "no-mag profile");
            if (!imu_.setFilterProfile(nomag))
                LOG_WARN << "calibration: IMU did not confirm the no-mag profile (continuing)";
        } else if (nomag < 0) {
            LOG_WARN << "calibration: no no-mag XKF profile available on this device (continuing)";
        }
    }
    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }

    // 7. Hold still: correct yaw drift (no-rotation gyro-bias update) then reset
    //    the IMU heading so this pose is yaw 0 — done right before the yaw sweep so
    //    the drift accumulated during the long pitch sweep doesn't bias the fit.
    LOG_INFO << "calibration: correcting yaw drift (no-rotation " << params_.norotation_s
             << " s, holding still)...";
    setProgress('Y', 0, N, "yaw drift correction");
    imu_.noRotation(params_.norotation_s);
    for (int waited = 0; waited < params_.norotation_s * 1000 + 500 && !cancel_; waited += 100)
        msleep(100);
    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }
    LOG_INFO << "calibration: resetting IMU heading (this pose is now yaw 0)";
    setProgress('Y', 0, N, "yaw heading reset");
    imu_.headingReset();
    msleep(params_.reset_settle_ms);
    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }

    // 8. Calibrate YAW from this (now zero-heading) first position.
    if (!sweepAxis('Y', result.yaw, ya, yb, ytol)) { return; }  // cancelled

    if (cancel_) { LOG_WARN << "calibration cancelled"; return; }

    std::string path = paths::writeLogFile(paths::timestampedLogName("calib"), log.str());
    if (!path.empty()) LOG_INFO << "calibration log written: " << path;

    report.valid  = true;
    report.ts_ms  = nowMs();
    report.all_ok = !report.axes.empty();
    for (const auto& ax : report.axes) report.all_ok = report.all_ok && ax.ok;
    {
        std::lock_guard<std::mutex> lk(result_mutex_);
        result_ = result;
        report_ = report;
    }
    LOG_INFO << "=== gimbal calibration complete (applied to session) ===";
}

}  // namespace fgc
