#include "imu.h"

#include "board_config.h"
#include "touch.h"

#include "driver/i2c_master.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include <math.h>

static const char *TAG = "imu";

#if BOARD_IMU_ENABLED

#define QMI8658_ADDR_SDO_LOW     0x6A
#define QMI8658_ADDR_SDO_HIGH    0x6B
#define QMI8658_WHO_AM_I         0x00
#define QMI8658_WHO_AM_I_ID      0x05
#define QMI8658_CTRL1            0x02
#define QMI8658_CTRL2            0x03
#define QMI8658_CTRL3            0x04
#define QMI8658_CTRL5            0x06
#define QMI8658_CTRL6            0x07
#define QMI8658_CTRL7            0x08
#define QMI8658_RESET            0x60
#define QMI8658_ACC_X_L          0x35

#define QMI8658_RESET_VALUE      0xB0
/* Little-endian + address auto-increment (matches int16 parse below) */
#define QMI8658_CTRL1_VALUE      0x40
/* ±8g, 1000 Hz ODR */
#define QMI8658_CTRL2_VALUE      0x23
/* Gyro ±512 dps, 1000 Hz — required even when only reading accel */
#define QMI8658_CTRL3_VALUE      0x53
#define QMI8658_CTRL5_VALUE      0x11
#define QMI8658_CTRL6_VALUE      0x00
#define QMI8658_CTRL7_VALUE      0x03

#define QMI8658_ACCEL_SENS_8G    4096.0f

#define IMU_SAMPLE_PERIOD_MS     10
#define IMU_DEBUG_LOG_INTERVAL_MS 1000

static i2c_master_bus_handle_t s_imu_bus;
static i2c_master_dev_handle_t s_imu_dev;
static imu_sample_cb_t s_sample_cb;

static esp_err_t imu_read_raw(int16_t *rx, int16_t *ry, int16_t *rz);

static int16_t imu_raw16_le(const uint8_t *lo, const uint8_t *hi)
{
    return (int16_t)((*hi << 8) | *lo);
}

static esp_err_t imu_write_u8(uint8_t reg, uint8_t val)
{
    uint8_t buf[2] = {reg, val};
    return i2c_master_transmit(s_imu_dev, buf, sizeof(buf), 100);
}

static esp_err_t imu_read_u8(uint8_t reg, uint8_t *val)
{
    return i2c_master_transmit_receive(s_imu_dev, &reg, 1, val, 1, 100);
}

static esp_err_t imu_read_bytes(uint8_t reg, uint8_t *data, size_t len)
{
    return i2c_master_transmit_receive(s_imu_dev, &reg, 1, data, len, 100);
}

static esp_err_t imu_probe_device(uint8_t addr)
{
    i2c_device_config_t dev_cfg = {
        .dev_addr_length = I2C_ADDR_BIT_LEN_7,
        .device_address = addr,
        .scl_speed_hz = BOARD_IMU_I2C_FREQ_HZ,
    };

    esp_err_t err = i2c_master_bus_add_device(s_imu_bus, &dev_cfg, &s_imu_dev);
    if (err != ESP_OK) {
        return err;
    }

    uint8_t who_am_i = 0;
    err = imu_read_u8(QMI8658_WHO_AM_I, &who_am_i);
    if (err != ESP_OK || who_am_i != QMI8658_WHO_AM_I_ID) {
        i2c_master_bus_rm_device(s_imu_dev);
        s_imu_dev = NULL;
        if (err == ESP_OK) {
            ESP_LOGW(TAG, "WHO_AM_I mismatch at 0x%02X: 0x%02X", addr, who_am_i);
            return ESP_ERR_NOT_FOUND;
        }
        return err;
    }

    ESP_LOGI(TAG, "QMI8658 found at 0x%02X (WHO_AM_I=0x%02X)", addr, who_am_i);
    return ESP_OK;
}

static esp_err_t imu_configure(void)
{
    esp_err_t err = imu_write_u8(QMI8658_RESET, QMI8658_RESET_VALUE);
    if (err != ESP_OK) {
        return err;
    }
    vTaskDelay(pdMS_TO_TICKS(50));

    err = imu_write_u8(QMI8658_CTRL1, QMI8658_CTRL1_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    err = imu_write_u8(QMI8658_CTRL2, QMI8658_CTRL2_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    err = imu_write_u8(QMI8658_CTRL3, QMI8658_CTRL3_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    err = imu_write_u8(QMI8658_CTRL5, QMI8658_CTRL5_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    err = imu_write_u8(QMI8658_CTRL6, QMI8658_CTRL6_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    err = imu_write_u8(QMI8658_CTRL7, QMI8658_CTRL7_VALUE);
    if (err != ESP_OK) {
        return err;
    }

    vTaskDelay(pdMS_TO_TICKS(50));
    return ESP_OK;
}

static void imu_log_sample(const char *label, float ax, float ay, float az,
                           int16_t rx, int16_t ry, int16_t rz)
{
    float mag = sqrtf(ax * ax + ay * ay + az * az);
    ESP_LOGI(TAG, "%s raw=(%d,%d,%d) g=(%.3f,%.3f,%.3f) |mag|=%.3f",
             label, rx, ry, rz, ax, ay, az, mag);
}

static void imu_task(void *arg)
{
    (void)arg;

    int64_t last_log_us = esp_timer_get_time();

    while (1) {
        float ax = 0.0f;
        float ay = 0.0f;
        float az = 0.0f;
        int16_t rx = 0;
        int16_t ry = 0;
        int16_t rz = 0;

        esp_err_t err = imu_read_raw(&rx, &ry, &rz);
        if (err == ESP_OK) {
            ax = (float)rx / QMI8658_ACCEL_SENS_8G;
            ay = (float)ry / QMI8658_ACCEL_SENS_8G;
            az = (float)rz / QMI8658_ACCEL_SENS_8G;

            int64_t now_us = esp_timer_get_time();
            if ((now_us - last_log_us) >= (int64_t)IMU_DEBUG_LOG_INTERVAL_MS * 1000) {
                imu_log_sample("sample", ax, ay, az, rx, ry, rz);
                last_log_us = now_us;
            }

            if (s_sample_cb != NULL) {
                s_sample_cb(ax, ay, az);
            }
        }

        vTaskDelay(pdMS_TO_TICKS(IMU_SAMPLE_PERIOD_MS));
    }
}

esp_err_t imu_init(void)
{
#if BOARD_TOUCH_ENABLED
    s_imu_bus = touch_get_i2c_bus();
    if (s_imu_bus == NULL) {
        ESP_LOGE(TAG, "Shared I2C bus not ready (touch not initialized?)");
        return ESP_ERR_INVALID_STATE;
    }
    ESP_LOGI(TAG, "Using shared I2C bus (SDA=%d SCL=%d)", BOARD_PIN_IMU_SDA,
             BOARD_PIN_IMU_SCL);
#else
    ESP_LOGI(TAG, "Initialize dedicated I2C for IMU");
    i2c_master_bus_config_t bus_config = {
        .clk_source = I2C_CLK_SRC_DEFAULT,
        .i2c_port = BOARD_IMU_I2C_NUM,
        .sda_io_num = BOARD_PIN_IMU_SDA,
        .scl_io_num = BOARD_PIN_IMU_SCL,
        .glitch_ignore_cnt = 7,
        .flags.enable_internal_pullup = true,
    };
    esp_err_t bus_err = i2c_new_master_bus(&bus_config, &s_imu_bus);
    if (bus_err != ESP_OK) {
        return bus_err;
    }
#endif

    esp_err_t err = imu_probe_device(QMI8658_ADDR_SDO_LOW);
    if (err != ESP_OK) {
        err = imu_probe_device(QMI8658_ADDR_SDO_HIGH);
    }
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "QMI8658 not found");
        return err;
    }

    err = imu_configure();
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "IMU configure failed: %s", esp_err_to_name(err));
        return err;
    }

    float ax = 0.0f;
    float ay = 0.0f;
    float az = 0.0f;
    int16_t rx = 0;
    int16_t ry = 0;
    int16_t rz = 0;
    if (imu_read_accel(&ax, &ay, &az) == ESP_OK) {
        imu_read_raw(&rx, &ry, &rz);
        imu_log_sample("init check", ax, ay, az, rx, ry, rz);
    } else {
        ESP_LOGW(TAG, "Init check read failed");
    }

    ESP_LOGI(TAG, "IMU initialized (±8g, 1000 Hz ODR)");
    return ESP_OK;
}

static esp_err_t imu_read_raw(int16_t *rx, int16_t *ry, int16_t *rz)
{
    if (s_imu_dev == NULL || rx == NULL || ry == NULL || rz == NULL) {
        return ESP_ERR_INVALID_STATE;
    }

    uint8_t raw[6];
    esp_err_t err = imu_read_bytes(QMI8658_ACC_X_L, raw, sizeof(raw));
    if (err != ESP_OK) {
        return err;
    }

    *rx = imu_raw16_le(&raw[0], &raw[1]);
    *ry = imu_raw16_le(&raw[2], &raw[3]);
    *rz = imu_raw16_le(&raw[4], &raw[5]);
    return ESP_OK;
}

esp_err_t imu_read_accel(float *ax, float *ay, float *az)
{
    int16_t rx = 0;
    int16_t ry = 0;
    int16_t rz = 0;

    esp_err_t err = imu_read_raw(&rx, &ry, &rz);
    if (err != ESP_OK) {
        return err;
    }

    *ax = (float)rx / QMI8658_ACCEL_SENS_8G;
    *ay = (float)ry / QMI8658_ACCEL_SENS_8G;
    *az = (float)rz / QMI8658_ACCEL_SENS_8G;
    return ESP_OK;
}

void imu_start_task(imu_sample_cb_t cb)
{
    s_sample_cb = cb;
    xTaskCreate(imu_task, "imu_task", 3072, NULL, 4, NULL);
    ESP_LOGI(TAG, "IMU sample task started");
}

#else

esp_err_t imu_init(void)
{
    return ESP_OK;
}

esp_err_t imu_read_accel(float *ax, float *ay, float *az)
{
    (void)ax;
    (void)ay;
    (void)az;
    return ESP_ERR_NOT_SUPPORTED;
}

void imu_start_task(imu_sample_cb_t cb)
{
    (void)cb;
}

#endif