#include "wifi_handler.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_netif.h"
#include "freertos/event_groups.h"
#include "esp_log.h"
#include "freertos/semphr.h"
#include "common/semaphore_guard.h"
#include "cJSON.h"

#define TAG "WifiHandler"
#define WIFI_CRED_FILE_PATH "wifi_credentials.json"

/*
 * WiFi Credentials JSON Structure:
 * {
 *   "current_ssid": "MyWiFi",
 *   "credentials": {
 *     "MyWiFi": {
 *       "password": "mypassword123"
 *     },
 *     "OtherNetwork": {
 *       "password": "otherpass456"
 *     }
 *   }
 * }
 */

WifiHandler::WifiHandler(
  std::shared_ptr<LittleFSHandler> fs_handler_
) : fs_handler_(std::move(fs_handler_)) {
  this->s_wifi_event_group = xEventGroupCreate();
  if (!this->s_wifi_event_group) {
    ESP_LOGE(TAG, "Failed to create WiFi event group");
  }
  this->scan_mutex = xSemaphoreCreateMutex();
  if (!this->scan_mutex) {
    ESP_LOGE(TAG, "Failed to create scan mutex");
  }
  this->connection_mutex = xSemaphoreCreateMutex();
  if (!this->connection_mutex) {
    ESP_LOGE(TAG, "Failed to create connection mutex");
  }
  this->credential_mutex = xSemaphoreCreateMutex();
  if (!this->credential_mutex) {
    ESP_LOGE(TAG, "Failed to create credential mutex");
  }
  if (this->fs_handler_ == nullptr) {
    ESP_LOGW(TAG, "FSHandler is null, WiFi credentials will not be stored");
  } else {
    esp_err_t err = this->fs_handler_->init("/littlefs");
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Failed to initialize FSHandler: %s", esp_err_to_name(err));
    }
  }
}

WifiHandler::~WifiHandler() {
  if (this->initialized) {
    esp_wifi_stop();
    // Check if it should be called
    esp_wifi_deinit();
    vEventGroupDelete(this->s_wifi_event_group);
    if (!this->current_ssid.empty()) {
      this->current_ssid.clear();
    }
    vSemaphoreDelete(this->scan_mutex);
    vSemaphoreDelete(this->connection_mutex);
    esp_event_handler_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, &WifiHandler::wifi_event_handler);
    esp_event_handler_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, &WifiHandler::wifi_event_handler);
    this->initialized = false;
    //
    this->fs_handler_ = nullptr;
  }
}

esp_err_t WifiHandler::init() {
  if (this->initialized) {
    ESP_LOGW(TAG, "Already initialized, skipping");
    return ESP_OK;
  }
  esp_err_t err;

  // initialize TCP/IP stack and default event loop
  err = esp_netif_init();
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_netif_init failed: %s", esp_err_to_name(err));
    return err;
  }

  err = esp_event_loop_create_default();
  if (err != ESP_OK && err != ESP_ERR_INVALID_STATE) {
    ESP_LOGE(TAG, "esp_event_loop_create_default failed: %s", esp_err_to_name(err));
    return err;
  }

  // create default WiFi station
  esp_netif_create_default_wifi_sta();

  // init WiFi driver
  wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
  err = esp_wifi_init(&cfg);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_wifi_init failed: %s", esp_err_to_name(err));
    return err;
  }

  // register event handlers for WiFi and IP events
  err = esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &WifiHandler::wifi_event_handler, this);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_event_handler_register failed: %s", esp_err_to_name(err));
    return err;
  }
  err = esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &WifiHandler::wifi_event_handler, this);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_event_handler_register failed: %s", esp_err_to_name(err));
    return err;
  }

  err = esp_wifi_set_mode(WIFI_MODE_STA);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_wifi_set_mode failed: %s", esp_err_to_name(err));
    return err;
  }

  err = esp_wifi_start();
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "esp_wifi_start failed: %s", esp_err_to_name(err));
    return err;
  }

  // get WiFi credentials from KV storage if available
  std::string ssid;
  std::string password;
  this->get_wifi_credentials(ssid, password);

  // If KV storage didn't provide credentials, allow build-time injected values
  // via compile-time defines BUILD_WIFI_SSID and BUILD_WIFI_PASSWORD.
#if defined(BUILD_WIFI_SSID) and defined(BUILD_WIFI_PASSWORD)
  if (ssid.empty()) {
    ssid = std::string(BUILD_WIFI_SSID);
    ESP_LOGI(TAG, "Using build-time injected WiFi SSID");
  }
  if (password.empty()) {
    password = std::string(BUILD_WIFI_PASSWORD);
    ESP_LOGI(TAG, "Using build-time injected WiFi password");
  }
#endif

  if (!ssid.empty() && !password.empty()) {
    ESP_LOGI(TAG, "Found stored WiFi credentials, connecting to SSID: %s", ssid.c_str());
    err = this->connect(ssid, password);
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Failed to connect to stored WiFi credentials: %s", esp_err_to_name(err));
    }
  } else {
    ESP_LOGI(TAG, "No stored WiFi credentials found, not connecting");
  }

  initialized = true;
  return ESP_OK;
}

esp_err_t WifiHandler::connect(const std::string& ssid, const std::string& password) {
  SemaphoreGuard guard(this->connection_mutex);
  // wait up to 5 seconds to take the mutex
  if (!guard.take(5000 / portTICK_PERIOD_MS)) {
    ESP_LOGE(TAG, "Failed to take connection mutex");
    return ESP_FAIL;
  }

  expect_disconnected = false;
  if (!this->current_ssid.empty()) {
    this->current_ssid.clear();
  }
  this->current_ssid = ssid;
  this->current_password = password;

  //
  wifi_config_t wifi_config = {};
  strncpy((char*)wifi_config.sta.ssid, this->current_ssid.c_str(), sizeof(wifi_config.sta.ssid));
  wifi_config.sta.ssid[sizeof(wifi_config.sta.ssid) - 1] = '\0';
  strncpy((char*)wifi_config.sta.password, password.c_str(), sizeof(wifi_config.sta.password));
  wifi_config.sta.password[sizeof(wifi_config.sta.password) - 1] = '\0';
  // set auth mode to WPA2_PSK minimum
  wifi_config.sta.threshold.authmode = WIFI_AUTH_WPA2_PSK;

  ESP_LOGI(TAG, "Connecting to SSID: %s", this->current_ssid.c_str());
  esp_err_t err = esp_wifi_set_config(wifi_interface_t::WIFI_IF_STA, &wifi_config);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to set WiFi config: %s", esp_err_to_name(err));
    return err;
  }
  err = esp_wifi_connect();
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to initiate WiFi connection: %s", esp_err_to_name(err));
    return err;
  }

  // Note: Credentials will be stored in the event handler after successful connection
  // to avoid storing credentials for failed connection attempts

  return ESP_OK;
}

esp_err_t WifiHandler::connect(const std::string& ssid) {
  std::string stored_ssid;
  std::string stored_password;
  this->get_wifi_credentials(stored_ssid, stored_password);
  if (stored_ssid.empty() || stored_ssid != ssid) {
    ESP_LOGE(TAG, "No stored credentials for SSID: %s", ssid.c_str());
    return ESP_FAIL;
  }
  esp_err_t err = this->connect(stored_ssid, stored_password);
  return err;
}

esp_err_t WifiHandler::reconnect() {
  if (this->current_ssid.empty()) {
    ESP_LOGE(TAG, "No current SSID set, cannot reconnect");
    return ESP_FAIL;
  }
  return this->connect(this->current_ssid);
}

void WifiHandler::disconnect() {
  SemaphoreGuard guard(this->connection_mutex);
  // wait up to 5 seconds to take the mutex
  if (!guard.take(5000 / portTICK_PERIOD_MS)) {
    ESP_LOGE(TAG, "Failed to take connection mutex");
    return;
  }

  expect_disconnected = true;
  esp_wifi_disconnect();
  xEventGroupClearBits(
    this->s_wifi_event_group,
    WIFI_CONNECTED_BIT
  );
}

esp_err_t WifiHandler::scan_networks(
  wifi_ap_record_t*& ap_records,
  uint16_t& ap_count
) {
  SemaphoreGuard guard(this->scan_mutex);
  // wait up to 5 seconds to take the mutex
  if (!guard.take(5000 / portTICK_PERIOD_MS)) {
    ESP_LOGE(TAG, "Failed to take scan mutex");
    return ESP_FAIL;
  }
  ap_records = nullptr;
  ap_count = 0;
  // start scan

  esp_err_t err = esp_wifi_scan_start(nullptr, true); // block until done
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to start WiFi scan: %s", esp_err_to_name(err));
    return err;
  }

  // get number of APs found
  uint16_t ap_count_local = 0;
  err = esp_wifi_scan_get_ap_num(&ap_count_local);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to get number of APs found: %s", esp_err_to_name(err));
    return err;
  }

  wifi_ap_record_t* ap_records_local = new wifi_ap_record_t[ap_count_local];
  err = esp_wifi_scan_get_ap_records(&ap_count_local, ap_records_local);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to get AP records: %s", esp_err_to_name(err));
    delete[] ap_records_local;
    return err;
  }

  ap_records = ap_records_local;
  ap_count = ap_count_local;
  return ESP_OK;
}

void WifiHandler::wifi_event_handler(void* arg, esp_event_base_t event_base, int32_t event_id, void* event_data) {
  WifiHandler* self = static_cast<WifiHandler*>(arg);
  if (self == nullptr) {
    ESP_LOGE(TAG, "wifi_event_handler received null WifiHandler pointer");
    return;
  }

  switch (event_id) {
    case WIFI_EVENT_STA_START:
      // When the station starts, attempt to connect
      ESP_LOGI(TAG, "WIFI_EVENT_STA_START");
      if (!self->expect_disconnected && !self->current_ssid.empty()) {
        ESP_LOGI(TAG, "Station started, attempting to connect to SSID: %s", self->current_ssid.c_str());
        self->reconnect();
      }
      // set the event bit to indicate started
      xEventGroupSetBits(
        self->s_wifi_event_group,
        WIFI_STARTED_BIT
      );
      break;
    case WIFI_EVENT_STA_DISCONNECTED:
      ESP_LOGI(TAG, "WIFI_EVENT_STA_DISCONNECTED");
      if (!self->expect_disconnected) {
        ESP_LOGI(TAG, "Unexpected disconnection, attempting to reconnect");
        self->reconnect();
      }
      xEventGroupClearBits(
        self->s_wifi_event_group,
        WIFI_CONNECTED_BIT
      );
      break;
    case IP_EVENT_STA_GOT_IP:
    {
      ip_event_got_ip_t* event = static_cast<ip_event_got_ip_t*>(event_data);
      ESP_LOGI(TAG, "WIFI_EVENT_STA_GOT_IP: " IPSTR, IP2STR(&event->ip_info.ip));
      xEventGroupSetBits(
        self->s_wifi_event_group,
        WIFI_CONNECTED_BIT
      );
      // Store credentials only after successful connection
      if (!self->current_ssid.empty() && !self->current_password.empty()) {
        self->store_wifi_credentials(self->current_ssid, self->current_password);
      }
      break;
    }
    default:
      ESP_LOGW(TAG, "Unhandled WiFi event: %d", event_id);
      break;
  }
}

//
// private methods
//

void WifiHandler::store_wifi_credentials(const std::string& ssid, const std::string& password) {
  if (!fs_handler_) {
    ESP_LOGW(TAG, "FSHandler not set, cannot store WiFi credentials");
    return;
  }
  SemaphoreGuard guard(this->credential_mutex);
  // wait up to 5 seconds to take the mutex
  if (!guard.take(5000 / portTICK_PERIOD_MS)) {
    ESP_LOGE(TAG, "Failed to take credential mutex");
    return;
  }

  cJSON* json = nullptr;

  // Try to read existing credentials file
  if (fs_handler_->file_exists(WIFI_CRED_FILE_PATH)) {
    // Read existing file
    size_t file_size = 0;
    esp_err_t err = fs_handler_->get_file_size(WIFI_CRED_FILE_PATH, file_size);
    if (err == ESP_OK && file_size > 0) {
      std::vector<uint8_t> file_data(file_size + 1); // +1 for null terminator
      size_t bytes_read = 0;
      err = fs_handler_->read_file(WIFI_CRED_FILE_PATH, file_size, file_data.data(), bytes_read);
      if (err == ESP_OK) {
        file_data[bytes_read] = '\0'; // Null terminate
        json = cJSON_Parse(reinterpret_cast<const char*>(file_data.data()));
        if (json == nullptr) {
          ESP_LOGE(TAG, "Failed to parse existing WiFi credentials JSON, creating new");
        }
      }
    }
  }

  // Create new JSON if parsing failed or file doesn't exist
  if (json == nullptr) {
    json = cJSON_CreateObject();
  }

  // Set current SSID
  cJSON* current_ssid_item = cJSON_GetObjectItem(json, "current_ssid");
  if (current_ssid_item != nullptr) {
    cJSON_SetValuestring(current_ssid_item, ssid.c_str());
  } else {
    cJSON_AddStringToObject(json, "current_ssid", ssid.c_str());
  }

  // Get or create credentials object
  cJSON* credentials = cJSON_GetObjectItem(json, "credentials");
  if (credentials == nullptr || !cJSON_IsObject(credentials)) {
    credentials = cJSON_CreateObject();
    cJSON_AddItemToObject(json, "credentials", credentials);
  }

  // Limit stored credentials to prevent excessive file size (keep max 10 SSIDs)
  int credential_count = cJSON_GetArraySize(credentials);
  if (credential_count >= 10) {
    ESP_LOGW(TAG, "Too many stored credentials (%d), keeping only current SSID", credential_count);
    // Keep only the current SSID's credentials
    cJSON* new_credentials = cJSON_CreateObject();
    cJSON_ReplaceItemInObject(json, "credentials", new_credentials);
    credentials = new_credentials;
  }

  // Remove existing entry for this SSID to update it
  cJSON_DeleteItemFromObject(credentials, ssid.c_str());

  // Create SSID object with password
  cJSON* ssid_item = cJSON_CreateObject();
  cJSON_AddStringToObject(ssid_item, "password", password.c_str());
  cJSON_AddItemToObject(credentials, ssid.c_str(), ssid_item);

  // Serialize and write to file
  char* json_str = cJSON_PrintUnformatted(json);
  if (json_str) {
    size_t bytes_written = 0;
    esp_err_t err = fs_handler_->write_file(
      WIFI_CRED_FILE_PATH,
      reinterpret_cast<const uint8_t*>(json_str),
      strlen(json_str),
      bytes_written
    );
    if (err != ESP_OK) {
      ESP_LOGE(TAG, "Failed to write WiFi credentials to file: %s", esp_err_to_name(err));
    } else {
      ESP_LOGI(TAG, "Stored WiFi credentials for SSID: %s", ssid.c_str());
    }
    cJSON_free(json_str);
  } else {
    ESP_LOGE(TAG, "Failed to serialize WiFi credentials JSON");
  }
  cJSON_Delete(json);
}

void WifiHandler::get_wifi_credentials(std::string& out_ssid, std::string& out_password) {
  if (!fs_handler_) {
    ESP_LOGW(TAG, "FSHandler not set, cannot get WiFi credentials");
    out_ssid = "";
    out_password = "";
    return;
  }

  SemaphoreGuard guard(this->credential_mutex);
  // wait up to 5 seconds to take the mutex
  if (!guard.take(5000 / portTICK_PERIOD_MS)) {
    ESP_LOGE(TAG, "Failed to take credential mutex");
    out_ssid = "";
    out_password = "";
    return;
  }

  // Check if credentials file exists
  if (!fs_handler_->file_exists(WIFI_CRED_FILE_PATH)) {
    ESP_LOGD(TAG, "WiFi credentials file does not exist");
    out_ssid = "";
    out_password = "";
    return;
  }

  // Read credentials file
  size_t file_size = 0;
  esp_err_t err = fs_handler_->get_file_size(WIFI_CRED_FILE_PATH, file_size);
  if (err != ESP_OK || file_size == 0) {
    ESP_LOGE(TAG, "Failed to get WiFi credentials file size");
    out_ssid = "";
    out_password = "";
    return;
  }

  std::vector<uint8_t> file_data(file_size + 1); // +1 for null terminator
  size_t bytes_read = 0;
  err = fs_handler_->read_file(WIFI_CRED_FILE_PATH, file_size, file_data.data(), bytes_read);
  if (err != ESP_OK) {
    ESP_LOGE(TAG, "Failed to read WiFi credentials file: %s", esp_err_to_name(err));
    out_ssid = "";
    out_password = "";
    return;
  }
  file_data[bytes_read] = '\0'; // Null terminate

  // Parse JSON
  cJSON* json = cJSON_Parse(reinterpret_cast<const char*>(file_data.data()));
  if (json == nullptr) {
    ESP_LOGE(TAG, "Failed to parse WiFi credentials JSON");
    out_ssid = "";
    out_password = "";
    return;
  }

  // Get current SSID
  cJSON* current_ssid_item = cJSON_GetObjectItem(json, "current_ssid");
  if (current_ssid_item == nullptr || !cJSON_IsString(current_ssid_item)) {
    ESP_LOGE(TAG, "WiFi credentials JSON does not contain valid 'current_ssid' field");
    out_ssid = "";
    out_password = "";
    cJSON_Delete(json);
    return;
  }
  out_ssid = current_ssid_item->valuestring;

  // Get credentials object
  cJSON* credentials = cJSON_GetObjectItem(json, "credentials");
  if (credentials == nullptr || !cJSON_IsObject(credentials)) {
    ESP_LOGE(TAG, "WiFi credentials JSON does not contain valid 'credentials' object");
    out_ssid = "";
    out_password = "";
    cJSON_Delete(json);
    return;
  }

  // Get SSID entry
  cJSON* ssid_item = cJSON_GetObjectItem(credentials, out_ssid.c_str());
  if (ssid_item == nullptr || !cJSON_IsObject(ssid_item)) {
    ESP_LOGE(TAG, "WiFi credentials JSON does not contain entry for SSID: %s", out_ssid.c_str());
    out_ssid = "";
    out_password = "";
    cJSON_Delete(json);
    return;
  }

  // Get password
  cJSON* password_item = cJSON_GetObjectItem(ssid_item, "password");
  if (password_item == nullptr || !cJSON_IsString(password_item)) {
    ESP_LOGE(TAG, "WiFi credentials JSON does not contain valid 'password' field for SSID: %s", out_ssid.c_str());
    out_ssid = "";
    out_password = "";
    cJSON_Delete(json);
    return;
  }
  out_password = password_item->valuestring;

  cJSON_Delete(json);
  ESP_LOGD(TAG, "Retrieved WiFi credentials for SSID: %s", out_ssid.c_str());
}

EventBits_t WifiHandler::wait_for_connection(TickType_t ticks_to_wait) {
  return xEventGroupWaitBits(
    s_wifi_event_group,
    WIFI_CONNECTED_BIT,
    pdFALSE,
    pdTRUE,
    ticks_to_wait
  );
}