#include <stdbool.h>
#include <stdio.h>
#include <stdatomic.h>

#include "adxl345.h"
#include "driver/gpio.h"
#include "driver/i2c_master.h"
#include "driver/usb_serial_jtag_vfs.h"
#include "esp_err.h"
#include "esp_log.h"
#include "esp_timer.h"
#include "freertos/FreeRTOS.h"
#include "freertos/queue.h"
#include "freertos/task.h"
#include "l3g4200d.h"
#include "trikke_protocol.h"
#include "trikke_transport.h"
#if CONFIG_TRIKKE_TRANSPORT_BLE
#include "trikke_ble_transport.h"
#else
#include "trikke_usb_transport.h"
#endif

// Seeed Studio XIAO ESP32-C3: D4/SDA = GPIO6, D5/SCL = GPIO7.
#define TRIKKE_I2C_PORT       I2C_NUM_0
#define TRIKKE_I2C_SDA_GPIO   GPIO_NUM_6
#define TRIKKE_I2C_SCL_GPIO   GPIO_NUM_7
#define TRIKKE_I2C_FREQ_HZ    400000
#define TRIKKE_SAMPLE_RATE_HZ 100
#define TRIKKE_SAMPLE_TICKS   pdMS_TO_TICKS(1000 / TRIKKE_SAMPLE_RATE_HZ)
#define TRIKKE_SAMPLE_QUEUE_DEPTH 512
#define TRIKKE_METADATA_INTERVAL_PACKETS 64
#define TRIKKE_STATUS_INTERVAL_PACKETS   64
#define TRIKKE_TRANSPORT_RETRY_DELAY_MS 10

// Software calibration from the 2026-08-17 enclosure six-face capture.
// Accelerometer coefficients are measured. Gyroscope scale is nominal; its
// zero-rate biases and all three axis polarities were measured.
#define TRIKKE_ACCEL_X_OFFSET_COUNTS (-1.417678f)
#define TRIKKE_ACCEL_Y_OFFSET_COUNTS (-4.400892f)
#define TRIKKE_ACCEL_Z_OFFSET_COUNTS (11.776132f)
#define TRIKKE_ACCEL_X_COUNTS_PER_G  (258.890661f)
#define TRIKKE_ACCEL_Y_COUNTS_PER_G  (259.825135f)
#define TRIKKE_ACCEL_Z_COUNTS_PER_G  (245.755573f)
#define TRIKKE_GYRO_X_BIAS_COUNTS    (9.0482f)
#define TRIKKE_GYRO_Y_BIAS_COUNTS    (-153.6198f)
#define TRIKKE_GYRO_Z_BIAS_COUNTS    (-7.0238f)
#define TRIKKE_GYRO_MDPS_PER_LSB     (17.5f)

static const char *TAG = "trikke";

typedef struct {
    int32_t x;
    int32_t y;
    int32_t z;
} trikke_axes_sample_t;

typedef struct {
    adxl345_t accelerometer;
    l3g4200d_t gyroscope;
    QueueHandle_t sample_queue;
    atomic_uint_least32_t sensor_read_failure_count;
    atomic_uint_least32_t queue_overflow_count;
    atomic_uint_least32_t loop_overrun_count;
    trikke_transport_t transport;
#if CONFIG_TRIKKE_TRANSPORT_BLE
    trikke_ble_transport_t ble_transport;
#else
    trikke_usb_transport_t usb_transport;
#endif
} trikke_context_t;

static trikke_context_t s_context;

static const trikke_wire_metadata_t TRIKKE_METADATA = {
    .sample_rate_hz = TRIKKE_SAMPLE_RATE_HZ,
    .accel_range_g = 8,
    .gyro_range_dps = 500,
    .flags = TRIKKE_METADATA_FLAG_ACCEL_Y_NEGX_Z |
             TRIKKE_METADATA_FLAG_GYRO_IDENTITY |
             TRIKKE_METADATA_FLAG_GYRO_POLARITY |
             TRIKKE_METADATA_FLAG_GYRO_SCALE_NOMINAL,
    .accel_offset_counts = {
        TRIKKE_ACCEL_X_OFFSET_COUNTS,
        TRIKKE_ACCEL_Y_OFFSET_COUNTS,
        TRIKKE_ACCEL_Z_OFFSET_COUNTS,
    },
    .accel_counts_per_g = {
        TRIKKE_ACCEL_X_COUNTS_PER_G,
        TRIKKE_ACCEL_Y_COUNTS_PER_G,
        TRIKKE_ACCEL_Z_COUNTS_PER_G,
    },
    .gyro_bias_counts = {
        TRIKKE_GYRO_X_BIAS_COUNTS,
        TRIKKE_GYRO_Y_BIAS_COUNTS,
        TRIKKE_GYRO_Z_BIAS_COUNTS,
    },
    .gyro_mdps_per_lsb = TRIKKE_GYRO_MDPS_PER_LSB,
};

static trikke_axes_sample_t map_accel_to_enclosure(const adxl345_sample_t *native)
{
    // Enclosure frame: +X right, +Y toward the top, +Z toward the cover.
    // Mounted ADXL345: native +Y right, native +X down, native +Z toward cover.
    return (trikke_axes_sample_t) {
        .x = native->y,
        .y = -(int32_t)native->x,
        .z = native->z,
    };
}

static trikke_axes_sample_t map_gyro_to_enclosure(const l3g4200d_sample_t *native)
{
    // The mounted L3G4200D axes already match the enclosure frame.
    return (trikke_axes_sample_t) {
        .x = native->x,
        .y = native->y,
        .z = native->z,
    };
}

static uint32_t dropped_sample_count(const trikke_context_t *context)
{
    return atomic_load(&context->sensor_read_failure_count) +
           atomic_load(&context->queue_overflow_count);
}

static trikke_wire_status_t status_snapshot(
    trikke_context_t *context,
    const trikke_transport_sender_t *sender)
{
    trikke_wire_status_t status = {
        .sensor_read_failure_count =
            atomic_load(&context->sensor_read_failure_count),
        .queue_overflow_count = atomic_load(&context->queue_overflow_count),
        .transport_begin_retry_count = sender->begin_retry_count,
    };
#if CONFIG_TRIKKE_TRANSPORT_BLE
    trikke_ble_transport_counters_t ble_counters = {0};
    trikke_ble_transport_get_counters(
        &context->ble_transport, &ble_counters);
    status.transport_disconnect_count = ble_counters.disconnect_count;
    status.transport_send_failure_count = ble_counters.send_failure_count;
    status.transport_replay_count = ble_counters.replay_count;
    status.transport_invalid_ack_count = ble_counters.invalid_ack_count;
#endif
    return status;
}

static void acquisition_task(void *argument)
{
    trikke_context_t *context = argument;
    ulTaskNotifyTake(pdTRUE, portMAX_DELAY);

    uint32_t sequence = 0;
    TickType_t last_wake = xTaskGetTickCount();
    while (true) {
        adxl345_sample_t accel = {0};
        l3g4200d_sample_t gyro = {0};
        uint8_t accel_status = 0;
        uint8_t gyro_status = 0;
        const int64_t timestamp_us = esp_timer_get_time();

        const esp_err_t accel_err =
            adxl345_read_raw(&context->accelerometer, &accel, &accel_status);
        const esp_err_t gyro_err =
            l3g4200d_read_raw(&context->gyroscope, &gyro, &gyro_status);

        if (accel_err == ESP_OK && gyro_err == ESP_OK) {
            const trikke_axes_sample_t enclosure_accel =
                map_accel_to_enclosure(&accel);
            const trikke_axes_sample_t enclosure_gyro =
                map_gyro_to_enclosure(&gyro);
            const trikke_wire_sample_t sample = {
                .sequence = sequence,
                .timestamp_us = timestamp_us,
                .accel_x = (int16_t)enclosure_accel.x,
                .accel_y = (int16_t)enclosure_accel.y,
                .accel_z = (int16_t)enclosure_accel.z,
                .gyro_x = (int16_t)enclosure_gyro.x,
                .gyro_y = (int16_t)enclosure_gyro.y,
                .gyro_z = (int16_t)enclosure_gyro.z,
                .accel_status = accel_status,
                .gyro_status = gyro_status,
            };
            if (xQueueSend(context->sample_queue, &sample, 0) != pdPASS) {
                atomic_fetch_add(&context->queue_overflow_count, 1);
            }
        } else {
            atomic_fetch_add(&context->sensor_read_failure_count, 1);
        }

        ++sequence;
        if (xTaskDelayUntil(&last_wake, TRIKKE_SAMPLE_TICKS) == pdFALSE) {
            atomic_fetch_add(&context->loop_overrun_count, 1);
            last_wake = xTaskGetTickCount();
            vTaskDelay(1);
        }
    }
}

static void write_binary_packet_until_sent(
    trikke_context_t *context,
    trikke_transport_sender_t *sender,
    const uint8_t *packet,
    size_t packet_size)
{
    while (true) {
        const bool was_pending = sender->pending;
        const trikke_transport_status_t status =
            trikke_transport_sender_step(
                sender, &context->transport, packet, packet_size);
        if (status == TRIKKE_TRANSPORT_COMPLETE) {
            return;
        }
        if (status == TRIKKE_TRANSPORT_FATAL) {
            // The binary stream is terminal at this point, so a final text
            // diagnostic cannot corrupt later frames. Preserve the packet and
            // stop consuming the sample queue.
            esp_log_level_set(TAG, ESP_LOG_ERROR);
            ESP_LOGE(TAG, "fatal transport invariant; output task suspended");
            while (true) {
                vTaskSuspend(NULL);
            }
        }
        if (status == TRIKKE_TRANSPORT_PENDING && !was_pending) {
            // Poll once immediately after acceptance. Subsequent pending polls
            // are paced so a future nonblocking backend cannot busy-spin.
            continue;
        }
        vTaskDelay(pdMS_TO_TICKS(TRIKKE_TRANSPORT_RETRY_DELAY_MS));
    }
}

static void output_task(void *argument)
{
    trikke_context_t *context = argument;
    ulTaskNotifyTake(pdTRUE, portMAX_DELAY);

    uint8_t packet[TRIKKE_WIRE_MAX_PACKET_SIZE] = {0};
    uint32_t packet_sequence = 0;
    uint32_t sample_packet_count = 0;
    trikke_transport_sender_t sender;
    trikke_transport_sender_init(&sender);

    size_t packet_size = trikke_encode_metadata_packet(
        packet, sizeof(packet), packet_sequence++, esp_timer_get_time(),
        dropped_sample_count(context),
        atomic_load(&context->loop_overrun_count), &TRIKKE_METADATA);
    write_binary_packet_until_sent(context, &sender, packet, packet_size);

    trikke_wire_status_t status = status_snapshot(context, &sender);
    packet_size = trikke_encode_status_packet(
        packet, sizeof(packet), packet_sequence++, esp_timer_get_time(),
        dropped_sample_count(context),
        atomic_load(&context->loop_overrun_count), &status);
    write_binary_packet_until_sent(context, &sender, packet, packet_size);

    while (true) {
        trikke_wire_sample_t samples[TRIKKE_WIRE_MAX_RECORDS] = {0};
        size_t sample_count = 0;
        if (xQueueReceive(context->sample_queue, &samples[sample_count],
                          portMAX_DELAY) != pdPASS) {
            continue;
        }
        ++sample_count;

        // output_task is the queue's sole consumer. This peek-then-receive
        // sequence relies on that invariant; transports must not dequeue here.
        while (sample_count < TRIKKE_WIRE_MAX_RECORDS) {
            trikke_wire_sample_t next_sample = {0};
            if (xQueuePeek(context->sample_queue, &next_sample,
                           pdMS_TO_TICKS(15)) != pdPASS) {
                break;
            }
            if (!trikke_wire_timestamp_delta_fits(
                    samples[sample_count - 1].timestamp_us,
                    next_sample.timestamp_us)) {
                break;
            }
            if (xQueueReceive(context->sample_queue, &samples[sample_count], 0) !=
                pdPASS) {
                break;
            }
            ++sample_count;
        }

        if (sample_packet_count > 0 &&
            sample_packet_count % TRIKKE_METADATA_INTERVAL_PACKETS == 0) {
            packet_size = trikke_encode_metadata_packet(
                packet, sizeof(packet), packet_sequence++, esp_timer_get_time(),
                dropped_sample_count(context),
                atomic_load(&context->loop_overrun_count), &TRIKKE_METADATA);
            write_binary_packet_until_sent(
                context, &sender, packet, packet_size);
        }

        if (sample_packet_count > 0 &&
            sample_packet_count % TRIKKE_STATUS_INTERVAL_PACKETS == 0) {
            status = status_snapshot(context, &sender);
            packet_size = trikke_encode_status_packet(
                packet, sizeof(packet), packet_sequence++, esp_timer_get_time(),
                dropped_sample_count(context),
                atomic_load(&context->loop_overrun_count), &status);
            write_binary_packet_until_sent(
                context, &sender, packet, packet_size);
        }

        packet_size = trikke_encode_sample_packet(
            packet, sizeof(packet), packet_sequence++,
            dropped_sample_count(context),
            atomic_load(&context->loop_overrun_count), samples, sample_count);
        write_binary_packet_until_sent(context, &sender, packet, packet_size);
        ++sample_packet_count;
    }
}

static esp_err_t init_i2c(i2c_master_bus_handle_t *bus)
{
    const i2c_master_bus_config_t config = {
        .i2c_port = TRIKKE_I2C_PORT,
        .sda_io_num = TRIKKE_I2C_SDA_GPIO,
        .scl_io_num = TRIKKE_I2C_SCL_GPIO,
        .clk_source = I2C_CLK_SRC_DEFAULT,
        .glitch_ignore_cnt = 7,
        .flags.enable_internal_pullup = true,
    };

    return i2c_new_master_bus(&config, bus);
}

void app_main(void)
{
    // Keep the byte stream unbuffered so a write result describes the complete
    // packet rather than bytes still retained inside stdio.
    if (setvbuf(stdout, NULL, _IONBF, 0) != 0) {
        ESP_LOGE(TAG, "failed to configure unbuffered telemetry output");
        return;
    }

    ESP_LOGI(TAG, "Trikke motion telemetry prototype v0");
    ESP_LOGI(TAG, "I2C: SDA=GPIO%d, SCL=GPIO%d, clock=%d Hz",
             TRIKKE_I2C_SDA_GPIO, TRIKKE_I2C_SCL_GPIO, TRIKKE_I2C_FREQ_HZ);

    i2c_master_bus_handle_t bus = NULL;
    esp_err_t err = init_i2c(&bus);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "I2C initialization failed: %s", esp_err_to_name(err));
        return;
    }

    err = adxl345_init(&s_context.accelerometer, bus, TRIKKE_I2C_FREQ_HZ);
    if (err != ESP_OK) {
        ESP_LOGE(TAG,
                 "ADXL345 not found at 0x53 or 0x1D (expected DEVID 0xE5): %s",
                 esp_err_to_name(err));
        i2c_del_master_bus(bus);
        return;
    }
    ESP_LOGI(TAG, "ADXL345 detected at 0x%02X; 100 Hz, +/-8 g, full resolution",
             adxl345_address(&s_context.accelerometer));

    err = l3g4200d_init(&s_context.gyroscope, bus, TRIKKE_I2C_FREQ_HZ);
    if (err != ESP_OK) {
        ESP_LOGE(TAG,
                 "L3G4200D not found at 0x69 or 0x68 (expected WHO_AM_I 0xD3): %s",
                 esp_err_to_name(err));
        adxl345_deinit(&s_context.accelerometer);
        i2c_del_master_bus(bus);
        return;
    }
    ESP_LOGI(TAG, "L3G4200D detected at 0x%02X; 100 Hz, 25 Hz BW, +/-500 dps",
             l3g4200d_address(&s_context.gyroscope));

    // Discard the gyroscope's visible startup transient before beginning the stream.
    vTaskDelay(pdMS_TO_TICKS(500));

    printf("# format=trikke_calibrated_v3\n");
    printf("# accel_calibration=offset_counts:(%.6f,%.6f,%.6f),"
           "counts_per_g:(%.6f,%.6f,%.6f)\n",
           TRIKKE_ACCEL_X_OFFSET_COUNTS, TRIKKE_ACCEL_Y_OFFSET_COUNTS,
           TRIKKE_ACCEL_Z_OFFSET_COUNTS, TRIKKE_ACCEL_X_COUNTS_PER_G,
           TRIKKE_ACCEL_Y_COUNTS_PER_G, TRIKKE_ACCEL_Z_COUNTS_PER_G);
    printf("# gyro_calibration=bias_counts:(%.4f,%.4f,%.4f),"
           "nominal_mdps_per_lsb:%.1f,polarity:verified\n",
           TRIKKE_GYRO_X_BIAS_COUNTS, TRIKKE_GYRO_Y_BIAS_COUNTS,
           TRIKKE_GYRO_Z_BIAS_COUNTS, TRIKKE_GYRO_MDPS_PER_LSB);
    printf("# enclosure_axes=+x:right,+y:top,+z:toward_cover\n");
    printf("# mapping=accel(x,y,z)=(native_y,-native_x,native_z);"
           "gyro(x,y,z)=(native_x,native_y,native_z)\n");
    printf("# status_masks=accel_data_ready:0x80,accel_overrun:0x01,"
           "gyro_data_ready:0x08,gyro_overrun:0x80\n");
    printf("# binary_stream=TRK1,wire_version=%d,record_size=%d,"
           "records_per_packet=%d\n",
           TRIKKE_WIRE_VERSION, TRIKKE_WIRE_SAMPLE_RECORD_SIZE,
           TRIKKE_WIRE_MAX_RECORDS);
#if CONFIG_TRIKKE_TRANSPORT_BLE
    printf("# transport=ble,device_name=TrikkeSensor\n");
#else
    printf("# transport=usb_serial_jtag\n");
#endif
    fflush(stdout);

#if !CONFIG_TRIKKE_TRANSPORT_BLE
    // The console defaults to CRLF conversion, which would insert bytes into
    // binary frames whenever a payload byte equals LF.
    usb_serial_jtag_vfs_set_tx_line_endings(ESP_LINE_ENDINGS_LF);
#endif

    s_context.sample_queue =
        xQueueCreate(TRIKKE_SAMPLE_QUEUE_DEPTH, sizeof(trikke_wire_sample_t));
    if (s_context.sample_queue == NULL) {
        ESP_LOGE(TAG, "sample queue allocation failed");
        l3g4200d_deinit(&s_context.gyroscope);
        adxl345_deinit(&s_context.accelerometer);
        i2c_del_master_bus(bus);
        return;
    }

    TaskHandle_t output_task_handle = NULL;
    TaskHandle_t acquisition_task_handle = NULL;
    if (xTaskCreate(output_task, "trikke_output", 4096, &s_context, 5,
                    &output_task_handle) != pdPASS ||
        xTaskCreate(acquisition_task, "trikke_acquire", 4096, &s_context, 10,
                    &acquisition_task_handle) != pdPASS) {
        if (output_task_handle != NULL) {
            vTaskDelete(output_task_handle);
        }
        if (acquisition_task_handle != NULL) {
            vTaskDelete(acquisition_task_handle);
        }
        vQueueDelete(s_context.sample_queue);
        ESP_LOGE(TAG, "telemetry task creation failed");
        l3g4200d_deinit(&s_context.gyroscope);
        adxl345_deinit(&s_context.accelerometer);
        i2c_del_master_bus(bus);
        return;
    }

#if CONFIG_TRIKKE_TRANSPORT_BLE
    err = trikke_ble_transport_init(
        &s_context.ble_transport, &s_context.transport);
#else
    err = trikke_usb_transport_init(
        &s_context.usb_transport, &s_context.transport);
#endif
    if (err != ESP_OK) {
        vTaskDelete(output_task_handle);
        vTaskDelete(acquisition_task_handle);
        vQueueDelete(s_context.sample_queue);
#if CONFIG_TRIKKE_TRANSPORT_BLE
        ESP_LOGE(TAG, "BLE transport initialization failed: %s",
                 esp_err_to_name(err));
#else
        ESP_LOGE(TAG, "USB transport initialization failed: %s",
                 esp_err_to_name(err));
#endif
        l3g4200d_deinit(&s_context.gyroscope);
        adxl345_deinit(&s_context.accelerometer);
        i2c_del_master_bus(bus);
        return;
    }

    // No text may share the byte stream once framed binary output begins.
    esp_log_level_set("*", ESP_LOG_NONE);
    xTaskNotifyGive(output_task_handle);
    xTaskNotifyGive(acquisition_task_handle);
}