#include "wifi_transport.h"
#include "protocol.h"
#include <string.h>
#include <stdio.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/queue.h"
#include "freertos/semphr.h"
#include "esp_wifi.h"
#include "esp_wifi_ap_get_sta_list.h"
#include "esp_netif.h"
#include "esp_event.h"
#include "lwip/sockets.h"

#define WIFI_SSID     "EIS4"
#define WIFI_PASS     "eis4data"
#define WIFI_CHANNEL  1
#define WIFI_MAX_CONN 10

#define UDP_PORT           5941
#define UDP_BUF_SIZE       128
#define UDP_CLIENTS_MAX    16
#define REAP_THRESHOLD     10
#define REAP_WINDOW_MS     200
#define REAP_INTERVAL_MS   5000

static int udp_sock = -1;
static esp_netif_t *ap_netif;

static struct {
    struct sockaddr_in addr;
    uint8_t mac[6];
    uint32_t last_touch_ms;
    bool active;
} clients[UDP_CLIENTS_MAX];

static int client_count;
static SemaphoreHandle_t client_mutex;

static void client_touch(const struct sockaddr_in *addr)
{
    xSemaphoreTake(client_mutex, portMAX_DELAY);

    for (int i = 0; i < client_count; i++) {
        if (clients[i].addr.sin_addr.s_addr == addr->sin_addr.s_addr &&
            clients[i].addr.sin_port == addr->sin_port) {
            clients[i].last_touch_ms = xTaskGetTickCount() * portTICK_PERIOD_MS;
            xSemaphoreGive(client_mutex);
            return;
        }
    }

    if (client_count < UDP_CLIENTS_MAX) {
        clients[client_count].addr = *addr;
        clients[client_count].active = true;
        memset(clients[client_count].mac, 0, 6);

        wifi_sta_list_t sta_list;
        wifi_sta_mac_ip_list_t ip_list;
        if (esp_wifi_ap_get_sta_list(&sta_list) == ESP_OK &&
            esp_wifi_ap_get_sta_list_with_ip(&sta_list, &ip_list) == ESP_OK) {
            for (int j = 0; j < ip_list.num; j++) {
                if (ip_list.sta[j].ip.addr == addr->sin_addr.s_addr) {
                    memcpy(clients[client_count].mac, ip_list.sta[j].mac, 6);
                    break;
                }
            }
        }

        clients[client_count].last_touch_ms = xTaskGetTickCount() * portTICK_PERIOD_MS;
        client_count++;
        printf("UDP: client added (%d)\n", client_count);
    }

    xSemaphoreGive(client_mutex);
}

static void client_remove_by_mac(const uint8_t *mac)
{
    xSemaphoreTake(client_mutex, portMAX_DELAY);
    for (int i = 0; i < client_count; ) {
        if (memcmp(clients[i].mac, mac, 6) == 0) {
            clients[i] = clients[--client_count];
            printf("UDP: client removed by MAC (%d)\n", client_count);
        } else {
            i++;
        }
    }
    xSemaphoreGive(client_mutex);
}

/* caller must hold client_mutex */
static void client_remove_by_index(int idx)
{
    if (idx < 0 || idx >= client_count) return;
    printf("REAP: removing client %d\n", idx);
    clients[idx] = clients[--client_count];
}

static void parse_udp_sysex(const uint8_t *data, uint16_t len)
{
    if (len < 3 || data[0] != 0xF0 || data[1] != 0x7D)
        return;

    uint16_t end = 0;
    for (uint16_t i = 2; i < len; i++) {
        if (data[i] == 0xF7) { end = i; break; }
    }
    if (!end) return;

    Command cmd;
    memset(&cmd, 0, sizeof(cmd));
    cmd.type = data[2];

    switch (cmd.type) {
    case CMD_SET_SWEEP:
        if (len < 16) return;
        cmd.sweep.freq_start = decode_float(&data[3]);
        cmd.sweep.freq_stop  = decode_float(&data[8]);
        cmd.sweep.ppd        = decode_u16(&data[13]);
        break;
    case CMD_SET_RTIA:
        if (len < 4) return;
        cmd.rtia = data[3];
        break;
    case CMD_SET_RCAL:
        if (len < 4) return;
        cmd.rcal = data[3];
        break;
    case CMD_SET_ELECTRODE:
        if (len < 4) return;
        cmd.electrode = data[3];
        break;
    case CMD_START_LSV:
        if (len < 19) return;
        cmd.lsv.v_start   = decode_float(&data[3]);
        cmd.lsv.v_stop    = decode_float(&data[8]);
        cmd.lsv.scan_rate = decode_float(&data[13]);
        cmd.lsv.lp_rtia   = data[18];
        if (len >= 22)
            cmd.lsv.num_points = decode_u16(&data[19]);
        break;
    case CMD_START_AMP:
        if (len < 19) return;
        cmd.amp.v_hold       = decode_float(&data[3]);
        cmd.amp.interval_ms  = decode_float(&data[8]);
        cmd.amp.duration_s   = decode_float(&data[13]);
        cmd.amp.lp_rtia      = data[18];
        break;
    case CMD_START_CL:
        if (len < 34) return;
        cmd.cl.v_cond    = decode_float(&data[3]);
        cmd.cl.t_cond_ms = decode_float(&data[8]);
        cmd.cl.v_free    = decode_float(&data[13]);
        cmd.cl.v_total   = decode_float(&data[18]);
        cmd.cl.t_dep_ms  = decode_float(&data[23]);
        cmd.cl.t_meas_ms = decode_float(&data[28]);
        cmd.cl.lp_rtia   = data[33];
        break;
    case CMD_START_PH:
        if (len < 8) return;
        cmd.ph.stabilize_s = decode_float(&data[3]);
        break;
    case CMD_START_CLEAN:
        if (len < 13) return;
        cmd.clean.v_mv       = decode_float(&data[3]);
        cmd.clean.duration_s = decode_float(&data[8]);
        break;
    case CMD_SET_CELL_K:
        if (len < 8) return;
        cmd.cell_k = decode_float(&data[3]);
        break;
    case CMD_SET_CL_FACTOR:
        if (len < 8) return;
        cmd.cl_factor = decode_float(&data[3]);
        break;
    case CMD_SET_PH_CAL:
        if (len < 13) return;
        cmd.ph_cal.slope  = decode_float(&data[3]);
        cmd.ph_cal.offset = decode_float(&data[8]);
        break;
    case CMD_SESSION_CREATE:
        if (len < 5) return;
        cmd.session_create.name_len = data[3] & 0x7F;
        if (cmd.session_create.name_len > MAX_SESSION_NAME)
            cmd.session_create.name_len = MAX_SESSION_NAME;
        for (uint8_t i = 0; i < cmd.session_create.name_len && (4 + i) < end; i++)
            cmd.session_create.name[i] = data[4 + i];
        break;
    case CMD_SESSION_SWITCH:
        if (len < 4) return;
        cmd.session_switch.id = data[3] & 0x7F;
        break;
    case CMD_SESSION_RENAME:
        if (len < 6) return;
        cmd.session_rename.id = data[3] & 0x7F;
        cmd.session_rename.name_len = data[4] & 0x7F;
        if (cmd.session_rename.name_len > MAX_SESSION_NAME)
            cmd.session_rename.name_len = MAX_SESSION_NAME;
        for (uint8_t i = 0; i < cmd.session_rename.name_len && (5 + i) < end; i++)
            cmd.session_rename.name[i] = data[5 + i];
        break;
    case CMD_START_SWEEP:
    case CMD_GET_CONFIG:
    case CMD_STOP_AMP:
    case CMD_GET_TEMP:
    case CMD_GET_CELL_K:
    case CMD_GET_CL_FACTOR:
    case CMD_GET_PH_CAL:
    case CMD_START_REFS:
    case CMD_GET_REFS:
    case CMD_CLEAR_REFS:
    case CMD_OPEN_CAL:
    case CMD_CLEAR_OPEN_CAL:
    case CMD_SESSION_LIST:
    case CMD_HEARTBEAT:
        break;
    default:
        return;
    }

    protocol_push_command(&cmd);
}

static void udp_rx_task(void *param)
{
    (void)param;
    uint8_t buf[UDP_BUF_SIZE];
    struct sockaddr_in src;
    socklen_t slen = sizeof(src);

    for (;;) {
        int n = recvfrom(udp_sock, buf, sizeof(buf), 0,
                         (struct sockaddr *)&src, &slen);
        if (n <= 0) continue;

        client_touch(&src);
        parse_udp_sysex(buf, (uint16_t)n);
    }
}

static void udp_reaper_task(void *arg)
{
    (void)arg;
    for (;;) {
        vTaskDelay(pdMS_TO_TICKS(REAP_INTERVAL_MS));

        xSemaphoreTake(client_mutex, portMAX_DELAY);
        if (client_count < REAP_THRESHOLD) {
            xSemaphoreGive(client_mutex);
            continue;
        }

        uint32_t cutoff = xTaskGetTickCount() * portTICK_PERIOD_MS;
        printf("REAP: cycle start, %d clients\n", client_count);
        xSemaphoreGive(client_mutex);

        send_keepalive();
        vTaskDelay(pdMS_TO_TICKS(REAP_WINDOW_MS));

        xSemaphoreTake(client_mutex, portMAX_DELAY);
        int reaped = 0;
        for (int i = client_count - 1; i >= 0; i--) {
            if (clients[i].last_touch_ms < cutoff) {
                client_remove_by_index(i);
                reaped++;
            }
        }
        xSemaphoreGive(client_mutex);

        if (reaped) printf("REAP: removed %d zombie(s), %d remain\n", reaped, client_count);
    }
}

int wifi_send_sysex(const uint8_t *sysex, uint16_t len)
{
    if (udp_sock < 0 || client_count == 0)
        return -1;

    xSemaphoreTake(client_mutex, portMAX_DELAY);
    int sent = 0;
    for (int i = 0; i < client_count; i++) {
        int r = sendto(udp_sock, sysex, len, 0,
                       (struct sockaddr *)&clients[i].addr,
                       sizeof(clients[i].addr));
        if (r > 0) sent++;
    }
    xSemaphoreGive(client_mutex);
    return sent > 0 ? 0 : -1;
}

int wifi_get_client_count(void)
{
    return client_count;
}

static void wifi_event_handler(void *arg, esp_event_base_t base,
                               int32_t id, void *data)
{
    (void)arg;
    if (base == WIFI_EVENT) {
        if (id == WIFI_EVENT_AP_STACONNECTED) {
            printf("WiFi: station connected\n");
        } else if (id == WIFI_EVENT_AP_STADISCONNECTED) {
            wifi_event_ap_stadisconnected_t *evt = data;
            printf("WiFi: station disconnected\n");
            client_remove_by_mac(evt->mac);
        }
    }
}

#ifndef STA_SSID
#define STA_SSID ""
#endif
#ifndef STA_PASS
#define STA_PASS ""
#endif

static void sta_event_handler(void *arg, esp_event_base_t base,
                              int32_t id, void *data)
{
    (void)arg;
    if (base == WIFI_EVENT && id == WIFI_EVENT_STA_DISCONNECTED) {
        printf("WiFi: STA disconnected, reconnecting...\n");
        esp_wifi_connect();
    } else if (base == IP_EVENT && id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t *evt = (ip_event_got_ip_t *)data;
        printf("WiFi: STA connected, IP " IPSTR "\n", IP2STR(&evt->ip_info.ip));
    }
}

static int wifi_ap_init(void)
{
    ap_netif = esp_netif_create_default_wifi_ap();
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t wifi_cfg = WIFI_INIT_CONFIG_DEFAULT();
    esp_err_t err = esp_wifi_init(&wifi_cfg);
    if (err) return err;

    esp_event_handler_instance_t inst;
    esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID,
                                        wifi_event_handler, NULL, &inst);
    esp_event_handler_instance_register(WIFI_EVENT, WIFI_EVENT_STA_DISCONNECTED,
                                        sta_event_handler, NULL, &inst);
    esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP,
                                        sta_event_handler, NULL, &inst);

    wifi_config_t ap_cfg = {
        .ap = {
            .ssid = WIFI_SSID,
            .ssid_len = sizeof(WIFI_SSID) - 1,
            .channel = WIFI_CHANNEL,
            .password = WIFI_PASS,
            .max_connection = WIFI_MAX_CONN,
            .authmode = WIFI_AUTH_WPA2_PSK,
        },
    };

    esp_wifi_set_mode(WIFI_MODE_APSTA);
    esp_wifi_set_config(WIFI_IF_AP, &ap_cfg);

    if (strlen(STA_SSID) > 0) {
        wifi_config_t sta_cfg = {0};
        strncpy((char *)sta_cfg.sta.ssid, STA_SSID, sizeof(sta_cfg.sta.ssid) - 1);
        strncpy((char *)sta_cfg.sta.password, STA_PASS, sizeof(sta_cfg.sta.password) - 1);
        sta_cfg.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
        esp_wifi_set_config(WIFI_IF_STA, &sta_cfg);
    }

    err = esp_wifi_start();
    if (err) return err;

    esp_wifi_set_inactive_time(WIFI_IF_AP, 120);

    if (strlen(STA_SSID) > 0) {
        esp_wifi_connect();
        printf("WiFi: STA connecting to \"%s\"\n", STA_SSID);
    }

    printf("WiFi: AP \"%s\" on channel %d\n", WIFI_SSID, WIFI_CHANNEL);
    return 0;
}

static int udp_init(void)
{
    udp_sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (udp_sock < 0) return -1;

    struct sockaddr_in bind_addr = {
        .sin_family = AF_INET,
        .sin_port = htons(UDP_PORT),
        .sin_addr.s_addr = htonl(INADDR_ANY),
    };

    if (bind(udp_sock, (struct sockaddr *)&bind_addr, sizeof(bind_addr)) < 0) {
        close(udp_sock);
        udp_sock = -1;
        return -1;
    }

    printf("UDP: listening on port %d\n", UDP_PORT);
    return 0;
}

int wifi_transport_init(void)
{
    client_mutex = xSemaphoreCreateMutex();

    int rc = wifi_ap_init();
    if (rc) {
        printf("WiFi: AP init failed: %d\n", rc);
        return rc;
    }

    rc = udp_init();
    if (rc) {
        printf("UDP: init failed\n");
        return rc;
    }

    xTaskCreate(udp_rx_task, "udp_rx", 4096, NULL, 5, NULL);
    xTaskCreate(udp_reaper_task, "reaper", 2048, NULL, 4, NULL);
    return 0;
}