#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_log.h"
#include "esp_websocket_client.h"
#include "esp_crt_bundle.h"
#include "esp_wifi.h"
#include "esp_timer.h"

#include "socket_client.h"
#include "audio_driver.h"
#include "wifi_connect.h"
#include "state_manager.h"
#include "device_config.h"
#include "ota_manager.h"

static const char *TAG = "socket_client";

static EventGroupHandle_t s_ws_event_group;
#define WS_CONNECTED_BIT BIT0

static esp_websocket_client_handle_t s_client = NULL;

static void send_telemetry(void)
{
    if (s_client == NULL || !(xEventGroupGetBits(s_ws_event_group) & WS_CONNECTED_BIT)) {
        return;
    }

    device_state_t st;
    state_get_current(&st);

    int free_heap = esp_get_free_heap_size();
    int uptime = esp_timer_get_time() / 1000000; // seconds
    int rssi = 0;
    
    wifi_ap_record_t ap_info;
    if (esp_wifi_sta_get_ap_info(&ap_info) == ESP_OK) {
        rssi = ap_info.rssi;
    }

    char payload[512];
    int len = snprintf(payload, sizeof(payload),
        "{\"type\":\"state_report\",\"system\":{\"model\":\"%s\"},\"health\":{\"free_heap_bytes\":%d,\"uptime_seconds\":%d},\"network\":{\"wifi_rssi\":%d},\"peripherals\":{\"audio\":{\"volume\":%d},\"led_ring\":{\"state_id\":%d,\"brightness_coef\":%.2f},\"haptic_intensity\":%d,\"ambient_light\":%d,\"touch_active\":%s}}",
        DEVICE_MODEL, free_heap, uptime, rssi, st.volume, st.led_state, st.led_brightness,
        st.haptic_intensity, st.ambient_light, st.touch_active ? "true" : "false");

    if (len > 0 && len < sizeof(payload)) {
        ESP_LOGI(TAG, "Sending state report: %s", payload);
        esp_websocket_client_send_text(s_client, payload, len, portMAX_DELAY);
    }
}

void trigger_telemetry_broadcast(void)
{
    send_telemetry();
}

static char *extract_json_string(const char *json, const char *key)
{
    char *key_pos = strstr(json, key);
    if (!key_pos) return NULL;
    
    // Find the colon after key
    char *colon = strchr(key_pos + strlen(key), ':');
    if (!colon) return NULL;
    
    // Find the opening quote of the string value
    char *start_quote = strchr(colon, '"');
    if (!start_quote) return NULL;
    
    // Find the closing quote, making sure we handle escaped quotes \"
    char *end_ptr = start_quote + 1;
    while (*end_ptr) {
        if (*end_ptr == '"' && *(end_ptr - 1) != '\\') {
            break;
        }
        end_ptr++;
    }
    if (*end_ptr != '"') return NULL;
    
    size_t raw_len = end_ptr - start_quote - 1;
    char *out = malloc(raw_len + 1);
    if (!out) return NULL;
    
    // Unescape the string while copying
    size_t out_idx = 0;
    for (size_t i = 0; i < raw_len; i++) {
        if (start_quote[1 + i] == '\\' && i + 1 < raw_len) {
            char next_char = start_quote[1 + i + 1];
            if (next_char == 'n') {
                out[out_idx++] = '\n';
                i++;
            } else if (next_char == 'r') {
                out[out_idx++] = '\r';
                i++;
            } else if (next_char == 't') {
                out[out_idx++] = '\t';
                i++;
            } else if (next_char == '\\' || next_char == '"' || next_char == '/') {
                out[out_idx++] = next_char;
                i++;
            } else {
                out[out_idx++] = '\\';
            }
        } else {
            out[out_idx++] = start_quote[1 + i];
        }
    }
    out[out_idx] = '\0';
    return out;
}

static void websocket_event_handler(void *handler_args, esp_event_base_t base, int32_t event_id, void *event_data)
{
    esp_websocket_event_data_t *data = (esp_websocket_event_data_t *)event_data;
    switch (event_id) {
        case WEBSOCKET_EVENT_CONNECTED:
            ESP_LOGI(TAG, "Websocket connected to server!");
            xEventGroupSetBits(s_ws_event_group, WS_CONNECTED_BIT);
            state_set_websocket_connected(true);
            break;
        case WEBSOCKET_EVENT_DISCONNECTED:
            ESP_LOGI(TAG, "Websocket disconnected!");
            xEventGroupClearBits(s_ws_event_group, WS_CONNECTED_BIT);
            state_set_websocket_connected(false);
            break;
        case WEBSOCKET_EVENT_DATA:
            if (data->data_ptr != NULL && data->data_len > 0) {
                ESP_LOGI(TAG, "Received message: %.*s", data->data_len, (char *)data->data_ptr);
                
                char *msg_str = malloc(data->data_len + 1);
                if (msg_str != NULL) {
                    memcpy(msg_str, data->data_ptr, data->data_len);
                    msg_str[data->data_len] = '\0';
                    
                    if (strstr(msg_str, "\"set_volume\"") != NULL) {
                        char *val_ptr = strstr(msg_str, "\"value\"");
                        if (val_ptr != NULL) {
                            int level = -1;
                            if (sscanf(val_ptr, "\"value\"%*[^0-9]%d", &level) == 1) {
                                if (level >= 0 && level <= 100) {
                                    ESP_LOGI(TAG, "Parsed set_volume command. Mutating state to: %d", level);
                                    state_set_volume((uint8_t)level);
                                }
                            }
                        }
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"set_led\"") != NULL) {
                        char *val_ptr = strstr(msg_str, "\"value\"");
                        if (val_ptr != NULL) {
                            if (strstr(val_ptr, "\"rainbow\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed set_led command: rainbow");
                                state_set_led(LED_STATE_CONNECTED);
                            } else if (strstr(val_ptr, "\"connecting\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed set_led command: connecting");
                                state_set_led(LED_STATE_SERVER_CONNECTING);
                            } else if (strstr(val_ptr, "\"failed\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed set_led command: failed");
                                state_set_led(LED_STATE_SERVER_FAILED);
                            } else if (strstr(val_ptr, "\"idle\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed set_led command: idle");
                                state_set_led(LED_STATE_OFF);
                            } else if (strstr(val_ptr, "\"updating\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed set_led command: updating");
                                state_set_led(LED_STATE_UPDATING);
                            }
                        }
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"play_audio\"") != NULL) {
                        char *val_ptr = strstr(msg_str, "\"value\"");
                        if (val_ptr != NULL) {
                            if (strstr(val_ptr, "\"startup\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed play_audio command: startup");
                                Audio_Play_Music_To_End("file://spiffs/startup.mp3");
                            } else if (strstr(val_ptr, "\"wifi_connecting\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed play_audio command: wifi_connecting");
                                Audio_Play_Music_To_End("file://spiffs/connecting_wifi.mp3");
                            } else if (strstr(val_ptr, "\"wifi_connected\"") != NULL) {
                                ESP_LOGI(TAG, "Parsed play_audio command: wifi_connected");
                                Audio_Play_Music_To_End("file://spiffs/connected_wifi.mp3");
                            }
                        }
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"set_brightness\"") != NULL) {
                        char *val_ptr = strstr(msg_str, "\"value\"");
                        if (val_ptr != NULL) {
                            float brightness = -1.0f;
                            if (sscanf(val_ptr, "\"value\"%*[^0-9.-]%f", &brightness) == 1) {
                                if (brightness >= 0.0f && brightness <= 1.0f) {
                                    ESP_LOGI(TAG, "Parsed set_brightness command. Mutating state to: %f", brightness);
                                    state_set_brightness(brightness);
                                }
                            }
                        }
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"set_haptics\"") != NULL) {
                        char *val_ptr = strstr(msg_str, "\"value\"");
                        if (val_ptr != NULL) {
                            int intensity = -1;
                            if (sscanf(val_ptr, "\"value\"%*[^0-9]%d", &intensity) == 1) {
                                if (intensity >= 0 && intensity <= 100) {
                                    ESP_LOGI(TAG, "Parsed set_haptics command. Mutating state to: %d", intensity);
                                    state_set_haptic_intensity((uint8_t)intensity);
                                }
                            }
                        }
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"query_state\"") != NULL) {
                        ESP_LOGI(TAG, "Parsed query_state command. Dispatching telemetry update...");
                        send_telemetry();
                    }
                    else if (strstr(msg_str, "\"start_ota\"") != NULL) {
                        char *url = extract_json_string(msg_str, "\"url\"");
                        
                        if (url != NULL) {
                            ESP_LOGI(TAG, "Parsed start_ota command. URL: %s", url);
                            ota_start(url);
                            free(url);
                        } else {
                            ESP_LOGE(TAG, "start_ota: missing url field");
                        }
                    }
                    free(msg_str);
                }
            }
            break;
        case WEBSOCKET_EVENT_ERROR:
            ESP_LOGE(TAG, "Websocket error occurred!");
            break;
    }
}

static void socket_client_task(void *pvParameters)
{
    bool was_connected = false;

    while (1) {
        // Wait until Wi-Fi is connected
        while (!wifi_is_connected()) {
            state_set_wifi_connected(false);
            vTaskDelay(pdMS_TO_TICKS(1000));
        }
        state_set_wifi_connected(true);

        EventBits_t bits = xEventGroupGetBits(s_ws_event_group);
        if (!(bits & WS_CONNECTED_BIT)) {
            // We are not connected!
            if (was_connected) {
                // We lost connection!
                ESP_LOGW(TAG, "WebSocket connection lost! Playing alert...");
                state_set_led(LED_STATE_SERVER_CONNECTING);
                Audio_Play_Music_To_End("file://spiffs/disconnected_server.mp3");
                was_connected = false;
            }

            ESP_LOGI(TAG, "Attempting WebSocket connection...");
            state_set_led(LED_STATE_SERVER_CONNECTING);
            
            // Play connecting audio: "Connecting to the dev server"
            Audio_Play_Music_To_End("file://spiffs/connecting_server.mp3");

            // Stop (if running) and start connection
            esp_websocket_client_stop(s_client);
            esp_websocket_client_start(s_client);

            // Wait up to 10 seconds for connection
            ESP_LOGI(TAG, "Waiting for websocket connection...");
            bits = xEventGroupWaitBits(s_ws_event_group, WS_CONNECTED_BIT, pdFALSE, pdFALSE, pdMS_TO_TICKS(10000));

            if (bits & WS_CONNECTED_BIT) {
                ESP_LOGI(TAG, "WebSocket connected successfully.");
                was_connected = true;
                state_set_led(LED_STATE_CONNECTED);
                Audio_Play_Music_To_End("file://spiffs/connected_server.mp3");
            } else {
                ESP_LOGE(TAG, "Failed to connect to dev server within 10 seconds!");
                state_set_led(LED_STATE_SERVER_FAILED);
                esp_websocket_client_stop(s_client);
                Audio_Play_Music_To_End("file://spiffs/connect_server_fail.mp3");

                // Play "Still not working, let me try again in a moment"
                Audio_Play_Music_To_End("file://spiffs/retry_later.mp3");

                // Wait 10 seconds before next connection attempt
                vTaskDelay(pdMS_TO_TICKS(10000));
            }
        } else {
            // Already connected! Keep monitoring
            vTaskDelay(pdMS_TO_TICKS(1000));
        }
    }
}

void socket_client_init(void)
{
    s_ws_event_group = xEventGroupCreate();

    // Load WebSocket configuration from SPIFFS
    char url_buf[128] = {0};
    FILE *f = fopen("/spiffs/server_url.txt", "r");
    if (f != NULL) {
        if (fgets(url_buf, sizeof(url_buf), f) != NULL) {
            // Strip trailing whitespace/newlines
            size_t len = strlen(url_buf);
            while (len > 0 && (url_buf[len - 1] == '\n' || url_buf[len - 1] == '\r' || url_buf[len - 1] == ' ')) {
                url_buf[len - 1] = '\0';
                len--;
            }
            ESP_LOGI(TAG, "Loaded WebSocket URI from SPIFFS: %s", url_buf);
        }
        fclose(f);
    }

    const char *ws_uri = "wss://example.com"; // Replace with your deployment endpoint
    if (strlen(url_buf) > 0) {
        ws_uri = url_buf;
    } else {
        ESP_LOGW(TAG, "No server_url.txt configuration found or empty, using default: %s", ws_uri);
    }

    esp_websocket_client_config_t websocket_cfg = {
        .uri = ws_uri,
        .buffer_size = 4096,
    };

    ESP_LOGI(TAG, "Initializing WebSocket client...");
    s_client = esp_websocket_client_init(&websocket_cfg);

    // Register event handler
    ESP_ERROR_CHECK(esp_websocket_register_events(s_client, WEBSOCKET_EVENT_ANY, websocket_event_handler, NULL));

    // Create the background connection task
    ESP_LOGI(TAG, "Starting WebSocket monitor task...");
    xTaskCreate(socket_client_task, "socket_client_task", 4096, NULL, 5, NULL);
}

/* ------------------------------------------------------------------
 * OTA WebSocket message helpers
 * Called from ota_manager.c to report OTA state back to the server.
 * ------------------------------------------------------------------ */

void socket_send_ota_progress(int percent, int bytes_written, int bytes_total)
{
    if (s_client == NULL || !(xEventGroupGetBits(s_ws_event_group) & WS_CONNECTED_BIT)) {
        return;
    }
    char buf[128];
    int len = snprintf(buf, sizeof(buf),
        "{\"type\":\"ota_progress\",\"percent\":%d,\"bytes_written\":%d,\"bytes_total\":%d}",
        percent, bytes_written, bytes_total);
    if (len > 0 && len < (int)sizeof(buf)) {
        esp_websocket_client_send_text(s_client, buf, len, portMAX_DELAY);
    }
}

void socket_send_ota_complete(void)
{
    if (s_client == NULL || !(xEventGroupGetBits(s_ws_event_group) & WS_CONNECTED_BIT)) {
        return;
    }
    const char *msg = "{\"type\":\"ota_complete\"}";
    esp_websocket_client_send_text(s_client, msg, strlen(msg), portMAX_DELAY);
}

void socket_send_ota_error(const char *reason)
{
    if (s_client == NULL || !(xEventGroupGetBits(s_ws_event_group) & WS_CONNECTED_BIT)) {
        return;
    }
    char buf[192];
    int len = snprintf(buf, sizeof(buf),
        "{\"type\":\"ota_error\",\"reason\":\"%s\"}", reason ? reason : "unknown");
    if (len > 0 && len < (int)sizeof(buf)) {
        esp_websocket_client_send_text(s_client, buf, len, portMAX_DELAY);
    }
}