#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <time.h>
#include <pthread.h>
#include <sqlite3.h>
#include <libwebsockets.h>

// Include nostr_core_lib for Nostr functionality
#include "../nostr_core_lib/cjson/cJSON.h"
#include "../nostr_core_lib/nostr_core/nostr_core.h"

// Configuration
#define DEFAULT_PORT 8888
#define DEFAULT_HOST "127.0.0.1"
#define DATABASE_PATH "db/c_nostr_relay.db"
#define MAX_CLIENTS 100

// Global state
static sqlite3* g_db = NULL;
static int g_server_running = 1;

// Color constants for logging
#define RED     "\033[31m"
#define GREEN   "\033[32m" 
#define YELLOW  "\033[33m"
#define BLUE    "\033[34m"
#define BOLD    "\033[1m"
#define RESET   "\033[0m"

// Logging functions
void log_info(const char* message) {
    printf(BLUE "[INFO]" RESET " %s\n", message);
    fflush(stdout);
}

void log_success(const char* message) {
    printf(GREEN "[SUCCESS]" RESET " %s\n", message);
    fflush(stdout);
}

void log_error(const char* message) {
    printf(RED "[ERROR]" RESET " %s\n", message);
    fflush(stdout);
}

void log_warning(const char* message) {
    printf(YELLOW "[WARNING]" RESET " %s\n", message);
    fflush(stdout);
}

// Signal handler for graceful shutdown
void signal_handler(int sig) {
    if (sig == SIGINT || sig == SIGTERM) {
        log_info("Received shutdown signal");
        g_server_running = 0;
    }
}

// Initialize database connection
int init_database() {
    int rc = sqlite3_open(DATABASE_PATH, &g_db);
    if (rc != SQLITE_OK) {
        log_error("Cannot open database");
        return -1;
    }
    
    log_success("Database connection established");
    return 0;
}

// Close database connection
void close_database() {
    if (g_db) {
        sqlite3_close(g_db);
        g_db = NULL;
        log_info("Database connection closed");
    }
}

// Store event in database
int store_event(cJSON* event) {
    if (!g_db || !event) {
        return -1;
    }
    
    // Extract event fields
    cJSON* id = cJSON_GetObjectItem(event, "id");
    cJSON* pubkey = cJSON_GetObjectItem(event, "pubkey");
    cJSON* created_at = cJSON_GetObjectItem(event, "created_at");
    cJSON* kind = cJSON_GetObjectItem(event, "kind");
    cJSON* content = cJSON_GetObjectItem(event, "content");
    cJSON* sig = cJSON_GetObjectItem(event, "sig");
    cJSON* tags = cJSON_GetObjectItem(event, "tags");
    
    if (!id || !pubkey || !created_at || !kind || !content || !sig) {
        log_error("Invalid event - missing required fields");
        return -1;
    }
    
    // Prepare SQL statement for event insertion
    const char* sql =
        "INSERT INTO event (id, pubkey, created_at, kind, content, sig) "
        "VALUES (?, ?, ?, ?, ?, ?)";
    
    sqlite3_stmt* stmt;
    int rc = sqlite3_prepare_v2(g_db, sql, -1, &stmt, NULL);
    if (rc != SQLITE_OK) {
        log_error("Failed to prepare event insert statement");
        return -1;
    }
    
    // Bind parameters
    sqlite3_bind_text(stmt, 1, cJSON_GetStringValue(id), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 2, cJSON_GetStringValue(pubkey), -1, SQLITE_STATIC);
    sqlite3_bind_int64(stmt, 3, (sqlite3_int64)cJSON_GetNumberValue(created_at));
    sqlite3_bind_int(stmt, 4, (int)cJSON_GetNumberValue(kind));
    sqlite3_bind_text(stmt, 5, cJSON_GetStringValue(content), -1, SQLITE_STATIC);
    sqlite3_bind_text(stmt, 6, cJSON_GetStringValue(sig), -1, SQLITE_STATIC);
    
    // Execute statement
    rc = sqlite3_step(stmt);
    sqlite3_finalize(stmt);
    
    if (rc != SQLITE_DONE) {
        if (rc == SQLITE_CONSTRAINT) {
            log_warning("Event already exists in database");
            return 0; // Not an error, just duplicate
        }
        char error_msg[256];
        snprintf(error_msg, sizeof(error_msg), "Failed to insert event: %s", sqlite3_errmsg(g_db));
        log_error(error_msg);
        return -1;
    }
    
    // Insert tags if present
    if (tags && cJSON_IsArray(tags)) {
        const char* event_id = cJSON_GetStringValue(id);
        cJSON* tag;
        cJSON_ArrayForEach(tag, tags) {
            if (cJSON_IsArray(tag) && cJSON_GetArraySize(tag) >= 2) {
                cJSON* tag_name = cJSON_GetArrayItem(tag, 0);
                cJSON* tag_value = cJSON_GetArrayItem(tag, 1);
                
                if (cJSON_IsString(tag_name) && cJSON_IsString(tag_value)) {
                    // Collect additional tag parameters if present
                    char* parameters = NULL;
                    if (cJSON_GetArraySize(tag) > 2) {
                        cJSON* params_array = cJSON_CreateArray();
                        for (int i = 2; i < cJSON_GetArraySize(tag); i++) {
                            cJSON_AddItemToArray(params_array, cJSON_Duplicate(cJSON_GetArrayItem(tag, i), 1));
                        }
                        parameters = cJSON_Print(params_array);
                        cJSON_Delete(params_array);
                    }
                    
                    const char* tag_sql =
                        "INSERT INTO tag (id, name, value, parameters) VALUES (?, ?, ?, ?)";
                    
                    sqlite3_stmt* tag_stmt;
                    rc = sqlite3_prepare_v2(g_db, tag_sql, -1, &tag_stmt, NULL);
                    if (rc == SQLITE_OK) {
                        sqlite3_bind_text(tag_stmt, 1, event_id, -1, SQLITE_STATIC);
                        sqlite3_bind_text(tag_stmt, 2, cJSON_GetStringValue(tag_name), -1, SQLITE_STATIC);
                        sqlite3_bind_text(tag_stmt, 3, cJSON_GetStringValue(tag_value), -1, SQLITE_STATIC);
                        sqlite3_bind_text(tag_stmt, 4, parameters, -1, SQLITE_TRANSIENT);
                        
                        sqlite3_step(tag_stmt);
                        sqlite3_finalize(tag_stmt);
                    }
                    
                    if (parameters) free(parameters);
                }
            }
        }
    }
    
    log_success("Event stored in database");
    return 0;
}

// Retrieve event from database
cJSON* retrieve_event(const char* event_id) {
    if (!g_db || !event_id) {
        return NULL;
    }
    
    const char* sql =
        "SELECT id, pubkey, created_at, kind, content, sig FROM event WHERE id = ?";
    
    sqlite3_stmt* stmt;
    int rc = sqlite3_prepare_v2(g_db, sql, -1, &stmt, NULL);
    if (rc != SQLITE_OK) {
        return NULL;
    }
    
    sqlite3_bind_text(stmt, 1, event_id, -1, SQLITE_STATIC);
    
    cJSON* event = NULL;
    if (sqlite3_step(stmt) == SQLITE_ROW) {
        event = cJSON_CreateObject();
        
        cJSON_AddStringToObject(event, "id", (char*)sqlite3_column_text(stmt, 0));
        cJSON_AddStringToObject(event, "pubkey", (char*)sqlite3_column_text(stmt, 1));
        cJSON_AddNumberToObject(event, "created_at", sqlite3_column_int64(stmt, 2));
        cJSON_AddNumberToObject(event, "kind", sqlite3_column_int(stmt, 3));
        cJSON_AddStringToObject(event, "content", (char*)sqlite3_column_text(stmt, 4));
        cJSON_AddStringToObject(event, "sig", (char*)sqlite3_column_text(stmt, 5));
        
        // Add tags array - retrieve from tag table
        cJSON* tags_array = cJSON_CreateArray();
        
        const char* tag_sql = "SELECT name, value, parameters FROM tag WHERE id = ?";
        sqlite3_stmt* tag_stmt;
        if (sqlite3_prepare_v2(g_db, tag_sql, -1, &tag_stmt, NULL) == SQLITE_OK) {
            sqlite3_bind_text(tag_stmt, 1, event_id, -1, SQLITE_STATIC);
            
            while (sqlite3_step(tag_stmt) == SQLITE_ROW) {
                cJSON* tag = cJSON_CreateArray();
                cJSON_AddItemToArray(tag, cJSON_CreateString((char*)sqlite3_column_text(tag_stmt, 0)));
                cJSON_AddItemToArray(tag, cJSON_CreateString((char*)sqlite3_column_text(tag_stmt, 1)));
                
                // Add parameters if they exist
                const char* parameters = (char*)sqlite3_column_text(tag_stmt, 2);
                if (parameters && strlen(parameters) > 0) {
                    cJSON* params = cJSON_Parse(parameters);
                    if (params && cJSON_IsArray(params)) {
                        int param_count = cJSON_GetArraySize(params);
                        for (int i = 0; i < param_count; i++) {
                            cJSON* param = cJSON_GetArrayItem(params, i);
                            cJSON_AddItemToArray(tag, cJSON_Duplicate(param, 1));
                        }
                    }
                    if (params) cJSON_Delete(params);
                }
                
                cJSON_AddItemToArray(tags_array, tag);
            }
            sqlite3_finalize(tag_stmt);
        }
        
        cJSON_AddItemToObject(event, "tags", tags_array);
    }
    
    sqlite3_finalize(stmt);
    return event;
}

// Handle REQ message (subscription)
int handle_req_message(const char* sub_id, cJSON* filters) {
    log_info("Handling REQ message");
    
    // For now, just handle simple event ID requests
    if (cJSON_IsArray(filters)) {
        cJSON* filter = cJSON_GetArrayItem(filters, 0);
        if (filter) {
            cJSON* ids = cJSON_GetObjectItem(filter, "ids");
            if (ids && cJSON_IsArray(ids)) {
                cJSON* event_id = cJSON_GetArrayItem(ids, 0);
                if (event_id && cJSON_IsString(event_id)) {
                    cJSON* event = retrieve_event(cJSON_GetStringValue(event_id));
                    if (event) {
                        log_success("Found event for subscription");
                        cJSON_Delete(event);
                        return 1; // Found event
                    }
                }
            }
        }
    }
    
    return 0; // No events found
}

// Handle EVENT message (publish)
int handle_event_message(cJSON* event) {
    log_info("Handling EVENT message");
    
    // Validate event structure (basic check)
    cJSON* id = cJSON_GetObjectItem(event, "id");
    if (!id || !cJSON_IsString(id)) {
        log_error("Invalid event - no ID");
        return -1;
    }
    
    // Store event in database
    if (store_event(event) == 0) {
        log_success("Event stored successfully");
        return 0;
    }
    
    return -1;
}

// Global WebSocket context
static struct lws_context *ws_context = NULL;

// Per-session data structure
struct per_session_data {
    int authenticated;
    char subscription_id[64];
};

// WebSocket callback function for Nostr relay protocol
static int nostr_relay_callback(struct lws *wsi, enum lws_callback_reasons reason,
                               void *user, void *in, size_t len) {
    struct per_session_data *pss = (struct per_session_data *)user;
    
    switch (reason) {
        case LWS_CALLBACK_ESTABLISHED:
            log_info("WebSocket connection established");
            memset(pss, 0, sizeof(*pss));
            break;
            
        case LWS_CALLBACK_RECEIVE:
            if (len > 0) {
                char *message = malloc(len + 1);
                if (message) {
                    memcpy(message, in, len);
                    message[len] = '\0';
                    
                    log_info("Received WebSocket message");
                    
                    // Parse JSON message
                    cJSON* json = cJSON_Parse(message);
                    if (json && cJSON_IsArray(json)) {
                        // Get message type
                        cJSON* type = cJSON_GetArrayItem(json, 0);
                        if (type && cJSON_IsString(type)) {
                            const char* msg_type = cJSON_GetStringValue(type);
                            
                            if (strcmp(msg_type, "EVENT") == 0) {
                                // Handle EVENT message
                                cJSON* event = cJSON_GetArrayItem(json, 1);
                                if (event && cJSON_IsObject(event)) {
                                    int result = handle_event_message(event);
                                    
                                    // Send OK response
                                    cJSON* event_id = cJSON_GetObjectItem(event, "id");
                                    if (event_id && cJSON_IsString(event_id)) {
                                        cJSON* response = cJSON_CreateArray();
                                        cJSON_AddItemToArray(response, cJSON_CreateString("OK"));
                                        cJSON_AddItemToArray(response, cJSON_CreateString(cJSON_GetStringValue(event_id)));
                                        cJSON_AddItemToArray(response, cJSON_CreateBool(result == 0));
                                        cJSON_AddItemToArray(response, cJSON_CreateString(result == 0 ? "" : "error: failed to store event"));
                                        
                                        char *response_str = cJSON_Print(response);
                                        if (response_str) {
                                            size_t response_len = strlen(response_str);
                                            unsigned char *buf = malloc(LWS_PRE + response_len);
                                            if (buf) {
                                                memcpy(buf + LWS_PRE, response_str, response_len);
                                                lws_write(wsi, buf + LWS_PRE, response_len, LWS_WRITE_TEXT);
                                                free(buf);
                                            }
                                            free(response_str);
                                        }
                                        cJSON_Delete(response);
                                    }
                                }
                            } else if (strcmp(msg_type, "REQ") == 0) {
                                // Handle REQ message
                                cJSON* sub_id = cJSON_GetArrayItem(json, 1);
                                cJSON* filters = cJSON_GetArrayItem(json, 2);
                                
                                if (sub_id && cJSON_IsString(sub_id)) {
                                    const char* subscription_id = cJSON_GetStringValue(sub_id);
                                    strncpy(pss->subscription_id, subscription_id, sizeof(pss->subscription_id) - 1);
                                    
                                    handle_req_message(subscription_id, filters);
                                    
                                    // Send EOSE (End of Stored Events)
                                    cJSON* eose_response = cJSON_CreateArray();
                                    cJSON_AddItemToArray(eose_response, cJSON_CreateString("EOSE"));
                                    cJSON_AddItemToArray(eose_response, cJSON_CreateString(subscription_id));
                                    
                                    char *eose_str = cJSON_Print(eose_response);
                                    if (eose_str) {
                                        size_t eose_len = strlen(eose_str);
                                        unsigned char *buf = malloc(LWS_PRE + eose_len);
                                        if (buf) {
                                            memcpy(buf + LWS_PRE, eose_str, eose_len);
                                            lws_write(wsi, buf + LWS_PRE, eose_len, LWS_WRITE_TEXT);
                                            free(buf);
                                        }
                                        free(eose_str);
                                    }
                                    cJSON_Delete(eose_response);
                                }
                            } else if (strcmp(msg_type, "CLOSE") == 0) {
                                // Handle CLOSE message
                                log_info("Subscription closed");
                            }
                        }
                    }
                    
                    if (json) cJSON_Delete(json);
                    free(message);
                }
            }
            break;
            
        case LWS_CALLBACK_CLOSED:
            log_info("WebSocket connection closed");
            break;
            
        default:
            break;
    }
    
    return 0;
}

// WebSocket protocol definition
static struct lws_protocols protocols[] = {
    {
        "nostr-relay-protocol",
        nostr_relay_callback,
        sizeof(struct per_session_data),
        4096, // rx buffer size
        0, NULL, 0
    },
    { NULL, NULL, 0, 0, 0, NULL, 0 } // terminator
};

// Start libwebsockets-based WebSocket Nostr relay server
int start_websocket_relay() {
    struct lws_context_creation_info info;

    log_info("Starting libwebsockets-based Nostr relay server...");

    memset(&info, 0, sizeof(info));
    info.port = DEFAULT_PORT;
    info.protocols = protocols;
    info.gid = -1;
    info.uid = -1;

    // Minimal libwebsockets configuration
    info.options = LWS_SERVER_OPTION_VALIDATE_UTF8;

    // Remove interface restrictions - let system choose
    // info.vhost_name = NULL;
    // info.iface = NULL;

    // Increase max connections for relay usage
    info.max_http_header_pool = 16;
    info.timeout_secs = 10;

    // Max payload size for Nostr events
    info.max_http_header_data = 4096;

    ws_context = lws_create_context(&info);
    if (!ws_context) {
        log_error("Failed to create libwebsockets context");
        perror("libwebsockets creation error");
        return -1;
    }

    log_success("WebSocket relay started on ws://127.0.0.1:8888");

    // Main event loop with proper signal handling
    fd_set rfds;
    struct timeval tv;

    while (g_server_running) {
        FD_ZERO(&rfds);
        tv.tv_sec = 1;
        tv.tv_usec = 0;

        int result = lws_service(ws_context, 1000);

        if (result < 0) {
            log_error("libwebsockets service error");
            break;
        }
    }

    log_info("Shutting down WebSocket server...");
    lws_context_destroy(ws_context);
    ws_context = NULL;

    log_success("WebSocket relay shut down cleanly");
    return 0;
}

// Print usage information
void print_usage(const char* program_name) {
    printf("Usage: %s [OPTIONS]\n", program_name);
    printf("\n");
    printf("C Nostr Relay Server\n");
    printf("\n");
    printf("Options:\n");
    printf("  -p, --port PORT    Listen port (default: %d)\n", DEFAULT_PORT);
    printf("  -h, --help         Show this help message\n");
    printf("\n");
}

int main(int argc, char* argv[]) {
    int port = DEFAULT_PORT;
    
    // Parse command line arguments
    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) {
            print_usage(argv[0]);
            return 0;
        } else if (strcmp(argv[i], "-p") == 0 || strcmp(argv[i], "--port") == 0) {
            if (i + 1 < argc) {
                port = atoi(argv[++i]);
                if (port <= 0 || port > 65535) {
                    log_error("Invalid port number");
                    return 1;
                }
            } else {
                log_error("Port argument requires a value");
                return 1;
            }
        } else {
            log_error("Unknown argument");
            print_usage(argv[0]);
            return 1;
        }
    }
    
    // Set up signal handlers
    signal(SIGINT, signal_handler);
    signal(SIGTERM, signal_handler);
    
    printf(BLUE BOLD "=== C Nostr Relay Server ===" RESET "\n");
    
    // Initialize database
    if (init_database() != 0) {
        log_error("Failed to initialize database");
        return 1;
    }
    
    // Initialize nostr library
    if (nostr_init() != 0) {
        log_error("Failed to initialize nostr library");
        close_database();
        return 1;
    }
    
    log_info("Starting relay server...");
    
    // Start WebSocket Nostr relay server
    int result = start_websocket_relay();
    
    // Cleanup
    nostr_cleanup();
    close_database();
    
    if (result == 0) {
        log_success("Server shutdown complete");
    } else {
        log_error("Server shutdown with errors");
    }
    
    return result;
}