#include "maze32_network.h"

#include <stdio.h>
#include <string.h>

#include "esp_log.h"
#include "esp_mac.h"
#include "esp_wifi.h"
#include "maze32_config_store.h"

#define MAZE32_CONNECTED_BIT BIT0
#define MAZE32_FAILED_BIT BIT1
#define MAZE32_STATION_RETRY_LIMIT 3U
#define MAZE32_STATION_TIMEOUT_MS 12000U

static const char *TAG = "maze32_network";
static maze32_network_t *active_network;

static void set_state(maze32_network_t *network, maze32_network_state_t state)
{
    network->state = state;
    ESP_LOGI(TAG, "Network state: %s", maze32_network_state_name(state));
}

static void network_event(void *argument,
                          esp_event_base_t event_base,
                          int32_t event_id,
                          void *event_data)
{
    (void)event_data;
    maze32_network_t *network = argument;
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START &&
        network->state == MAZE32_NETWORK_STATION_CONNECTING) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED &&
               network->state == MAZE32_NETWORK_STATION_CONNECTING) {
        if (network->retry_count < MAZE32_STATION_RETRY_LIMIT) {
            ++network->retry_count;
            ESP_LOGW(TAG, "Station connection failed; retry %u/%u",
                     network->retry_count, MAZE32_STATION_RETRY_LIMIT);
            esp_wifi_connect();
        } else {
            xEventGroupSetBits(network->events, MAZE32_FAILED_BIT);
        }
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        xEventGroupSetBits(network->events, MAZE32_CONNECTED_BIT);
    }
}

static esp_err_t stop_wifi(void)
{
    esp_err_t result = esp_wifi_stop();
    if (result == ESP_ERR_WIFI_NOT_STARTED || result == ESP_ERR_WIFI_NOT_INIT) {
        return ESP_OK;
    }
    return result;
}

static esp_err_t start_provisioning(maze32_network_t *network)
{
    set_state(network, MAZE32_NETWORK_PROVISIONING);
    esp_err_t result = stop_wifi();
    if (result != ESP_OK) {
        return result;
    }

    wifi_config_t access_point = {0};
    size_t ssid_length = strlen(network->access_point_ssid);
    memcpy(access_point.ap.ssid, network->access_point_ssid, ssid_length);
    access_point.ap.ssid_len = ssid_length;
    access_point.ap.channel = 1U;
    access_point.ap.max_connection = 4U;
    access_point.ap.authmode = WIFI_AUTH_OPEN;
    access_point.ap.pmf_cfg.required = false;

    result = esp_wifi_set_mode(WIFI_MODE_APSTA);
    if (result == ESP_OK) {
        result = esp_wifi_set_config(WIFI_IF_AP, &access_point);
    }
    if (result == ESP_OK) {
        result = esp_wifi_start();
    }
    if (result == ESP_OK) {
        ESP_LOGI(TAG, "Fallback access point ready: SSID=%s, URL=http://192.168.4.1/provision",
                 network->access_point_ssid);
    } else {
        set_state(network, MAZE32_NETWORK_ERROR);
    }
    return result;
}

static esp_err_t connect_station(maze32_network_t *network,
                                 const maze32_wifi_credentials_t *credentials)
{
    esp_err_t result = stop_wifi();
    if (result != ESP_OK) {
        return result;
    }

    wifi_config_t station = {0};
    memcpy(station.sta.ssid, credentials->ssid, strlen(credentials->ssid));
    memcpy(station.sta.password, credentials->password, strlen(credentials->password));
    station.sta.threshold.authmode = WIFI_AUTH_OPEN;
    station.sta.sae_pwe_h2e = WPA3_SAE_PWE_BOTH;

    network->retry_count = 0U;
    xEventGroupClearBits(network->events, MAZE32_CONNECTED_BIT | MAZE32_FAILED_BIT);
    set_state(network, MAZE32_NETWORK_STATION_CONNECTING);
    ESP_LOGI(TAG, "Connecting to configured station SSID: %s", credentials->ssid);

    result = esp_wifi_set_mode(WIFI_MODE_STA);
    if (result == ESP_OK) {
        result = esp_wifi_set_config(WIFI_IF_STA, &station);
    }
    if (result == ESP_OK) {
        result = esp_wifi_start();
    }
    if (result != ESP_OK) {
        ESP_LOGE(TAG, "Station startup failed: %s", esp_err_to_name(result));
        return start_provisioning(network);
    }

    EventBits_t bits = xEventGroupWaitBits(
        network->events,
        MAZE32_CONNECTED_BIT | MAZE32_FAILED_BIT,
        pdTRUE,
        pdFALSE,
        pdMS_TO_TICKS(MAZE32_STATION_TIMEOUT_MS));
    if ((bits & MAZE32_CONNECTED_BIT) != 0U) {
        set_state(network, MAZE32_NETWORK_STATION_CONNECTED);
        char address[16];
        if (maze32_network_ip_address(network, address, sizeof(address)) == ESP_OK) {
            ESP_LOGI(TAG, "Station connected: IP=%s", address);
        }
        return ESP_OK;
    }

    ESP_LOGW(TAG, "Station connection timed out or exhausted retries; enabling fallback");
    return start_provisioning(network);
}

static int save_credentials(const maze32_wifi_credentials_t *credentials, void *context)
{
    (void)context;
    return maze32_config_store_save(credentials) == ESP_OK ? 0 : -1;
}

static int clear_credentials(void *context)
{
    (void)context;
    return maze32_config_store_clear() == ESP_OK ? 0 : -1;
}

esp_err_t maze32_network_init(maze32_network_t *network,
                              maze32_reconfiguration_guard_fn guard,
                              void *guard_context)
{
    if (network == NULL || guard == NULL || active_network != NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    memset(network, 0, sizeof(*network));
    network->state = MAZE32_NETWORK_BOOTING;
    network->reconfiguration_guard = guard;
    network->guard_context = guard_context;
    network->events = xEventGroupCreate();
    network->operation_mutex = xSemaphoreCreateMutex();
    if (network->events == NULL || network->operation_mutex == NULL) {
        return ESP_ERR_NO_MEM;
    }

    uint8_t mac[6];
    esp_err_t result = esp_read_mac(mac, ESP_MAC_WIFI_SOFTAP);
    if (result != ESP_OK) {
        return result;
    }
    snprintf(network->access_point_ssid, sizeof(network->access_point_ssid),
             "maze32-%02X%02X", mac[4], mac[5]);

    wifi_init_config_t wifi_configuration = WIFI_INIT_CONFIG_DEFAULT();
    result = esp_wifi_init(&wifi_configuration);
    if (result != ESP_OK) {
        return result;
    }
    network->station_netif = esp_netif_create_default_wifi_sta();
    network->access_point_netif = esp_netif_create_default_wifi_ap();
    if (network->station_netif == NULL || network->access_point_netif == NULL) {
        return ESP_ERR_NO_MEM;
    }
    result = esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, network_event, network);
    if (result == ESP_OK) {
        result = esp_event_handler_register(IP_EVENT, IP_EVENT_STA_GOT_IP, network_event, network);
    }
    if (result == ESP_OK) {
        active_network = network;
    }
    return result;
}

esp_err_t maze32_network_start(maze32_network_t *network)
{
    if (network == NULL || xSemaphoreTake(network->operation_mutex, portMAX_DELAY) != pdTRUE) {
        return ESP_ERR_INVALID_ARG;
    }
    maze32_wifi_credentials_t credentials = {0};
    esp_err_t result = maze32_config_store_load(&credentials);
    if (result == ESP_OK) {
        result = connect_station(network, &credentials);
    } else if (result == ESP_ERR_NOT_FOUND || result == ESP_ERR_INVALID_VERSION ||
               result == ESP_ERR_INVALID_STATE) {
        ESP_LOGI(TAG, "No usable saved station configuration; enabling fallback");
        result = start_provisioning(network);
    }
    xSemaphoreGive(network->operation_mutex);
    return result;
}

esp_err_t maze32_network_apply_credentials(maze32_network_t *network,
                                           const maze32_wifi_credentials_t *credentials)
{
    if (network == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    int result = maze32_wifi_config_apply(credentials,
                                          network->reconfiguration_guard,
                                          save_credentials,
                                          network->guard_context);
    if (result == -1) {
        return ESP_ERR_INVALID_ARG;
    }
    if (result == -2) {
        return ESP_ERR_INVALID_STATE;
    }
    return result == 0 ? ESP_OK : ESP_FAIL;
}

esp_err_t maze32_network_reset_credentials(maze32_network_t *network)
{
    if (network == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    int result = maze32_wifi_config_reset(network->reconfiguration_guard,
                                          clear_credentials,
                                          network->guard_context);
    if (result == -2) {
        return ESP_ERR_INVALID_STATE;
    }
    return result == 0 ? ESP_OK : ESP_FAIL;
}

void maze32_network_retry_saved_credentials(void *context)
{
    maze32_network_t *network = context;
    if (network == NULL || xSemaphoreTake(network->operation_mutex, portMAX_DELAY) != pdTRUE) {
        return;
    }
    maze32_wifi_credentials_t credentials = {0};
    esp_err_t result = maze32_config_store_load(&credentials);
    if (result == ESP_OK) {
        result = connect_station(network, &credentials);
    }
    if (result != ESP_OK) {
        ESP_LOGE(TAG, "Saved station retry failed: %s", esp_err_to_name(result));
        start_provisioning(network);
    }
    xSemaphoreGive(network->operation_mutex);
}

void maze32_network_return_to_provisioning(void *context)
{
    maze32_network_t *network = context;
    if (network == NULL || xSemaphoreTake(network->operation_mutex, portMAX_DELAY) != pdTRUE) {
        return;
    }
    esp_err_t result = start_provisioning(network);
    if (result != ESP_OK) {
        ESP_LOGE(TAG, "Failed to restore provisioning access point: %s", esp_err_to_name(result));
    }
    xSemaphoreGive(network->operation_mutex);
}

esp_err_t maze32_network_scan(maze32_network_t *network,
                              maze32_wifi_scan_result_t *results,
                              size_t capacity,
                              size_t *count)
{
    if (network == NULL || results == NULL || count == NULL || capacity == 0U ||
        capacity > MAZE32_WIFI_SCAN_LIMIT ||
        xSemaphoreTake(network->operation_mutex, pdMS_TO_TICKS(1000U)) != pdTRUE) {
        return ESP_ERR_INVALID_ARG;
    }
    if (network->state != MAZE32_NETWORK_PROVISIONING) {
        xSemaphoreGive(network->operation_mutex);
        return ESP_ERR_INVALID_STATE;
    }

    wifi_scan_config_t scan_configuration = {0};
    esp_err_t result = esp_wifi_scan_start(&scan_configuration, true);
    uint16_t found = (uint16_t)capacity;
    wifi_ap_record_t records[MAZE32_WIFI_SCAN_LIMIT];
    if (result == ESP_OK) {
        result = esp_wifi_scan_get_ap_records(&found, records);
    }
    if (result == ESP_OK) {
        *count = found;
        for (size_t index = 0U; index < found; ++index) {
            memcpy(results[index].ssid, records[index].ssid, MAZE32_WIFI_SSID_MAX_BYTES);
            results[index].ssid[MAZE32_WIFI_SSID_MAX_BYTES] = '\0';
            results[index].rssi = records[index].rssi;
            results[index].auth_mode = (uint8_t)records[index].authmode;
        }
    }
    xSemaphoreGive(network->operation_mutex);
    return result;
}

maze32_network_state_t maze32_network_state(const maze32_network_t *network)
{
    return network == NULL ? MAZE32_NETWORK_ERROR : network->state;
}

const char *maze32_network_state_name(maze32_network_state_t state)
{
    switch (state) {
        case MAZE32_NETWORK_BOOTING: return "booting";
        case MAZE32_NETWORK_STATION_CONNECTING: return "station_connecting";
        case MAZE32_NETWORK_STATION_CONNECTED: return "station_connected";
        case MAZE32_NETWORK_PROVISIONING: return "provisioning";
        case MAZE32_NETWORK_ERROR: return "error";
        default: return "unknown";
    }
}

const char *maze32_network_access_point_ssid(const maze32_network_t *network)
{
    return network == NULL ? "" : network->access_point_ssid;
}

esp_err_t maze32_network_ip_address(const maze32_network_t *network,
                                    char *output,
                                    size_t output_size)
{
    if (network == NULL || output == NULL || output_size < 16U ||
        network->station_netif == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    esp_netif_ip_info_t information;
    esp_err_t result = esp_netif_get_ip_info(network->station_netif, &information);
    if (result == ESP_OK) {
        snprintf(output, output_size, IPSTR, IP2STR(&information.ip));
    }
    return result;
}