#include "zh_rf_gateway.h"

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

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/semphr.h"
#include "driver/spi_master.h"
#include "driver/gpio.h"
#include "esp_log.h"
#include "esp_rom_sys.h"
#include "nvs_flash.h"
#include "nvs.h"

ESP_EVENT_DEFINE_BASE(ZH_NRF_GATEWAY_EVENTS);

static const char *TAG = "NRF_GW";

/* ================== nRF24 КОМАНДЫ ================== */
#define NRF_CMD_R_REGISTER 0x00
#define NRF_CMD_W_REGISTER 0x20
#define NRF_CMD_R_RX_PAYLOAD 0x61
#define NRF_CMD_W_TX_PAYLOAD 0xA0
#define NRF_CMD_FLUSH_TX 0xE1
#define NRF_CMD_FLUSH_RX 0xE2
#define NRF_CMD_NOP 0xFF
#define NRF_CMD_ACTIVATE 0x50
#define NRF_CMD_R_RX_PL_WID 0x60

/* ================== nRF24 РЕГИСТРЫ ================== */
#define NRF_REG_CONFIG 0x00
#define NRF_REG_EN_AA 0x01
#define NRF_REG_EN_RXADDR 0x02
#define NRF_REG_SETUP_AW 0x03
#define NRF_REG_SETUP_RETR 0x04
#define NRF_REG_RF_CH 0x05
#define NRF_REG_RF_SETUP 0x06
#define NRF_REG_STATUS 0x07
#define NRF_REG_RX_ADDR_P0 0x0A
#define NRF_REG_TX_ADDR 0x10
#define NRF_REG_FIFO_STATUS 0x17
#define NRF_REG_DYNPD 0x1C
#define NRF_REG_FEATURE 0x1D

/* ================== nRF24 БИТЫ ================== */
#define NRF_CONFIG_MASK_RX_DR 0x40
#define NRF_CONFIG_MASK_TX_DS 0x20
#define NRF_CONFIG_MASK_MAX_RT 0x10
#define NRF_CONFIG_EN_CRC 0x08
#define NRF_CONFIG_PWR_UP 0x02
#define NRF_CONFIG_PRIM_RX 0x01

#define NRF_CONFIG_ALL_MASKED                        \
    (NRF_CONFIG_MASK_RX_DR | NRF_CONFIG_MASK_TX_DS | \
     NRF_CONFIG_MASK_MAX_RT | NRF_CONFIG_EN_CRC)

#define NRF_CONFIG_RX_UNMASKED \
    (NRF_CONFIG_EN_CRC | NRF_CONFIG_PWR_UP | NRF_CONFIG_PRIM_RX)

#define NRF_STATUS_RX_DR 0x40
#define NRF_STATUS_TX_DS 0x20
#define NRF_STATUS_MAX_RT 0x10

#define NRF_FIFO_RX_EMPTY 0x01

#define NRF_FEATURE_EN_DPL 0x04
#define NRF_DYNPD_DPL_P0 0x01

/* NVS namespace and key */
#define NVS_NAMESPACE "zh_gw"
#define NVS_KEY_NEXT_ID "next_id"

/* ================== СОСТОЯНИЕ КОМПОНЕНТА ================== */
static struct
{
    zh_nrf_gateway_config_t cfg;
    spi_device_handle_t spi;
    SemaphoreHandle_t irq_sem;
    TaskHandle_t task;
    bool initialized;
} s_ctx = {0};

/* ============================================================
 *                     SPI / nRF24 HELPERS
 * ============================================================ */
static void nrf_csn_low(void) { gpio_set_level(s_ctx.cfg.pin_csn, 0); }
static void nrf_csn_high(void) { gpio_set_level(s_ctx.cfg.pin_csn, 1); }
static void nrf_ce_low(void) { gpio_set_level(s_ctx.cfg.pin_ce, 0); }
static void nrf_ce_high(void) { gpio_set_level(s_ctx.cfg.pin_ce, 1); }

static uint8_t nrf_read_register(uint8_t reg)
{
    uint8_t tx[2] = {NRF_CMD_R_REGISTER | (reg & 0x1F), 0xFF};
    uint8_t rx[2] = {0};
    spi_transaction_t t = {.length = 16, .tx_buffer = tx, .rx_buffer = rx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
    return rx[1];
}

static void nrf_write_register(uint8_t reg, uint8_t val)
{
    uint8_t tx[2] = {NRF_CMD_W_REGISTER | (reg & 0x1F), val};
    spi_transaction_t t = {.length = 16, .tx_buffer = tx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
}

static void nrf_write_buffer(uint8_t reg, const uint8_t *buf, uint8_t len)
{
    uint8_t tx[32 + 1];
    tx[0] = NRF_CMD_W_REGISTER | (reg & 0x1F);
    memcpy(&tx[1], buf, len);
    spi_transaction_t t = {.length = (len + 1) * 8, .tx_buffer = tx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
}

static uint8_t nrf_send_command(uint8_t cmd)
{
    uint8_t tx[1] = {cmd};
    uint8_t rx[1] = {0};
    spi_transaction_t t = {.length = 8, .tx_buffer = tx, .rx_buffer = rx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
    return rx[0];
}

static void nrf_read_payload(uint8_t *buf, uint8_t len)
{
    uint8_t cmd = NRF_CMD_R_RX_PAYLOAD;
    nrf_csn_low();
    spi_transaction_t t_cmd = {.length = 8, .tx_buffer = &cmd};
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t_cmd));
    spi_transaction_t t_data = {.length = len * 8, .rx_buffer = buf};
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t_data));
    nrf_csn_high();
}

static uint8_t nrf_read_payload_width(void)
{
    uint8_t tx[2] = {NRF_CMD_R_RX_PL_WID, 0xFF};
    uint8_t rx[2] = {0};
    spi_transaction_t t = {.length = 16, .tx_buffer = tx, .rx_buffer = rx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
    return rx[1];
}

static void nrf_activate(void)
{
    uint8_t tx[2] = {NRF_CMD_ACTIVATE, 0x73};
    spi_transaction_t t = {.length = 16, .tx_buffer = tx};
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();
}

/* ============================================================
 *                     IRQ (GPIO ISR)
 * ============================================================ */
static void IRAM_ATTR nrf_irq_isr(void *arg)
{
    BaseType_t woken = pdFALSE;
    xSemaphoreGiveFromISR(s_ctx.irq_sem, &woken);
    if (woken)
        portYIELD_FROM_ISR();
}

static esp_err_t nrf_irq_init(void)
{
    s_ctx.irq_sem = xSemaphoreCreateBinary();
    if (!s_ctx.irq_sem)
        return ESP_ERR_NO_MEM;

    gpio_config_t io_conf = {
        .pin_bit_mask = (1ULL << s_ctx.cfg.pin_irq),
        .mode = GPIO_MODE_INPUT,
        .pull_up_en = GPIO_PULLUP_ENABLE,
        .pull_down_en = GPIO_PULLDOWN_DISABLE,
        .intr_type = GPIO_INTR_NEGEDGE,
    };
    esp_err_t err = gpio_config(&io_conf);
    if (err != ESP_OK)
        return err;

    err = gpio_install_isr_service(0);
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE)
        return err;

    return gpio_isr_handler_add(s_ctx.cfg.pin_irq, nrf_irq_isr, NULL);
}

/* ============================================================
 *                     NVS: ВЫДАЧА ID
 * ============================================================ */
static uint8_t allocate_device_id(void)
{
    nvs_handle_t h;
    uint8_t id = 0;

    if (nvs_open(NVS_NAMESPACE, NVS_READWRITE, &h) != ESP_OK)
    {
        ESP_LOGW(TAG, "NVS open failed, using fallback id=1");
        return 1;
    }

    if (nvs_get_u8(h, NVS_KEY_NEXT_ID, &id) != ESP_OK || id == 0 || id == 0xFF)
    {
        id = 1;
    }

    uint8_t next = id + 1;
    if (next == 0 || next == 0xFF)
        next = 1;

    nvs_set_u8(h, NVS_KEY_NEXT_ID, next);
    nvs_commit(h);
    nvs_close(h);

    ESP_LOGI(TAG, "Pairing: issued device_id=%u (next=%u)", id, next);
    return id;
}

/* ============================================================
 *                     ОТПРАВКА PAIRING-ACK
 * ============================================================ */
static void send_pairing_ack(uint8_t req_type, uint16_t token, uint8_t new_id)
{
    uint8_t pkt[ZH_RF_PAIRING_PACKET_SIZE] = {
        req_type,
        ZH_RF_PAIRING_ACK_MAGIC,
        (uint8_t)(token >> 8),
        (uint8_t)(token & 0xFF),
        new_id};

    /* 1. Заглушить ВСЕ IRQ + перейти в PTX standby */
    nrf_ce_low();
    nrf_write_register(NRF_REG_CONFIG, NRF_CONFIG_ALL_MASKED | NRF_CONFIG_PWR_UP);
    vTaskDelay(pdMS_TO_TICKS(2));

    /* 2. Сбросить только TX-флаги */
    nrf_write_register(NRF_REG_STATUS, NRF_STATUS_TX_DS | NRF_STATUS_MAX_RT);
    nrf_send_command(NRF_CMD_FLUSH_TX);

    /* 3. Записать payload и запустить TX */
    uint8_t tx[1 + ZH_RF_PAIRING_PACKET_SIZE];
    tx[0] = NRF_CMD_W_TX_PAYLOAD;
    memcpy(&tx[1], pkt, ZH_RF_PAIRING_PACKET_SIZE);
    spi_transaction_t t = {
        .length = (1 + ZH_RF_PAIRING_PACKET_SIZE) * 8,
        .tx_buffer = tx,
    };
    nrf_csn_low();
    ESP_ERROR_CHECK(spi_device_transmit(s_ctx.spi, &t));
    nrf_csn_high();

    nrf_ce_high();
    esp_rom_delay_us(20);
    nrf_ce_low();

    /* 4. Опрос TX_DS / MAX_RT */
    uint8_t status = 0;
    for (int i = 0; i < 100; i++)
    {
        status = nrf_send_command(NRF_CMD_NOP);
        if (status & (NRF_STATUS_TX_DS | NRF_STATUS_MAX_RT))
            break;
        vTaskDelay(pdMS_TO_TICKS(1));
    }

    if (status & NRF_STATUS_TX_DS)
    {
        ESP_LOGI(TAG, "Pairing ACK sent (type=%u, id=%u)", req_type, new_id);
    }
    else
    {
        ESP_LOGW(TAG, "Pairing ACK NOT delivered (status=0x%02X)", status);
    }

    /* 5. Сбросить TX-флаги */
    nrf_write_register(NRF_REG_STATUS, NRF_STATUS_TX_DS | NRF_STATUS_MAX_RT);
    nrf_send_command(NRF_CMD_FLUSH_TX);

    /* 6. Возврат в RX, снимаем маску */
    nrf_write_register(NRF_REG_CONFIG, NRF_CONFIG_RX_UNMASKED);
    vTaskDelay(pdMS_TO_TICKS(5));
    nrf_ce_high();
    vTaskDelay(pdMS_TO_TICKS(1));
}

/* ============================================================
 *                     ОБРАБОТКА ПАКЕТОВ
 * ============================================================ */

/* Pairing request: {type, 0xA5, token_hi, token_lo, 0}. Internal only. */
static void handle_pairing_request(const uint8_t *data)
{
    uint8_t req_type = data[0];
    uint16_t token = ((uint16_t)data[2] << 8) | data[3];

    ESP_LOGI(TAG, "Pairing request: type=%u (%s), token=0x%04X",
             req_type,
             zh_get_rf_sensor_type_name((zh_rf_sensor_type_t)req_type),
             token);

    uint8_t new_id = allocate_device_id();
    send_pairing_ack(req_type, token, new_id);
}

/* Data packet: parse and post to event loop. */
static void handle_sensor_data(const uint8_t *buf, uint8_t len)
{
    if (len < ZH_RF_DATA_HEADER_SIZE)
    {
        ESP_LOGW(TAG, "Too short data packet (%u)", len);
        return;
    }

    zh_rf_sensor_transmitted_data_t d = {0};
    d.id = zh_rf_sensor_read_i16_le_at(buf, 0);
    d.type = zh_rf_sensor_read_i16_le_at(buf, 2);
    d.battery = zh_rf_sensor_read_i16_le_at(buf, 4);

    uint8_t expected = zh_get_rf_sensor_payload_size(d.type);
    if (expected == 0)
    {
        ESP_LOGW(TAG, "Unknown sensor type %d, dropping", d.type);
        return;
    }
    if (len != expected)
    {
        ESP_LOGW(TAG, "Length %u != expected %u for type %d (%s)",
                 len, expected, d.type,
                 zh_get_rf_sensor_type_name((zh_rf_sensor_type_t)d.type));
    }

    switch (d.type)
    {
    case ZH_RF_SENSOR_TOUCH_SWITCH:
        break;

    case ZH_RF_SENSOR_WATER_LEAKAGE:
    case ZH_RF_SENSOR_OPEN_CLOSE:
    case ZH_RF_SENSOR_PLANT_HUMIDITY:
        if (len >= ZH_RF_DATA_HEADER_SIZE + ZH_RF_DATA_FIELD_SIZE)
            d.value_1 = zh_rf_sensor_read_i16_le_at(buf, 6);
        break;

    case ZH_RF_SENSOR_BMP280:
        if (len >= ZH_RF_DATA_HEADER_SIZE + 2 * ZH_RF_DATA_FIELD_SIZE)
        {
            d.value_1 = zh_rf_sensor_read_i16_le_at(buf, 6);
            d.value_2 = zh_rf_sensor_read_i16_le_at(buf, 8);
        }
        break;

    case ZH_RF_SENSOR_BME280:
        if (len >= ZH_RF_DATA_HEADER_SIZE + 3 * ZH_RF_DATA_FIELD_SIZE)
        {
            d.value_1 = zh_rf_sensor_read_i16_le_at(buf, 6);
            d.value_2 = zh_rf_sensor_read_i16_le_at(buf, 8);
            d.value_3 = zh_rf_sensor_read_i16_le_at(buf, 10);
        }
        break;

    case ZH_RF_SENSOR_BME680:
        if (len >= ZH_RF_DATA_HEADER_SIZE + 4 * ZH_RF_DATA_FIELD_SIZE)
        {
            d.value_1 = zh_rf_sensor_read_i16_le_at(buf, 6);
            d.value_2 = zh_rf_sensor_read_i16_le_at(buf, 8);
            d.value_3 = zh_rf_sensor_read_i16_le_at(buf, 10);
            d.value_4 = zh_rf_sensor_read_i16_le_at(buf, 12);
        }
        break;

    default:
        break;
    }

    esp_err_t err = esp_event_post(ZH_NRF_GATEWAY_EVENTS,
                                   ZH_NRF_GATEWAY_EVENT_SENSOR_DATA,
                                   &d, sizeof(d),
                                   pdMS_TO_TICKS(100));
    if (err != ESP_OK)
    {
        ESP_LOGW(TAG, "esp_event_post failed: %s", esp_err_to_name(err));
    }
}

/* One packet from RX FIFO. Returns false if FIFO is empty. */
static bool process_one_packet(void)
{
    uint8_t fifo = nrf_read_register(NRF_REG_FIFO_STATUS);
    if (fifo & NRF_FIFO_RX_EMPTY)
        return false;

    uint8_t width = nrf_read_payload_width();
    if (width > ZH_RF_MAX_PAYLOAD_SIZE)
    {
        ESP_LOGW(TAG, "Corrupt payload width %u, FLUSH_RX", width);
        nrf_send_command(NRF_CMD_FLUSH_RX);
        nrf_write_register(NRF_REG_STATUS, NRF_STATUS_RX_DR);
        return true;
    }

    uint8_t buf[ZH_RF_MAX_PAYLOAD_SIZE] = {0};
    nrf_read_payload(buf, width);
    nrf_write_register(NRF_REG_STATUS, NRF_STATUS_RX_DR);

    /* Pairing: 5 bytes, magic @ offset 1. Data: >=6 bytes, id @ offset 0. */
    if (width == ZH_RF_PAIRING_PACKET_SIZE && buf[1] == ZH_RF_PAIRING_REQ_MAGIC)
    {
        handle_pairing_request(buf);
    }
    else
    {
        handle_sensor_data(buf, width);
    }
    return true;
}

static void process_all_pending(void)
{
    while (process_one_packet())
    {
    }
}

/* ============================================================
 *                     RECEIVER TASK
 * ============================================================ */
static void gateway_task(void *arg)
{
    while (1)
    {
        if (xSemaphoreTake(s_ctx.irq_sem, pdMS_TO_TICKS(1000)) == pdTRUE)
        {
            process_all_pending();
        }
        else
        {
            /* Fallback polling once per second */
            process_all_pending();
        }
    }
}

/* ============================================================
 *                     PUBLIC API
 * ============================================================ */
esp_err_t zh_nrf_gateway_init(const zh_nrf_gateway_config_t *config)
{
    if (!config)
        return ESP_ERR_INVALID_ARG;
    if (s_ctx.initialized)
        return ESP_ERR_INVALID_STATE;

    s_ctx.cfg = *config;
    if (s_ctx.cfg.spi_clock_hz == 0)
        s_ctx.cfg.spi_clock_hz = 4 * 1000 * 1000;
    if (s_ctx.cfg.channel > 125)
        s_ctx.cfg.channel = 120;

    /* NVS - init if needed */
    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND)
    {
        ESP_LOGW(TAG, "NVS flash needs erase, erasing");
        ESP_ERROR_CHECK(nvs_flash_erase());
        err = nvs_flash_init();
    }
    if (err != ESP_OK)
        return err;

    /* Default event loop - create if needed */
    err = esp_event_loop_create_default();
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE)
        return err;

    /* SPI bus */
    spi_bus_config_t buscfg = {
        .miso_io_num = s_ctx.cfg.pin_miso,
        .mosi_io_num = s_ctx.cfg.pin_mosi,
        .sclk_io_num = s_ctx.cfg.pin_sck,
        .quadwp_io_num = -1,
        .quadhd_io_num = -1,
        .max_transfer_sz = ZH_RF_MAX_PAYLOAD_SIZE,
    };
    err = spi_bus_initialize(s_ctx.cfg.spi_host, &buscfg, SPI_DMA_CH_AUTO);
    if (err != ESP_OK)
        return err;

    spi_device_interface_config_t devcfg = {
        .clock_speed_hz = (int)s_ctx.cfg.spi_clock_hz,
        .mode = 0,
        .spics_io_num = s_ctx.cfg.pin_csn,
        .queue_size = 1,
    };
    err = spi_bus_add_device(s_ctx.cfg.spi_host, &devcfg, &s_ctx.spi);
    if (err != ESP_OK)
        return err;

    /* CE pin */
    gpio_config_t ce_conf = {
        .pin_bit_mask = (1ULL << s_ctx.cfg.pin_ce),
        .mode = GPIO_MODE_OUTPUT,
    };
    err = gpio_config(&ce_conf);
    if (err != ESP_OK)
        return err;

    nrf_ce_low();
    vTaskDelay(pdMS_TO_TICKS(100));

    /* nRF24 config */
    nrf_write_register(NRF_REG_CONFIG, 0x00);
    nrf_write_register(NRF_REG_EN_AA, 0x01);
    nrf_write_register(NRF_REG_EN_RXADDR, 0x01);
    nrf_write_register(NRF_REG_SETUP_AW, 0x01); /* 3-byte address */
    nrf_write_register(NRF_REG_SETUP_RETR, 0xFF);
    nrf_write_register(NRF_REG_RF_CH, s_ctx.cfg.channel);
    nrf_write_register(NRF_REG_RF_SETUP, 0x26); /* 1 Mbps, 0 dBm */

    uint8_t addr_bytes[3] = {
        (uint8_t)(s_ctx.cfg.address >> 0),
        (uint8_t)(s_ctx.cfg.address >> 8),
        (uint8_t)(s_ctx.cfg.address >> 16),
    };
    nrf_write_buffer(NRF_REG_TX_ADDR, addr_bytes, 3);
    nrf_write_buffer(NRF_REG_RX_ADDR_P0, addr_bytes, 3);

    nrf_write_register(NRF_REG_STATUS,
                       NRF_STATUS_RX_DR | NRF_STATUS_TX_DS | NRF_STATUS_MAX_RT);

    /* Enable Dynamic Payload Length */
    nrf_activate();
    nrf_write_register(NRF_REG_FEATURE, NRF_FEATURE_EN_DPL);
    nrf_write_register(NRF_REG_DYNPD, NRF_DYNPD_DPL_P0);

    /* Enter RX */
    nrf_write_register(NRF_REG_CONFIG, NRF_CONFIG_RX_UNMASKED);
    vTaskDelay(pdMS_TO_TICKS(5));
    nrf_ce_high();
    vTaskDelay(pdMS_TO_TICKS(1));

    /* IRQ GPIO */
    err = nrf_irq_init();
    if (err != ESP_OK)
        return err;

    xSemaphoreTake(s_ctx.irq_sem, 0);

    s_ctx.initialized = true;
    ESP_LOGI(TAG, "Initialized: ch=%u addr=0x%06X DPL on, RX mode",
             s_ctx.cfg.channel, (unsigned)s_ctx.cfg.address & 0xFFFFFF);
    return ESP_OK;
}

esp_err_t zh_nrf_gateway_start(void)
{
    if (!s_ctx.initialized)
        return ESP_ERR_INVALID_STATE;
    if (s_ctx.task)
        return ESP_OK;

    BaseType_t r = xTaskCreatePinnedToCore(gateway_task, "nrf_gateway",
                                           4096, NULL, 5,
                                           &s_ctx.task, tskNO_AFFINITY);
    if (r != pdPASS)
    {
        s_ctx.task = NULL;
        return ESP_ERR_NO_MEM;
    }
    ESP_LOGI(TAG, "Receiver task started");
    return ESP_OK;
}