#include "esp_event.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "nvs_flash.h"

static const char *TAG = "WPS_EXAMPLE";

// Function to start WPS (e.g., called by button interrupt)
void start_wps_process()
{
    ESP_LOGI(TAG, "Starting WPS...");
    // esp_wifi_wps_config_t wps_config = WPS_CONFIG_INIT_DEFAULT(WPS_TYPE_PBC);
    // ESP_ERROR_CHECK(esp_wifi_wps_enable(&wps_config));
    // ESP_ERROR_CHECK(esp_wifi_wps_start(0)); // Start WPS indefinitely (handle timeout via event)
}

static void wifi_event_handler(void *arg, esp_event_base_t event_base, int32_t event_id, void *event_data)
{
    if (event_base == WIFI_EVENT)
    {
        switch (event_id)
        {
        case WIFI_EVENT_STA_START:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_START: Wi-Fi station started.");
            // Option 1: Start WPS directly here
            // start_wps_process();
            // Option 2: Wait for external trigger (e.g., button), see start_wps_process()
            // Optional: Try to connect normally if credentials are stored
            // esp_wifi_connect();
            break;

        case WIFI_EVENT_STA_CONNECTED:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_CONNECTED");
            // If connected (e.g., from flash), disable WPS if it was still running
            // esp_wifi_wps_disable();
            break;

        case WIFI_EVENT_STA_DISCONNECTED:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_DISCONNECTED");
            // Possibly start normal connection attempt or wait for WPS trigger
            // esp_wifi_connect(); // Or: Wait for button for WPS
            break;

        case WIFI_EVENT_STA_WPS_ER_SUCCESS:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_SUCCESS");
            // WPS successful, credentials were received.
            // The Wi-Fi stack now automatically tries to connect.
            // Disable WPS.
            ESP_ERROR_CHECK(esp_wifi_wps_disable());
            // Optional: Read and save credentials if needed (usually not required as the stack uses them internally)
            break;

        case WIFI_EVENT_STA_WPS_ER_FAILED:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_FAILED");
            ESP_ERROR_CHECK(esp_wifi_wps_disable());
            // Here error handling: Info to user, possibly allow restart
            break;

        case WIFI_EVENT_STA_WPS_ER_TIMEOUT:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_TIMEOUT");
            ESP_ERROR_CHECK(esp_wifi_wps_disable());
            // Here error handling: Info to user, possibly allow restart
            break;

        case WIFI_EVENT_STA_WPS_ER_PIN:
            ESP_LOGI(TAG, "WIFI_EVENT_STA_WPS_ER_PIN");
            // Relevant for PIN mode
            // wifi_event_sta_wps_er_pin_t* event = (wifi_event_sta_wps_er_pin_t*) event_data;
            // ESP_LOGI(TAG, "WPS PIN: %.*s", event->pin_code_len, event->pin_code); // Example: ESP32 generates PIN
            break;

        default:
            break;
        }
    }
    else if (event_base == IP_EVENT && 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 address: " IPSTR, IP2STR(&event->ip_info.ip));
        // Connection successfully established
    }
}

void wifi_init_sta_wps(void)
{
    // 1. Initialize NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND)
    {
        ESP_ERROR_CHECK(nvs_flash_erase());
        ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    // 2. Initialize TCP/IP stack
    ESP_ERROR_CHECK(esp_netif_init());

    // 3. Create event loop
    ESP_ERROR_CHECK(esp_event_loop_create_default());

    // 4. Create default Wi-Fi Station Netif
    esp_netif_create_default_wifi_sta();

    // 5. Initialize Wi-Fi driver
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // 6. Register event handler
    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL));
    ESP_ERROR_CHECK(esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &wifi_event_handler, NULL));

    // 7. Set Wi-Fi mode to station
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));

    // 8. Start Wi-Fi (triggers WIFI_EVENT_STA_START)
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG, "Wi-Fi initialization completed.");
    ESP_LOGI(TAG, "Ready for WPS trigger (e.g., button press or automatic start in STA_START).");
}