// socket_backend.cpp (Native Sockets Implementation)
#include "wren.h"
#include <iostream>
#include <string>
#include <vector>
#include <thread>
#include <mutex>
#include <condition_variable>
#include <queue>
#include <atomic>
#include <chrono>
#include <string.h>
#ifdef _WIN32
  #include <winsock2.h>
  #include <ws2tcpip.h>
  #pragma comment(lib, "ws2_32.lib")
  typedef SOCKET socket_t;
  #define IS_SOCKET_VALID(s) ((s) != INVALID_SOCKET)
  #define CLOSE_SOCKET(s) closesocket(s)
#else
  #include <sys/socket.h>
  #include <netinet/in.h>
  #include <arpa/inet.h>
  #include <unistd.h>
  #include <fcntl.h>
  #include <netdb.h>
  #include <cerrno>
  typedef int socket_t;
  #define INVALID_SOCKET -1
  #define IS_SOCKET_VALID(s) ((s) >= 0)
  #define CLOSE_SOCKET(s) close(s)
#endif


// --- Thread-Safe Queue for Asynchronous Operations ---

template <typename T>
class ThreadSafeQueue {
public:
    void push(T item) {
        std::lock_guard<std::mutex> lock(mutex_);
        queue_.push(item);
        cond_.notify_one();
    }

    T pop() {
        std::unique_lock<std::mutex> lock(mutex_);
        cond_.wait(lock, [this] { return !queue_.empty(); });
        T item = queue_.front();
        queue_.pop();
        return item;
    }

    bool empty() {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.empty();
    }

private:
    std::queue<T> queue_;
    std::mutex mutex_;
    std::condition_variable cond_;
};

// --- Context for Asynchronous Socket Operations ---

enum class SocketOp {
    CONNECT,
    ACCEPT,
    READ,
    WRITE
};

struct SocketData {
    socket_t sock;
    bool isListener;
};

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

    // Operation-specific data
    std::string host;
    int port;
    std::string data;
    int bytesToRead;

    // Result data
    bool success;
    std::string resultData;
    std::string errorMessage;
    socket_t newSocket; // For accepted connections
};


// --- Asynchronous Socket Manager ---

class AsyncSocketManager {
public:
    AsyncSocketManager(WrenVM* vm) : vm_(vm), running_(true) {
        for (int i = 0; i < 4; ++i) {
            threads_.emplace_back([this] { workerThread(); });
        }
    }

    ~AsyncSocketManager() {
        running_ = false;
        for (size_t i = 0; i < threads_.size(); ++i) {
            requestQueue_.push(nullptr);
        }
        for (auto& thread : threads_) {
            thread.join();
        }
    }

    void submit(SocketContext* context) {
        requestQueue_.push(context);
    }
    
    bool completionQueueEmpty() {
        return completionQueue_.empty();
    }

    void processCompletions() {
        while (!completionQueue_.empty()) {
            SocketContext* context = completionQueue_.pop();
            
            WrenHandle* callHandle = wrenMakeCallHandle(vm_, "call(_,_)");
            wrenEnsureSlots(vm_, 3);
            wrenSetSlotHandle(vm_, 0, context->callback);

            if (context->success) {
                wrenSetSlotNull(vm_, 1); // No error
                switch (context->operation) {
                    case SocketOp::ACCEPT: {
                        wrenGetVariable(vm_, "socket", "Socket", 2);
                        void* foreign = wrenSetSlotNewForeign(vm_, 2, 2, sizeof(SocketData));
                        SocketData* clientData = (SocketData*)foreign;
                        clientData->sock = context->newSocket;
                        clientData->isListener = false;
                        break;
                    }
                    case SocketOp::READ:
                        wrenSetSlotBytes(vm_, 2, context->resultData.c_str(), context->resultData.length());
                        break;
                    default:
                        wrenSetSlotNull(vm_, 2);
                        break;
                }
            } else {
                wrenSetSlotString(vm_, 1, context->errorMessage.c_str());
                wrenSetSlotNull(vm_, 2);
            }

            wrenCall(vm_, callHandle);

            wrenReleaseHandle(vm_, context->socketHandle);
            wrenReleaseHandle(vm_, context->callback);
            wrenReleaseHandle(vm_, callHandle);
            delete context;
        }
    }

private:
    void workerThread() {
        while (running_) {
            SocketContext* context = requestQueue_.pop();
            if (!context || !running_) break;

            wrenEnsureSlots(context->vm, 1);
            wrenSetSlotHandle(context->vm, 0, context->socketHandle);
            SocketData* socketData = (wrenGetSlotType(context->vm, 0) == WREN_TYPE_FOREIGN)
                ? (SocketData*)wrenGetSlotForeign(context->vm, 0)
                : nullptr;

            if (!socketData) {
                context->success = false;
                context->errorMessage = "Invalid socket object.";
                completionQueue_.push(context);
                continue;
            }

            switch (context->operation) {
                case SocketOp::CONNECT: {
                    addrinfo hints = {}, *addrs;
                    hints.ai_family = AF_UNSPEC;
                    hints.ai_socktype = SOCK_STREAM;
                    getaddrinfo(context->host.c_str(), std::to_string(context->port).c_str(), &hints, &addrs);
                    
                    socket_t sock = INVALID_SOCKET;
                    for (addrinfo* addr = addrs; addr; addr = addr->ai_next) {
                        sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
                        if (!IS_SOCKET_VALID(sock)) continue;
                        if (connect(sock, addr->ai_addr, (int)addr->ai_addrlen) == 0) break;
                        CLOSE_SOCKET(sock);
                        sock = INVALID_SOCKET;
                    }
                    freeaddrinfo(addrs);

                    if (IS_SOCKET_VALID(sock)) {
                        socketData->sock = sock;
                        socketData->isListener = false;
                        context->success = true;
                    } else {
                        context->success = false;
                        context->errorMessage = "Connection failed.";
                    }
                    break;
                }
                case SocketOp::ACCEPT: {
                    if (socketData->isListener) {
                        context->newSocket = accept(socketData->sock, nullptr, nullptr);
                        context->success = IS_SOCKET_VALID(context->newSocket);
                        if (!context->success) context->errorMessage = "Accept failed.";
                    } else {
                        context->success = false;
                        context->errorMessage = "Cannot accept on a non-listening socket.";
                    }
                    break;
                }
                case SocketOp::READ: {
                    if (!socketData->isListener) {
                        char buf[4096];
                        ssize_t len = recv(socketData->sock, buf, sizeof(buf), 0);
                        if (len > 0) {
                            context->resultData.assign(buf, len);
                            context->success = true;
                        } else {
                             context->success = false;
                             context->errorMessage = "Read failed or connection closed.";
                        }
                    }
                    break;
                }
                case SocketOp::WRITE: {
                    if (!socketData->isListener) {
                       ssize_t written = send(socketData->sock, context->data.c_str(), context->data.length(), 0);
                       context->success = (written == (ssize_t)context->data.length());
                       if(!context->success) context->errorMessage = "Write failed.";
                    }
                    break;
                }
            }
            completionQueue_.push(context);
        }
    }

    WrenVM* vm_;
    std::atomic<bool> running_;
    std::vector<std::thread> threads_;
    ThreadSafeQueue<SocketContext*> requestQueue_;
    ThreadSafeQueue<SocketContext*> completionQueue_;
};

static AsyncSocketManager* socketManager = nullptr;

// --- Socket Foreign Class/Methods ---

void socketAllocate(WrenVM* vm) {
    SocketData* data = (SocketData*)wrenSetSlotNewForeign(vm, 0, 0, sizeof(SocketData));
    data->sock = INVALID_SOCKET;
    data->isListener = false;
}

void socketConnect(WrenVM* vm) {
    SocketContext* context = new SocketContext();
    context->operation = SocketOp::CONNECT;
    context->vm = vm;
    context->socketHandle = wrenGetSlotHandle(vm, 0);
    context->host = wrenGetSlotString(vm, 1);
    context->port = (int)wrenGetSlotDouble(vm, 2);
    context->callback = wrenGetSlotHandle(vm, 3);
    socketManager->submit(context);
    wrenSetSlotNull(vm, 0);
}

void socketListen(WrenVM* vm) {
    const char* host = wrenGetSlotString(vm, 1);
    int port = (int)wrenGetSlotDouble(vm, 2);
    int backlog = (int)wrenGetSlotDouble(vm, 3);
    
    addrinfo hints = {}, *addrs;
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags = AI_PASSIVE;
    getaddrinfo(host, std::to_string(port).c_str(), &hints, &addrs);

    socket_t sock = INVALID_SOCKET;
    for (addrinfo* addr = addrs; addr; addr = addr->ai_next) {
        sock = socket(addr->ai_family, addr->ai_socktype, addr->ai_protocol);
        if (!IS_SOCKET_VALID(sock)) continue;
        
        int yes = 1;
        setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (const char*)&yes, sizeof(yes));

        if (bind(sock, addr->ai_addr, (int)addr->ai_addrlen) == 0) break;
        
        CLOSE_SOCKET(sock);
        sock = INVALID_SOCKET;
    }
    freeaddrinfo(addrs);
    
    if (IS_SOCKET_VALID(sock) && listen(sock, backlog) == 0) {
        SocketData* data = (SocketData*)wrenGetSlotForeign(vm, 0);
        data->sock = sock;
        data->isListener = true;
        wrenSetSlotBool(vm, 0, true);
    } else {
        if(IS_SOCKET_VALID(sock)) CLOSE_SOCKET(sock);
        wrenSetSlotBool(vm, 0, false);
    }
}

void socketAccept(WrenVM* vm) {
    SocketContext* context = new SocketContext();
    context->operation = SocketOp::ACCEPT;
    context->vm = vm;
    context->socketHandle = wrenGetSlotHandle(vm, 0);
    context->callback = wrenGetSlotHandle(vm, 1);
    socketManager->submit(context);
    wrenSetSlotNull(vm, 0);
}

void socketRead(WrenVM* vm) {
    SocketContext* context = new SocketContext();
    context->operation = SocketOp::READ;
    context->vm = vm;
    context->socketHandle = wrenGetSlotHandle(vm, 0);
    context->bytesToRead = (int)wrenGetSlotDouble(vm, 1);
    context->callback = wrenGetSlotHandle(vm, 2);
    socketManager->submit(context);
    wrenSetSlotNull(vm, 0);
}

void socketWrite(WrenVM* vm) {
    SocketContext* context = new SocketContext();
    context->operation = SocketOp::WRITE;
    context->vm = vm;
    context->socketHandle = wrenGetSlotHandle(vm, 0);
    int len;
    const char* bytes = wrenGetSlotBytes(vm, 1, &len);
    context->data.assign(bytes, len);
    context->callback = wrenGetSlotHandle(vm, 2);
    socketManager->submit(context);
    wrenSetSlotNull(vm, 0);
}

void socketClose(WrenVM* vm) {
    SocketData* data = (SocketData*)wrenGetSlotForeign(vm, 0);
    if (IS_SOCKET_VALID(data->sock)) {
        CLOSE_SOCKET(data->sock);
        data->sock = INVALID_SOCKET;
    }
}

void socketIsOpen(WrenVM* vm) {
    SocketData* data = (SocketData*)wrenGetSlotForeign(vm, 0);
    wrenSetSlotBool(vm, 0, IS_SOCKET_VALID(data->sock));
}

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, "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;
        if (strcmp(signature, "close()") == 0) return socketClose;
        if (strcmp(signature, "isOpen") == 0) return socketIsOpen;
    }
    return NULL;
}

WrenForeignClassMethods bindSocketForeignClass(WrenVM* vm, const char* module, const char* className) {
    WrenForeignClassMethods methods = {0};
    if (strcmp(module, "socket") == 0 && strcmp(className, "Socket") == 0) {
        methods.allocate = socketAllocate;
    }
    return methods;
}