#include "fgc/PitchSetup.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

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

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

bool PitchSetupRoutine::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 PitchSetupRoutine::cancel() { cancel_ = true; }

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

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

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

void PitchSetupRoutine::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 mean measured pitch in degrees.
// A plain arithmetic mean is right here: pitch is a gravity-referenced elevation
// over a ~40 deg range, so it cannot wrap the way a heading can.
double dwellAndMeasure(IImuSource& imu, 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(s->pitch_deg);
        msleep(sample_ms);
    }
    if (vals.empty()) return 0.0;
    double m = 0;
    for (double v : vals) m += v;
    return m / vals.size();
}

}  // namespace

void PitchSetupRoutine::run() {
    struct Done {
        PitchSetupRoutine* 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 pitch setup starting ===";
    setProgress('?', 0, N, "starting");
    CalibReport report;

    // 1. Preconditions.
    if (!imu_.connected()) {
        LOG_WARN << "pitch setup aborted: IMU not connected (declare [Hardware] imu)";
        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 << "pitch setup: gimbal not homed; homing first...";
        setProgress('?', 0, N, "homing");
        motor_.sendCommand("HOME");
        if (!waitHomed(motor_, params_.home_timeout_ms, cancel_)) {
            if (cancel_) { LOG_WARN << "pitch setup cancelled"; return; }
            LOG_WARN << "pitch setup 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 << "pitch setup: 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 << "pitch setup aborted: could not read soft limits";
        return;
    }

    Geometry result = initial_;
    std::ostringstream log;
    log << "gimbal pitch setup\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 << "measuring 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)
        for (int i = 0; i < N; ++i) {
            if (cancel_) { LOG_WARN << "pitch setup 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");
            const double deg = dwellAndMeasure(imu_, params_.dwell_ms, params_.sample_ms, cancel_);
            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 double old_dac = map.deg_at_center;
        // The fit gives counts as a function of degrees; the map is anchored on
        // the homed centre, so what we store is the angle that centre sits at.
        map.counts_per_deg = fit.slope;
        map.deg_at_center =
            (static_cast<double>(map.center_count) - fit.intercept) / fit.slope;
        LOG_INFO << "  fit: counts_per_deg " << old_cpd << " -> " << fit.slope
                 << ", deg_at_center " << old_dac << " -> " << map.deg_at_center
                 << "  (R^2=" << fit.r2 << ")";
        log << "  fit: counts_per_deg=" << fit.slope << " deg_at_center=" << map.deg_at_center
            << " R2=" << fit.r2 << " (was cpd=" << old_cpd << " dac=" << old_dac << ")\n";
        report.axes.push_back({axis, true, fit.slope, map.deg_at_center, fit.r2, fit.n});
        return true;
    };

    // Sweep pitch across its inset travel, dwelling at each step to record the
    // gravity-referenced IMU pitch, and fit. Yaw is NOT touched: it is tied to
    // north by `gimbal heading setup` against real landmarks, which needs no IMU
    // and does not inherit the MTi's heading drift.
    if (!pitch_present) {
        LOG_WARN << "pitch setup aborted: this rig reports no pitch axis";
        return;
    }
    long pa = 0, pb = 0, ptol = 0;
    if (!axisSpan('P', pa, pb, ptol)) {
        LOG_WARN << "pitch setup aborted: no pitch soft limits in dump";
        return;
    }
    if (!sweepAxis('P', result.pitch, pa, pb, ptol)) return;  // cancelled

    // Park level, using the map just fitted - a visible confirmation that it is
    // right, and a sane place to leave the axis.
    const long plevel = result.pitch.toCounts(0.0);
    LOG_INFO << "pitch setup: moving to 0 deg / level (counts=" << plevel << ")";
    setProgress('P', N, N, "pitch->level");
    motor_.sendCommand(std::string("MOVE P ") + std::to_string(plevel));
    waitSettle(motor_, 'P', plevel, ptol, params_.settle_timeout_ms, cancel_);
    if (cancel_) { LOG_WARN << "pitch setup cancelled"; return; }

    std::string path = paths::writeLogFile(paths::timestampedLogName("pitch_setup"), log.str());
    if (!path.empty()) LOG_INFO << "pitch setup 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 pitch setup complete (applied to session) ===";
}

}  // namespace fgc