#include <groundext.h>
#include <curl/curl.h>
#include <groundvm.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

bool curl_init = false;

typedef struct {
    char* data;
    size_t size;
} ResponseBuffer;

typedef struct {
    CURL* handle;
    bool connected;
    char* url;
} WebSocketConnection;

// Global storage for WebSocket connections (simple implementation)
#define MAX_WS_CONNECTIONS 10
WebSocketConnection ws_connections[MAX_WS_CONNECTIONS] = {0};

void curl_setup() {
    if (!curl_init) {
        curl_global_init(CURL_GLOBAL_ALL);
        curl_init = true;
    }
}

// Proper write callback that accumulates data
size_t write_callback(void* ptr, size_t size, size_t nmemb, void* userdata) {
    size_t total_size = size * nmemb;
    ResponseBuffer* buffer = (ResponseBuffer*)userdata;
    
    char* new_data = realloc(buffer->data, buffer->size + total_size + 1);
    if (new_data == NULL) {
        printf("Failed to allocate memory\n");
        return 0;
    }
    
    buffer->data = new_data;
    memcpy(buffer->data + buffer->size, ptr, total_size);
    buffer->size += total_size;
    buffer->data[buffer->size] = '\0';
    
    return total_size;
}

// GET request
GroundValue get_request(GroundScope* scope, List args) {
    curl_setup();
    CURL* handle = curl_easy_init();
    
    if (!handle) {
        ERROR("CURL failed to init", "GenericRequestError");
    }
    
    ResponseBuffer buffer = {0};
    buffer.data = malloc(1);
    buffer.data[0] = '\0';
    buffer.size = 0;
    
    curl_easy_setopt(handle, CURLOPT_URL, args.values[0].data.stringVal);
    curl_easy_setopt(handle, CURLOPT_HTTPGET, 1L);
    curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(handle, CURLOPT_WRITEDATA, &buffer);
    curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 1L);
    
    CURLcode result = curl_easy_perform(handle);
    
    if (result != CURLE_OK) {
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "Curl request failed: %s\n", curl_easy_strerror(result));
        ERROR(buf, "ActionRequestError");
    }
    
    curl_easy_cleanup(handle);
    
    GroundValue ret = groundCreateValue(STRING, buffer.data);
    return ret;
}

// POST request
GroundValue post_request(GroundScope* scope, List args) {
    curl_setup();
    CURL* handle = curl_easy_init();
    
    if (!handle) {
        ERROR("CURL failed to init", "GenericRequestError");
    }
    
    ResponseBuffer buffer = {0};
    buffer.data = malloc(1);
    buffer.data[0] = '\0';
    buffer.size = 0;
    
    curl_easy_setopt(handle, CURLOPT_URL, args.values[0].data.stringVal);
    curl_easy_setopt(handle, CURLOPT_POST, 1L);
    curl_easy_setopt(handle, CURLOPT_POSTFIELDS, args.values[1].data.stringVal);
    curl_easy_setopt(handle, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(handle, CURLOPT_WRITEDATA, &buffer);
    curl_easy_setopt(handle, CURLOPT_FOLLOWLOCATION, 1L);
    
    CURLcode result = curl_easy_perform(handle);
    
    if (result != CURLE_OK) {
        free(buffer.data);
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "Curl request failed: %s\n", curl_easy_strerror(result));
        curl_easy_cleanup(handle);
        ERROR(buf, "ActionRequestError");
    }
    
    curl_easy_cleanup(handle);
    
    GroundValue ret = groundCreateValue(STRING, buffer.data);
    return ret;
}

// ============== WebSocket Functions ==============

// Find an empty slot for a new WebSocket connection
int find_ws_slot() {
    for (int i = 0; i < MAX_WS_CONNECTIONS; i++) {
        if (!ws_connections[i].connected) {
            return i;
        }
    }
    return -1;
}

// WebSocket Connect - returns connection ID as INT
GroundValue ws_connect(GroundScope* scope, List args) {
    curl_setup();
    
    int slot = find_ws_slot();
    if (slot == -1) {
        ERROR("Maximum WebSocket connections reached", "OverflowError");
    }
    
    CURL* handle = curl_easy_init();
    if (!handle) {
        ERROR("CURL failed to init", "GenericWSError");
    }
    
    curl_easy_setopt(handle, CURLOPT_URL, args.values[0].data.stringVal);
    curl_easy_setopt(handle, CURLOPT_CONNECT_ONLY, 2L); // WebSocket mode
    curl_easy_setopt(handle, CURLOPT_TIMEOUT, 10L);
    
    CURLcode result = curl_easy_perform(handle);
    
    if (result != CURLE_OK) {
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "Curl WebSocket failed: %s\n", curl_easy_strerror(result));
        curl_easy_cleanup(handle);
        ERROR(buf, "ActionWSError");
    }
    
    ws_connections[slot].handle = handle;
    ws_connections[slot].connected = true;
    ws_connections[slot].url = strdup(args.values[0].data.stringVal);
    
    return groundCreateValue(INT, (int64_t)slot);
}

// WebSocket Send - sends a text message, returns bytes sent as INT
GroundValue ws_send(GroundScope* scope, List args) {
    int conn_id = (int)args.values[0].data.intVal;
    if (conn_id < 0 || conn_id >= MAX_WS_CONNECTIONS || !ws_connections[conn_id].connected) {
        ERROR("Invalid WebSocket connection ID", "GenericWSError");
    }
    
    CURL* handle = ws_connections[conn_id].handle;
    const char* message = args.values[1].data.stringVal;
    size_t sent;
    
    CURLcode result = curl_ws_send(handle, message, strlen(message), &sent, 0, CURLWS_TEXT);
    
    if (result != CURLE_OK) {
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "WebSocket send failed: %s\n", curl_easy_strerror(result));
        ERROR(buf, "ActionWSError");
    }
    
    return groundCreateValue(INT, (int64_t)sent);
}

// WebSocket Receive - receives a message (blocking)
GroundValue ws_receive(GroundScope* scope, List args) {
    int conn_id = (int)args.values[0].data.intVal;
    if (conn_id < 0 || conn_id >= MAX_WS_CONNECTIONS || !ws_connections[conn_id].connected) {
        ERROR("Invalid WebSocket connection ID", "GenericWSError");
    }
    
    CURL* handle = ws_connections[conn_id].handle;
    char buffer[4096];
    size_t received;
    const struct curl_ws_frame* meta;
    
    CURLcode result = curl_ws_recv(handle, buffer, sizeof(buffer) - 1, &received, &meta);
    
    if (result != CURLE_OK) {
        if (result == CURLE_AGAIN) {
            // No data available right now
            return groundCreateValue(STRING, "");
        }
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "WebSocket receive failed: %s\n", curl_easy_strerror(result));
        ERROR(buf, "ActionWSError");
    }
    
    buffer[received] = '\0';
    
    // Handle different frame types
    if (meta->flags & CURLWS_CLOSE) {
        printf("WebSocket close frame received\n");
        ws_connections[conn_id].connected = false;
        return groundCreateValue(STRING, "[CLOSED]");
    }
    
    if (meta->flags & CURLWS_PING) {
        // Automatically respond to ping with pong
        curl_ws_send(handle, "", 0, &received, 0, CURLWS_PONG);
        return groundCreateValue(STRING, "[PING]");
    }
    
    char* data = malloc(received + 1);
    memcpy(data, buffer, received);
    data[received] = '\0';
    
    return groundCreateValue(STRING, data);
}

// WebSocket Receive with timeout (non-blocking version)
GroundValue ws_receive_timeout(GroundScope* scope, List args) {
    int conn_id = (int)args.values[0].data.intVal;
    long timeout_ms = (long)args.values[1].data.intVal;
    
    if (conn_id < 0 || conn_id >= MAX_WS_CONNECTIONS || !ws_connections[conn_id].connected) {
        ERROR("Invalid WebSocket connection ID", "GenericWSError");
    }
    
    CURL* handle = ws_connections[conn_id].handle;
    
    // Set socket to non-blocking mode
    curl_easy_setopt(handle, CURLOPT_TIMEOUT_MS, timeout_ms);
    
    char buffer[4096];
    size_t received;
    const struct curl_ws_frame* meta;
    
    CURLcode result = curl_ws_recv(handle, buffer, sizeof(buffer) - 1, &received, &meta);
    
    if (result == CURLE_AGAIN || result == CURLE_OPERATION_TIMEDOUT) {
        return groundCreateValue(STRING, "");
    }
    
    if (result != CURLE_OK) {
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "WebSocket receive failed: %s\n", curl_easy_strerror(result));
        ERROR(buf, "ActionWSError");
    }
    
    buffer[received] = '\0';
    
    if (meta->flags & CURLWS_CLOSE) {
        ws_connections[conn_id].connected = false;
        return groundCreateValue(STRING, "[CLOSED]");
    }
    
    if (meta->flags & CURLWS_PING) {
        curl_ws_send(handle, "", 0, &received, 0, CURLWS_PONG);
        return groundCreateValue(STRING, "[PING]");
    }
    
    char* data = malloc(received + 1);
    memcpy(data, buffer, received);
    data[received] = '\0';
    
    return groundCreateValue(STRING, data);
}

// WebSocket Close - properly close a connection, returns BOOL success
GroundValue ws_close(GroundScope* scope, List args) {
    int conn_id = (int)args.values[0].data.intVal;
    if (conn_id < 0 || conn_id >= MAX_WS_CONNECTIONS || !ws_connections[conn_id].connected) {
        ERROR("Invalid or already connected websocket", "GenericWSError");
    }
    
    CURL* handle = ws_connections[conn_id].handle;
    size_t sent;
    
    // Send close frame
    curl_ws_send(handle, "", 0, &sent, 0, CURLWS_CLOSE);
    
    // Cleanup
    curl_easy_cleanup(handle);
    free(ws_connections[conn_id].url);
    ws_connections[conn_id].handle = NULL;
    ws_connections[conn_id].connected = false;
    ws_connections[conn_id].url = NULL;
    
    return groundCreateValue(BOOL, true);
}

// WebSocket Send Binary data, returns bytes sent as INT
GroundValue ws_send_binary(GroundScope* scope, List args) {
    int conn_id = (int)args.values[0].data.intVal;
    if (conn_id < 0 || conn_id >= MAX_WS_CONNECTIONS || !ws_connections[conn_id].connected) {
        ERROR("Invalid WebSocket connection ID", "GenericWSError");
    }
    
    CURL* handle = ws_connections[conn_id].handle;
    const char* data = args.values[1].data.stringVal;
    size_t sent;
    
    CURLcode result = curl_ws_send(handle, data, strlen(data), &sent, 0, CURLWS_BINARY);
    
    if (result != CURLE_OK) {
        char buf[256];
        snprintf(buf, sizeof(buf) - 1, "WebSocket binary send failed: %s\n", curl_easy_strerror(result));
        ERROR(buf, "ActionWSError");
    }
    
    return groundCreateValue(INT, (int64_t)sent);
}

void ground_init(GroundScope* scope) {
    // HTTP Functions
    groundAddNativeFunction(scope, "request_Get", get_request, STRING, 1, STRING, "url");
    groundAddNativeFunction(scope, "request_Post", post_request, STRING, 2, STRING, "url", STRING, "data");
    
    // WebSocket Functions
    groundAddNativeFunction(scope, "ws_Connect", ws_connect, INT, 1, STRING, "url");
    groundAddNativeFunction(scope, "ws_Send", ws_send, INT, 2, INT, "conn_id", STRING, "message");
    groundAddNativeFunction(scope, "ws_Receive", ws_receive, STRING, 1, INT, "conn_id");
    groundAddNativeFunction(scope, "ws_ReceiveTimeout", ws_receive_timeout, STRING, 2, INT, "conn_id", INT, "timeout_ms");
    groundAddNativeFunction(scope, "ws_Close", ws_close, BOOL, 1, INT, "conn_id");
    groundAddNativeFunction(scope, "ws_SendBinary", ws_send_binary, INT, 2, INT, "conn_id", STRING, "data");
}