// Copyright (2025) Beijing Volcano Engine Technology Ltd.
// SPDX-License-Identifier: Apache-2.0

#include "network.h"

#include <freertos/FreeRTOS.h>
#include <freertos/event_groups.h>
#include <esp_err.h>
#include <esp_event.h>
#include <esp_wifi.h>
#include <esp_log.h>
#include <nvs.h>
#include <nvs_flash.h>

#include "configuration_ap.h"
#include "platform_client.h"

#define TAG "NETWORK"

#define WIFI_EVENT_CONNECTED BIT0
#define WIFI_EVENT_FAILED BIT1
#define WIFI_EVENT_STARTED BIT2
#define MAX_RECONNECT_COUNT 5

static network_t network;

int8_t get_rssi()
{
  // Get station info
  wifi_ap_record_t ap_info;
  if (esp_wifi_sta_get_ap_info(&ap_info) != ESP_OK) return 0;
  return ap_info.rssi;
}

uint8_t get_channel()
{
  // Get station info
  wifi_ap_record_t ap_info;
  ESP_ERROR_CHECK(esp_wifi_sta_get_ap_info(&ap_info));
  return ap_info.primary;
}

void handler_event_wifi_st(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
  if (event_id == WIFI_EVENT_STA_START)
  {
    xEventGroupSetBits(network.event_group, WIFI_EVENT_STARTED);
  }
  else if (event_id == WIFI_EVENT_STA_DISCONNECTED)
  {
    xEventGroupClearBits(network.event_group, WIFI_EVENT_CONNECTED);
    if (network.reconnect_count < MAX_RECONNECT_COUNT)
    {
      /* Retry the selected AP first, but allow another AP if it is no longer
       * reachable. Never permanently bind this device to one mesh node. */
      if (network.reconnect_count == 2) {
        wifi_config_t config;
        if (esp_wifi_get_config(WIFI_IF_STA, &config) == ESP_OK && config.sta.bssid_set) {
          config.sta.bssid_set = false;
          config.sta.channel = 0;
          esp_wifi_set_config(WIFI_IF_STA, &config);
          ESP_LOGW(TAG, "Preferred AP unavailable; falling back to SSID scan");
        }
      }
      esp_wifi_connect();
      network.reconnect_count++;
      ESP_LOGI(TAG, "Reconnecting WiFi (attempt %d)", network.reconnect_count);
    }
    else
    {
      xEventGroupSetBits(network.event_group, WIFI_EVENT_FAILED);
      ESP_LOGI(TAG, "WiFi connection failed");
    }
  }
}

void handler_event_ip_st(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
  ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
  network.reconnect_count = 0;

  esp_ip4addr_ntoa(&event->ip_info.ip, network.ip_address, sizeof(network.ip_address));
  ESP_LOGI(TAG, "Got IP: %s", network.ip_address);
  xEventGroupSetBits(network.event_group, WIFI_EVENT_CONNECTED);
}

bool is_connected()
{
  return xEventGroupGetBits(network.event_group) & WIFI_EVENT_CONNECTED;
}

bool configure_network()
{

  // Create the event group
  network.ssid[0] = '\0';
  network.password[0] = '\0';
  network.reconnect_count = 0;
  network.ip_address[0] = '\0';
  network.event_group = xEventGroupCreate();

  // Get ssid and password from NVS
  nvs_handle_t nvs_handle;
  uint8_t setup_requested = 0;
  esp_err_t ret = nvs_open("wifi", NVS_READONLY, &nvs_handle);
  if (ret == ESP_OK)
  {
    nvs_get_u8(nvs_handle, "setup_requested", &setup_requested);
    size_t length = sizeof(network.ssid);
    nvs_get_str(nvs_handle, "ssid", network.ssid, &length);
    length = sizeof(network.password);
    nvs_get_str(nvs_handle, "password", network.password, &length);
    nvs_close(nvs_handle);
  }
  if (setup_requested || !platform_identity_is_configured())
  {
    ESP_LOGI(TAG, "Starting setup: requested or platform identity missing");
    configure_ap();
    return false;
  }
  // Try to connect to WiFi, if failed, launch the WiFi configuration AP
  if (network.ssid[0] != '\0')
  {
    ESP_LOGI(TAG, "Connecting to WiFi %s", network.ssid);
    // Initialize the TCP/IP stack
    ESP_ERROR_CHECK(esp_netif_init());

    esp_event_handler_instance_t instance_any_id;
    esp_event_handler_instance_t instance_got_ip;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &handler_event_wifi_st,
                                                        NULL,
                                                        &instance_any_id));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &handler_event_ip_st,
                                                        NULL,
                                                        &instance_got_ip));

    // Create the default event loop
    esp_netif_t *station_netif = esp_netif_create_default_wifi_sta();

    // Initialize the WiFi stack in station mode
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_storage(WIFI_STORAGE_RAM));

    ESP_LOGI(TAG, "Connecting to WiFi ssid=%s password=***", network.ssid);
    wifi_config_t wifi_config;
    bzero(&wifi_config, sizeof(wifi_config));
    strncpy((char *)wifi_config.sta.ssid, network.ssid, sizeof(wifi_config.sta.ssid));
    strncpy((char *)wifi_config.sta.password, network.password, sizeof(wifi_config.sta.password));
    wifi_config.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;
    wifi_config.sta.sort_method = WIFI_CONNECT_AP_BY_SIGNAL;
    wifi_config.sta.threshold.authmode = WIFI_AUTH_OPEN;
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));

    // Start the WiFi stack
    ESP_ERROR_CHECK(esp_wifi_start());
    ESP_ERROR_CHECK(esp_wifi_set_ps(WIFI_PS_NONE));
    xEventGroupWaitBits(network.event_group, WIFI_EVENT_STARTED, pdFALSE, pdFALSE, pdMS_TO_TICKS(5000));
    wifi_scan_config_t scan = {.ssid = (uint8_t *)network.ssid};
    if (esp_wifi_scan_start(&scan, true) == ESP_OK) {
      wifi_ap_record_t records[16];
      uint16_t count = 16;
      if (esp_wifi_scan_get_ap_records(&count, records) == ESP_OK && count > 0) {
        uint16_t best = 0;
        for (uint16_t i = 1; i < count; ++i) {
          if (records[i].rssi > records[best].rssi) best = i;
        }
        wifi_config.sta.bssid_set = true;
        memcpy(wifi_config.sta.bssid, records[best].bssid, sizeof(wifi_config.sta.bssid));
        wifi_config.sta.channel = records[best].primary;
        ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
        ESP_LOGI(TAG, "Selected strongest AP among %u: rssi=%d channel=%u bssid=%02x:%02x:%02x:%02x:%02x:%02x",
                 count, records[best].rssi, records[best].primary,
                 records[best].bssid[0], records[best].bssid[1], records[best].bssid[2],
                 records[best].bssid[3], records[best].bssid[4], records[best].bssid[5]);
      }
    }
    ESP_ERROR_CHECK(esp_wifi_connect());

    // Wait for the WiFi stack to start
    EventBits_t bits = xEventGroupWaitBits(network.event_group, WIFI_EVENT_CONNECTED | WIFI_EVENT_FAILED, pdFALSE, pdFALSE, pdMS_TO_TICKS(30 * 1000));
    if (!(bits & WIFI_EVENT_CONNECTED))
    {
      ESP_LOGE(TAG, "WifiStation failed");
      // Reset the WiFi stack

      // 取消注册事件处理程序
      ESP_ERROR_CHECK(esp_event_handler_instance_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, instance_any_id));
      ESP_ERROR_CHECK(esp_event_handler_instance_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, instance_got_ip));
      ESP_ERROR_CHECK(esp_wifi_stop());
      ESP_ERROR_CHECK(esp_wifi_deinit());
      esp_netif_destroy_default_wifi(station_netif);
      xEventGroupClearBits(network.event_group, WIFI_EVENT_CONNECTED | WIFI_EVENT_FAILED);
      // nvs_flash_erase();
      // return false;
    } else {
        ESP_LOGI(TAG, "Connected to %s rssi=%d channel=%d", network.ssid, get_rssi(), get_channel());
    }
  }

  if (!is_connected())
  {
    configure_ap();
    ESP_LOGI(TAG, "WiFi not connected, starting AP");
    return false;
  }

  return true;
}