// 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 <esp_heap_caps.h>
#include <esp_timer.h>
#include <nvs.h>
#include <nvs_flash.h>

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

#define TAG "NETWORK"

#define WIFI_EVENT_CONNECTED BIT0
#define WIFI_EVENT_FAILED BIT1
#define WIFI_EVENT_STARTED BIT2
#define MAX_RECONNECT_COUNT 5
#define RECONNECT_DELAY_US (10 * 1000 * 1000)

static network_t network;
static esp_err_t network_error;
static esp_timer_handle_t reconnect_timer;
static bool reconnect_enabled;
esp_err_t network_last_error(void) { return network_error; }

static void reconnect_saved_network(void *arg)
{
  (void)arg;
  if (!reconnect_enabled ||
      (xEventGroupGetBits(network.event_group) & WIFI_EVENT_CONNECTED)) return;
  esp_err_t err = esp_wifi_connect();
  if (err != ESP_OK) {
    ESP_LOGW(TAG, "WiFi reconnect deferred: %s", esp_err_to_name(err));
    esp_timer_start_once(reconnect_timer, RECONNECT_DELAY_US);
  }
}

static void schedule_reconnect(void)
{
  // Fast retries recover a short interruption; slower retries continue for as
  // long as the saved network is unavailable, including a phone hotspot.
  if (reconnect_enabled && reconnect_timer && !esp_timer_is_active(reconnect_timer))
    esp_timer_start_once(reconnect_timer, RECONNECT_DELAY_US);
}

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);
    network.ip_address[0] = '\0';
    if (!reconnect_enabled) return;
    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_err_t err = esp_wifi_connect();
      network.reconnect_count++;
      ESP_LOGI(TAG, "Reconnecting WiFi (attempt %d)", network.reconnect_count);
      if (err != ESP_OK) schedule_reconnect();
    }
    else
    {
      xEventGroupSetBits(network.event_group, WIFI_EVENT_FAILED);
      schedule_reconnect();
      ESP_LOGI(TAG, "WiFi unavailable; retrying every 10 seconds");
    }
  }
  else if (event_id == WIFI_EVENT_STA_STOP)
  {
    reconnect_enabled = false;
    if (reconnect_timer) esp_timer_stop(reconnect_timer);
    xEventGroupClearBits(network.event_group, WIFI_EVENT_CONNECTED | WIFI_EVENT_STARTED);
    network.ip_address[0] = '\0';
  }
}

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;
  if (reconnect_timer) esp_timer_stop(reconnect_timer);

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

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

bool configure_network()
{
  network_error=ESP_OK;

  // 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();
  if(!network.event_group){network_error=ESP_ERR_NO_MEM;return false;}
  if (!reconnect_timer) {
    const esp_timer_create_args_t timer_config = {
      .callback = reconnect_saved_network,
      .name = "wifi_reconnect"
    };
    network_error = esp_timer_create(&timer_config, &reconnect_timer);
    if (network_error != ESP_OK) return false;
  }

  // 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();
    // RTC and TLS share internal RAM with the Wi-Fi DMA buffers. Six static
    // RX/TX frames retain streaming capacity without reserving 32 DMA frames.
    cfg.static_rx_buf_num=6;
    cfg.static_tx_buf_num=6;
    // Dynamic RX buffers prefer PSRAM; match the 32KiB TCP receive window.
    cfg.dynamic_rx_buf_num=32;
    cfg.cache_tx_buf_num=8;
    cfg.rx_ba_win=6;
    ESP_LOGI(TAG,"WiFi init internal_free=%u largest=%u rx=%d tx=%d",
        (unsigned)heap_caps_get_free_size(MALLOC_CAP_INTERNAL),
        (unsigned)heap_caps_get_largest_free_block(MALLOC_CAP_INTERNAL),cfg.static_rx_buf_num,cfg.static_tx_buf_num);
    network_error=esp_wifi_init(&cfg);
    if(network_error!=ESP_OK){
        ESP_LOGE(TAG,"WiFi initialization unavailable: %s",esp_err_to_name(network_error));
        esp_event_handler_instance_unregister(WIFI_EVENT,ESP_EVENT_ANY_ID,instance_any_id);
        esp_event_handler_instance_unregister(IP_EVENT,IP_EVENT_STA_GOT_IP,instance_got_ip);
        esp_netif_destroy_default_wifi(station_netif);
        return false; // No abort/reboot loop on insufficient internal RAM.
    }
    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
    reconnect_enabled = true;
    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

      // 取消注册事件处理程序
      reconnect_enabled = false;
      esp_timer_stop(reconnect_timer);
      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())
  {
    if (ble_cloud_configured()) {
      // Completed BLE devices need the same phone flow when taken away from
      // their old Wi-Fi. Reboot to release normal-mode buffers before NimBLE.
      network_error = nvs_open("wifi", NVS_READWRITE, &nvs_handle);
      if (network_error == ESP_OK) {
        network_error = nvs_set_u8(nvs_handle, "setup_requested", 1);
        if (network_error == ESP_OK) network_error = nvs_commit(nvs_handle);
        nvs_close(nvs_handle);
      }
      ESP_LOGI(TAG, "Saved WiFi unavailable; requesting Bluetooth setup");
      return false;
    }
    configure_ap();
    ESP_LOGI(TAG, "WiFi not connected, starting AP");
    return false;
  }

  return true;
}