#include <stdio.h>
#include <string.h>
#include <stdbool.h>
#include <stdint.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/spi_master.h"
#include "driver/gpio.h"
#include "driver/pulse_cnt.h"
#include "esp_adc/adc_oneshot.h"
#include "esp_system.h"
#include "esp_log.h"

// ==========================================
// VERIFIED HARDWARE PIN MAPPINGS (ESP32-S3)
// ==========================================
#define PIN_NUM_MOSI    11
#define PIN_NUM_MISO    13
#define PIN_NUM_CLK     12

#define NRF_PIN_CS      14
#define NRF_PIN_CE      21

#define OLED_PIN_CS     10
#define OLED_PIN_DC     9
#define OLED_PIN_RES    8

#define PIN_JOY_VRX     4   // ADC1 Channel 3 (GPIO 4)
#define PIN_JOY_VRY     5   // ADC1 Channel 4 (GPIO 5)
#define JOY_ADC_CH_X    ADC_CHANNEL_3
#define JOY_ADC_CH_Y    ADC_CHANNEL_4

#define PIN_ENCODER_CLK 15
#define PIN_ENCODER_DT  16
#define PIN_ENCODER_SW  17

#define NUM_MENU_ITEMS  4

static const char *TAG = "MFC_CORE";
spi_device_handle_t oled_spi;
spi_device_handle_t nrf_spi;
adc_oneshot_unit_handle_t adc1_handle;

// Global Radio & Calibration State
int current_radio_channel = 76;
int joy_center_x = 2048;
int joy_center_y = 2048;

typedef struct __attribute__((packed)) {
    uint16_t joy_x;
    uint16_t joy_y;
    uint8_t  flags;
    uint8_t  channel;
} rc_packet_t;

typedef enum {
    SCREEN_MAIN_MENU = 0,
    SCREEN_DRIVE,
    SCREEN_RADIO_CONFIG,
    SCREEN_CALIBRATE,
    SCREEN_SYSTEM_INFO
} screen_state_t;

// ==========================================
// 5x8 BARE-METAL FONT LIBRARY
// ==========================================
const uint8_t font_AZ[26][5] = {
    {0x7E, 0x11, 0x11, 0x11, 0x7E}, {0x7F, 0x49, 0x49, 0x49, 0x36}, {0x3E, 0x41, 0x41, 0x41, 0x22},
    {0x7F, 0x41, 0x41, 0x22, 0x1C}, {0x7F, 0x49, 0x49, 0x49, 0x41}, {0x7F, 0x09, 0x09, 0x09, 0x01},
    {0x3E, 0x41, 0x49, 0x49, 0x7A}, {0x7F, 0x08, 0x08, 0x08, 0x7F}, {0x00, 0x41, 0x7F, 0x41, 0x00},
    {0x20, 0x40, 0x41, 0x3F, 0x01}, {0x7F, 0x08, 0x14, 0x22, 0x41}, {0x7F, 0x40, 0x40, 0x40, 0x40},
    {0x7F, 0x02, 0x0C, 0x02, 0x7F}, {0x7F, 0x04, 0x08, 0x10, 0x7F}, {0x3E, 0x41, 0x41, 0x41, 0x3E},
    {0x7F, 0x09, 0x09, 0x09, 0x06}, {0x3E, 0x41, 0x51, 0x21, 0x5E}, {0x7F, 0x09, 0x19, 0x29, 0x46},
    {0x46, 0x49, 0x49, 0x49, 0x31}, {0x01, 0x01, 0x7F, 0x01, 0x01}, {0x3F, 0x40, 0x40, 0x40, 0x3F},
    {0x1F, 0x20, 0x40, 0x20, 0x1F}, {0x3F, 0x40, 0x38, 0x40, 0x3F}, {0x63, 0x14, 0x08, 0x14, 0x63},
    {0x03, 0x04, 0x78, 0x04, 0x03}, {0x61, 0x51, 0x49, 0x45, 0x43}
};

const uint8_t font_nums[10][5] = {
    {0x3E, 0x51, 0x49, 0x45, 0x3E}, {0x00, 0x42, 0x7F, 0x40, 0x00}, {0x42, 0x61, 0x51, 0x49, 0x46},
    {0x21, 0x41, 0x45, 0x4B, 0x31}, {0x18, 0x14, 0x12, 0x7F, 0x10}, {0x27, 0x45, 0x45, 0x45, 0x39},
    {0x3C, 0x4A, 0x49, 0x49, 0x30}, {0x01, 0x71, 0x09, 0x05, 0x03}, {0x36, 0x49, 0x49, 0x49, 0x36},
    {0x06, 0x49, 0x49, 0x29, 0x1E}
};

const uint8_t* get_font_data(char c) {
    static const uint8_t f_spc[5]   = {0x00, 0x00, 0x00, 0x00, 0x00};
    static const uint8_t f_gt[5]    = {0x08, 0x04, 0x7E, 0x04, 0x08}; // '>'
    static const uint8_t f_lt[5]    = {0x08, 0x14, 0x22, 0x41, 0x00}; // '<'
    static const uint8_t f_lbr[5]   = {0x00, 0x7F, 0x41, 0x41, 0x00}; // '['
    static const uint8_t f_rbr[5]   = {0x00, 0x41, 0x41, 0x7F, 0x00}; // ']'
    static const uint8_t f_lpar[5]  = {0x00, 0x1C, 0x22, 0x41, 0x00}; // '('
    static const uint8_t f_rpar[5]  = {0x00, 0x41, 0x22, 0x1C, 0x00}; // ')'
    static const uint8_t f_slash[5] = {0x20, 0x10, 0x08, 0x04, 0x02}; // '/'
    static const uint8_t f_dash[5]  = {0x08, 0x08, 0x08, 0x08, 0x08}; // '-'
    static const uint8_t f_col[5]   = {0x00, 0x36, 0x36, 0x00, 0x00}; // ':'
    static const uint8_t f_dot[5]   = {0x00, 0x60, 0x60, 0x00, 0x00}; // '.'

    if (c >= 'A' && c <= 'Z') { return font_AZ[c - 'A']; }
    if (c >= 'a' && c <= 'z') { return font_AZ[c - 'a']; }
    if (c >= '0' && c <= '9') { return font_nums[c - '0']; }
    if (c == '>') { return f_gt; }
    if (c == '<') { return f_lt; }
    if (c == '[') { return f_lbr; }
    if (c == ']') { return f_rbr; }
    if (c == '(') { return f_lpar; }
    if (c == ')') { return f_rpar; }
    if (c == '/') { return f_slash; }
    if (c == '-') { return f_dash; }
    if (c == ':') { return f_col; }
    if (c == '.') { return f_dot; }
    return f_spc;
}

// ==========================================
// OLED & NRF SPI DRIVERS
// ==========================================
void oled_cmd(uint8_t cmd) {
    gpio_set_level(OLED_PIN_DC, 0);
    spi_transaction_t t = { .length = 8, .tx_buffer = &cmd };
    spi_device_polling_transmit(oled_spi, &t);
}

void oled_data(uint8_t data) {
    gpio_set_level(OLED_PIN_DC, 1);
    spi_transaction_t t = { .length = 8, .tx_buffer = &data };
    spi_device_polling_transmit(oled_spi, &t);
}

void oled_draw_string(uint8_t page, uint8_t col, const char* str) {
    oled_cmd(0xB0 + page);
    uint8_t phys_col = col + 2; // SH1106 132x64 internal RAM 2-px offset
    oled_cmd(0x00 | (phys_col & 0x0F));
    oled_cmd(0x10 | ((phys_col >> 4) & 0x0F));
    while (*str) {
        const uint8_t* pixels = get_font_data(*str++);
        for (int i = 0; i < 5; i++) { oled_data(pixels[i]); }
        oled_data(0x00);
    }
}

void oled_clear(void) {
    for (uint8_t p = 0; p < 8; p++) {
        oled_cmd(0xB0 + p); oled_cmd(0x02); oled_cmd(0x10);
        for (uint8_t c = 0; c < 128; c++) { oled_data(0x00); }
    }
}

void oled_init(void) {
    gpio_set_level(OLED_PIN_RES, 0); vTaskDelay(pdMS_TO_TICKS(50));
    gpio_set_level(OLED_PIN_RES, 1); vTaskDelay(pdMS_TO_TICKS(50));
    uint8_t init_cmds[] = {
        0xAE, 0xD5, 0x80, 0xA8, 0x3F, 0xD3, 0x00, 0x40,
        0x8D, 0x14, 0xA1, 0xC8, 0xDA, 0x12, 0x81, 0xCF,
        0xD9, 0xF1, 0xDB, 0x40, 0xA4, 0xA6, 0xAF
    };
    for (size_t i = 0; i < sizeof(init_cmds); i++) { oled_cmd(init_cmds[i]); }
}

static void nrf_write_reg(uint8_t reg, uint8_t val) {
    uint8_t tx[2] = { (uint8_t)(0x20 | (reg & 0x1F)), val };
    spi_transaction_t t = { .length = 16, .tx_buffer = tx };
    spi_device_polling_transmit(nrf_spi, &t);
}

bool nrf_check_connection(void) {
    uint8_t tx[2] = { 0x03, 0xFF }; // Read SETUP_AW register (0x03)
    uint8_t rx[2] = { 0 };
    spi_transaction_t t = { .length = 16, .tx_buffer = tx, .rx_buffer = rx };
    spi_device_polling_transmit(nrf_spi, &t);
    return (rx[1] >= 0x01 && rx[1] <= 0x03);
}

void set_nrf_channel(uint8_t channel) {
    nrf_write_reg(0x05, channel); // Write RF_CH register (0x05)
}

void nrf_init_tx(uint8_t channel) {
    gpio_set_level(NRF_PIN_CE, 0);
    nrf_write_reg(0x00, 0x0E); // CONFIG: PWR_UP, EN_CRC (2-byte), PTX mode
    nrf_write_reg(0x01, 0x00); // EN_AA: Disable auto-ack for low-latency broadcast
    nrf_write_reg(0x02, 0x01); // EN_RXADDR: Pipe 0
    nrf_write_reg(0x03, 0x03); // SETUP_AW: 5-byte address width
    nrf_write_reg(0x04, 0x00); // SETUP_RETR: Retransmit disabled
    nrf_write_reg(0x05, channel); // RF_CH: Set active 2.4 GHz channel
    nrf_write_reg(0x06, 0x06); // RF_SETUP: 1 Mbps, 0 dBm
    nrf_write_reg(0x07, 0x70); // STATUS: Clear IRQ flags
}

void nrf_send_packet(const rc_packet_t *pkt) {
    uint8_t flush_cmd = 0xE1; // FLUSH_TX
    spi_transaction_t tf = { .length = 8, .tx_buffer = &flush_cmd };
    spi_device_polling_transmit(nrf_spi, &tf);

    uint8_t tx_buf[1 + sizeof(rc_packet_t)];
    tx_buf[0] = 0xA0; // W_TX_PAYLOAD
    memcpy(&tx_buf[1], pkt, sizeof(rc_packet_t));
    spi_transaction_t tp = { .length = (1 + sizeof(rc_packet_t)) * 8, .tx_buffer = tx_buf };
    spi_device_polling_transmit(nrf_spi, &tp);

    gpio_set_level(NRF_PIN_CE, 1);
    esp_rom_delay_us(15);
    gpio_set_level(NRF_PIN_CE, 0);
}

static void read_joystick(int *raw_x, int *raw_y) {
    if (adc_oneshot_read(adc1_handle, JOY_ADC_CH_X, raw_x) != ESP_OK) { *raw_x = 0; }
    if (adc_oneshot_read(adc1_handle, JOY_ADC_CH_Y, raw_y) != ESP_OK) { *raw_y = 0; }
}

// ==========================================
// UI SCREEN RENDERING
// ==========================================
void draw_menu(int selected_index) {
    const char* items[] = {"DRIVE VEHICLE", "RADIO CONFIG", "CALIBRATE", "SYSTEM INFO"};
    char buffer[32];

    oled_draw_string(0, 20, "- MAIN MENU -");

    for (int i = 0; i < NUM_MENU_ITEMS; i++) {
        if (i == selected_index) {
            snprintf(buffer, sizeof(buffer), "> %-14s", items[i]);
        } else {
            snprintf(buffer, sizeof(buffer), "  %-14s", items[i]);
        }
        oled_draw_string(2 + (i * 2), 5, buffer);
    }
}

void draw_drive_screen(int raw_x, int raw_y, int channel) {
    char buffer[32];
    oled_draw_string(0, 10, "- DRIVE VEHICLE -");

    snprintf(buffer, sizeof(buffer), "VRX: %-4d VRY: %-4d", raw_x, raw_y);
    oled_draw_string(2, 2, buffer);

    snprintf(buffer, sizeof(buffer), "TX CH: %03d [ACTIVE]", channel);
    oled_draw_string(4, 2, buffer);

    oled_draw_string(6, 10, "[CLICK TO BACK]");
}

void draw_calibrate_screen(int raw_x, int raw_y) {
    char buffer[32];
    oled_draw_string(0, 10, "- CALIBRATION -");

    snprintf(buffer, sizeof(buffer), "CUR: %04d / %04d", raw_x, raw_y);
    oled_draw_string(2, 2, buffer);

    snprintf(buffer, sizeof(buffer), "CTR: %04d / %04d", joy_center_x, joy_center_y);
    oled_draw_string(4, 2, buffer);

    oled_draw_string(6, 5, "[CLICK SET CTR]");
}

void draw_system_info(void) {
    char buffer[32];
    oled_draw_string(0, 15, "- SYSTEM INFO -");

    snprintf(buffer, sizeof(buffer), "CHIP: ESP32-S3");
    oled_draw_string(2, 5, buffer);

    snprintf(buffer, sizeof(buffer), "HEAP: %lu KB", (unsigned long)(esp_get_free_heap_size() / 1024));
    oled_draw_string(4, 5, buffer);

    oled_draw_string(6, 10, "[CLICK TO BACK]");
}

void draw_radio_config_screen(int channel) {
    char buffer[32];
    oled_draw_string(0, 5, "- RADIO CONFIG -");

    if (nrf_check_connection()) {
        snprintf(buffer, sizeof(buffer), "CH: <%03d> (2.4G)", channel);
        oled_draw_string(2, 2, buffer);
        oled_draw_string(4, 2, "TURN DIAL TO TUNE");
    } else {
        oled_draw_string(2, 2, "STATUS: FAILED   ");
        oled_draw_string(4, 2, "CHECK WIRING     ");
    }

    oled_draw_string(6, 5, "[CLICK TO SAVE]");
}

// ==========================================
// MAIN APPLICATION ENTRY
// ==========================================
void app_main(void) {
    // 1. Initialize Shared SPI2_HOST Bus (FSPI)
    spi_bus_config_t buscfg = {
        .miso_io_num = PIN_NUM_MISO,
        .mosi_io_num = PIN_NUM_MOSI,
        .sclk_io_num = PIN_NUM_CLK,
        .quadwp_io_num = -1,
        .quadhd_io_num = -1,
        .max_transfer_sz = 4096
    };
    ESP_ERROR_CHECK(spi_bus_initialize(SPI2_HOST, &buscfg, SPI_DMA_CH_AUTO));

    // 2. Attach 1.3" SH1106 OLED Display (10 MHz SPI)
    spi_device_interface_config_t oled_cfg = {
        .clock_speed_hz = 10 * 1000 * 1000,
        .mode = 0,
        .spics_io_num = OLED_PIN_CS,
        .queue_size = 1
    };
    ESP_ERROR_CHECK(spi_bus_add_device(SPI2_HOST, &oled_cfg, &oled_spi));
    gpio_reset_pin(OLED_PIN_DC);  gpio_set_direction(OLED_PIN_DC, GPIO_MODE_OUTPUT);
    gpio_reset_pin(OLED_PIN_RES); gpio_set_direction(OLED_PIN_RES, GPIO_MODE_OUTPUT);

    oled_init();
    oled_clear();

    // 3. Attach nRF24L01+ Radio Module (1 MHz SPI)
    gpio_reset_pin(NRF_PIN_CE);
    gpio_set_direction(NRF_PIN_CE, GPIO_MODE_OUTPUT);
    gpio_set_level(NRF_PIN_CE, 0);

    spi_device_interface_config_t nrf_cfg = {
        .clock_speed_hz = 1 * 1000 * 1000,
        .mode = 0,
        .spics_io_num = NRF_PIN_CS,
        .queue_size = 1
    };
    ESP_ERROR_CHECK(spi_bus_add_device(SPI2_HOST, &nrf_cfg, &nrf_spi));
    nrf_init_tx((uint8_t)current_radio_channel);

    // 4. Initialize One-Shot ADC1 for Hall-Effect / Analog Joystick (GPIO 4 & GPIO 5)
    adc_oneshot_unit_init_cfg_t init_config1 = {
        .unit_id = ADC_UNIT_1,
    };
    ESP_ERROR_CHECK(adc_oneshot_new_unit(&init_config1, &adc1_handle));

    adc_oneshot_chan_cfg_t adc_cfg = {
        .bitwidth = ADC_BITWIDTH_DEFAULT,
        .atten = ADC_ATTEN_DB_12,
    };
    ESP_ERROR_CHECK(adc_oneshot_config_channel(adc1_handle, JOY_ADC_CH_X, &adc_cfg));
    ESP_ERROR_CHECK(adc_oneshot_config_channel(adc1_handle, JOY_ADC_CH_Y, &adc_cfg));

    // 5. Initialize Rotary Encoder Push Switch (GPIO 17)
    gpio_config_t sw_cfg = {
        .pin_bit_mask = (1ULL << PIN_ENCODER_SW),
        .mode = GPIO_MODE_INPUT,
        .pull_up_en = GPIO_PULLUP_ENABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_DISABLE
    };
    gpio_config(&sw_cfg);

    // 6. Initialize PCNT Hardware Quadrature Decoder (GPIO 15 & GPIO 16)
    pcnt_unit_handle_t pcnt_unit = NULL;
    pcnt_unit_config_t unit_config = { .high_limit = 10000, .low_limit = -10000 };
    ESP_ERROR_CHECK(pcnt_new_unit(&unit_config, &pcnt_unit));

    pcnt_glitch_filter_config_t filter_config = { .max_glitch_ns = 1000 };
    ESP_ERROR_CHECK(pcnt_unit_set_glitch_filter(pcnt_unit, &filter_config));

    pcnt_channel_handle_t pcnt_chan_a = NULL, pcnt_chan_b = NULL;
    pcnt_chan_config_t chan_a_config = { .edge_gpio_num = PIN_ENCODER_CLK, .level_gpio_num = PIN_ENCODER_DT };
    ESP_ERROR_CHECK(pcnt_new_channel(pcnt_unit, &chan_a_config, &pcnt_chan_a));
    pcnt_chan_config_t chan_b_config = { .edge_gpio_num = PIN_ENCODER_DT, .level_gpio_num = PIN_ENCODER_CLK };
    ESP_ERROR_CHECK(pcnt_new_channel(pcnt_unit, &chan_b_config, &pcnt_chan_b));

    ESP_ERROR_CHECK(pcnt_channel_set_edge_action(pcnt_chan_a, PCNT_CHANNEL_EDGE_ACTION_DECREASE, PCNT_CHANNEL_EDGE_ACTION_INCREASE));
    ESP_ERROR_CHECK(pcnt_channel_set_level_action(pcnt_chan_a, PCNT_CHANNEL_LEVEL_ACTION_KEEP, PCNT_CHANNEL_LEVEL_ACTION_INVERSE));
    ESP_ERROR_CHECK(pcnt_channel_set_edge_action(pcnt_chan_b, PCNT_CHANNEL_EDGE_ACTION_INCREASE, PCNT_CHANNEL_EDGE_ACTION_DECREASE));
    ESP_ERROR_CHECK(pcnt_channel_set_level_action(pcnt_chan_b, PCNT_CHANNEL_LEVEL_ACTION_KEEP, PCNT_CHANNEL_LEVEL_ACTION_INVERSE));

    ESP_ERROR_CHECK(pcnt_unit_enable(pcnt_unit));
    ESP_ERROR_CHECK(pcnt_unit_clear_count(pcnt_unit));
    ESP_ERROR_CHECK(pcnt_unit_start(pcnt_unit));

    ESP_LOGI(TAG, "MFC Multi-Tasking OS State Machine Online (ESP32-S3).");

    // 7. Master UI & Telemetry Loop
    int count = 0;
    int last_menu_index = -1;
    int last_sw_state = 1;
    TickType_t last_button_press = 0;

    screen_state_t current_screen = SCREEN_MAIN_MENU;
    bool screen_needs_redraw = true;

    int radio_start_count = 0;
    int last_tuned_channel = current_radio_channel;

    while (1) {
        pcnt_unit_get_count(pcnt_unit, &count);
        int sw_state = gpio_get_level(PIN_ENCODER_SW);

        if (current_screen == SCREEN_MAIN_MENU) {
            int raw_steps = count / 4;
            int current_menu_index = ((raw_steps % NUM_MENU_ITEMS) + NUM_MENU_ITEMS) % NUM_MENU_ITEMS;

            if (current_menu_index != last_menu_index || screen_needs_redraw) {
                oled_clear();
                draw_menu(current_menu_index);
                last_menu_index = current_menu_index;
                screen_needs_redraw = false;
            }

            // Click to Select with 300ms Debounce
            if (sw_state == 0 && last_sw_state == 1) {
                if ((xTaskGetTickCount() - last_button_press) > pdMS_TO_TICKS(300)) {
                    current_screen = (screen_state_t)(current_menu_index + 1);
                    screen_needs_redraw = true;

                    if (current_screen == SCREEN_RADIO_CONFIG) {
                        radio_start_count = count;
                        last_tuned_channel = current_radio_channel;
                    }

                    last_button_press = xTaskGetTickCount();
                }
            }
        } else {
            // Sub-screen handling
            if (current_screen == SCREEN_RADIO_CONFIG) {
                int encoder_delta = (count - radio_start_count) / 4;
                int target_channel = current_radio_channel + encoder_delta;

                if (target_channel < 0) target_channel = 0;
                if (target_channel > 125) target_channel = 125;

                if (target_channel != last_tuned_channel || screen_needs_redraw) {
                    current_radio_channel = target_channel;
                    set_nrf_channel((uint8_t)current_radio_channel);

                    oled_clear();
                    draw_radio_config_screen(current_radio_channel);
                    last_tuned_channel = target_channel;
                    screen_needs_redraw = false;
                }
            } else if (current_screen == SCREEN_DRIVE) {
                int raw_x = 0, raw_y = 0;
                read_joystick(&raw_x, &raw_y);

                rc_packet_t pkt = {
                    .joy_x = (uint16_t)raw_x,
                    .joy_y = (uint16_t)raw_y,
                    .flags = 0x01,
                    .channel = (uint8_t)current_radio_channel
                };
                nrf_send_packet(&pkt);

                if (screen_needs_redraw) {
                    oled_clear();
                    screen_needs_redraw = false;
                }
                draw_drive_screen(raw_x, raw_y, current_radio_channel);
            } else if (current_screen == SCREEN_CALIBRATE) {
                int raw_x = 0, raw_y = 0;
                read_joystick(&raw_x, &raw_y);

                if (screen_needs_redraw) {
                    oled_clear();
                    screen_needs_redraw = false;
                }
                draw_calibrate_screen(raw_x, raw_y);
            } else if (current_screen == SCREEN_SYSTEM_INFO) {
                if (screen_needs_redraw) {
                    oled_clear();
                    draw_system_info();
                    screen_needs_redraw = false;
                }
            }

            // Click to return to Main Menu (and save calibration if in CALIBRATE) with 300ms Debounce
            if (sw_state == 0 && last_sw_state == 1) {
                if ((xTaskGetTickCount() - last_button_press) > pdMS_TO_TICKS(300)) {
                    if (current_screen == SCREEN_CALIBRATE) {
                        read_joystick(&joy_center_x, &joy_center_y);
                    }
                    current_screen = SCREEN_MAIN_MENU;
                    screen_needs_redraw = true;
                    last_button_press = xTaskGetTickCount();
                }
            }
        }

        last_sw_state = sw_state;
        vTaskDelay(pdMS_TO_TICKS(30));
    }
}