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

#ifdef _WIN32
  #include <windows.h>
  typedef HANDLE thread_t;
  typedef CRITICAL_SECTION mutex_t;
  typedef CONDITION_VARIABLE cond_t;
#else
  #include <pthread.h>
  typedef pthread_t thread_t;
  typedef pthread_mutex_t mutex_t;
  typedef pthread_cond_t cond_t;
#endif

// --- Data Structures ---

typedef enum {
    DB_OP_OPEN,
    DB_OP_EXEC,
    DB_OP_QUERY,
    DB_OP_CLOSE
} DBOp;

typedef struct {
    sqlite3* db;
} DatabaseData;

// C-side representation of query results to pass from worker to main thread.
typedef struct DbValue {
    int type;
    union {
        double num;
        struct {
            char* text;
            int length;
        } str;
    } as;
} DbValue;

typedef struct DbRow {
    char** columns;
    DbValue* values;
    int count;
    struct DbRow* next;
} DbRow;

typedef struct DbContext {
    WrenVM* vm;
    DBOp operation;
    WrenHandle* callback;
    WrenHandle* dbHandle;
    char* path;
    sqlite3* newDb; 
    char* sql;
    sqlite3* db; 
    bool success;
    char* errorMessage;
    DbRow* resultRows; 
    struct DbContext* next;
} DbContext;

// --- Thread-Safe Queue ---

typedef struct {
    DbContext *head, *tail;
    mutex_t mutex;
    cond_t cond;
} DbThreadSafeQueue;

void db_queue_init(DbThreadSafeQueue* 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 db_queue_destroy(DbThreadSafeQueue* q) {
    #ifdef _WIN32
        DeleteCriticalSection(&q->mutex);
    #else
        pthread_mutex_destroy(&q->mutex);
        pthread_cond_destroy(&q->cond);
    #endif
}

void db_queue_push(DbThreadSafeQueue* q, DbContext* 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
}

DbContext* db_queue_pop(DbThreadSafeQueue* 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
    DbContext* 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 db_queue_empty(DbThreadSafeQueue* 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 DB Manager ---

typedef struct {
    WrenVM* vm;
    volatile bool running;
    thread_t threads[2]; 
    DbThreadSafeQueue requestQueue;
    DbThreadSafeQueue completionQueue;
} AsyncDbManager;

static AsyncDbManager* dbManager = NULL;

void free_db_result_rows(DbRow* rows) {
    while (rows) {
        DbRow* next = rows->next;
        if (rows->columns) {
            for (int i = 0; i < rows->count; i++) {
                free(rows->columns[i]);
            }
            free(rows->columns);
        }
        if (rows->values) {
            for (int i = 0; i < rows->count; i++) {
                if (rows->values[i].type == SQLITE_TEXT || rows->values[i].type == SQLITE_BLOB) {
                    free(rows->values[i].as.str.text);
                }
            }
            free(rows->values);
        }
        free(rows);
        rows = next;
    }
}

void free_db_context(DbContext* context) {
    if (context == NULL) return;
    free(context->path);
    free(context->sql);
    free(context->errorMessage);
    if (context->dbHandle) wrenReleaseHandle(context->vm, context->dbHandle);
    if (context->callback) wrenReleaseHandle(context->vm, context->callback);
    if (context->resultRows) free_db_result_rows(context->resultRows);
    free(context);
}

static void set_context_error(DbContext* context, const char* message) {
    if (context == NULL) return;
    context->success = false;
    if (context->errorMessage) free(context->errorMessage);
    context->errorMessage = message ? strdup(message) : strdup("An unknown database error occurred.");
}

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

void dbManager_create(WrenVM* vm) {
    if (dbManager != NULL) return;
    dbManager = (AsyncDbManager*)malloc(sizeof(AsyncDbManager));
    if (dbManager == NULL) return;
    dbManager->vm = vm;
    dbManager->running = true;
    db_queue_init(&dbManager->requestQueue);
    db_queue_init(&dbManager->completionQueue);
    for (int i = 0; i < 2; ++i) {
        #ifdef _WIN32
            dbManager->threads[i] = CreateThread(NULL, 0, dbWorkerThread, dbManager, 0, NULL);
        #else
            pthread_create(&dbManager->threads[i], NULL, dbWorkerThread, dbManager);
        #endif
    }
}

void dbManager_destroy() {
    if (!dbManager) return;
    dbManager->running = false;
    for (int i = 0; i < 2; ++i) db_queue_push(&dbManager->requestQueue, NULL);
    for (int i = 0; i < 2; ++i) {
        #ifdef _WIN32
            WaitForSingleObject(dbManager->threads[i], INFINITE);
            CloseHandle(dbManager->threads[i]);
        #else
            pthread_join(dbManager->threads[i], NULL);
        #endif
    }
    while(!db_queue_empty(&dbManager->requestQueue)) free_db_context(db_queue_pop(&dbManager->requestQueue));
    while(!db_queue_empty(&dbManager->completionQueue)) free_db_context(db_queue_pop(&dbManager->completionQueue));
    db_queue_destroy(&dbManager->requestQueue);
    db_queue_destroy(&dbManager->completionQueue);
    free(dbManager);
    dbManager = NULL;
}

void dbManager_processCompletions() {
    if (!dbManager || !dbManager->vm || db_queue_empty(&dbManager->completionQueue)) return;
    
    while (!db_queue_empty(&dbManager->completionQueue)) {
        DbContext* context = db_queue_pop(&dbManager->completionQueue);
        if (context == NULL) continue;

        if (context->success && context->dbHandle) {
            wrenEnsureSlots(dbManager->vm, 1);
            wrenSetSlotHandle(dbManager->vm, 0, context->dbHandle);
            DatabaseData* dbData = (DatabaseData*)wrenGetSlotForeign(dbManager->vm, 0);
            if (dbData) {
                if (context->operation == DB_OP_OPEN) dbData->db = context->newDb;
                else if (context->operation == DB_OP_CLOSE) dbData->db = NULL;
            }
        }

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

        WrenHandle* callHandle = NULL;
        int numArgs = 0;
        if (context->operation == DB_OP_QUERY) {
            callHandle = wrenMakeCallHandle(dbManager->vm, "call(_,_)");
            numArgs = 2;
        } else {
            callHandle = wrenMakeCallHandle(dbManager->vm, "call(_)");
            numArgs = 1;
        }

        if (callHandle == NULL) {
            free_db_context(context);
            continue;
        }

        // Ensure enough slots for callback, args, and temp work.
        // Slots 0, 1, 2 are for the callback and its arguments.
        // Slots 3, 4, 5 are for temporary work building maps.
        wrenEnsureSlots(dbManager->vm, 6); 
        wrenSetSlotHandle(dbManager->vm, 0, context->callback);

        if (context->success) {
            wrenSetSlotNull(dbManager->vm, 1); // error is null
            if (numArgs == 2) { // Query case
                if (context->resultRows) {
                    wrenSetSlotNewList(dbManager->vm, 2); // Result list in slot 2
                    DbRow* row = context->resultRows;
                    while(row) {
                        wrenSetSlotNewMap(dbManager->vm, 3); // Temp map for row in slot 3
                        for (int i = 0; i < row->count; i++) {
                            // Use slots 4 and 5 for key/value to avoid conflicts
                            wrenSetSlotString(dbManager->vm, 4, row->columns[i]);
                            DbValue* val = &row->values[i];
                            switch (val->type) {
                                case SQLITE_INTEGER: case SQLITE_FLOAT:
                                    wrenSetSlotDouble(dbManager->vm, 5, val->as.num); break;
                                case SQLITE_TEXT: case SQLITE_BLOB:
                                    wrenSetSlotBytes(dbManager->vm, 5, val->as.str.text, val->as.str.length); break;
                                case SQLITE_NULL:
                                    wrenSetSlotNull(dbManager->vm, 5); break;
                            }
                            wrenSetMapValue(dbManager->vm, 3, 4, 5); // map=3, key=4, val=5
                        }
                        wrenInsertInList(dbManager->vm, 2, -1, 3); // list=2, element=3
                        row = row->next;
                    }
                } else {
                    wrenSetSlotNewList(dbManager->vm, 2); // Return empty list for success with no rows
                }
            }
        } else {
            wrenSetSlotString(dbManager->vm, 1, context->errorMessage ? context->errorMessage : "Unknown error.");
            if (numArgs == 2) wrenSetSlotNull(dbManager->vm, 2);
        }

        wrenCall(dbManager->vm, callHandle);
        wrenReleaseHandle(dbManager->vm, callHandle);
        free_db_context(context);
    }
}

// --- Worker Thread ---
#ifdef _WIN32
DWORD WINAPI dbWorkerThread(LPVOID arg) {
#else
void* dbWorkerThread(void* arg) {
#endif
    AsyncDbManager* manager = (AsyncDbManager*)arg;
    while (manager->running) {
        DbContext* context = db_queue_pop(&manager->requestQueue);
        if (!context || !manager->running) {
            if (context) free_db_context(context);
            break;
        }
        switch (context->operation) {
            case DB_OP_OPEN: {
                int rc = sqlite3_open_v2(context->path, &context->newDb, 
                                         SQLITE_OPEN_READWRITE | SQLITE_OPEN_CREATE | SQLITE_OPEN_FULLMUTEX, 
                                         NULL);
                if (rc != SQLITE_OK) {
                    set_context_error(context, sqlite3_errmsg(context->newDb));
                    sqlite3_close(context->newDb);
                    context->newDb = NULL;
                } else {
                    context->success = true;
                }
                break;
            }
            case DB_OP_EXEC: {
                if (!context->db) { set_context_error(context, "Database is not open."); break; }
                char* err = NULL;
                int rc = sqlite3_exec(context->db, context->sql, 0, 0, &err);
                if (rc != SQLITE_OK) {
                    set_context_error(context, err);
                    sqlite3_free(err);
                } else {
                    context->success = true;
                }
                break;
            }
            case DB_OP_QUERY: {
                if (!context->db) { set_context_error(context, "Database is not open."); break; }
                sqlite3_stmt* stmt;
                int rc = sqlite3_prepare_v2(context->db, context->sql, -1, &stmt, 0);
                if (rc != SQLITE_OK) { set_context_error(context, sqlite3_errmsg(context->db)); break; }

                int colCount = sqlite3_column_count(stmt);
                DbRow* head = NULL, *tail = NULL;
                bool oom = false;
                while ((rc = sqlite3_step(stmt)) == SQLITE_ROW) {
                    DbRow* row = (DbRow*)calloc(1, sizeof(DbRow));
                    if (!row) { oom = true; break; }
                    row->count = colCount;
                    row->columns = (char**)malloc(sizeof(char*) * colCount);
                    row->values = (DbValue*)malloc(sizeof(DbValue) * colCount);
                    if (!row->columns || !row->values) { free(row->columns); free(row->values); free(row); oom = true; break; }
                    for (int i = 0; i < colCount; i++) {
                        const char* colName = sqlite3_column_name(stmt, i);
                        row->columns[i] = colName ? strdup(colName) : strdup("");
                        if (!row->columns[i]) { for (int j=0; j<i; j++) free(row->columns[j]); free(row->columns); free(row->values); free(row); oom = true; goto query_loop_end; }
                        DbValue* val = &row->values[i];
                        val->type = sqlite3_column_type(stmt, i);
                        switch (val->type) {
                            case SQLITE_INTEGER: case SQLITE_FLOAT: val->as.num = sqlite3_column_double(stmt, i); break;
                            case SQLITE_TEXT: case SQLITE_BLOB: {
                                const void* blob = sqlite3_column_blob(stmt, i);
                                int len = sqlite3_column_bytes(stmt, i);
                                val->as.str.text = (char*)malloc(len);
                                if (!val->as.str.text) { for (int j=0; j<=i; j++) free(row->columns[j]); free(row->columns); free(row->values); free(row); oom = true; goto query_loop_end; }
                                memcpy(val->as.str.text, blob, len);
                                val->as.str.length = len;
                                break;
                            }
                            case SQLITE_NULL: break;
                        }
                    }
                    if (!head) head = tail = row; else { tail->next = row; tail = row; }
                }
                query_loop_end:;
                if (oom) { set_context_error(context, "Out of memory during query."); free_db_result_rows(head); }
                else if (rc != SQLITE_DONE) { set_context_error(context, sqlite3_errmsg(context->db)); free_db_result_rows(head); }
                else { context->success = true; context->resultRows = head; }
                sqlite3_finalize(stmt);
                break;
            }
            case DB_OP_CLOSE: {
                if (context->db) sqlite3_close(context->db);
                context->success = true;
                break;
            }
        }
        db_queue_push(&manager->completionQueue, context);
    }
    return 0;
}

// --- Wren FFI ---

void dbAllocate(WrenVM* vm) {
    DatabaseData* data = (DatabaseData*)wrenSetSlotNewForeign(vm, 0, 0, sizeof(DatabaseData));
    if (data) data->db = NULL;
}

void dbFinalize(void* data) {
    DatabaseData* dbData = (DatabaseData*)data;
    if (dbData && dbData->db) sqlite3_close(dbData->db);
}

static void create_db_context(WrenVM* vm, DBOp op, int sqlSlot, int cbSlot) {
    DbContext* context = (DbContext*)calloc(1, sizeof(DbContext));
    if (!context) { wrenSetSlotString(vm, 0, "Out of memory."); wrenAbortFiber(vm, 0); return; }
    context->vm = vm;
    context->operation = op;
    context->dbHandle = wrenGetSlotHandle(vm, 0);
    context->callback = wrenGetSlotHandle(vm, cbSlot);
    if (sqlSlot != -1) {
        if (wrenGetSlotType(vm, sqlSlot) != WREN_TYPE_STRING) {
            wrenSetSlotString(vm, 0, "SQL argument must be a string.");
            wrenAbortFiber(vm, 0);
            free_db_context(context);
            return;
        }
        const char* sql_str = wrenGetSlotString(vm, sqlSlot);
        if (sql_str) context->sql = strdup(sql_str);
        if (!context->sql) { set_context_error(context, "Out of memory."); db_queue_push(&dbManager->requestQueue, context); return; }
    }
    DatabaseData* dbData = (DatabaseData*)wrenGetSlotForeign(vm, 0);
    if (!dbData) { set_context_error(context, "Invalid database object."); db_queue_push(&dbManager->requestQueue, context); return; }
    context->db = dbData->db;
    db_queue_push(&dbManager->requestQueue, context);
}

void dbOpen(WrenVM* vm) {
    DbContext* context = (DbContext*)calloc(1, sizeof(DbContext));
    if (!context) { wrenSetSlotString(vm, 0, "Out of memory."); wrenAbortFiber(vm, 0); return; }
    context->vm = vm;
    context->operation = DB_OP_OPEN;
    const char* path_str = wrenGetSlotString(vm, 1);
    if (path_str) context->path = strdup(path_str);
    if (!context->path) { free(context); wrenSetSlotString(vm, 0, "Out of memory."); wrenAbortFiber(vm, 0); return; }
    context->dbHandle = wrenGetSlotHandle(vm, 0);
    context->callback = wrenGetSlotHandle(vm, 2);
    db_queue_push(&dbManager->requestQueue, context);
}

void dbExec(WrenVM* vm) { create_db_context(vm, DB_OP_EXEC, 1, 2); }
void dbQuery(WrenVM* vm) { create_db_context(vm, DB_OP_QUERY, 1, 2); }
void dbClose(WrenVM* vm) { create_db_context(vm, DB_OP_CLOSE, -1, 1); }

WrenForeignMethodFn bindSqliteForeignMethod(WrenVM* vm, const char* module, const char* className, bool isStatic, const char* signature) {
    if (strcmp(module, "sqlite") != 0) return NULL;
    if (strcmp(className, "Database") == 0 && !isStatic) {
        if (strcmp(signature, "open_(_,_)") == 0) return dbOpen;
        if (strcmp(signature, "exec_(_,_)") == 0) return dbExec;
        if (strcmp(signature, "query_(_,_)") == 0) return dbQuery;
        if (strcmp(signature, "close_(_)") == 0) return dbClose;
    }
    return NULL;
}

WrenForeignClassMethods bindSqliteForeignClass(WrenVM* vm, const char* module, const char* className) {
    if (strcmp(module, "sqlite") == 0 && strcmp(className, "Database") == 0) {
        WrenForeignClassMethods methods = {dbAllocate, dbFinalize};
        return methods;
    }
    WrenForeignClassMethods methods = {0, 0};
    return methods;
}