#include <Adafruit_BNO08x.h>
#include <Wire.h>
#include <math.h>
#include "imu.h"
#include "sensor_node_config.h"

namespace {
Adafruit_BNO08x bno086(-1);
sh2_SensorValue_t value;
bool initialized = false;
unsigned long lastInitAttemptMs = 0;
float radiansToDegrees(float r) { return r * 180.0f / PI; }
bool enableReports() {
    bool rotationEnabled = bno086.enableReport(SH2_ROTATION_VECTOR, IMU_REPORT_INTERVAL_US);
    bno086.enableReport(SH2_LINEAR_ACCELERATION, IMU_REPORT_INTERVAL_US);
    bno086.enableReport(SH2_GYROSCOPE_CALIBRATED, IMU_REPORT_INTERVAL_US);
    bno086.enableReport(SH2_GRAVITY, IMU_REPORT_INTERVAL_US);
    bno086.enableReport(SH2_STABILITY_CLASSIFIER, IMU_STABILITY_INTERVAL_US);
    return rotationEnabled;
}
bool tryInitImu() {
    lastInitAttemptMs = millis();
    initialized = bno086.begin_I2C(BNO086_I2C_ADDRESS, &Wire);
    if (initialized) initialized = enableReports();
    return initialized;
}
void updateEuler(ImuState& s) {
    float sinr = 2 * (s.real * s.i + s.j * s.k);
    float cosr = 1 - 2 * (s.i * s.i + s.j * s.j);
    s.rollDegrees = radiansToDegrees(atan2f(sinr, cosr));
    float sinp = 2 * (s.real * s.j - s.k * s.i);
    s.pitchDegrees = radiansToDegrees(fabsf(sinp) >= 1 ? copysignf(PI / 2, sinp) : asinf(sinp));
    float siny = 2 * (s.real * s.k + s.i * s.j);
    float cosy = 1 - 2 * (s.j * s.j + s.k * s.k);
    s.yawDegrees = radiansToDegrees(atan2f(siny, cosy));
}
String stabilityName(uint8_t value) {
    switch (value) {
        case STABILITY_CLASSIFIER_ON_TABLE: return "on_table";
        case STABILITY_CLASSIFIER_STATIONARY: return "stationary";
        case STABILITY_CLASSIFIER_STABLE: return "stable";
        case STABILITY_CLASSIFIER_MOTION: return "motion";
        default: return "unknown";
    }
}
}

bool initImu() {
    Wire.begin(BNO086_SDA_PIN, BNO086_SCL_PIN);
    Wire.setClock(BNO086_I2C_CLOCK_HZ);
    delay(100);
    Serial.print("Starting BNO086 at 0x");
    Serial.println(BNO086_I2C_ADDRESS, HEX);
    for (int attempt = 1; attempt <= BNO086_INIT_ATTEMPTS; attempt++) {
        if (tryInitImu()) return true;
        Serial.print("BNO086 init attempt failed: ");
        Serial.println(attempt);
        delay(BNO086_INIT_RETRY_DELAY_MS);
    }
    return false;
}

void maintainImu(ImuState& s) {
    if (initialized || millis() - lastInitAttemptMs < BNO086_RUNTIME_RETRY_MS) return;
    Serial.println("Retrying BNO086 initialization");
    s.online = tryInitImu();
    if (s.online) Serial.println("BNO086 recovered");
}

void updateImu(ImuState& s) {
    s.online = initialized;
    if (initialized && bno086.wasReset()) {
        initialized = enableReports();
        s.online = initialized;
        if (!initialized) return;
    }
    if (!initialized || !bno086.getSensorEvent(&value)) return;
    unsigned long now = millis();
    Vector3State* vector = nullptr;
    switch (value.sensorId) {
        case SH2_ROTATION_VECTOR:
            s.i = value.un.rotationVector.i; s.j = value.un.rotationVector.j;
            s.k = value.un.rotationVector.k; s.real = value.un.rotationVector.real;
            s.accuracy = value.status; s.valid = true; s.lastUpdateMs = now; updateEuler(s); return;
        case SH2_LINEAR_ACCELERATION:
            vector = &s.linearAcceleration;
            vector->x = value.un.linearAcceleration.x; vector->y = value.un.linearAcceleration.y; vector->z = value.un.linearAcceleration.z; break;
        case SH2_GYROSCOPE_CALIBRATED:
            vector = &s.angularVelocity;
            vector->x = value.un.gyroscope.x; vector->y = value.un.gyroscope.y; vector->z = value.un.gyroscope.z; break;
        case SH2_GRAVITY:
            vector = &s.gravity;
            vector->x = value.un.gravity.x; vector->y = value.un.gravity.y; vector->z = value.un.gravity.z; break;
        case SH2_STABILITY_CLASSIFIER:
            s.stabilityCode = value.un.stabilityClassifier.classification;
            s.stability = stabilityName(s.stabilityCode); s.stabilityLastUpdateMs = now; return;
        default: return;
    }
    vector->accuracy = value.status; vector->lastUpdateMs = now;
}