/**
 * @file nrf24l01.c
 * @brief NRF24L01+ driver implementation using SPI
 */
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "driver/spi_master.h"
#include "driver/gpio.h"
#include "esp_log.h"
#include "nrf24l01.h"

static const char *TAG = "nrf24";

/* SPI device handle */
static spi_device_handle_t nrf_spi = NULL;
static int nrf_ce_gpio = -1;

/* Default 5-byte address */
static const uint8_t default_addr[5] = {0xE7, 0xE7, 0xE7, 0xE7, 0xE7};

/* SPI transaction helper */
static uint8_t nrf_spi_transfer(uint8_t byte)
{
    spi_transaction_t t = {
        .length = 8,
        .tx_buffer = &byte,
        .rx_buffer = &byte,
    };
    spi_device_transmit(nrf_spi, &t);
    return byte;
}

/* Write a register */
static void nrf_write_reg(uint8_t reg, uint8_t value)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_W_REGISTER | reg);
    nrf_spi_transfer(value);
}

/* Read a register */
static uint8_t nrf_read_reg(uint8_t reg)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_R_REGISTER | reg);
    return nrf_spi_transfer(NRF_CMD_NOP);
}

/* Write multiple bytes to a register (for addresses) */
static void nrf_write_multi(uint8_t reg, const uint8_t *data, uint8_t len)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_W_REGISTER | reg);
    for (int i = 0; i < len; i++) {
        nrf_spi_transfer(data[i]);
    }
}

/* Read multiple bytes from a register */
static void nrf_read_multi(uint8_t reg, uint8_t *data, uint8_t len)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_R_REGISTER | reg);
    for (int i = 0; i < len; i++) {
        data[i] = nrf_spi_transfer(NRF_CMD_NOP);
    }
}

esp_err_t nrf24_init(int mosi, int miso, int sclk, int cs, int ce, int irq)
{
    nrf_ce_gpio = ce;
    
    /* Configure SPI bus */
    spi_bus_config_t bus_cfg = {
        .mosi_io_num = mosi,
        .miso_io_num = miso,
        .sclk_io_num = sclk,
        .quadwp_io_num = -1,
        .quadhd_io_num = -1,
        .max_transfer_sz = 32,
    };
    
    spi_device_interface_config_t dev_cfg = {
        .mode = 0,              /* SPI mode 0: CPOL=0, CPHA=0 */
        .clock_speed_hz = 10 * 1000 * 1000,  /* 10 MHz */
        .spics_io_num = cs,
        .queue_size = 1,
        .flags = SPI_DEVICE_HALFDUPLEX,
    };
    
    ESP_ERROR_CHECK(spi_bus_initialize(SPI2_HOST, &bus_cfg, SPI_DMA_CH_AUTO));
    ESP_ERROR_CHECK(spi_bus_add_device(SPI2_HOST, &dev_cfg, &nrf_spi));
    
    /* Configure CE pin */
    gpio_config_t ce_conf = {
        .pin_bit_mask = (1ULL << ce),
        .mode = GPIO_MODE_OUTPUT,
        .pull_up_en = GPIO_PULLUP_DISABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_DISABLE,
    };
    gpio_config(&ce_conf);
    gpio_set_level(ce, 0);
    
    /* Wait for power-on */
    vTaskDelay(pdMS_TO_TICKS(5));
    
    /* Configure NRF24L01 */
    /* Disable auto-ack on all pipes (we use ACK payload) */
    nrf_write_reg(NRF_REG_EN_AA, 0x01);  /* Only pipe 0 auto-ack */
    nrf_write_reg(NRF_REG_EN_RXADDR, 0x01);  /* Enable pipe 0 */
    nrf_write_reg(NRF_REG_SETUP_AW, 0x03);   /* 5-byte address */
    nrf_write_reg(NRF_REG_SETUP_RETR, 0x1A); /* 250us retry, 10 retransmits */
    nrf_write_reg(NRF_REG_RF_CH, 100);        /* Default channel */
    nrf_write_reg(NRF_REG_RF_SETUP, 0x06);    /* 1Mbps, 0dBm */
    nrf_write_reg(NRF_REG_STATUS, 0x70);      /* Clear interrupts */
    
    /* Set default address */
    nrf_write_multi(NRF_REG_TX_ADDR, default_addr, 5);
    nrf_write_multi(NRF_REG_RX_ADDR_P0, default_addr, 5);
    
    /* Set payload width for pipe 0 */
    nrf_write_reg(NRF_REG_RX_PW_P0, 32);
    
    /* Enable dynamic payload length and ACK payload */
    nrf_write_reg(NRF_REG_FEATURE, 0x06);  /* EN_DPL | EN_ACK_PAY */
    nrf_write_reg(NRF_REG_DYNPD, 0x01);    /* DPL for pipe 0 */
    
    /* Power up in RX mode by default */
    nrf_write_reg(NRF_REG_CONFIG, NRF_CONFIG_EN_CRC | NRF_CONFIG_CRCO | NRF_CONFIG_PWR_UP | NRF_CONFIG_PRIM_RX);
    gpio_set_level(ce, 1);
    
    vTaskDelay(pdMS_TO_TICKS(2));
    
    ESP_LOGI(TAG, "NRF24L01 initialized");
    return ESP_OK;
}

void nrf24_set_channel(uint8_t channel)
{
    if (channel > 125) channel = 125;
    nrf_write_reg(NRF_REG_RF_CH, channel);
}

void nrf24_set_data_rate(nrf_data_rate_t rate)
{
    uint8_t rf_setup = nrf_read_reg(NRF_REG_RF_SETUP);
    rf_setup &= 0xD7;  /* Clear RF_DR bits */
    
    switch (rate) {
        case NRF_1MBPS:
            rf_setup |= (0 << 3);  /* RF_DR_LOW=0, RF_DR_HIGH=0 */
            break;
        case NRF_2MBPS:
            rf_setup |= (1 << 3);  /* RF_DR_HIGH=1 */
            break;
        case NRF_250KBPS:
            rf_setup |= (1 << 5) | (0 << 3);  /* RF_DR_LOW=1 */
            break;
    }
    nrf_write_reg(NRF_REG_RF_SETUP, rf_setup);
}

void nrf24_set_power(nrf_pa_power_t power)
{
    uint8_t rf_setup = nrf_read_reg(NRF_REG_RF_SETUP);
    rf_setup &= 0xF9;  /* Clear PWR bits */
    rf_setup |= (power << 1);
    nrf_write_reg(NRF_REG_RF_SETUP, rf_setup);
}

void nrf24_set_tx_addr(const uint8_t *addr)
{
    nrf_write_multi(NRF_REG_TX_ADDR, addr, 5);
}

void nrf24_set_rx_addr(const uint8_t *addr, uint8_t pipe)
{
    if (pipe == 0) {
        nrf_write_multi(NRF_REG_RX_ADDR_P0, addr, 5);
    } else if (pipe <= 5) {
        nrf_write_reg(NRF_REG_RX_ADDR_P0 + pipe, addr[4]);  /* Only LSB */
    }
}

void nrf24_power_up_tx(void)
{
    uint8_t config = nrf_read_reg(NRF_REG_CONFIG);
    config &= ~NRF_CONFIG_PRIM_RX;  /* Set PTX */
    config |= NRF_CONFIG_PWR_UP;
    nrf_write_reg(NRF_REG_CONFIG, config);
    gpio_set_level(nrf_ce_gpio, 0);
    vTaskDelay(pdMS_TO_TICKS(2));
}

void nrf24_power_up_rx(void)
{
    uint8_t config = nrf_read_reg(NRF_REG_CONFIG);
    config |= NRF_CONFIG_PWR_UP | NRF_CONFIG_PRIM_RX;
    nrf_write_reg(NRF_REG_CONFIG, config);
    gpio_set_level(nrf_ce_gpio, 1);
    vTaskDelay(pdMS_TO_TICKS(2));
}

void nrf24_power_down(void)
{
    uint8_t config = nrf_read_reg(NRF_REG_CONFIG);
    config &= ~NRF_CONFIG_PWR_UP;
    nrf_write_reg(NRF_REG_CONFIG, config);
    gpio_set_level(nrf_ce_gpio, 0);
}

esp_err_t nrf24_tx_payload(const uint8_t *data, uint8_t len)
{
    if (len > 32) len = 32;
    
    nrf24_power_up_tx();
    
    /* Write payload */
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_W_TX_PAYLOAD);
    for (int i = 0; i < len; i++) {
        nrf_spi_transfer(data[i]);
    }
    
    /* Pulse CE to start transmission */
    gpio_set_level(nrf_ce_gpio, 1);
    vTaskDelay(pdMS_TO_TICKS(1));
    gpio_set_level(nrf_ce_gpio, 0);
    
    /* Wait for TX_DS or MAX_RT */
    int timeout = 100;
    while (timeout--) {
        uint8_t status = nrf24_get_status();
        if (status & (1 << 5)) {  /* TX_DS */
            nrf_write_reg(NRF_REG_STATUS, 0x20);  /* Clear TX_DS */
            return ESP_OK;
        }
        if (status & (1 << 4)) {  /* MAX_RT */
            nrf_write_reg(NRF_REG_STATUS, 0x10);  /* Clear MAX_RT */
            nrf24_flush_tx();
            return ESP_FAIL;
        }
        vTaskDelay(pdMS_TO_TICKS(1));
    }
    
    return ESP_ERR_TIMEOUT;
}

esp_err_t nrf24_rx_payload(uint8_t *data, uint8_t *len)
{
    if (!nrf24_data_available()) {
        return ESP_ERR_NOT_FOUND;
    }
    
    /* Get payload width */
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_R_RX_PL_WID);
    *len = nrf_spi_transfer(NRF_CMD_NOP);
    
    if (*len > 32 || *len == 0) {
        nrf24_flush_rx();
        return ESP_ERR_INVALID_SIZE;
    }
    
    /* Read payload */
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_R_RX_PAYLOAD);
    for (int i = 0; i < *len; i++) {
        data[i] = nrf_spi_transfer(NRF_CMD_NOP);
    }
    
    /* Clear RX_DR */
    nrf_write_reg(NRF_REG_STATUS, 0x40);
    
    return ESP_OK;
}

uint8_t nrf24_data_available(void)
{
    uint8_t status = nrf24_get_status();
    return (status & (1 << 6)) ? 1 : 0;  /* RX_DR */
}

uint8_t nrf24_get_status(void)
{
    return nrf_spi_transfer(NRF_CMD_NOP);
}

void nrf24_flush_tx(void)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_FLUSH_TX);
}

void nrf24_flush_rx(void)
{
    gpio_set_level(nrf_ce_gpio, 0);
    nrf_spi_transfer(NRF_CMD_FLUSH_RX);
}