#include "platform.h"

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

#include "esp_event.h"
#include "esp_log.h"
#include "esp_netif.h"
#include "esp_wifi.h"
#include "esp_wifi_default.h"
#include "freertos/FreeRTOS.h"
#include "freertos/event_groups.h"

#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT BIT1

static const char *TAG = "wifi_manager";

static EventGroupHandle_t s_wifi_event_group;
static int s_retry_num;
static bool s_ready;
static bool s_started;
static esp_netif_t *s_sta_netif;
static esp_event_handler_instance_t s_wifi_event_inst;
static esp_event_handler_instance_t s_ip_event_inst;

static void wifi_manager_apply_sta_throughput_profile(void) {
    const uint8_t protocol = WIFI_PROTOCOL_11B | WIFI_PROTOCOL_11G | WIFI_PROTOCOL_11N;
    esp_err_t proto_rc = esp_wifi_set_protocol(WIFI_IF_STA, protocol);
    if (proto_rc != ESP_OK) {
        ESP_LOGW(TAG, "failed to set 11bgn protocol mask, rc=0x%x", (unsigned)proto_rc);
    } else {
        ESP_LOGI(TAG, "Wi-Fi protocol mask set to 11b/11g/11n");
    }

    esp_err_t bw_rc = esp_wifi_set_bandwidth(WIFI_IF_STA, WIFI_BW_HT40);
    if (bw_rc == ESP_OK) {
        ESP_LOGI(TAG, "Wi-Fi bandwidth set to HT40");
        return;
    }

    ESP_LOGW(TAG, "failed to set HT40, fallback to HT20, rc=0x%x", (unsigned)bw_rc);
    bw_rc = esp_wifi_set_bandwidth(WIFI_IF_STA, WIFI_BW_HT20);
    if (bw_rc != ESP_OK) {
        ESP_LOGW(TAG, "failed to set HT20 fallback, rc=0x%x", (unsigned)bw_rc);
    } else {
        ESP_LOGI(TAG, "Wi-Fi bandwidth set to HT20");
    }
}

static void wifi_event_handler(void *arg,
                               esp_event_base_t event_base,
                               int32_t event_id,
                               void *event_data) {
    (void)arg;
    (void)event_data;

    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        if (!s_started) {
            return;
        }
        (void)esp_wifi_connect();
        return;
    }

    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        s_ready = false;
        if (!s_started) {
            return;
        }
        if (s_retry_num < CONFIG_TQ_WIFI_MAXIMUM_RETRY) {
            (void)esp_wifi_connect();
            s_retry_num++;
            ESP_LOGW(TAG, "retry to connect to the AP (%d/%d)", s_retry_num, CONFIG_TQ_WIFI_MAXIMUM_RETRY);
        } else if (s_wifi_event_group != NULL) {
            xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT);
        }
        return;
    }

    if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        if (!s_started) {
            return;
        }
        ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
        ESP_LOGI(TAG, "got ip:" IPSTR, IP2STR(&event->ip_info.ip));
        s_retry_num = 0;
        s_ready = true;
        if (s_wifi_event_group != NULL) {
            xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
        }
        return;
    }
}

static esp_err_t start_sta_mode(void) {
    wifi_config_t wifi_config = {0};
    if (strlen(CONFIG_TQ_WIFI_SSID) == 0) {
        ESP_LOGE(TAG, "CONFIG_TQ_WIFI_SSID is empty; STA-only mode requires valid SSID");
        return ESP_ERR_INVALID_STATE;
    }

    strlcpy((char *)wifi_config.sta.ssid, CONFIG_TQ_WIFI_SSID, sizeof(wifi_config.sta.ssid));
    strlcpy((char *)wifi_config.sta.password, CONFIG_TQ_WIFI_PASSWORD, sizeof(wifi_config.sta.password));
    wifi_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;
    wifi_config.sta.pmf_cfg.capable = true;
    wifi_config.sta.pmf_cfg.required = false;

    esp_err_t rc = esp_wifi_set_mode(WIFI_MODE_STA);
    if (rc != ESP_OK) {
        return rc;
    }
    rc = esp_wifi_set_config(WIFI_IF_STA, &wifi_config);
    if (rc != ESP_OK) {
        return rc;
    }
    rc = esp_wifi_start();
    if (rc != ESP_OK && rc != ESP_ERR_INVALID_STATE) {
        return rc;
    }

    ESP_LOGI(TAG, "wifi_init_sta finished");

    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group,
                                           WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
                                           pdFALSE,
                                           pdFALSE,
                                           pdMS_TO_TICKS(runtime_policy_wifi_connect_timeout_ms()));

    if (bits & WIFI_CONNECTED_BIT) {
        ESP_LOGI(TAG, "connected to AP SSID:%s", CONFIG_TQ_WIFI_SSID);
        wifi_manager_apply_sta_throughput_profile();
        esp_err_t ps_rc = esp_wifi_set_ps(WIFI_PS_NONE);
        if (ps_rc == ESP_OK) {
            ESP_LOGI(TAG, "Wi-Fi power save disabled (WIFI_PS_NONE)");
        } else {
            ESP_LOGW(TAG, "failed to set WIFI_PS_NONE, rc=0x%x", (unsigned)ps_rc);
        }
        s_ready = true;
        return ESP_OK;
    }

    if (bits & WIFI_FAIL_BIT) {
        ESP_LOGW(TAG, "Failed to connect to SSID:%s", CONFIG_TQ_WIFI_SSID);
        return ESP_FAIL;
    }

    ESP_LOGW(TAG, "Wi-Fi connect timeout");
    return ESP_ERR_TIMEOUT;
}

esp_err_t wifi_manager_start(void) {
    if (s_started) {
        return s_ready ? ESP_OK : ESP_ERR_INVALID_STATE;
    }

    s_retry_num = 0;
    s_ready = false;

    if (s_wifi_event_group == NULL) {
        s_wifi_event_group = xEventGroupCreate();
        if (s_wifi_event_group == NULL) {
            return ESP_ERR_NO_MEM;
        }
    }
    xEventGroupClearBits(s_wifi_event_group, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT);

    esp_err_t err = esp_netif_init();
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
        runtime_diag_record_error("wifi_start", err, "esp_netif_init failed");
        return err;
    }

    err = esp_event_loop_create_default();
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
        runtime_diag_record_error("wifi_start", err, "event loop init failed");
        return err;
    }

    if (s_sta_netif == NULL) {
        s_sta_netif = esp_netif_create_default_wifi_sta();
        if (s_sta_netif == NULL) {
            runtime_diag_record_error("wifi_start", ESP_ERR_NO_MEM, "create sta netif failed");
            return ESP_ERR_NO_MEM;
        }
    }

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    err = esp_wifi_init(&cfg);
    if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
        runtime_diag_record_error("wifi_start", err, "esp_wifi_init failed");
        return err;
    }

    if (s_wifi_event_inst == NULL) {
        err = esp_event_handler_instance_register(WIFI_EVENT,
                                                  ESP_EVENT_ANY_ID,
                                                  &wifi_event_handler,
                                                  NULL,
                                                  &s_wifi_event_inst);
        if (err != ESP_OK) {
            runtime_diag_record_error("wifi_start", err, "register WIFI event failed");
            return err;
        }
    }

    if (s_ip_event_inst == NULL) {
        err = esp_event_handler_instance_register(IP_EVENT,
                                                  IP_EVENT_STA_GOT_IP,
                                                  &wifi_event_handler,
                                                  NULL,
                                                  &s_ip_event_inst);
        if (err != ESP_OK) {
            runtime_diag_record_error("wifi_start", err, "register IP event failed");
            return err;
        }
    }

    s_started = true;
    err = start_sta_mode();
    if (err == ESP_OK) {
        runtime_diag_counter_add(RUNTIME_DIAG_COUNTER_WIFI_CONNECT_SUCCESS, 1);
    } else {
        if (err == ESP_ERR_TIMEOUT) {
            runtime_diag_counter_add(RUNTIME_DIAG_COUNTER_WIFI_CONNECT_TIMEOUT, 1);
        } else {
            runtime_diag_counter_add(RUNTIME_DIAG_COUNTER_WIFI_CONNECT_FAILED, 1);
        }
        runtime_diag_record_error("wifi_connect", err, "STA connect failed");
        (void)wifi_manager_stop();
    }
    return err;
}

esp_err_t wifi_manager_stop(void) {
    if (!s_started) {
        return ESP_OK;
    }

    esp_err_t first_err = ESP_OK;

    s_ready = false;
    s_retry_num = 0;

    if (s_wifi_event_inst != NULL) {
        esp_err_t err = esp_event_handler_instance_unregister(WIFI_EVENT,
                                                              ESP_EVENT_ANY_ID,
                                                              s_wifi_event_inst);
        if (err != ESP_OK && first_err == ESP_OK) {
            first_err = err;
        }
        s_wifi_event_inst = NULL;
    }

    if (s_ip_event_inst != NULL) {
        esp_err_t err = esp_event_handler_instance_unregister(IP_EVENT,
                                                              IP_EVENT_STA_GOT_IP,
                                                              s_ip_event_inst);
        if (err != ESP_OK && first_err == ESP_OK) {
            first_err = err;
        }
        s_ip_event_inst = NULL;
    }

    esp_err_t err = esp_wifi_disconnect();
    if (err != ESP_OK && err != ESP_ERR_WIFI_NOT_STARTED && err != ESP_ERR_WIFI_CONN) {
        if (first_err == ESP_OK) {
            first_err = err;
        }
    }

    err = esp_wifi_stop();
    if (err != ESP_OK && err != ESP_ERR_WIFI_NOT_INIT && err != ESP_ERR_WIFI_NOT_STARTED) {
        if (first_err == ESP_OK) {
            first_err = err;
        }
    }

    err = esp_wifi_deinit();
    if (err != ESP_OK && err != ESP_ERR_WIFI_NOT_INIT) {
        if (first_err == ESP_OK) {
            first_err = err;
        }
    }

    if (s_sta_netif != NULL) {
        esp_netif_destroy_default_wifi(s_sta_netif);
        s_sta_netif = NULL;
    }

    if (s_wifi_event_group != NULL) {
        vEventGroupDelete(s_wifi_event_group);
        s_wifi_event_group = NULL;
    }

    s_started = false;

    if (first_err != ESP_OK) {
        runtime_diag_record_error("wifi_stop", first_err, "wifi stop failed");
    }

    return first_err;
}

bool wifi_manager_is_ready(void) {
    return s_ready;
}

void wifi_manager_get_ip(char *buf, size_t buf_len) {
    if (buf == NULL || buf_len == 0) {
        return;
    }

    buf[0] = '\0';

    esp_netif_ip_info_t ip_info;
    memset(&ip_info, 0, sizeof(ip_info));

    if (s_sta_netif != NULL && esp_netif_get_ip_info(s_sta_netif, &ip_info) == ESP_OK) {
        snprintf(buf, buf_len, IPSTR, IP2STR(&ip_info.ip));
        return;
    }

    strlcpy(buf, "0.0.0.0", buf_len);
}