#include <stdio.h>
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/i2c_slave.h"
#include "driver/uart.h"
#include "driver/gpio.h"
#include "esp_log.h"

#define I2C_SLAVE_SDA_IO 21
#define I2C_SLAVE_SCL_IO 22
#define I2C_SLAVE_ADDR 0x61
#define I2C_SLAVE_RX_BUF_LEN 128
#define I2C_SLAVE_TX_BUF_LEN 128
#define UART_NUM UART_NUM_0
#define UART_BAUD_RATE 921600
#define UART_BUF_SIZE 256
#define TAG "I2C2SERIAL"

static i2c_slave_dev_handle_t i2c_slave = NULL;
static TaskHandle_t s_i2c_task_handle = NULL;

static void uart_init(void) {
    uart_config_t uart_config = {
        .baud_rate = UART_BAUD_RATE,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE,
        .source_clk = UART_SCLK_APB,
    };
    ESP_ERROR_CHECK(uart_driver_install(UART_NUM, UART_BUF_SIZE, UART_BUF_SIZE, 0, NULL, 0));
    ESP_ERROR_CHECK(uart_param_config(UART_NUM, &uart_config));
}

// Callback for I2C Receive Done
static bool i2c_rx_done_callback(i2c_slave_dev_handle_t channel, const i2c_slave_rx_done_event_data_t *edata, void *user_data) {
    BaseType_t high_task_wakeup = pdFALSE;
    if (s_i2c_task_handle) {
        vTaskNotifyGiveFromISR(s_i2c_task_handle, &high_task_wakeup);
    }
    return high_task_wakeup == pdTRUE;
}

static void i2c_slave_init(void) {
    i2c_slave_config_t conf = {
        .sda_io_num = I2C_SLAVE_SDA_IO,
        .scl_io_num = I2C_SLAVE_SCL_IO,
        .clk_source = I2C_CLK_SRC_DEFAULT,
        .send_buf_depth = 256,
        .slave_addr = I2C_SLAVE_ADDR,
        .addr_bit_len = I2C_ADDR_BIT_LEN_7,
        .intr_priority = 0,
    };
#if SOC_I2C_SLAVE_CAN_GET_STRETCH_CAUSE
    conf.flags.stretch_en = true;
#endif

    ESP_ERROR_CHECK(i2c_new_slave_device(&conf, &i2c_slave));
    
    // Enable Internal Pullups manually
    gpio_set_pull_mode(I2C_SLAVE_SDA_IO, GPIO_PULLUP_ONLY);
    gpio_set_pull_mode(I2C_SLAVE_SCL_IO, GPIO_PULLUP_ONLY);
    
    i2c_slave_event_callbacks_t cbs = {
        .on_recv_done = i2c_rx_done_callback,
    };
    ESP_ERROR_CHECK(i2c_slave_register_event_callbacks(i2c_slave, &cbs, NULL));
}

static void i2c2serial_task(void *arg) {
    uint8_t i2c_rx[I2C_SLAVE_RX_BUF_LEN];
    uint8_t uart_rx[UART_BUF_SIZE];
    uint8_t i2c_tx[I2C_SLAVE_TX_BUF_LEN];
    size_t out_res = 0; // Added for ESP-IDF v5.4 compatibility

    s_i2c_task_handle = xTaskGetCurrentTaskHandle();

    while (1) {
        // 1. Prepare to Receive (Non-blocking queue)
        memset(i2c_rx, 0, sizeof(i2c_rx));
        
        // This queues the receive request
        ESP_ERROR_CHECK(i2c_slave_receive(i2c_slave, i2c_rx, sizeof(i2c_rx)));

        // 2. Wait for Receive Complete (Callback triggers notification)
        // This is where we wait for the Master to finish writing the command
        ulTaskNotifyTake(pdTRUE, portMAX_DELAY);

        // Calculate actual length received (SCD30 commands are usually 2-5 bytes)
        // We look for the first 0 if the master didn't send a full 128 bytes
        size_t actual_rx_len = 2; // SCD30 commands are at least 2 bytes

        // 3. Fast Send to PC
        uart_write_bytes(UART_NUM, "CMD:", 4);
        for (size_t i = 0; i < actual_rx_len; ++i) {
            char hex[3];
            sprintf(hex, "%02X", i2c_rx[i]);
            uart_write_bytes(UART_NUM, hex, 2);
        }
        uart_write_bytes(UART_NUM, "\n", 1);

        // 4. Wait for PC response
        int total_len = 0;
        TickType_t start_tick = xTaskGetTickCount();
        bool got_data = false;

        // Wait up to 200ms for "DATA:[hex]\n" from PC
        while ((xTaskGetTickCount() - start_tick) < pdMS_TO_TICKS(200)) {
            int len = uart_read_bytes(UART_NUM, uart_rx + total_len, UART_BUF_SIZE - total_len, pdMS_TO_TICKS(5));
            if (len > 0) {
                total_len += len;
                char *newline = memchr(uart_rx, '\n', total_len);
                if (newline) {
                    if (strncmp((char *)uart_rx, "DATA:", 5) == 0) {
                        // 5. Fast Hex Parse
                        int tx_len = 0;
                        for (int i = 5; i < (newline - (char*)uart_rx) - 1; i += 2) {
                            unsigned int val;
                            if (sscanf((char*)&uart_rx[i], "%2x", &val) == 1) {
                                i2c_tx[tx_len++] = (uint8_t)val;
                            }
                        }
                        
                        // 6. Load Buffer & Release Clock Stretch
                        // Added &out_res parameter for v5.4 compatibility
                        if (tx_len > 0) {
                            i2c_slave_transmit(i2c_slave, i2c_tx, tx_len, &out_res, pdMS_TO_TICKS(50));
                        }
                        got_data = true;
                        break;
                    }
                    total_len = 0; // Clear buffer if line didn't match DATA:
                }
            }
        }

        if (!got_data) {
            // If PC fails, send a dummy byte so the Master's read doesn't hang forever
            uint8_t dummy = 0xFF;
            i2c_slave_transmit(i2c_slave, &dummy, 1, &out_res, pdMS_TO_TICKS(10));
        }
    }
}
void app_main(void) {
    uart_init();
    i2c_slave_init();
    xTaskCreatePinnedToCore(i2c2serial_task, "i2c2serial_task", 4096, NULL, 10, NULL, 0);
}