#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;
  #define GET_THREAD_ID() GetCurrentThreadId()
#else
  #include <pthread.h>
  typedef pthread_t thread_t;
  typedef pthread_mutex_t mutex_t;
  typedef pthread_cond_t cond_t;
  #define GET_THREAD_ID() pthread_self()
#endif

#define TRACE() printf("[TRACE] %s:%d\n", __FUNCTION__, __LINE__)

// --- 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) {
    TRACE();
    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
    TRACE();
}

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

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

DbContext* db_queue_pop(DbThreadSafeQueue* q) {
    TRACE();
    #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
    TRACE();
    DbContext* context = q->head;
    TRACE();
    q->head = q->head->next;
    TRACE();
    if (q->head == NULL) q->tail = NULL;
    TRACE();
    #ifdef _WIN32
        LeaveCriticalSection(&q->mutex);
    #else
        pthread_mutex_unlock(&q->mutex);
    #endif
    TRACE();
    return context;
}

bool db_queue_empty(DbThreadSafeQueue* q) {
    TRACE();
    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
    TRACE();
    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) {
    TRACE();
    while (rows) {
        TRACE();
        DbRow* next = rows->next;
        if (rows->columns) {
            TRACE();
            for (int i = 0; i < rows->count; i++) {
                free(rows->columns[i]);
            }
            free(rows->columns);
        }
        if (rows->values) {
            TRACE();
            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;
    }
    TRACE();
}

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

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

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

void dbManager_create(WrenVM* vm) {
    TRACE();
    if (dbManager != NULL) { TRACE(); return; }
    TRACE();
    dbManager = (AsyncDbManager*)malloc(sizeof(AsyncDbManager));
    TRACE();
    if (dbManager == NULL) { TRACE(); return; }
    TRACE();
    dbManager->vm = vm;
    TRACE();
    dbManager->running = true;
    TRACE();
    db_queue_init(&dbManager->requestQueue);
    TRACE();
    db_queue_init(&dbManager->completionQueue);
    TRACE();
    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
    }
    TRACE();
}

void dbManager_destroy() {
    TRACE();
    if (!dbManager) { TRACE(); return; }
    TRACE();
    dbManager->running = false;
    TRACE();
    for (int i = 0; i < 2; ++i) {
        db_queue_push(&dbManager->requestQueue, NULL);
    }
    TRACE();
    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
    }
    TRACE();
    while(!db_queue_empty(&dbManager->requestQueue)) free_db_context(db_queue_pop(&dbManager->requestQueue));
    TRACE();
    while(!db_queue_empty(&dbManager->completionQueue)) free_db_context(db_queue_pop(&dbManager->completionQueue));
    TRACE();
    db_queue_destroy(&dbManager->requestQueue);
    TRACE();
    db_queue_destroy(&dbManager->completionQueue);
    TRACE();
    free(dbManager);
    TRACE();
    dbManager = NULL;
    TRACE();
}

void dbManager_processCompletions() {
    TRACE();
    if (!dbManager || !dbManager->vm || db_queue_empty(&dbManager->completionQueue)) { TRACE(); return; }
    TRACE();
    WrenHandle* callHandle = wrenMakeCallHandle(dbManager->vm, "call(_,_)");
    TRACE();
    if (callHandle == NULL) { TRACE(); return; }
    TRACE();
    while (!db_queue_empty(&dbManager->completionQueue)) {
        TRACE();
        DbContext* context = db_queue_pop(&dbManager->completionQueue);
        TRACE();
        if (context == NULL) { TRACE(); continue; }
        TRACE();
        if (context->success && context->dbHandle) {
            TRACE();
            wrenEnsureSlots(dbManager->vm, 1);
            TRACE();
            wrenSetSlotHandle(dbManager->vm, 0, context->dbHandle);
            TRACE();
            DatabaseData* dbData = (DatabaseData*)wrenGetSlotForeign(dbManager->vm, 0);
            TRACE();
            if (dbData) {
                TRACE();
                if (context->operation == DB_OP_OPEN) {
                    dbData->db = context->newDb;
                } else if (context->operation == DB_OP_CLOSE) {
                    dbData->db = NULL;
                }
            }
        }
        TRACE();
        if (context->callback == NULL) {
            TRACE();
            free_db_context(context);
            continue;
        }
        TRACE();
        wrenEnsureSlots(dbManager->vm, 3);
        TRACE();
        wrenSetSlotHandle(dbManager->vm, 0, context->callback);
        TRACE();
        if (context->success) {
            TRACE();
            wrenSetSlotNull(dbManager->vm, 1); // error
            TRACE();
            if (context->resultRows) {
                TRACE();
                wrenSetSlotNewList(dbManager->vm, 2);
                DbRow* row = context->resultRows;
                while(row) {
                    wrenSetSlotNewMap(dbManager->vm, 1);
                    for (int i = 0; i < row->count; i++) {
                        wrenSetSlotString(dbManager->vm, 0, row->columns[i]); // key
                        DbValue* val = &row->values[i];
                        switch (val->type) {
                            case SQLITE_INTEGER: case SQLITE_FLOAT:
                                wrenSetSlotDouble(dbManager->vm, 1, val->as.num); break;
                            case SQLITE_TEXT:
                                wrenSetSlotBytes(dbManager->vm, 1, val->as.str.text, val->as.str.length); break;
                            case SQLITE_BLOB:
                                wrenSetSlotBytes(dbManager->vm, 1, val->as.str.text, val->as.str.length); break;
                            case SQLITE_NULL:
                                wrenSetSlotNull(dbManager->vm, 1); break;
                        }
                        wrenSetMapValue(dbManager->vm, 1, 0, 1);
                    }
                    wrenInsertInList(dbManager->vm, 2, -1, 1);
                    row = row->next;
                }
            } else {
                TRACE();
                wrenSetSlotNull(dbManager->vm, 2);
            }
        } else {
            TRACE();
            wrenSetSlotString(dbManager->vm, 1, context->errorMessage ? context->errorMessage : "Unknown error.");
            TRACE();
            wrenSetSlotNull(dbManager->vm, 2);
        }
        TRACE();
        wrenCall(dbManager->vm, callHandle);
        TRACE();
        free_db_context(context);
    }
    TRACE();
    wrenReleaseHandle(dbManager->vm, callHandle);
    TRACE();
}

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

                    if (row->columns == NULL || row->values == NULL) {
                        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] == NULL) {
                           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 == NULL) {
                                   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 == NULL) head = tail = row;
                    else { tail->next = row; tail = row; }
                }
                query_loop_end:;
                TRACE();
                if (oom) {
                    set_context_error(context, "Memory allocation failed 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;
                }
                TRACE();
                sqlite3_finalize(stmt);
                break;
            }
            case DB_OP_CLOSE: {
                TRACE();
                if (context->db) {
                    sqlite3_close(context->db);
                }
                TRACE();
                context->success = true;
                break;
            }
        }
        TRACE();
        db_queue_push(&manager->completionQueue, context);
    }
    TRACE();
    return 0;
}

// --- Wren FFI ---

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

void dbFinalize(void* data) {
    TRACE();
    DatabaseData* dbData = (DatabaseData*)data;
    if (dbData && dbData->db) {
        sqlite3_close(dbData->db);
    }
    TRACE();
}
static void ensureDbManager(WrenVM* vm)
{
    if (!dbManager) dbManager_create(vm);
}

static void create_db_context(WrenVM* vm,
                              DBOp  op,
                              int   sqlSlot,
                              int   cbSlot)
{
  TRACE();

  DbContext* context = calloc(1, sizeof(*context));
  if (!context)
  {
    wrenSetSlotString(vm, 0, "Out of memory creating DbContext.");
    wrenAbortFiber(vm, 0);
    return;
  }

  context->vm        = vm;
  context->operation = op;

  /* ------------------------------------------------------------------ */
  /* 1.  Grab the SQL bytes *first* (before any API call that might GC) */
  /* ------------------------------------------------------------------ */
  if (sqlSlot != -1)
  {
    if (wrenGetSlotType(vm, sqlSlot) != WREN_TYPE_STRING)
    {
      wrenSetSlotString(vm, 0, "SQL argument must be a string.");
      wrenAbortFiber(vm, 0);
      free(context);
      return;
    }

    int   len   = 0;
    const char* bytes = wrenGetSlotBytes(vm, sqlSlot, &len);

    context->sql = malloc((size_t)len + 1);
    if (!context->sql)
    {
      wrenSetSlotString(vm, 0, "Out of memory copying SQL string.");
      wrenAbortFiber(vm, 0);
      free(context);
      return;
    }

    memcpy(context->sql, bytes, (size_t)len);
    context->sql[len] = '\0';          /* NUL‑terminate for SQLite */
  }

  /* -------------------------------------------------------------- */
  /* 2.  Now take the handles – these *may* allocate / trigger GC   */
  /* -------------------------------------------------------------- */
  context->dbHandle = wrenGetSlotHandle(vm, 0);
  context->callback = wrenGetSlotHandle(vm, cbSlot);

  /* -------------------------------------------------------------- */
  /* 3.  Stash live DB pointer and queue the request                */
  /* -------------------------------------------------------------- */
  DatabaseData* dbData = wrenGetSlotForeign(vm, 0);
  if (!dbData)
  {
    set_context_error(context, "Internal error: bad Database object.");
    db_queue_push(&dbManager->requestQueue, context);
    return;
  }
  context->db = dbData->db;

  db_queue_push(&dbManager->requestQueue, context);
  TRACE();
}




void dbOpen(WrenVM* vm) {
    
    ensureDbManager(vm);
    TRACE();
    DbContext* context = (DbContext*)calloc(1, sizeof(DbContext));
    if (context == NULL) {
        wrenSetSlotString(vm, 0, "Failed to allocate memory for database operation.");
        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 == NULL) {
        free(context);
        wrenSetSlotString(vm, 0, "Failed to allocate memory for database path.");
        wrenAbortFiber(vm, 0);
        return;
    }
    context->dbHandle = wrenGetSlotHandle(vm, 0);
    context->callback = wrenGetSlotHandle(vm, 2);
    db_queue_push(&dbManager->requestQueue, context);
    TRACE();
}

void dbExec(WrenVM* vm) {
    ensureDbManager(vm);
        TRACE();
    create_db_context(vm, DB_OP_EXEC, 1, 2);
    TRACE();
}

void dbQuery(WrenVM* vm) {

    ensureDbManager(vm);
        TRACE();
    TRACE();
    create_db_context(vm, DB_OP_QUERY, 1, 2);
    TRACE();
}

void dbClose(WrenVM* vm) {
    
    ensureDbManager(vm);
        TRACE();
    TRACE();
    create_db_context(vm, DB_OP_CLOSE, -1, 1);
    TRACE();
}

WrenForeignMethodFn bindSqliteForeignMethod(WrenVM* vm, const char* module, const char* className, bool isStatic, const char* signature) {
    TRACE();
    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) {
    TRACE();
    if (strcmp(module, "sqlite") == 0 && strcmp(className, "Database") == 0) {
        WrenForeignClassMethods methods = {dbAllocate, dbFinalize};
        return methods;
    }
    WrenForeignClassMethods methods = {0, 0};
    return methods;
}