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

#include "configuration_ap.h"

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

#include <lwip/ip_addr.h>

#include <stdio.h>
#include <stdlib.h>
#include <ctype.h>

#include "platform_client.h"
#include "cJSON.h"
#include "freertos/task.h"

#define TAG "CONFIGURATION_AP"

#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT BIT1

#define WEB_SERVER_URL "http://192.168.4.1"

extern const char index_html_start[] asm("_binary_wifi_configuration_ap_html_start");
ap_t ap;
typedef enum { SETUP_IDLE, SETUP_CONNECTING, SETUP_SAVED, SETUP_WIFI_FAILED, SETUP_SAVE_FAILED } setup_state_t;
static setup_state_t setup_state;
static portMUX_TYPE setup_mux = portMUX_INITIALIZER_UNLOCKED;
static setup_state_t get_setup_state(void) {
  portENTER_CRITICAL(&setup_mux);
  setup_state_t state = setup_state;
  portEXIT_CRITICAL(&setup_mux);
  return state;
}
static void set_setup_state(setup_state_t state) {
  portENTER_CRITICAL(&setup_mux);
  setup_state = state;
  portEXIT_CRITICAL(&setup_mux);
}
static esp_err_t send_json(httpd_req_t *r, const char *status, const char *body) {
  httpd_resp_set_status(r, status);
  httpd_resp_set_type(r, "application/json; charset=utf-8");
  httpd_resp_set_hdr(r, "Cache-Control", "no-store");
  return httpd_resp_send(r, body, HTTPD_RESP_USE_STRLEN);
}


void start_access_point()
{
  // Get the SSID
  uint8_t mac[6];
  esp_read_mac(mac, ESP_MAC_WIFI_STA);
  snprintf(ap.ssid, 32, "%s-%02X%02X%02X", "VolcConvAI", mac[3], mac[4], mac[5]);
  // Initialize the TCP/IP stack
  ESP_ERROR_CHECK(esp_netif_init());

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

  // Set the router IP address to 192.168.4.1
  esp_netif_ip_info_t ip_info;
  IP4_ADDR(&ip_info.ip, 192, 168, 4, 1);
  IP4_ADDR(&ip_info.gw, 192, 168, 4, 1);
  IP4_ADDR(&ip_info.netmask, 255, 255, 255, 0);
  esp_netif_dhcps_stop(netif);
  esp_netif_set_ip_info(netif, &ip_info);
  esp_netif_dhcps_start(netif);

  // Initialize the WiFi stack in Access Point mode
  wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
  ESP_ERROR_CHECK(esp_wifi_init(&cfg));

  // Set the WiFi configuration
  wifi_config_t wifi_config = {};
  strncpy((char *)wifi_config.ap.ssid, ap.ssid, sizeof(wifi_config.ap.ssid));
  wifi_config.ap.ssid_len = strlen(ap.ssid);
  wifi_config.ap.max_connection = 4;
  wifi_config.ap.authmode = WIFI_AUTH_WPA2_PSK;
  strncpy((char *)wifi_config.ap.password, ap.ssid, sizeof(wifi_config.ap.password) - 1);
  wifi_config.ap.password[sizeof(wifi_config.ap.password) - 1] = '\0';

  // Start the WiFi Access Point
  ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_APSTA));
  ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
  ESP_ERROR_CHECK(esp_wifi_start());

  ESP_LOGI(TAG, "Access Point started with SSID %s", ap.ssid);
}

esp_err_t handler_index_html(httpd_req_t *r)
{
  httpd_resp_set_type(r, "text/html; charset=utf-8");
  httpd_resp_set_hdr(r, "Cache-Control", "no-store");
  httpd_resp_send(r, index_html_start, strlen(index_html_start));
  return ESP_OK;
}

esp_err_t handler_scan(httpd_req_t *r)
{
  if (get_setup_state() == SETUP_CONNECTING || get_setup_state() == SETUP_SAVED)
    return send_json(r, "409 Conflict", "{\"error\":\"busy\"}");
  if (esp_wifi_scan_start(NULL, true) != ESP_OK)
    return send_json(r, "503 Service Unavailable", "{\"error\":\"scan_failed\"}");
  uint16_t count = 0;
  esp_wifi_scan_get_ap_num(&count);
  if (count > 32) count = 32;
  wifi_ap_record_t *records = count ? calloc(count, sizeof(*records)) : NULL;
  if (count && !records) {
    esp_wifi_clear_ap_list();
    return send_json(r, "503 Service Unavailable", "{\"error\":\"no_memory\"}");
  }
  if (count && esp_wifi_scan_get_ap_records(&count, records) != ESP_OK) {
    free(records);
    esp_wifi_clear_ap_list();
    return send_json(r, "503 Service Unavailable", "{\"error\":\"scan_failed\"}");
  }
  esp_wifi_clear_ap_list();
  cJSON *list = cJSON_CreateArray();
  for (int i = 0; list && i < count; i++) {
    cJSON *item = cJSON_CreateObject();
    if (!item) break;
    cJSON_AddStringToObject(item, "ssid", (const char *)records[i].ssid);
    cJSON_AddNumberToObject(item, "rssi", records[i].rssi);
    cJSON_AddNumberToObject(item, "authmode", records[i].authmode);
    cJSON_AddItemToArray(list, item);
  }
  free(records);
  char *body = list ? cJSON_PrintUnformatted(list) : NULL;
  cJSON_Delete(list);
  esp_err_t ret = body ? send_json(r, "200 OK", body) :
    send_json(r, "503 Service Unavailable", "{\"error\":\"no_memory\"}");
  free(body);
  return ret;
}

bool connect2wifi(char *ssid, char *password, EventGroupHandle_t event_group)
{
  xEventGroupClearBits(event_group, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT);
  if (esp_wifi_disconnect() == ESP_OK) {
    xEventGroupWaitBits(event_group, WIFI_FAIL_BIT, pdTRUE, pdFALSE, pdMS_TO_TICKS(1000));
  }
  xEventGroupClearBits(event_group, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT);
  // auto esp_netif = esp_netif_create_default_wifi_sta();

  wifi_config_t wifi_config;
  bzero(&wifi_config, sizeof(wifi_config));
  strncpy((char *)wifi_config.sta.ssid, ssid, sizeof(wifi_config.sta.ssid));
  strncpy((char *)wifi_config.sta.password, password, sizeof(wifi_config.sta.password));
  wifi_config.sta.scan_method = WIFI_ALL_CHANNEL_SCAN;
  wifi_config.sta.failure_retry_cnt = 1;

  if (esp_wifi_set_config(WIFI_IF_STA, &wifi_config) != ESP_OK) return false;
  esp_err_t ret = esp_wifi_connect();
  if (ret != ESP_OK)
  {
    ESP_LOGE(TAG, "Failed to connect to WiFi: %d", ret);
    return false;
  }
  ESP_LOGI(TAG, "Connecting to WiFi %s", ssid);

  // Require an IP address, with a bounded wait so the form can be retried.
  EventBits_t bits = xEventGroupWaitBits(event_group, WIFI_CONNECTED_BIT | WIFI_FAIL_BIT, pdTRUE, pdFALSE, pdMS_TO_TICKS(30000));
  if (bits & WIFI_CONNECTED_BIT)
  {
    ESP_LOGI(TAG, "Connected to WiFi %s", ssid);
    return true;
  }
  else
  {
    ESP_LOGE(TAG, "Failed to connect to WiFi %s", ssid);
    return false;
  }
}

typedef struct {
  char ssid[33];
  char password[65];
  char device_id[37];
  char device_secret[129];
} setup_input_t;

static esp_err_t save_setup(const setup_input_t *input)
{
  nvs_handle_t nvs;
  esp_err_t ret = nvs_open("wifi", NVS_READWRITE, &nvs);
  if (ret != ESP_OK) return ret;
  // Keep recovery available if persisting either namespace fails.
  ret = nvs_set_u8(nvs, "setup_requested", 1);
  if (ret == ESP_OK) ret = nvs_commit(nvs);
  if (ret == ESP_OK) ret = platform_identity_save(input->device_id, input->device_secret);
  if (ret == ESP_OK) ret = nvs_set_str(nvs, "ssid", input->ssid);
  if (ret == ESP_OK) ret = nvs_set_str(nvs, "password", input->password);
  if (ret == ESP_OK) ret = nvs_set_u8(nvs, "setup_requested", 0);
  if (ret == ESP_OK) ret = nvs_commit(nvs);
  nvs_close(nvs);
  return ret;
}

static void setup_worker(void *arg)
{
  setup_input_t *input = arg;
  if (!connect2wifi(input->ssid, input->password, ap.event_group)) {
    set_setup_state(SETUP_WIFI_FAILED);
  } else {
    esp_err_t ret = save_setup(input);
    ESP_LOGI(TAG, "Setup save result: %s", esp_err_to_name(ret));
    set_setup_state(ret == ESP_OK ? SETUP_SAVED : SETUP_SAVE_FAILED);
  }
  memset(input, 0, sizeof(*input));
  free(input);
  if (get_setup_state() == SETUP_SAVED) {
    // Leave time for the browser to retrieve the result before AP shutdown.
    vTaskDelay(pdMS_TO_TICKS(15000));
    esp_restart();
  }
  vTaskDelete(NULL);
}

static bool read_field(cJSON *json, const char *name, char *out, size_t capacity, size_t min)
{
  cJSON *value = cJSON_GetObjectItemCaseSensitive(json, name);
  if (!cJSON_IsString(value) || !value->valuestring) return false;
  size_t length = strlen(value->valuestring);
  if (length < min || length >= capacity) return false;
  memcpy(out, value->valuestring, length + 1);
  return true;
}

esp_err_t handler_setup_status(httpd_req_t *r)
{
  const char *states[] = {"idle", "connecting", "saved", "wifi_failed", "save_failed"};
  char body[64];
  snprintf(body, sizeof(body), "{\"state\":\"%s\"}", states[get_setup_state()]);
  return send_json(r, "200 OK", body);
}

esp_err_t handler_form_submission(httpd_req_t *r)
{
  if (get_setup_state() == SETUP_CONNECTING || get_setup_state() == SETUP_SAVED)
    return send_json(r, "409 Conflict", "{\"error\":\"busy\"}");
  if (r->content_len <= 0 || r->content_len > 1536)
    return send_json(r, "400 Bad Request", "{\"error\":\"invalid_fields\"}");
  char *body = calloc(1, r->content_len + 1);
  if (!body) return send_json(r, "503 Service Unavailable", "{\"error\":\"no_memory\"}");
  size_t received = 0;
  while (received < r->content_len) {
    int count = httpd_req_recv(r, body + received, r->content_len - received);
    if (count <= 0) {
      memset(body, 0, r->content_len);
      free(body);
      return send_json(r, "408 Request Timeout", "{\"error\":\"receive_timeout\"}");
    }
    received += count;
  }
  cJSON *json = cJSON_Parse(body);
  memset(body, 0, r->content_len);
  free(body);
  setup_input_t *input = calloc(1, sizeof(*input));
  bool valid = input && cJSON_IsObject(json) &&
    read_field(json, "ssid", input->ssid, sizeof(input->ssid), 1) &&
    read_field(json, "password", input->password, sizeof(input->password), 0) &&
    read_field(json, "device_id", input->device_id, sizeof(input->device_id), 36) &&
    read_field(json, "device_secret", input->device_secret, sizeof(input->device_secret), 32);
  cJSON_Delete(json);
  if (!valid) {
    if (input) { memset(input, 0, sizeof(*input)); free(input); }
    ESP_LOGW(TAG, "Setup submission rejected: invalid fields");
    return send_json(r, "400 Bad Request", "{\"error\":\"invalid_fields\"}");
  }
  set_setup_state(SETUP_CONNECTING);
  if (xTaskCreate(setup_worker, "setup_connect", 6144, input, 5, NULL) != pdPASS) {
    memset(input, 0, sizeof(*input));
    free(input);
    set_setup_state(SETUP_IDLE);
    return send_json(r, "503 Service Unavailable", "{\"error\":\"no_memory\"}");
  }
  ESP_LOGI(TAG, "Setup submission accepted; connecting in background");
  return send_json(r, "202 Accepted", "{\"state\":\"connecting\"}");
}

void start_web_server()
{
  // Start the web server
  httpd_config_t config = HTTPD_DEFAULT_CONFIG();
  config.uri_match_fn = httpd_uri_match_wildcard;
  ESP_ERROR_CHECK(httpd_start(&(ap.server), &config));

  // Register the index.html file
  httpd_uri_t index_html = {
      .uri = "/",
      .method = HTTP_GET,
      .handler = handler_index_html,
      .user_ctx = NULL};
  ESP_ERROR_CHECK(httpd_register_uri_handler(ap.server, &index_html));

  // Register the /scan URI
  httpd_uri_t scan = {
      .uri = "/scan",
      .method = HTTP_GET,
      .handler = handler_scan,
      .user_ctx = NULL};
  ESP_ERROR_CHECK(httpd_register_uri_handler(ap.server, &scan));

  // Register the form submission
  httpd_uri_t form_submit = {
      .uri = "/submit",
      .method = HTTP_POST,
      .handler = handler_form_submission,
      .user_ctx = NULL};
  ESP_ERROR_CHECK(httpd_register_uri_handler(ap.server, &form_submit));

  httpd_uri_t status = { .uri = "/setup-status", .method = HTTP_GET,
      .handler = handler_setup_status };
  ESP_ERROR_CHECK(httpd_register_uri_handler(ap.server, &status));
  ESP_LOGI(TAG, "Web server started");
}

void handler_event_wifi_ap(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
  if (event_id == WIFI_EVENT_AP_STACONNECTED)
  {
    wifi_event_ap_staconnected_t *event = (wifi_event_ap_staconnected_t *)event_data;
    ESP_LOGI(TAG, "Station " MACSTR " joined, AID=%d", MAC2STR(event->mac), event->aid);
  }
  else if (event_id == WIFI_EVENT_AP_STADISCONNECTED)
  {
    wifi_event_ap_stadisconnected_t *event = (wifi_event_ap_stadisconnected_t *)event_data;
    ESP_LOGI(TAG, "Station " MACSTR " left, AID=%d", MAC2STR(event->mac), event->aid);
  }
  else if (event_id == WIFI_EVENT_STA_CONNECTED)
  {
    // Association alone does not mean DHCP has completed.
  }
  else if (event_id == WIFI_EVENT_STA_DISCONNECTED)
  {
    xEventGroupSetBits(ap.event_group, WIFI_FAIL_BIT);
  }
}

void handler_event_ip_ap(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
  if (event_id == IP_EVENT_STA_GOT_IP)
  {
    ip_event_got_ip_t *event = (ip_event_got_ip_t *)event_data;
    ESP_LOGI(TAG, "Got IP:" IPSTR, IP2STR(&event->ip_info.ip));
    xEventGroupSetBits(ap.event_group, WIFI_CONNECTED_BIT);
  }
}

void configure_ap()
{
  ap.ssid[0] = '\0';
  ap.event_group = xEventGroupCreate();
  // Register event handlers
  ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                      ESP_EVENT_ANY_ID,
                                                      &handler_event_wifi_ap,
                                                      NULL,
                                                      &(ap.instance_any_id)));

  ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                      IP_EVENT_STA_GOT_IP,
                                                      &handler_event_ip_ap,
                                                      NULL,
                                                      &(ap.instance_got_ip)));

  start_access_point();
  start_web_server();
  ESP_LOGI(TAG, "Setup portal ready at " WEB_SERVER_URL);
}