#include <stdio.h>
#include <string.h>

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "esp_log.h"

#include "esp_lcd_panel_io.h"
#include "esp_lcd_panel_ops.h"
#include "esp_lcd_panel_vendor.h"

#include "driver/i2c_master.h"
#include "driver/uart.h"

static const char *TAG = "PMME-Device";

#define I2C_HOST            0
#define I2C_SDA_PIN         21
#define I2C_SCL_PIN         22
#define I2C_CLOCK_HZ        400000    // Reduced from 400000 to a safer 100KHz
#define SSD1306_ADDR        0x3C
#define SSD1306_WIDTH       128
#define SSD1306_HEIGHT      64

#define UART_NUM UART_NUM_2
#define BUF_SIZE 32
#define UART_BUF_SIZE 128
#define TXD_PIN 17
#define RXD_PIN 16

#define SWAP_BYTES(x) ((x >> 8) | (x << 8))

typedef struct __attribute__((packed)) {
    // PM concentrations (CF=1, standard particle) in μg/m³
    uint16_t pm1_0_cf1;
    uint16_t pm2_5_cf1;
    uint16_t pm10_cf1;

    // PM concentrations (atmospheric environment) in μg/m³
    uint16_t pm1_0_atm;
    uint16_t pm2_5_atm;
    uint16_t pm10_atm;

    // Particle counts (particles per 0.1L air)
    uint16_t particles_0_3um;
    uint16_t particles_0_5um;
    uint16_t particles_1_0um;
    uint16_t particles_2_5um;
    uint16_t particles_5_0um;
    uint16_t particles_10um;
} pms7003_data_t;

void send_pms7003_command(uint8_t cmd, uint8_t datah, uint8_t datal) {
    uint8_t command[7] = {0x42, 0x4D, cmd, datah, datal, 0, 0};

    uint16_t checksum = 0;
    for (int i = 0; i < 5; i++) {
        checksum += command[i];
    }
    command[5] = (checksum >> 8) & 0xFF;
    command[6] = checksum & 0xFF;
    uart_write_bytes(UART_NUM, command, 7);
}

void read_pm_data(uint8_t *buf) {
    uint16_t checksum = 0;
    for (int i = 0; i <= 29; i++) {
        checksum += buf[i];
    }
    uint16_t c = (((uint16_t)buf[30]) << 8) | ((uint16_t)buf[31]);
    if (checksum != c) {
        ESP_LOGW(TAG, "Checksum from READ does not match, data corrupted");
        return;
    }
    ESP_LOGI(TAG, "SIZEOF pms7003_data_t: %d", sizeof(pms7003_data_t));
    if (sizeof(pms7003_data_t) == 24) {
        pms7003_data_t sensor_data = {0};
        memcpy(&sensor_data, &buf[4], sizeof(pms7003_data_t));
        uint16_t *sd = (uint16_t*)&sensor_data;
        for (int i = 0; i < sizeof(pms7003_data_t) / sizeof(uint16_t); i++) {
          sd[i] = SWAP_BYTES(sd[i]);
        }
        ESP_LOGI(TAG, "PM Readings:\nPM  1.0: %d\nPM  2.5: %d\nPM 10.0:%d",
                 sensor_data.pm1_0_atm,
                 sensor_data.pm2_5_atm,
                 sensor_data.pm10_atm);

    } else {
        ESP_LOGW(TAG, "Misaligned struct, ignoring memcpy");
    }
}

void pm_sensor_task(void *pvParameters) {
    ESP_LOGI(TAG, "Initialize UART, Wait 1 second");
    vTaskDelay(pdMS_TO_TICKS(1000));
    ESP_LOGI(TAG, "Initialize UART, ok now go...");
    uart_config_t uart_config = {
        .baud_rate = 9600,
        .data_bits = UART_DATA_8_BITS,
        .parity = UART_PARITY_DISABLE,
        .stop_bits = UART_STOP_BITS_1,
        .flow_ctrl = UART_HW_FLOWCTRL_DISABLE,
        // .rx_flow_ctrl_thresh = 122,
    };
    ESP_ERROR_CHECK(uart_param_config(UART_NUM, &uart_config));
    ESP_ERROR_CHECK(uart_set_pin(UART_NUM, TXD_PIN, RXD_PIN, UART_PIN_NO_CHANGE, UART_PIN_NO_CHANGE));
    ESP_ERROR_CHECK(uart_driver_install(UART_NUM, UART_BUF_SIZE * 2, 0, 0, NULL, 0));

    vTaskDelay(pdMS_TO_TICKS(100));
    send_pms7003_command(0xE1, 0x00, 0x00);
    vTaskDelay(pdMS_TO_TICKS(100));


    ESP_LOGI(TAG, "Alloc %d buffer for UART data", BUF_SIZE);
    uint8_t uart_data[BUF_SIZE];
    while (1) {
        send_pms7003_command(0xE2, 0x00, 0x00);

        int len = uart_read_bytes(UART_NUM, uart_data, BUF_SIZE, 100 / portTICK_PERIOD_MS);
        ESP_LOGI(TAG, "Received %d bytes from sensor", len);
        if (len > 0) {
            if (len == 32 && uart_data[0] != 0x42 && uart_data[1] != 0x4D) {
                ESP_LOGW(TAG, "Frame Start does not match 0x42, 0x4D, instead got %x %x", uart_data[0], uart_data[1]);
            } else {
                read_pm_data(uart_data);
                // for(int i = 0; i < len && i < BUF_SIZE; i++) {
                //     ESP_LOGI(TAG, "0x%02X", uart_data[i]);
                // }
            }
        }
        uart_flush(UART_NUM);
        vTaskDelay(pdMS_TO_TICKS(5000));
    }
}

void oled_display_task(void *pvParameters) {
    ESP_LOGI(TAG, "Initialize I2C bus");
    i2c_master_bus_handle_t i2c_bus = NULL;
    i2c_master_bus_config_t bus_config = {
        .clk_source = I2C_CLK_SRC_DEFAULT,
        .i2c_port = I2C_HOST,
        .sda_io_num = I2C_SDA_PIN,
        .scl_io_num = I2C_SCL_PIN,
        .flags.enable_internal_pullup = true,
    };
    ESP_ERROR_CHECK(i2c_new_master_bus(&bus_config, &i2c_bus));

    ESP_LOGI(TAG, "Install panel IO");
    esp_lcd_panel_io_handle_t io_handle = NULL;
    esp_lcd_panel_io_i2c_config_t io_config = {
        .dev_addr = SSD1306_ADDR,
        .scl_speed_hz = I2C_CLOCK_HZ,    // Added this line to set the clock speed
        .control_phase_bytes = 1,
        .lcd_cmd_bits = 8,
        .lcd_param_bits = 8,
        .dc_bit_offset = 6,
    };
    ESP_ERROR_CHECK(esp_lcd_new_panel_io_i2c(i2c_bus, &io_config, &io_handle));

    ESP_LOGI(TAG, "Install SSD1306 panel driver");
    esp_lcd_panel_handle_t panel_handle = NULL;
    esp_lcd_panel_ssd1306_config_t vendor_config = {
        .height = SSD1306_HEIGHT,
    };
    esp_lcd_panel_dev_config_t panel_config = {
        .bits_per_pixel = 1,
        .reset_gpio_num = -1,
        .vendor_config = &vendor_config,
    };
    ESP_ERROR_CHECK(esp_lcd_new_panel_ssd1306(io_handle, &panel_config, &panel_handle));

    ESP_LOGI(TAG, "Initialize the display");
    ESP_ERROR_CHECK(esp_lcd_panel_reset(panel_handle));
    ESP_ERROR_CHECK(esp_lcd_panel_init(panel_handle));
    ESP_ERROR_CHECK(esp_lcd_panel_disp_on_off(panel_handle, true));
    ESP_LOGI(TAG, "Create a simple pattern");

    uint8_t *fb = heap_caps_malloc(SSD1306_WIDTH * SSD1306_HEIGHT / 8, MALLOC_CAP_DMA);
    int idx = 0;
    while (1) {
        memset(fb, 0, SSD1306_WIDTH * SSD1306_HEIGHT / 8);
        taskYIELD();
        fb[idx] = 0xFF;

        esp_lcd_panel_draw_bitmap(panel_handle, 0, 0, SSD1306_WIDTH, SSD1306_HEIGHT, fb);
        idx = (idx + 4) % (SSD1306_WIDTH * SSD1306_HEIGHT / 8);
        // shift = (shift + 1) % 8;
        vTaskDelay(pdMS_TO_TICKS(100));
    }
    memset(fb, 0, SSD1306_WIDTH * SSD1306_HEIGHT / 8);
    free(fb);
}

void app_main(void) {
    xTaskCreate(pm_sensor_task, "pm_sensor_task", 4096, NULL, 5, NULL);
    // xTaskCreate(oled_display_task, "oled_display_task", 4096, NULL, 10, NULL);
}