#include <ctype.h>
#include <string.h>
#include <stdlib.h>
#include <sys/param.h>

#include "esp_event.h"
#include "esp_log.h"
#include "esp_mac.h"
#include "esp_random.h"

#include "nvs_flash.h"
#include "esp_wifi.h"
#include "esp_netif.h"
#include "lwip/inet.h"

#include "esp_http_server.h"
#include "dns_server.h"

#include "freertos/FreeRTOS.h"
#include "freertos/task.h"

#include "epd7in3e.h"
#include "qr_onboarding.h"
#include "wifi_provisioning.h"

#define NVS_NAMESPACE "frame_cfg"

extern const char root_start[] asm("_binary_root_html_start");
extern const char root_end[] asm("_binary_root_html_end");

static const char *TAG = "wifi_provisioning";

/* Uppercase letters + digits, excluding visually ambiguous characters
 * (0/O, 1/I/L) so a human can read the password off the panel and type it
 * on a desktop/laptop that can't scan the QR code. */
static const char AP_PASSWORD_CHARSET[] = "ABCDEFGHJKMNPQRSTUVWXYZ23456789";

/* ------------------------------------------------------------------------
 * NVS-backed config
 * ---------------------------------------------------------------------- */

esp_err_t frame_config_load(frame_config_t *out)
{
    memset(out, 0, sizeof(*out));

    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return err;
    }

    size_t len = sizeof(out->sta_ssid);
    err = nvs_get_str(handle, "sta_ssid", out->sta_ssid, &len);
    if (err != ESP_OK) {
        nvs_close(handle);
        return err;
    }

    len = sizeof(out->toolsserver);
    err = nvs_get_str(handle, "toolsserver", out->toolsserver, &len);
    if (err != ESP_OK) {
        nvs_close(handle);
        return err;
    }

    /* Password is allowed to be empty (open home network), so a missing
     * key just means "no password" rather than "not provisioned". */
    len = sizeof(out->sta_password);
    esp_err_t pass_err = nvs_get_str(handle, "sta_pass", out->sta_password, &len);
    if (pass_err != ESP_OK && pass_err != ESP_ERR_NVS_NOT_FOUND) {
        nvs_close(handle);
        return pass_err;
    }

    nvs_close(handle);
    return ESP_OK;
}

esp_err_t frame_config_save(const frame_config_t *cfg)
{
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle);
    if (err != ESP_OK) {
        return err;
    }

    err = nvs_set_str(handle, "sta_ssid", cfg->sta_ssid);
    if (err == ESP_OK) {
        err = nvs_set_str(handle, "sta_pass", cfg->sta_password);
    }
    if (err == ESP_OK) {
        err = nvs_set_str(handle, "toolsserver", cfg->toolsserver);
    }
    if (err == ESP_OK) {
        /* Fresh (re)provisioning -- the next successful connection should
         * show the status screen again. */
        err = nvs_set_u8(handle, "connected_once", 0);
    }
    if (err == ESP_OK) {
        err = nvs_commit(handle);
    }

    nvs_close(handle);
    return err;
}

bool frame_config_has_connected_once(void)
{
    nvs_handle_t handle;
    esp_err_t err = nvs_open(NVS_NAMESPACE, NVS_READONLY, &handle);
    if (err != ESP_OK) {
        return false;
    }

    uint8_t value = 0;
    err = nvs_get_u8(handle, "connected_once", &value);
    nvs_close(handle);

    return err == ESP_OK && value != 0;
}

void frame_config_mark_connected_once(void)
{
    nvs_handle_t handle;
    if (nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle) != ESP_OK) {
        return;
    }
    nvs_set_u8(handle, "connected_once", 1);
    nvs_commit(handle);
    nvs_close(handle);
}

void frame_config_clear(void)
{
    nvs_handle_t handle;
    if (nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle) != ESP_OK) {
        return;
    }
    nvs_erase_key(handle, "sta_ssid");
    nvs_erase_key(handle, "sta_pass");
    nvs_erase_key(handle, "toolsserver");
    nvs_erase_key(handle, "connected_once");
    nvs_commit(handle);
    nvs_close(handle);
}

static void generate_ap_password(char *out, size_t out_size)
{
    size_t len = MIN(FRAME_AP_PASSWORD_LEN, out_size - 1);
    uint8_t random_bytes[FRAME_AP_PASSWORD_LEN];
    esp_fill_random(random_bytes, len);

    size_t charset_len = strlen(AP_PASSWORD_CHARSET);
    for (size_t i = 0; i < len; i++) {
        out[i] = AP_PASSWORD_CHARSET[random_bytes[i] % charset_len];
    }
    out[len] = '\0';
}

void ap_identity_get(char *ssid_out, size_t ssid_len, char *pass_out, size_t pass_len)
{
    /* Suffix with the last 3 bytes of the station MAC so two frames on the
     * same network never advertise the same SSID. esp_read_mac() works
     * before the WiFi driver is started. */
    uint8_t mac[6];
    ESP_ERROR_CHECK(esp_read_mac(mac, ESP_MAC_WIFI_STA));
    snprintf(ssid_out, ssid_len, "%s_%02X%02X%02X", CONFIG_ESP_AP_SSID, mac[3], mac[4], mac[5]);

    nvs_handle_t handle;
    ESP_ERROR_CHECK(nvs_open(NVS_NAMESPACE, NVS_READWRITE, &handle));

    size_t len = pass_len;
    esp_err_t err = nvs_get_str(handle, "ap_pass", pass_out, &len);
    if (err == ESP_ERR_NVS_NOT_FOUND) {
        generate_ap_password(pass_out, pass_len);
        ESP_ERROR_CHECK(nvs_set_str(handle, "ap_pass", pass_out));
        ESP_ERROR_CHECK(nvs_commit(handle));
    } else {
        ESP_ERROR_CHECK(err);
    }

    nvs_close(handle);
}

/* ------------------------------------------------------------------------
 * Tiny application/x-www-form-urlencoded parser for the provisioning form
 * ---------------------------------------------------------------------- */

static void url_decode(char *dst, const char *src, size_t dst_size)
{
    size_t di = 0;
    for (size_t si = 0; src[si] != '\0' && di + 1 < dst_size; si++) {
        char c = src[si];
        if (c == '+') {
            dst[di++] = ' ';
        } else if (c == '%' && isxdigit((unsigned char)src[si + 1]) && isxdigit((unsigned char)src[si + 2])) {
            char hex[3] = { src[si + 1], src[si + 2], '\0' };
            dst[di++] = (char) strtol(hex, NULL, 16);
            si += 2;
        } else {
            dst[di++] = c;
        }
    }
    dst[di] = '\0';
}

static void extract_form_value(const char *body, const char *key, char *out, size_t out_size)
{
    out[0] = '\0';
    size_t key_len = strlen(key);
    const char *p = body;
    while (p != NULL && *p != '\0') {
        const char *amp = strchr(p, '&');
        size_t token_len = amp ? (size_t)(amp - p) : strlen(p);
        if (token_len > key_len && p[key_len] == '=' && strncmp(p, key, key_len) == 0) {
            char raw[FRAME_CFG_SERVER_MAX_LEN + 1] = {0};
            size_t value_len = token_len - key_len - 1;
            if (value_len > sizeof(raw) - 1) {
                value_len = sizeof(raw) - 1;
            }
            memcpy(raw, p + key_len + 1, value_len);
            raw[value_len] = '\0';
            url_decode(out, raw, out_size);
            return;
        }
        if (amp == NULL) {
            break;
        }
        p = amp + 1;
    }
}

/* ------------------------------------------------------------------------
 * HTTP handlers
 * ---------------------------------------------------------------------- */

static esp_err_t root_get_handler(httpd_req_t *req)
{
    const uint32_t root_len = root_end - root_start;

    ESP_LOGI(TAG, "Serve root");
    httpd_resp_set_type(req, "text/html");
    httpd_resp_send(req, root_start, root_len);

    return ESP_OK;
}

static const httpd_uri_t root = {
    .uri = "/",
    .method = HTTP_GET,
    .handler = root_get_handler
};

static esp_err_t save_config_post_handler(httpd_req_t *req)
{
    if (req->content_len <= 0 || req->content_len > 512) {
        httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "Form data too large");
        return ESP_FAIL;
    }

    char body[513];
    int received = 0;
    while (received < req->content_len) {
        int ret = httpd_req_recv(req, body + received, req->content_len - received);
        if (ret == HTTPD_SOCK_ERR_TIMEOUT) {
            continue;
        }
        if (ret <= 0) {
            httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "Failed to read body");
            return ESP_FAIL;
        }
        received += ret;
    }
    body[received] = '\0';

    frame_config_t cfg = {0};
    extract_form_value(body, "ssid", cfg.sta_ssid, sizeof(cfg.sta_ssid));
    extract_form_value(body, "password", cfg.sta_password, sizeof(cfg.sta_password));
    extract_form_value(body, "toolsserver", cfg.toolsserver, sizeof(cfg.toolsserver));

    if (strlen(cfg.sta_ssid) == 0 || strlen(cfg.toolsserver) == 0) {
        httpd_resp_send_err(req, HTTPD_400_BAD_REQUEST, "SSID and Tools Server are required");
        return ESP_FAIL;
    }

    esp_err_t err = frame_config_save(&cfg);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to save config: %s", esp_err_to_name(err));
        httpd_resp_send_500(req);
        return ESP_FAIL;
    }

    ESP_LOGI(TAG, "Saved config: ssid='%s' toolsserver='%s'", cfg.sta_ssid, cfg.toolsserver);

    static const char resp[] =
        "<html><body><h3>Saved. Restarting and connecting to your WiFi...</h3></body></html>";
    httpd_resp_set_type(req, "text/html");
    httpd_resp_send(req, resp, HTTPD_RESP_USE_STRLEN);

    /* Let the response flush to the client before rebooting into STA mode. */
    vTaskDelay(pdMS_TO_TICKS(1000));
    esp_restart();

    return ESP_OK;
}

static const httpd_uri_t save_config = {
    .uri = "/save_config",
    .method = HTTP_POST,
    .handler = save_config_post_handler
};

/* Redirects all unmatched requests to the captive portal root page. */
static esp_err_t http_404_error_handler(httpd_req_t *req, httpd_err_code_t err)
{
    httpd_resp_set_status(req, "303 See Other");
    httpd_resp_set_hdr(req, "Location", "/");
    /* iOS requires content in the response to detect a captive portal. */
    httpd_resp_send(req, "Redirect to the captive portal", HTTPD_RESP_USE_STRLEN);

    ESP_LOGI(TAG, "Redirecting to root");
    return ESP_OK;
}

static httpd_handle_t start_webserver(void)
{
    httpd_handle_t server = NULL;
    httpd_config_t config = HTTPD_DEFAULT_CONFIG();
    config.max_open_sockets = 13;
    config.lru_purge_enable = true;

    ESP_LOGI(TAG, "Starting server on port: '%d'", config.server_port);
    if (httpd_start(&server, &config) == ESP_OK) {
        ESP_LOGI(TAG, "Registering URI handlers");
        httpd_register_uri_handler(server, &root);
        httpd_register_uri_handler(server, &save_config);
        httpd_register_err_handler(server, HTTPD_404_NOT_FOUND, http_404_error_handler);
    }
    return server;
}

/* ------------------------------------------------------------------------
 * softAP + captive portal bring-up
 * ---------------------------------------------------------------------- */

static void wifi_event_handler(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 " join, 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 " leave, AID=%d, reason=%d",
                 MAC2STR(event->mac), event->aid, event->reason);
    }
}

static void wifi_init_softap(const char *ap_ssid, const char *ap_password)
{
    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    ESP_ERROR_CHECK(esp_event_handler_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &wifi_event_handler, NULL));

    wifi_config_t wifi_config = {
        .ap = {
            .max_connection = CONFIG_ESP_MAX_STA_CONN,
            .authmode = WIFI_AUTH_WPA_WPA2_PSK
        },
    };
    snprintf((char *)wifi_config.ap.ssid, sizeof(wifi_config.ap.ssid), "%s", ap_ssid);
    wifi_config.ap.ssid_len = strlen(ap_ssid);
    snprintf((char *)wifi_config.ap.password, sizeof(wifi_config.ap.password), "%s", ap_password);

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    esp_netif_ip_info_t ip_info;
    esp_netif_get_ip_info(esp_netif_get_handle_from_ifkey("WIFI_AP_DEF"), &ip_info);

    char ip_addr[16];
    inet_ntoa_r(ip_info.ip.addr, ip_addr, 16);
    ESP_LOGI(TAG, "Set up softAP with IP: %s", ip_addr);
    ESP_LOGI(TAG, "wifi_init_softap finished. SSID:'%s' password:'%s'", ap_ssid, ap_password);
}

#ifdef CONFIG_ESP_ENABLE_DHCP_CAPTIVEPORTAL
static void dhcp_set_captiveportal_url(void)
{
    esp_netif_ip_info_t ip_info;
    esp_netif_get_ip_info(esp_netif_get_handle_from_ifkey("WIFI_AP_DEF"), &ip_info);

    char ip_addr[16];
    inet_ntoa_r(ip_info.ip.addr, ip_addr, 16);

    char *captiveportal_uri = (char *) malloc(32 * sizeof(char));
    assert(captiveportal_uri && "Failed to allocate captiveportal_uri");
    strcpy(captiveportal_uri, "http://");
    strcat(captiveportal_uri, ip_addr);

    esp_netif_t *netif = esp_netif_get_handle_from_ifkey("WIFI_AP_DEF");

    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_netif_dhcps_stop(netif));
    ESP_ERROR_CHECK(esp_netif_dhcps_option(netif, ESP_NETIF_OP_SET, ESP_NETIF_CAPTIVEPORTAL_URI,
                                            captiveportal_uri, strlen(captiveportal_uri)));
    ESP_ERROR_CHECK_WITHOUT_ABORT(esp_netif_dhcps_start(netif));

    free(captiveportal_uri);
}
#endif // CONFIG_ESP_ENABLE_DHCP_CAPTIVEPORTAL

void wifi_provisioning_start(void)
{
    char ap_ssid[FRAME_CFG_SSID_MAX_LEN + 1];
    char ap_password[FRAME_AP_PASSWORD_LEN + 1];
    ap_identity_get(ap_ssid, sizeof(ap_ssid), ap_password, sizeof(ap_password));

    ESP_LOGI(TAG, "Provisioning AP: SSID='%s' password='%s'", ap_ssid, ap_password);

    /* Create (but don't yet start) the AP netif so its IP is known for the
     * config QR before the network is actually joinable -- the default AP
     * IP is fixed at netif creation, not at esp_wifi_start(). */
    esp_netif_create_default_wifi_ap();

    esp_netif_ip_info_t ap_ip_info;
    esp_netif_get_ip_info(esp_netif_get_handle_from_ifkey("WIFI_AP_DEF"), &ap_ip_info);
    char ap_ip_addr[16];
    inet_ntoa_r(ap_ip_info.ip.addr, ap_ip_addr, sizeof(ap_ip_addr));
    char config_url[32];
    snprintf(config_url, sizeof(config_url), "http://%s/", ap_ip_addr);

    /* Bring up the display and show the join QR/password + config QR
     * before the AP goes up, so the instructions are already on-screen by
     * the time the network is joinable. */
    esp_err_t epd_err = epd_init();
    if (epd_err == ESP_OK) {
        esp_err_t qr_err = qr_onboarding_show(ap_ssid, ap_password, config_url);
        if (qr_err != ESP_OK) {
            ESP_LOGW(TAG, "QR onboarding render failed (%s)", esp_err_to_name(qr_err));
        }
    } else {
        ESP_LOGW(TAG, "EPD init failed (%s), continuing without display", esp_err_to_name(epd_err));
    }

    wifi_init_softap(ap_ssid, ap_password);

#ifdef CONFIG_ESP_ENABLE_DHCP_CAPTIVEPORTAL
    dhcp_set_captiveportal_url();
#endif

    start_webserver();

    dns_server_config_t dns_config = DNS_SERVER_CONFIG_SINGLE("*" /* all A queries */, "WIFI_AP_DEF");
    start_dns_server(&dns_config);
}