#include <stdio.h>
#include <string.h>
#include <stdbool.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_system.h"
#include "esp_log.h"

// ==========================================
// VERIFIED HARDWARE PIN MAPPINGS
// ==========================================
#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_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;

// Global Radio State
int current_radio_channel = 76;

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]   = {0,0,0,0,0};
    static const uint8_t f_ptr[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 >= '0' && c <= '9') { return font_nums[c - '0']; }
    if (c == '>') { return f_ptr; }
    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 center 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() {
    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() {
    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(int i = 0; i < sizeof(init_cmds); i++) { oled_cmd(init_cmds[i]); }
}

bool nrf_check_connection() {
    uint8_t tx[2] = { 0x03, 0xFF }; 
    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] == 0x03);
}

void set_nrf_channel(uint8_t channel) {
    uint8_t tx[2] = { 0x25, channel }; // Write to RF_CH register (0x05)
    spi_transaction_t t = { .length = 16, .tx_buffer = tx };
    spi_device_polling_transmit(nrf_spi, &t);
}

// ==========================================
// UI SCREEN RENDERING
// ==========================================
void draw_menu(int selected_index) {
    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) {
            sprintf(buffer, "> %-14s", items[i]); 
        } else {
            sprintf(buffer, "  %-14s", items[i]);
        }
        oled_draw_string(2 + (i * 2), 5, buffer);
    }
}

void draw_system_info() {
    char buffer[32];
    oled_draw_string(0, 15, "- SYSTEM INFO -");
    
    sprintf(buffer, "CHIP: ESP32-S3");
    oled_draw_string(2, 5, buffer);
    
    sprintf(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()) {
        sprintf(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]");
}

void draw_placeholder_screen(const char* title) {
    oled_draw_string(0, 15, title);
    oled_draw_string(3, 20, "NOT CONNECTED");
    oled_draw_string(6, 10, "[CLICK TO BACK]");
}

// ==========================================
// MAIN APPLICATION ENTRY
// ==========================================
void app_main(void) {
    // 1. Initialize Shared SPI Bus
    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 
    };
    spi_bus_initialize(SPI2_HOST, &buscfg, SPI_DMA_CH_AUTO);

    // 2. Add OLED Device
    spi_device_interface_config_t oled_cfg = { 
        .clock_speed_hz = 10 * 1000 * 1000, 
        .mode = 0, 
        .spics_io_num = OLED_PIN_CS, 
        .queue_size = 1 
    };
    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. Add NRF24 Device
    spi_device_interface_config_t nrf_cfg = { 
        .clock_speed_hz = 1 * 1000 * 1000, 
        .mode = 0, 
        .spics_io_num = NRF_PIN_CS, 
        .queue_size = 1 
    };
    spi_bus_add_device(SPI2_HOST, &nrf_cfg, &nrf_spi);

    // 4. Init Rotary Encoder Switch
    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);

    // 5. Init PCNT Hardware Decoder for Quadrature Encoder
    pcnt_unit_handle_t pcnt_unit = NULL;
    pcnt_unit_config_t unit_config = { .high_limit = 10000, .low_limit = -10000 };
    pcnt_new_unit(&unit_config, &pcnt_unit);
    
    pcnt_glitch_filter_config_t filter_config = { .max_glitch_ns = 1000 };
    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 };
    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 };
    pcnt_new_channel(pcnt_unit, &chan_b_config, &pcnt_chan_b);

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

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

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

    // 6. Master UI State Machine 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(current_radio_channel);
                    
                    oled_clear();
                    draw_radio_config_screen(current_radio_channel);
                    last_tuned_channel = target_channel;
                    screen_needs_redraw = false;
                }
            } 
            else {
                if (screen_needs_redraw) {
                    oled_clear();
                    if (current_screen == SCREEN_SYSTEM_INFO) {
                        draw_system_info();
                    } else if (current_screen == SCREEN_DRIVE) {
                        draw_placeholder_screen("- DRIVE VEHICLE -");
                    } else if (current_screen == SCREEN_CALIBRATE) {
                        draw_placeholder_screen("- CALIBRATION -");
                    }
                    screen_needs_redraw = false;
                }
            }

            // Click to return to Main Menu with 300ms Debounce
            if (sw_state == 0 && last_sw_state == 1) {
                if ((xTaskGetTickCount() - last_button_press) > pdMS_TO_TICKS(300)) {
                    current_screen = SCREEN_MAIN_MENU;
                    screen_needs_redraw = true;
                    last_button_press = xTaskGetTickCount(); 
                }
            }
        }

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