#include "mpu6886.h"
#include <string.h>


#ifdef CONFIG_ENV_MQTT_ENABLED
static void wifi_init()
{
    ESP_ERROR_CHECK(nvs_flash_init());
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

#ifdef CONFIG_WIFI_AP_MODE
    esp_netif_create_default_wifi_ap();

    wifi_config_t wifi_config = { 0 };
    strncpy((char*)wifi_config.ap.ssid, CONFIG_WIFI_AP_SSID, sizeof(wifi_config.ap.ssid));
    wifi_config.ap.ssid_len = strlen(CONFIG_WIFI_AP_SSID);
    strncpy((char*)wifi_config.ap.password, CONFIG_WIFI_AP_PASSWORD, sizeof(wifi_config.ap.password));
    wifi_config.ap.max_connection = 4;
    if (strlen(CONFIG_WIFI_AP_PASSWORD) == 0) {
        wifi_config.ap.authmode = WIFI_AUTH_OPEN;
    } else {
        wifi_config.ap.authmode = WIFI_AUTH_WPA_WPA2_PSK;
    }

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_LOGI("WIFI", "AP started SSID:%s", CONFIG_WIFI_AP_SSID);
#else
    esp_netif_create_default_wifi_sta();

    wifi_config_t wifi_config = {
        .sta = {
            .ssid = CONFIG_WIFI_SSID,
            .password = CONFIG_WIFI_PASSWORD,
        },
    };

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_ERROR_CHECK(esp_wifi_connect());
    ESP_LOGI("WIFI", "STA started, connecting to: %s", CONFIG_WIFI_SSID);

#endif
}

static void mqtt_app_start(void)
{
    esp_mqtt_client_config_t mqtt_cfg = { 0 };

    mqtt_cfg.broker.address.uri = CONFIG_MQTT_BROKER_URI;

    s_mqtt_client = esp_mqtt_client_init(&mqtt_cfg);
    if (s_mqtt_client == NULL) {
        ESP_LOGW("MQTT", "failed to init mqtt client");
        return;
    }
    esp_mqtt_client_start(s_mqtt_client);
}
#endif // CONFIG_ENV_MQTT_ENABLED

static esp_err_t mpu6886_write_byte(mpu6886_t* dev, uint8_t reg, uint8_t data) {
    uint8_t tx[2] = { reg, data };
    return i2c_master_write_to_device(dev->i2c_port, dev->address, tx, sizeof(tx), pdMS_TO_TICKS(100));
}

static esp_err_t mpu6886_read_bytes(mpu6886_t* dev, uint8_t reg, uint8_t* data, size_t len) {
    return i2c_master_write_read_device(dev->i2c_port, dev->address, &reg, 1, data, len, pdMS_TO_TICKS(100));
}

static int16_t bytes_to_int16(uint8_t high, uint8_t low) {
    return (int16_t)((high << 8) | low);
}

static esp_err_t mpu6886_update_sensitivity(mpu6886_t* dev)
{
    uint8_t aconf = 0, gconf = 0;
    esp_err_t err;

    err = mpu6886_read_bytes(dev, MPU6886_ACCEL_CONFIG, &aconf, 1);
    if (err != ESP_OK) return err;
    err = mpu6886_read_bytes(dev, MPU6886_GYRO_CONFIG, &gconf, 1);
    if (err != ESP_OK) return err;

    uint8_t a_fs = (aconf >> 3) & 0x03; // AFS_SEL bits [4:3]
    switch (a_fs) {
        case 0: dev->accel_div = ACCEL_SO_2G; break;   // ±2g
        case 1: dev->accel_div = ACCEL_SO_4G; break;   // ±4g
        case 2: dev->accel_div = ACCEL_SO_8G; break;   // ±8g
        case 3: dev->accel_div = ACCEL_SO_16G; break;  // ±16g
        default: dev->accel_div = ACCEL_SO_2G; break;
    }

    uint8_t g_fs = (gconf >> 3) & 0x03; // FS_SEL bits [4:3]
    switch (g_fs) {
        case 0: dev->gyro_div = GYRO_SO_250DPS; break;   // ±250 dps
        case 1: dev->gyro_div = GYRO_SO_500DPS; break;   // ±500 dps
        case 2: dev->gyro_div = GYRO_SO_1000DPS; break;  // ±1000 dps
        case 3: dev->gyro_div = GYRO_SO_2000DPS; break;  // ±2000 dps
        default: dev->gyro_div = GYRO_SO_250DPS; break;
    }

    return ESP_OK;
}

static esp_err_t i2c_master_init(i2c_port_t i2c_num, gpio_num_t sda_io, gpio_num_t scl_io, uint32_t clk_speed_hz)
{
    i2c_config_t conf = {
        .mode = I2C_MODE_MASTER,
        .sda_io_num = sda_io,
        .sda_pullup_en = GPIO_PULLUP_ENABLE,
        .scl_io_num = scl_io,
        .scl_pullup_en = GPIO_PULLUP_ENABLE,
        .master.clk_speed = clk_speed_hz,
    };
    esp_err_t err = i2c_param_config(i2c_num, &conf);
    if (err != ESP_OK) return err;
    return i2c_driver_install(i2c_num, I2C_MODE_MASTER, 0, 0, 0);
}

esp_err_t mpu6886_init(mpu6886_t* dev, i2c_port_t i2c_port) {
    dev->i2c_port = i2c_port;
    dev->address = MPU6886_ADDR;

    uint8_t who_am_i = 0;
    if (mpu6886_read_bytes(dev, MPU6886_WHO_AM_I, &who_am_i, 1) != ESP_OK) {
        return ESP_FAIL;
    }
    if (who_am_i != 0x19) {
        return ESP_FAIL;
    }

    // Reset
    if (mpu6886_write_byte(dev, MPU6886_PWR_MGMT_1, 0x80) != ESP_OK) return ESP_FAIL;
    vTaskDelay(pdMS_TO_TICKS(100));

    // Auto select clock
    if (mpu6886_write_byte(dev, MPU6886_PWR_MGMT_1, 0x01) != ESP_OK) return ESP_FAIL;

    // Default config: set to ±2G accel, ±250DPS gyro (device may have different settings; we'll detect them)
    mpu6886_write_byte(dev, MPU6886_ACCEL_CONFIG, 0x00);
    mpu6886_write_byte(dev, MPU6886_GYRO_CONFIG, 0x00);

    dev->gyro_offset = (vec3_t) { 0, 0, 0 };
    dev->accel_offset = (vec3_t) { 0, 0, 0 };

    // detect actual sensitivities from device registers and set divisors
    if (mpu6886_update_sensitivity(dev) != ESP_OK) {
        // fallback to defaults if read fails
        dev->accel_div = ACCEL_SO_2G;
        dev->gyro_div = GYRO_SO_250DPS;
    }

    return ESP_OK;
}

esp_err_t mpu6886_read_accel(mpu6886_t* dev, vec3_t* accel) {
    uint8_t buf[6];
    esp_err_t ret = mpu6886_read_bytes(dev, MPU6886_ACCEL_XOUT_H, buf, 6);
    if (ret != ESP_OK) return ret;

    // convert raw -> m/s^2 and apply stored offsets
    accel->x = (float)bytes_to_int16(buf[0], buf[1]) / dev->accel_div * SF_M_S2 - dev->accel_offset.x;
    accel->y = (float)bytes_to_int16(buf[2], buf[3]) / dev->accel_div * SF_M_S2 - dev->accel_offset.y;
    accel->z = (float)bytes_to_int16(buf[4], buf[5]) / dev->accel_div * SF_M_S2 - dev->accel_offset.z;
    return ESP_OK;
}

esp_err_t mpu6886_read_gyro(mpu6886_t* dev, vec3_t* gyro) {
    uint8_t buf[6];
    esp_err_t ret = mpu6886_read_bytes(dev, MPU6886_GYRO_XOUT_H, buf, 6);
    if (ret != ESP_OK) return ret;

    gyro->x = (float)bytes_to_int16(buf[0], buf[1]) / dev->gyro_div * SF_RAD_S - dev->gyro_offset.x;
    gyro->y = (float)bytes_to_int16(buf[2], buf[3]) / dev->gyro_div * SF_RAD_S - dev->gyro_offset.y;
    gyro->z = (float)bytes_to_int16(buf[4], buf[5]) / dev->gyro_div * SF_RAD_S - dev->gyro_offset.z;
    return ESP_OK;
}

esp_err_t mpu6886_read_temp(mpu6886_t* dev, float* temp) {
    static uint8_t buf[2];
    esp_err_t ret = mpu6886_read_bytes(dev, MPU6886_TEMP_OUT_H, buf, 2);
    if (ret != ESP_OK) return ret;

    int16_t raw = bytes_to_int16(buf[0], buf[1]);
    *temp = ((float)raw / TEMP_SO) + TEMP_OFFSET;
    return ESP_OK;
}

esp_err_t mpu6886_calibrate_gyro(mpu6886_t* dev, int samples, int delay_ms) {
    vec3_t sum = { 0, 0, 0 }, g;
    for (int i = 0; i < samples; i++) {
        if (mpu6886_read_gyro(dev, &g) != ESP_OK) return ESP_FAIL;
        sum.x += g.x;
        sum.y += g.y;
        sum.z += g.z;
        vTaskDelay(pdMS_TO_TICKS(delay_ms));
    }
    dev->gyro_offset.x = sum.x / samples;
    dev->gyro_offset.y = sum.y / samples;
    dev->gyro_offset.z = sum.z / samples;
    return ESP_OK;
}

esp_err_t mpu6886_calibrate_accel(mpu6886_t* dev, int samples, int delay_ms)
{
    // Calibrate accelerometer offsets while the device is stationary.
    // The function will detect the device full-scale and compute offsets in m/s^2.
    vec3_t sum = { 0, 0, 0 }, a;
    for (int i = 0; i < samples; i++) {
        if (mpu6886_read_accel(dev, &a) != ESP_OK) {
            return ESP_FAIL;
        }
        sum.x += a.x;
        sum.y += a.y;
        sum.z += a.z;
        vTaskDelay(pdMS_TO_TICKS(delay_ms));
    }
    vec3_t avg = { sum.x / samples, sum.y / samples, sum.z / samples };

    // For X/Y expect ~0 m/s^2 when stationary; offset = measured average
    dev->accel_offset.x = avg.x;
    dev->accel_offset.y = avg.y;

    // For Z expect ~+1g (SF_M_S2) if +Z points up. If device is flipped you'll get -1g.
    // Compute expected gravity sign from measured avg.z magnitude and sign.
    float expected_g = (avg.z < 0) ? -SF_M_S2 : SF_M_S2;
    dev->accel_offset.z = avg.z - expected_g;

    return ESP_OK;
}

void app_main(void)
{
    esp_err_t err = i2c_master_init(I2C_NUM_0, CONFIG_I2C_MASTER_SDA, CONFIG_I2C_MASTER_SCL, CONFIG_I2C_MASTER_FREQUENCY); // adjust pins if needed
    if (err != ESP_OK) {
        ESP_LOGE("MPU6886", "I2C init failed: %d", err);
        return;
    }


#ifdef CONFIG_ENV_MQTT_ENABLED
    wifi_init();
    // Start MQTT client (will retry until network is available)
    mqtt_app_start();
#endif


    mpu6886_t mpu;
    mpu.i2c_port = I2C_NUM_0;
    mpu.address = MPU6886_ADDR;

    esp_err_t ret = mpu6886_init(&mpu, I2C_NUM_0);
    if (ret != ESP_OK) {
        ESP_LOGE("MPU6886", "init failed");
        return;
    }



    mpu6886_calibrate_gyro(&mpu, 100, 10);
    mpu6886_calibrate_accel(&mpu, 100, 10);
    vec3_t accel, gyro;
    float temp;

    while (1) {
        mpu6886_read_accel(&mpu, &accel);
        mpu6886_read_gyro(&mpu, &gyro);
        mpu6886_read_temp(&mpu, &temp);

#ifndef CONFIG_ENV_MQTT_ENABLED
        printf("{\"accel\":{\"x\":%8.3f,\"y\":%8.3f,\"z\":%8.3f},\"gyro\":{\"x\":%8.3f,\"y\":%8.3f,\"z\":%8.3f},\"Temp\":%8.2f}\n",
           accel.x, accel.y, accel.z,
           gyro.x, gyro.y, gyro.z,
           temp);
#else
        if (s_mqtt_client != NULL) {
            char payload[256];
            int n = snprintf(payload, sizeof(payload),
                "{\"accel\":{\"x\":%8.3f,\"y\":%8.3f,\"z\":%8.3f},\"gyro\":{\"x\":%8.3f,\"y\":%8.3f,\"z\":%8.3f},\"Temp\":%8.2f}",
                accel.x, accel.y, accel.z,
                gyro.x, gyro.y, gyro.z,
                temp);
            if (n > 0 && n < (int)sizeof(payload)) {
                int msg_id = esp_mqtt_client_publish(s_mqtt_client, CONFIG_MQTT_TOPIC, payload, 0, 1, 0);
                ESP_LOGD("MQTT", "published msg_id=%d payload_len=%d", msg_id, n);
            } else {
                ESP_LOGW("MQTT", "payload truncated or encoding error");
            }
        }
#endif
        vTaskDelay(pdMS_TO_TICKS(500));
    }
}