#include <string.h>

#include "esp_event.h"
#include "esp_log.h"
#include "esp_wifi.h"
#include "esp_wifi_default.h"
#include "esp_netif.h"
#include "esp_http_client.h"
#include "esp_sleep.h"

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

#include "epd7in3e.h"
#include "status_screen.h"
#include "manage_qr_overlay.h"

#include "frame_client.h"

static const char *TAG = "frame_client";

#define STA_CONNECTED_BIT BIT0
#define STA_FAILED_BIT BIT1

static EventGroupHandle_t s_sta_event_group;

/* wifi_sta_config_t's ssid/password fields are fixed-size byte arrays, not
 * necessarily null-terminated (a full 32-char SSID fills the field exactly).
 * snprintf() flags that as a possible truncation at -Werror, so copy by
 * hand instead. */
static void copy_wifi_field(uint8_t *dst, size_t dst_size, const char *src)
{
    size_t len = strnlen(src, dst_size);
    memcpy(dst, src, len);
    if (len < dst_size) {
        dst[len] = '\0';
    }
}

static void sta_event_handler(void *arg, esp_event_base_t event_base,
                               int32_t event_id, void *event_data)
{
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        ESP_LOGW(TAG, "Disconnected from home WiFi");
        xEventGroupSetBits(s_sta_event_group, STA_FAILED_BIT);
    } 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: " IPSTR, IP2STR(&event->ip_info.ip));
        xEventGroupSetBits(s_sta_event_group, STA_CONNECTED_BIT);
    }
}

esp_err_t frame_wifi_connect_sta(const frame_config_t *cfg)
{
    s_sta_event_group = xEventGroupCreate();

    esp_netif_t *sta_netif = esp_netif_create_default_wifi_sta();

    wifi_init_config_t init_cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&init_cfg));

    esp_event_handler_instance_t wifi_handler;
    esp_event_handler_instance_t ip_handler;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT, ESP_EVENT_ANY_ID, &sta_event_handler, NULL, &wifi_handler));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT, IP_EVENT_STA_GOT_IP, &sta_event_handler, NULL, &ip_handler));

    wifi_config_t wifi_config = {0};
    copy_wifi_field(wifi_config.sta.ssid, sizeof(wifi_config.sta.ssid), cfg->sta_ssid);
    copy_wifi_field(wifi_config.sta.password, sizeof(wifi_config.sta.password), cfg->sta_password);

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    esp_err_t result = ESP_FAIL;
    for (int attempt = 1; attempt <= CONFIG_FRAME_STA_CONNECT_MAX_RETRIES; attempt++) {
        ESP_LOGI(TAG, "Connecting to '%s' (attempt %d/%d)", cfg->sta_ssid, attempt,
                 CONFIG_FRAME_STA_CONNECT_MAX_RETRIES);

        xEventGroupClearBits(s_sta_event_group, STA_CONNECTED_BIT | STA_FAILED_BIT);
        esp_wifi_connect();

        EventBits_t bits = xEventGroupWaitBits(s_sta_event_group, STA_CONNECTED_BIT | STA_FAILED_BIT,
                                                pdTRUE, pdFALSE,
                                                pdMS_TO_TICKS(CONFIG_FRAME_STA_CONNECT_TIMEOUT_MS));

        if (bits & STA_CONNECTED_BIT) {
            result = ESP_OK;
            break;
        }
        ESP_LOGW(TAG, "Attempt %d/%d failed", attempt, CONFIG_FRAME_STA_CONNECT_MAX_RETRIES);
    }

    esp_event_handler_instance_unregister(WIFI_EVENT, ESP_EVENT_ANY_ID, wifi_handler);
    esp_event_handler_instance_unregister(IP_EVENT, IP_EVENT_STA_GOT_IP, ip_handler);
    vEventGroupDelete(s_sta_event_group);
    s_sta_event_group = NULL;

    if (result != ESP_OK) {
        /* Fully tear the WiFi driver back down on failure -- the caller
         * falls back to provisioning, which calls esp_wifi_init() again
         * for AP mode. Leaving the driver merely stopped (rather than
         * deinitialized) made that second esp_wifi_init() call fail with
         * ESP_ERR_INVALID_STATE and abort, confirmed on hardware. */
        esp_wifi_stop();
        esp_wifi_deinit();
        esp_netif_destroy_default_wifi(sta_netif);
    }

    return result;
}

typedef struct {
    bool reachable;
    uint32_t refresh_interval_s; /* CONFIG_FRAME_SLEEP_INTERVAL_S if absent/unparseable */
} frame_server_config_t;

/* Finds the first integer value associated with "key" in a small JSON
 * blob, e.g. 3600 in {"refresh_interval_s": 3600}. Not a general JSON
 * parser -- just enough for this project's small, flat config response,
 * to avoid pulling in a JSON library for one scalar field. */
static bool json_extract_uint(const char *json, const char *key, uint32_t *out)
{
    char needle[48];
    snprintf(needle, sizeof(needle), "\"%s\"", key);
    const char *pos = strstr(json, needle);
    if (pos == NULL) {
        return false;
    }
    pos = strchr(pos, ':');
    if (pos == NULL) {
        return false;
    }
    pos++;
    while (*pos == ' ') {
        pos++;
    }
    char *end;
    unsigned long value = strtoul(pos, &end, 10);
    if (end == pos) {
        return false;
    }
    *out = (uint32_t)value;
    return true;
}

/* GETs the server's /frame/config -- doubles as both the reachability
 * check (any completed HTTP response means the socket-level connection
 * succeeded) and the source of the server-configurable refresh interval. */
static frame_server_config_t fetch_frame_config(const char *toolsserver)
{
    frame_server_config_t result = {
        .reachable = false,
        .refresh_interval_s = CONFIG_FRAME_SLEEP_INTERVAL_S,
    };

    char url[160];
    snprintf(url, sizeof(url), "http://%s/frame/config", toolsserver);

    esp_http_client_config_t config = {
        .url = url,
        .method = HTTP_METHOD_GET,
        .timeout_ms = CONFIG_FRAME_SERVER_CHECK_TIMEOUT_MS,
    };
    esp_http_client_handle_t client = esp_http_client_init(&config);

    esp_err_t err = esp_http_client_open(client, 0);
    if (err != ESP_OK) {
        ESP_LOGW(TAG, "Server '%s' not reachable: %s", toolsserver, esp_err_to_name(err));
        esp_http_client_cleanup(client);
        return result;
    }

    esp_http_client_fetch_headers(client);
    result.reachable = true;

    char body[256];
    int total = 0;
    int n;
    while (total < (int)sizeof(body) - 1 &&
           (n = esp_http_client_read(client, body + total, sizeof(body) - 1 - total)) > 0) {
        total += n;
    }
    body[total] = '\0';

    esp_http_client_close(client);
    esp_http_client_cleanup(client);

    uint32_t interval;
    if (json_extract_uint(body, "refresh_interval_s", &interval)) {
        result.refresh_interval_s = interval;
    } else {
        ESP_LOGW(TAG, "'%s' response missing refresh_interval_s, using fallback %ds", url,
                 (int)result.refresh_interval_s);
    }

    return result;
}

typedef struct {
    esp_http_client_handle_t client;
    size_t stream_pos;                  /* running absolute offset into the frame, for overlay splicing */
    const manage_qr_overlay_t *overlay; /* NULL = no overlay this fetch */
} http_read_ctx_t;

/* Splices overlay pixels over the real photo bytes in chunk wherever
 * chunk's absolute byte range [chunk_start, chunk_start+chunk_len) within
 * the full frame intersects the overlay's rectangle. Rows/chunks outside
 * the overlay's footprint are left completely untouched. overlay->x0 is
 * always even (see manage_qr_overlay.h), so byte_x0 below is exact. */
static void splice_overlay(uint8_t *chunk, size_t chunk_len, size_t chunk_start, const manage_qr_overlay_t *overlay)
{
    int byte_x0 = overlay->x0 / 2;
    int byte_w = overlay->w / 2;
    size_t chunk_end = chunk_start + chunk_len;

    for (int row = overlay->y0; row < overlay->y0 + overlay->h; row++) {
        size_t row_start = (size_t)row * EPD_BYTES_PER_ROW + (size_t)byte_x0;
        size_t row_end = row_start + (size_t)byte_w;

        size_t lo = row_start > chunk_start ? row_start : chunk_start;
        size_t hi = row_end < chunk_end ? row_end : chunk_end;
        if (lo >= hi) {
            continue;
        }

        size_t overlay_row_offset = (size_t)(row - overlay->y0) * (size_t)byte_w + (lo - row_start);
        memcpy(chunk + (lo - chunk_start), overlay->buf + overlay_row_offset, hi - lo);
    }
}

/* Pulls the next chunk straight out of the in-progress HTTP response --
 * epd_write_frame() calls this to feed the panel without ever holding
 * the full ~192KB frame in RAM. Splices in ctx->overlay's pixels (if
 * set) as chunks pass through, so the panel driver never needs to know
 * an overlay exists at all. */
static size_t http_read_fn(uint8_t *chunk, size_t chunk_size, void *ctx_)
{
    http_read_ctx_t *ctx = (http_read_ctx_t *)ctx_;
    int n = esp_http_client_read(ctx->client, (char *)chunk, (int)chunk_size);
    if (n <= 0) {
        return 0;
    }

    if (ctx->overlay != NULL) {
        splice_overlay(chunk, (size_t)n, ctx->stream_pos, ctx->overlay);
    }
    ctx->stream_pos += (size_t)n;

    return (size_t)n;
}

/* GETs /frame/image (or, if force_advance, POSTs /frame/advance to skip
 * ahead immediately) and streams the response directly into the panel,
 * splicing in overlay's pixels (if non-NULL) as it streams. Returning
 * non-ESP_OK means the panel was never actually refreshed --
 * epd_display_stream() (see epd7in3e.c) refuses to trigger a physical
 * refresh on a short/wrong-size stream, so a failure here always leaves
 * the visible screen exactly as it was. */
static esp_err_t fetch_and_display(const frame_config_t *cfg, bool force_advance, const manage_qr_overlay_t *overlay)
{
    char url[160];
    snprintf(url, sizeof(url), "http://%s/%s", cfg->toolsserver, force_advance ? "frame/advance" : "frame/image");

    esp_http_client_config_t config = {
        .url = url,
        .method = force_advance ? HTTP_METHOD_POST : HTTP_METHOD_GET,
        .timeout_ms = CONFIG_FRAME_FETCH_TIMEOUT_MS,
    };
    esp_http_client_handle_t client = esp_http_client_init(&config);

    esp_err_t err = esp_http_client_open(client, 0);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Failed to open '%s': %s", url, esp_err_to_name(err));
        esp_http_client_cleanup(client);
        return err;
    }

    int content_length = esp_http_client_fetch_headers(client);
    int status = esp_http_client_get_status_code(client);
    if (status != 200) {
        ESP_LOGE(TAG, "'%s' returned HTTP %d", url, status);
        esp_http_client_close(client);
        esp_http_client_cleanup(client);
        return ESP_FAIL;
    }
    ESP_LOGI(TAG, "Fetching frame (%d bytes) from '%s'", content_length, url);

    http_read_ctx_t ctx = { .client = client, .overlay = overlay };
    uint32_t crc = 0;
    err = epd_write_frame(http_read_fn, &ctx, &crc);

    esp_http_client_close(client);
    esp_http_client_cleanup(client);

    if (err != ESP_OK) {
        return err;
    }

    uint32_t previous_crc;
    if (frame_config_get_last_display_crc32(&previous_crc) == ESP_OK && previous_crc == crc) {
        /* Same photo already on screen (e.g. redisplayed after a reboot,
         * before the refresh interval elapsed server-side) -- skip the
         * physical refresh, avoiding its visible flash and 15-30s
         * duration for no visual change. */
        ESP_LOGI(TAG, "Frame unchanged since last display, skipping refresh");
        return ESP_OK;
    }

    err = epd_turn_on_display();
    if (err == ESP_OK) {
        frame_config_set_last_display_crc32(crc);
    }
    return err;
}

/* Runs the appropriate fetch for this cycle: a plain fetch, or -- if
 * show_management_qr -- a fetch with the "scan to manage" QR overlay
 * spliced in, held on screen for 30s (device stays awake, doesn't sleep
 * the panel or the chip), then reverted with a second plain fetch.
 * Returns non-ESP_OK only if the FIRST fetch failed; a revert failure
 * afterward is logged but doesn't count as an overall failure -- the QR
 * itself displayed fine, which was the point of the button. */
static esp_err_t run_fetch_cycle(const frame_config_t *cfg, bool force_advance, bool show_management_qr)
{
    if (!show_management_qr) {
        return fetch_and_display(cfg, force_advance, NULL);
    }

    char management_url[160];
    snprintf(management_url, sizeof(management_url), "http://%s/", cfg->toolsserver);

    manage_qr_overlay_t overlay;
    esp_err_t overlay_err = manage_qr_overlay_render(management_url, &overlay);
    if (overlay_err != ESP_OK) {
        ESP_LOGW(TAG, "Could not render management QR overlay (%s), showing photo normally",
                 esp_err_to_name(overlay_err));
        return fetch_and_display(cfg, force_advance, NULL);
    }

    esp_err_t err = fetch_and_display(cfg, force_advance, &overlay);
    manage_qr_overlay_free(&overlay);
    if (err != ESP_OK) {
        return err;
    }

    ESP_LOGI(TAG, "Showing management QR for 30s");
    vTaskDelay(pdMS_TO_TICKS(30000));

    /* force_advance is always false here -- reverting shouldn't skip
     * ahead a second time. */
    esp_err_t revert_err = fetch_and_display(cfg, false, NULL);
    if (revert_err != ESP_OK) {
        ESP_LOGW(TAG, "Failed to revert management QR overlay (%s)", esp_err_to_name(revert_err));
    }

    return ESP_OK;
}

void frame_client_run(const frame_config_t *cfg, bool force_advance, bool show_management_qr)
{
    esp_err_t epd_err = epd_init();
    bool have_display = (epd_err == ESP_OK);
    if (!have_display) {
        ESP_LOGW(TAG, "EPD init failed (%s), continuing without display", esp_err_to_name(epd_err));
    }

    /* Always show the status screen on the first successful connection
     * after (re)provisioning, regardless of outcome -- confirms the
     * connection worked. Skipped on later wakes to save a refresh, except
     * when something's actually wrong (handled below). */
    bool first_connection = !frame_config_has_connected_once();
    if (first_connection) {
        frame_config_mark_connected_once();
        if (have_display) {
            status_screen_show(cfg->sta_ssid, STATUS_OK, cfg->toolsserver, STATUS_PENDING);
        }
    }

    /* The image fetch goes before the config check, not after. It has a
     * far more generous timeout (CONFIG_FRAME_FETCH_TIMEOUT_MS, 15s by
     * default, vs. the config check's 3s), so it comfortably absorbs the
     * extra connection-setup latency that's common on the very first
     * request after waking from a long deep sleep (stale ARP entries and
     * the like) -- confirmed on hardware: the config check's tight
     * timeout was intermittently tripping on exactly that latency while
     * it went first, even though the image fetch right after it (on an
     * already-warm connection) never had trouble. Trade-off: on a fully
     * down server, the device now waits up to the image fetch's longer
     * timeout to notice, instead of the config check's shorter one --
     * worth it to stop false-failing on the common case. */
    bool image_ok = true;
    if (have_display) {
        esp_err_t fetch_err = run_fetch_cycle(cfg, force_advance, show_management_qr);
        image_ok = (fetch_err == ESP_OK);
        if (!image_ok) {
            /* epd_display_stream() never triggers a physical refresh on a
             * failed/short/wrong-size stream (see epd7in3e.c), so the
             * visible screen is guaranteed untouched here -- always safe
             * to show what went wrong instead of leaving stale content
             * with no indication anything failed. */
            ESP_LOGW(TAG, "Fetch/display failed (%s), retrying sooner", esp_err_to_name(fetch_err));
            status_screen_show(cfg->sta_ssid, STATUS_OK, cfg->toolsserver, STATUS_FAILED);
        } else if (first_connection) {
            status_screen_show(cfg->sta_ssid, STATUS_OK, cfg->toolsserver, STATUS_OK);
        }
    }

    /* Only worth asking for the refresh interval if the image fetch
     * actually worked -- a failed fetch already means CONFIG_FRAME_RETRY_INTERVAL_S,
     * so there's nothing to gain from a config request whose result would
     * just be discarded. */
    uint32_t sleep_seconds = CONFIG_FRAME_RETRY_INTERVAL_S;
    if (image_ok) {
        frame_server_config_t server_cfg = fetch_frame_config(cfg->toolsserver);
        sleep_seconds = server_cfg.reachable ? server_cfg.refresh_interval_s : CONFIG_FRAME_RETRY_INTERVAL_S;
    }

    if (have_display) {
        epd_sleep();
    }

    ESP_LOGI(TAG, "Deep sleeping for %u seconds", (unsigned)sleep_seconds);
    esp_sleep_enable_timer_wakeup((uint64_t)sleep_seconds * 1000000ULL);
    esp_deep_sleep_start();
}