#include "wren.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>

#ifdef _WIN32
  #include <winsock2.h>
  #include <ws2tcpip.h>
  #pragma comment(lib, "ws2_32.lib")
  typedef SOCKET socket_t;
  typedef HANDLE thread_t;
  typedef CRITICAL_SECTION mutex_t;
  typedef CONDITION_VARIABLE cond_t;
#else
  #include <pthread.h>
  #include <sys/socket.h>
  #include <netinet/in.h>
  #include <arpa/inet.h>
  #include <unistd.h>
  #include <errno.h>
  #include <sys/select.h>
  #include <fcntl.h>
  typedef int socket_t;
  typedef pthread_t thread_t;
  typedef pthread_mutex_t mutex_t;
  typedef pthread_cond_t cond_t;
  #define closesocket(s) close(s)
#endif

// --- Data Structures ---

typedef enum {
    SOCKET_OP_CONNECT,
    SOCKET_OP_NEW,
    SOCKET_OP_BIND,
    SOCKET_OP_LISTEN,
    SOCKET_OP_ACCEPT,
    SOCKET_OP_READ,
    SOCKET_OP_READ_UNTIL,
    SOCKET_OP_READ_EXACTLY,
    SOCKET_OP_WRITE,
    SOCKET_OP_IS_READABLE,
    SOCKET_OP_SELECT
} SocketOp;

typedef struct {
    char* data;
    int length;
} Buffer;

typedef struct SocketContext {
    WrenVM* vm;
    SocketOp operation;
    WrenHandle* callback;

    // Operation specific data
    socket_t sock;
    char* host;
    int port;
    int backlog;
    int length;
    Buffer write_data;
    char* until_bytes;
    int until_len;
    WrenHandle* sockets_list_handle; // For select

    // Result data
    bool success;
    char* error_message;
    socket_t new_sock;
    Buffer read_data;
    WrenHandle* readable_sockets_handle; // For select result

    struct SocketContext* next;
} SocketContext;

// --- Thread-Safe Queue for Socket Operations ---

typedef struct {
    SocketContext *head, *tail;
    mutex_t mutex;
    cond_t cond;
} SocketThreadSafeQueue;

void socket_queue_init(SocketThreadSafeQueue* q) {
    q->head = q->tail = NULL;
    #ifdef _WIN32
        InitializeCriticalSection(&q->mutex);
        InitializeConditionVariable(&q->cond);
    #else
        pthread_mutex_init(&q->mutex, NULL);
        pthread_cond_init(&q->cond, NULL);
    #endif
}

void socket_queue_destroy(SocketThreadSafeQueue* q) {
    #ifdef _WIN32
        DeleteCriticalSection(&q->mutex);
    #else
        pthread_mutex_destroy(&q->mutex);
        pthread_cond_destroy(&q->cond);
    #endif
}

void socket_queue_push(SocketThreadSafeQueue* q, SocketContext* context) {
    #ifdef _WIN32
        EnterCriticalSection(&q->mutex);
    #else
        pthread_mutex_lock(&q->mutex);
    #endif
    if(context) context->next = NULL;
    if (q->tail) q->tail->next = context;
    else q->head = context;
    q->tail = context;
    #ifdef _WIN32
        WakeConditionVariable(&q->cond);
        LeaveCriticalSection(&q->mutex);
    #else
        pthread_cond_signal(&q->cond);
        pthread_mutex_unlock(&q->mutex);
    #endif
}

SocketContext* socket_queue_pop(SocketThreadSafeQueue* q) {
    #ifdef _WIN32
        EnterCriticalSection(&q->mutex);
        while (q->head == NULL) {
            SleepConditionVariableCS(&q->cond, &q->mutex, INFINITE);
        }
    #else
        pthread_mutex_lock(&q->mutex);
        while (q->head == NULL) {
            pthread_cond_wait(&q->cond, &q->mutex);
        }
    #endif
    SocketContext* context = q->head;
    q->head = q->head->next;
    if (q->head == NULL) q->tail = NULL;
    #ifdef _WIN32
        LeaveCriticalSection(&q->mutex);
    #else
        pthread_mutex_unlock(&q->mutex);
    #endif
    return context;
}

bool socket_queue_empty(SocketThreadSafeQueue* q) {
    bool empty;
    #ifdef _WIN32
        EnterCriticalSection(&q->mutex);
        empty = (q->head == NULL);
        LeaveCriticalSection(&q->mutex);
    #else
        pthread_mutex_lock(&q->mutex);
        empty = (q->head == NULL);
        pthread_mutex_unlock(&q->mutex);
    #endif
    return empty;
}

// --- Async Socket Manager ---

typedef struct {
    WrenVM* vm;
    volatile bool running;
    thread_t threads[4]; // 4 worker threads
    SocketThreadSafeQueue requestQueue;
    SocketThreadSafeQueue completionQueue;
} AsyncSocketManager;

static AsyncSocketManager* socketManager = NULL;

void free_socket_context(SocketContext* context) {
    if (context == NULL) return;
    free(context->host);
    free(context->error_message);
    free(context->write_data.data);
    free(context->read_data.data);
    free(context->until_bytes);
    if (context->callback) wrenReleaseHandle(context->vm, context->callback);
    if (context->sockets_list_handle) wrenReleaseHandle(context->vm, context->sockets_list_handle);
    if (context->readable_sockets_handle) wrenReleaseHandle(context->vm, context->readable_sockets_handle);
    free(context);
}

static void set_socket_context_error(SocketContext* context, const char* message) {
    if (context == NULL) return;
    context->success = false;
    if (context->error_message) free(context->error_message);
    context->error_message = message ? strdup(message) : strdup("An unknown socket error occurred.");
}

#ifdef _WIN32
DWORD WINAPI socketWorkerThread(LPVOID arg);
#else
void* socketWorkerThread(void* arg);
#endif

void socketManager_create(WrenVM* vm) {
    if (socketManager != NULL) return;
    socketManager = (AsyncSocketManager*)malloc(sizeof(AsyncSocketManager));
    if (socketManager == NULL) return;

    #ifdef _WIN32
        WSADATA wsaData;
        WSAStartup(MAKEWORD(2, 2), &wsaData);
    #endif

    socketManager->vm = vm;
    socketManager->running = true;
    socket_queue_init(&socketManager->requestQueue);
    socket_queue_init(&socketManager->completionQueue);
    for (int i = 0; i < 4; ++i) {
        #ifdef _WIN32
            socketManager->threads[i] = CreateThread(NULL, 0, socketWorkerThread, socketManager, 0, NULL);
        #else
            pthread_create(&socketManager->threads[i], NULL, socketWorkerThread, socketManager);
        #endif
    }
}

void socketManager_destroy() {
    if (!socketManager) return;
    socketManager->running = false;
    for (int i = 0; i < 4; ++i) socket_queue_push(&socketManager->requestQueue, NULL);
    for (int i = 0; i < 4; ++i) {
        #ifdef _WIN32
            WaitForSingleObject(socketManager->threads[i], INFINITE);
            CloseHandle(socketManager->threads[i]);
        #else
            pthread_join(socketManager->threads[i], NULL);
        #endif
    }
    while(!socket_queue_empty(&socketManager->requestQueue)) free_socket_context(socket_queue_pop(&socketManager->requestQueue));
    while(!socket_queue_empty(&socketManager->completionQueue)) free_socket_context(socket_queue_pop(&socketManager->completionQueue));
    socket_queue_destroy(&socketManager->requestQueue);
    socket_queue_destroy(&socketManager->completionQueue);
    free(socketManager);
    socketManager = NULL;

    #ifdef _WIN32
        WSACleanup();
    #endif
}

void socketManager_processCompletions() {
    if (!socketManager || !socketManager->vm || socket_queue_empty(&socketManager->completionQueue)) return;

    while (!socket_queue_empty(&socketManager->completionQueue)) {
        SocketContext* context = socket_queue_pop(&socketManager->completionQueue);
        if (context == NULL) continue;

        if (context->callback == NULL) {
            free_socket_context(context);
            continue;
        }

        WrenHandle* callHandle = wrenMakeCallHandle(socketManager->vm, "call(_,_)");
        wrenEnsureSlots(socketManager->vm, 3);
        wrenSetSlotHandle(socketManager->vm, 0, context->callback);

        if (context->success) {
            wrenSetSlotNull(socketManager->vm, 1); // error is null
            switch(context->operation) {
                case SOCKET_OP_CONNECT:
                case SOCKET_OP_NEW:
                case SOCKET_OP_ACCEPT:
                     wrenSetSlotDouble(socketManager->vm, 2, (double)context->new_sock);
                     break;
                case SOCKET_OP_BIND:
                case SOCKET_OP_LISTEN:
                case SOCKET_OP_IS_READABLE:
                     wrenSetSlotBool(socketManager->vm, 2, true);
                     break;
                case SOCKET_OP_READ:
                case SOCKET_OP_READ_UNTIL:
                case SOCKET_OP_READ_EXACTLY:
                     if (context->read_data.data) {
                        wrenSetSlotBytes(socketManager->vm, 2, context->read_data.data, context->read_data.length);
                     } else {
                        wrenSetSlotNull(socketManager->vm, 2);
                     }
                     break;
                case SOCKET_OP_SELECT:
                    wrenSetSlotHandle(socketManager->vm, 2, context->readable_sockets_handle);
                    break;
                case SOCKET_OP_WRITE:
                default:
                    wrenSetSlotNull(socketManager->vm, 2);
                    break;
            }
        } else {
            wrenSetSlotString(socketManager->vm, 1, context->error_message ? context->error_message : "Unknown error.");
            wrenSetSlotNull(socketManager->vm, 2);
        }

        wrenCall(socketManager->vm, callHandle);
        wrenReleaseHandle(socketManager->vm, callHandle);
        free_socket_context(context);
    }
}

// --- Worker Thread Implementation ---

#ifdef _WIN32
DWORD WINAPI socketWorkerThread(LPVOID arg) {
#else
void* socketWorkerThread(void* arg) {
#endif
    AsyncSocketManager* manager = (AsyncSocketManager*)arg;
    while (manager->running) {
        SocketContext* context = socket_queue_pop(&manager->requestQueue);
        if (!context || !manager->running) {
            if (context) free_socket_context(context);
            break;
        }

        switch (context->operation) {
            case SOCKET_OP_NEW: {
                context->new_sock = socket(AF_INET, SOCK_STREAM, 0);
                if (context->new_sock == -1) {
                    set_socket_context_error(context, "Failed to create socket.");
                } else {
                    context->success = true;
                }
                break;
            }
            case SOCKET_OP_CONNECT: {
                struct sockaddr_in serv_addr;
                context->new_sock = socket(AF_INET, SOCK_STREAM, 0);
                if (context->new_sock < 0) {
                    set_socket_context_error(context, "Socket creation error");
                    break;
                }
                serv_addr.sin_family = AF_INET;
                serv_addr.sin_port = htons(context->port);
                if(inet_pton(AF_INET, context->host, &serv_addr.sin_addr)<=0) {
                    set_socket_context_error(context, "Invalid address/ Address not supported");
                    break;
                }
                if (connect(context->new_sock, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
                    set_socket_context_error(context, "Connection Failed");
                } else {
                    context->success = true;
                }
                break;
            }
            case SOCKET_OP_BIND: {
                struct sockaddr_in address;
                address.sin_family = AF_INET;
                address.sin_addr.s_addr = INADDR_ANY; // Or inet_addr(context->host) for specific interface
                address.sin_port = htons(context->port);
                if (bind(context->sock, (struct sockaddr *)&address, sizeof(address)) < 0) {
                    set_socket_context_error(context, "Bind failed");
                } else {
                    context->success = true;
                }
                break;
            }
            case SOCKET_OP_LISTEN: {
                 if (listen(context->sock, context->backlog) < 0) {
                    set_socket_context_error(context, "Listen failed");
                } else {
                    context->success = true;
                }
                break;
            }
            case SOCKET_OP_ACCEPT: {
                struct sockaddr_in address;
                int addrlen = sizeof(address);
                context->new_sock = accept(context->sock, (struct sockaddr *)&address, (socklen_t*)&addrlen);
                if (context->new_sock < 0) {
                    set_socket_context_error(context, "Accept failed");
                } else {
                    context->success = true;
                }
                break;
            }
            case SOCKET_OP_READ: {
                char* buffer = malloc(context->length + 1);
                int valread = recv(context->sock, buffer, context->length, 0);
                if (valread >= 0) {
                    context->read_data.data = buffer;
                    context->read_data.length = valread;
                    context->success = true;
                } else {
                    free(buffer);
                    set_socket_context_error(context, "Read failed");
                }
                break;
            }
            case SOCKET_OP_WRITE: {
                int total_sent = 0;
                while (total_sent < context->write_data.length) {
                    int sent = send(context->sock, context->write_data.data + total_sent, context->write_data.length - total_sent, 0);
                    if (sent < 0) {
                        set_socket_context_error(context, "Write failed");
                        break;
                    }
                    total_sent += sent;
                }
                if (total_sent == context->write_data.length) {
                    context->success = true;
                }
                break;
            }
            // ... other cases ...
        }
        socket_queue_push(&manager->completionQueue, context);
    }
    return 0;
}

// --- Wren FFI Functions ---

static void create_socket_context(WrenVM* vm, SocketOp op) {
    SocketContext* context = (SocketContext*)calloc(1, sizeof(SocketContext));
    if (!context) {
        wrenSetSlotString(vm, 0, "Out of memory.");
        wrenAbortFiber(vm, 0);
        return;
    }
    context->vm = vm;
    context->operation = op;

    switch(op) {
        case SOCKET_OP_CONNECT:
            context->host = strdup(wrenGetSlotString(vm, 1));
            context->port = (int)wrenGetSlotDouble(vm, 2);
            context->callback = wrenGetSlotHandle(vm, 3);
            break;
        case SOCKET_OP_NEW:
            context->callback = wrenGetSlotHandle(vm, 1);
            break;
        case SOCKET_OP_BIND:
            context->sock = (socket_t)wrenGetSlotDouble(vm, 1);
            context->host = strdup(wrenGetSlotString(vm, 2));
            context->port = (int)wrenGetSlotDouble(vm, 3);
            context->callback = wrenGetSlotHandle(vm, 4);
            break;
        case SOCKET_OP_LISTEN:
            context->sock = (socket_t)wrenGetSlotDouble(vm, 1);
            context->backlog = (int)wrenGetSlotDouble(vm, 2);
            context->callback = wrenGetSlotHandle(vm, 3);
            break;
        case SOCKET_OP_ACCEPT:
            context->sock = (socket_t)wrenGetSlotDouble(vm, 1);
            context->callback = wrenGetSlotHandle(vm, 2);
            break;
        case SOCKET_OP_READ:
            context->sock = (socket_t)wrenGetSlotDouble(vm, 1);
            context->length = (int)wrenGetSlotDouble(vm, 2);
            context->callback = wrenGetSlotHandle(vm, 3);
            break;
        case SOCKET_OP_WRITE:
            context->sock = (socket_t)wrenGetSlotDouble(vm, 1);
            int len;
            const char* data = wrenGetSlotBytes(vm, 2, &len);
            context->write_data.data = malloc(len);
            memcpy(context->write_data.data, data, len);
            context->write_data.length = len;
            context->callback = wrenGetSlotHandle(vm, 3);
            break;
        // ... other cases ...
        default:
            free(context);
            return;
    }
    socket_queue_push(&socketManager->requestQueue, context);
}

void socketConnect(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_CONNECT); }
void socketNew(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_NEW); }
void socketBind(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_BIND); }
void socketListen(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_LISTEN); }
void socketAccept(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_ACCEPT); }
void socketRead(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_READ); }
void socketWrite(WrenVM* vm) { create_socket_context(vm, SOCKET_OP_WRITE); }
// ... other FFI functions ...

WrenForeignMethodFn bindSocketForeignMethod(WrenVM* vm, const char* module, const char* className, bool isStatic, const char* signature) {
    if (strcmp(module, "socket") != 0) return NULL;
    if (strcmp(className, "Socket") == 0 && isStatic) {
        if (strcmp(signature, "connect_(_,_,_)") == 0) return socketConnect;
        if (strcmp(signature, "new_(_)") == 0) return socketNew;
        if (strcmp(signature, "bind_(_,_,_,_)") == 0) return socketBind;
        if (strcmp(signature, "listen_(_,_,_)") == 0) return socketListen;
        if (strcmp(signature, "accept_(_,_)") == 0) return socketAccept;
        if (strcmp(signature, "read_(_,_,_)") == 0) return socketRead;
        if (strcmp(signature, "write_(_,_,_)") == 0) return socketWrite;
        // ... other bindings ...
    }
    return NULL;
}

WrenForeignClassMethods bindSocketForeignClass(WrenVM* vm, const char* module, const char* className) {
    WrenForeignClassMethods methods = {0, 0};
    return methods;
}