/*
 * NOSTR Core Library Implementation - Relay Pool Management
 * 
 * This file contains:
 * - Relay Pool Management
 * - Pool connection management
 * - Subscription handling
 * - Event processing and dispatching
 * - Statistics and latency tracking
 * - Multi-relay query and publish functions
 */

#define _GNU_SOURCE
#define _POSIX_C_SOURCE 200809L

#include "nostr_core.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

// Our production-ready WebSocket implementation
#include "../nostr_websocket/nostr_websocket_tls.h"

// cJSON for JSON handling
#include "../cjson/cJSON.h"

// NIP-42 Authentication
#include "nip042.h"

// =============================================================================
// TYPE DEFINITIONS FOR SYNCHRONOUS RELAY QUERIES
// =============================================================================

// Relay state enum (internal to this file)
typedef enum {
    RELAY_STATE_CONNECTING,
    RELAY_STATE_SUBSCRIBED,     // REQ sent successfully
    RELAY_STATE_ACTIVE,         // Receiving events
    RELAY_STATE_EOSE_RECEIVED,  // End of stored events
    RELAY_STATE_TIMED_OUT,      // No response within timeout
    RELAY_STATE_ERROR           // Connection or protocol error
} relay_state_enum_t;

// Internal relay connection structure
typedef struct {
    nostr_ws_client_t* client;
    char* url;
    time_t last_activity;        // Last message received
    time_t request_sent_time;    // When REQ was sent
    relay_state_enum_t state;
    int events_received;
    cJSON** events;              // Array of events from this relay
    int events_capacity;         // Allocated capacity
    char subscription_id[32];    // Unique subscription ID

    // NIP-42 Authentication fields
    nostr_auth_state_t auth_state;           // Current authentication state
    char auth_challenge[NOSTR_NIP42_MAX_CHALLENGE_LENGTH]; // Stored challenge
    time_t auth_challenge_time;              // When challenge was received
    int nip42_enabled;                       // Whether NIP-42 is enabled for this relay
} relay_connection_t;


// =============================================================================
// SYNCHRONOUS MULTI-RELAY QUERY WITH PROGRESS CALLBACKS
// =============================================================================
cJSON** synchronous_query_relays_with_progress(
    const char** relay_urls,
    int relay_count,
    cJSON* filter,
    relay_query_mode_t mode,
    int* result_count,
    int relay_timeout_seconds,
    relay_progress_callback_t callback,
    void* user_data,
    int nip42_enabled,
    const unsigned char* private_key) {
    
    if (!relay_urls || relay_count <= 0 || !filter || !result_count) {
        if (result_count) *result_count = 0;
        return NULL;
    }
    
    *result_count = 0;
    
    // Set default timeout if not specified
    if (relay_timeout_seconds <= 0) {
        relay_timeout_seconds = 2;  // Default 2 seconds
    }
    
    // Initialize relay connections
    relay_connection_t* relays = calloc(relay_count, sizeof(relay_connection_t));
    if (!relays) {
        return NULL;
    }
    
    // Setup connections
    int active_relays = 0;
    time_t start_time = time(NULL);
    
    for (int i = 0; i < relay_count; i++) {
        relays[i].url = strdup(relay_urls[i]);
        relays[i].state = RELAY_STATE_CONNECTING;
        relays[i].last_activity = start_time;
        relays[i].events_capacity = 10;
        relays[i].events = malloc(relays[i].events_capacity * sizeof(cJSON*));

        // Initialize NIP-42 authentication fields
        relays[i].auth_state = NOSTR_AUTH_STATE_NONE;
        memset(relays[i].auth_challenge, 0, sizeof(relays[i].auth_challenge));
        relays[i].auth_challenge_time = 0;
        relays[i].nip42_enabled = nip42_enabled;

        // Generate unique subscription ID
        snprintf(relays[i].subscription_id, sizeof(relays[i].subscription_id),
                 "sync_%d_%ld", i, start_time);

        if (callback) {
            callback(relays[i].url, "connecting", NULL, 0, relay_count, 0, user_data);
        }
        
        // Attempt connection
        relays[i].client = nostr_ws_connect(relays[i].url);
        if (relays[i].client) {
            relays[i].state = RELAY_STATE_SUBSCRIBED;
            relays[i].request_sent_time = time(NULL);
            relays[i].last_activity = relays[i].request_sent_time;
            
            // Send REQ message
            if (nostr_relay_send_req(relays[i].client, relays[i].subscription_id, filter) >= 0) {
                active_relays++;
                if (callback) {
                    callback(relays[i].url, "subscribed", NULL, 0, relay_count, 0, user_data);
                }
            } else {
                relays[i].state = RELAY_STATE_ERROR;
                if (callback) {
                    callback(relays[i].url, "error", NULL, 0, relay_count, 0, user_data);
                }
            }
        } else {
            relays[i].state = RELAY_STATE_ERROR;
            if (callback) {
                callback(relays[i].url, "error", NULL, 0, relay_count, 0, user_data);
            }
        }
    }
    
    if (active_relays == 0) {
        // Cleanup and return
        for (int i = 0; i < relay_count; i++) {
            free(relays[i].url);
            free(relays[i].events);
        }
        free(relays);
        return NULL;
    }
    
    // Main polling loop
    cJSON** result_array = NULL;
    int total_unique_events = 0;
    char seen_event_ids[1000][65]; // Simple deduplication
    int seen_count = 0;
    
    while (active_relays > 0) {
        time_t current_time = time(NULL);
        int completed_relays = relay_count - active_relays;
        
        for (int i = 0; i < relay_count; i++) {
            relay_connection_t* relay = &relays[i];
            
            // Skip finished relays
            if (relay->state == RELAY_STATE_EOSE_RECEIVED || 
                relay->state == RELAY_STATE_TIMED_OUT ||
                relay->state == RELAY_STATE_ERROR) {
                continue;
            }
            
            // Check for timeout
            time_t time_since_activity = current_time - relay->last_activity;
            if (time_since_activity > relay_timeout_seconds) {
                relay->state = RELAY_STATE_TIMED_OUT;
                active_relays--;
                if (callback) {
                    callback(relay->url, "timeout", NULL, relay->events_received, 
                            relay_count, completed_relays + 1, user_data);
                }
                if (relay->client) {
                    nostr_relay_send_close(relay->client, relay->subscription_id);
                    nostr_ws_close(relay->client);
                    relay->client = NULL;
                }
                continue;
            }
            
            // Try to receive message (short timeout to keep UI responsive)
            char buffer[8192];
            int len = nostr_ws_receive(relay->client, buffer, sizeof(buffer)-1, 50);
            
            if (len > 0) {
                buffer[len] = '\0';
                relay->last_activity = current_time;
                
                // Parse message
                char* msg_type = NULL;
                cJSON* parsed = NULL;
                if (nostr_parse_relay_message(buffer, &msg_type, &parsed) == 0) {
                    
                    if (msg_type && strcmp(msg_type, "AUTH") == 0) {
                        // Handle AUTH challenge message: ["AUTH", <challenge-string>]
                        if (relay->nip42_enabled && private_key && cJSON_IsArray(parsed) && cJSON_GetArraySize(parsed) >= 2) {
                            cJSON* challenge_json = cJSON_GetArrayItem(parsed, 1);
                            if (cJSON_IsString(challenge_json)) {
                                const char* challenge = cJSON_GetStringValue(challenge_json);

                                // Store challenge and attempt authentication
                                strncpy(relay->auth_challenge, challenge, sizeof(relay->auth_challenge) - 1);
                                relay->auth_challenge[sizeof(relay->auth_challenge) - 1] = '\0';
                                relay->auth_challenge_time = time(NULL);
                                relay->auth_state = NOSTR_AUTH_STATE_CHALLENGE_RECEIVED;

                                // Create and send authentication event
                                cJSON* auth_event = nostr_nip42_create_auth_event(challenge, relay->url, private_key, 0);
                                if (auth_event) {
                                    char* auth_message = nostr_nip42_create_auth_message(auth_event);
                                    if (auth_message) {
                                        if (nostr_ws_send_text(relay->client, auth_message) >= 0) {
                                            relay->auth_state = NOSTR_AUTH_STATE_AUTHENTICATING;
                                            if (callback) {
                                                callback(relay->url, "authenticating", NULL, 0, relay_count, completed_relays, user_data);
                                            }
                                        }
                                        free(auth_message);
                                    }
                                    cJSON_Delete(auth_event);
                                }
                            }
                        }

                    } else if (msg_type && strcmp(msg_type, "EVENT") == 0) {
                        // Handle EVENT message
                        if (cJSON_IsArray(parsed) && cJSON_GetArraySize(parsed) >= 3) {
                            cJSON* sub_id_json = cJSON_GetArrayItem(parsed, 1);
                            cJSON* event = cJSON_GetArrayItem(parsed, 2);

                            if (cJSON_IsString(sub_id_json) && event &&
                                strcmp(cJSON_GetStringValue(sub_id_json), relay->subscription_id) == 0) {

                                cJSON* event_id_json = cJSON_GetObjectItem(event, "id");
                                if (event_id_json && cJSON_IsString(event_id_json)) {
                                    const char* event_id = cJSON_GetStringValue(event_id_json);

                                    // Check for duplicate
                                    int is_duplicate = 0;
                                    for (int j = 0; j < seen_count; j++) {
                                        if (strcmp(seen_event_ids[j], event_id) == 0) {
                                            is_duplicate = 1;
                                            break;
                                        }
                                    }

                                    if (!is_duplicate && seen_count < 1000) {
                                        // New event - add to seen list
                                        strncpy(seen_event_ids[seen_count], event_id, 64);
                                        seen_event_ids[seen_count][64] = '\0';
                                        seen_count++;
                                        total_unique_events++;

                                        // Store event in relay's array
                                        if (relay->events_received >= relay->events_capacity) {
                                            relay->events_capacity *= 2;
                                            relay->events = realloc(relay->events,
                                                relay->events_capacity * sizeof(cJSON*));
                                        }

                                        relay->events[relay->events_received] = cJSON_Duplicate(event, 1);
                                        relay->events_received++;
                                        relay->state = RELAY_STATE_ACTIVE;

                                        if (callback) {
                                            callback(relay->url, "event_found", event_id,
                                                    relay->events_received, relay_count,
                                                    completed_relays, user_data);
                                        }

                                        // For FIRST_RESULT mode, return immediately
                                        if (mode == RELAY_QUERY_FIRST_RESULT) {
                                            result_array = malloc(sizeof(cJSON*));
                                            if (result_array) {
                                                result_array[0] = cJSON_Duplicate(event, 1);
                                                *result_count = 1;
                                                if (callback) {
                                                    callback(NULL, "first_result", event_id,
                                                            1, relay_count, completed_relays, user_data);
                                                }
                                            }
                                            goto cleanup; // Break out of all loops
                                        }
                                    }
                                }
                            }
                        }

                    } else if (msg_type && strcmp(msg_type, "EOSE") == 0) {
                        // Handle End of Stored Events
                        cJSON* sub_id_json = cJSON_GetArrayItem(parsed, 1);
                        if (cJSON_IsString(sub_id_json) &&
                            strcmp(cJSON_GetStringValue(sub_id_json), relay->subscription_id) == 0) {
                            
                            relay->state = RELAY_STATE_EOSE_RECEIVED;
                            active_relays--;
                            
                            if (callback) {
                                callback(relay->url, "eose", NULL, relay->events_received, 
                                        relay_count, completed_relays + 1, user_data);
                            }
                            
                            // Send CLOSE and cleanup
                            nostr_relay_send_close(relay->client, relay->subscription_id);
                            nostr_ws_close(relay->client);
                            relay->client = NULL;
                        }
                    }
                    
                    if (msg_type) free(msg_type);
                    if (parsed) cJSON_Delete(parsed);
                }
            }
        }
        
        // Small sleep to prevent busy waiting
        usleep(10000); // 10ms
    }
    
    // Handle different return modes
    if (mode == RELAY_QUERY_MOST_RECENT && total_unique_events > 0) {
        // Find the event with highest created_at
        time_t most_recent_time = 0;
        cJSON* most_recent_event = NULL;
        const char* most_recent_id = NULL;
        
        for (int i = 0; i < relay_count; i++) {
            for (int j = 0; j < relays[i].events_received; j++) {
                cJSON* event = relays[i].events[j];
                cJSON* created_at = cJSON_GetObjectItem(event, "created_at");
                if (created_at && cJSON_IsNumber(created_at)) {
                    time_t event_time = (time_t)cJSON_GetNumberValue(created_at);
                    if (event_time > most_recent_time) {
                        most_recent_time = event_time;
                        most_recent_event = event;
                        cJSON* id_json = cJSON_GetObjectItem(event, "id");
                        if (id_json && cJSON_IsString(id_json)) {
                            most_recent_id = cJSON_GetStringValue(id_json);
                        }
                    }
                }
            }
        }
        
        if (most_recent_event) {
            result_array = malloc(sizeof(cJSON*));
            if (result_array) {
                result_array[0] = cJSON_Duplicate(most_recent_event, 1);
                *result_count = 1;
                if (callback) {
                    callback(NULL, "all_complete", most_recent_id, 1, 
                            relay_count, relay_count, user_data);
                }
            }
        }
        
    } else if (mode == RELAY_QUERY_ALL_RESULTS && total_unique_events > 0) {
        // Return ALL unique events
        result_array = malloc(total_unique_events * sizeof(cJSON*));
        if (result_array) {
            int event_index = 0;
            
            // Collect all unique events from all relays
            for (int i = 0; i < relay_count; i++) {
                for (int j = 0; j < relays[i].events_received; j++) {
                    cJSON* event = relays[i].events[j];
                    cJSON* event_id_json = cJSON_GetObjectItem(event, "id");
                    
                    if (event_id_json && cJSON_IsString(event_id_json)) {
                        const char* event_id = cJSON_GetStringValue(event_id_json);
                        
                        // Check if we already added this event ID
                        int already_added = 0;
                        for (int k = 0; k < event_index; k++) {
                            cJSON* existing_id = cJSON_GetObjectItem(result_array[k], "id");
                            if (existing_id && cJSON_IsString(existing_id) &&
                                strcmp(cJSON_GetStringValue(existing_id), event_id) == 0) {
                                already_added = 1;
                                break;
                            }
                        }
                        
                        if (!already_added) {
                            result_array[event_index] = cJSON_Duplicate(event, 1);
                            event_index++;
                        }
                    }
                }
            }
            
            *result_count = event_index;
            
            if (callback) {
                callback(NULL, "all_complete", NULL, total_unique_events, 
                        relay_count, relay_count, user_data);
            }
        }
    }
    
cleanup:
    // Cleanup all relay connections and data
    for (int i = 0; i < relay_count; i++) {
        if (relays[i].client) {
            nostr_relay_send_close(relays[i].client, relays[i].subscription_id);
            nostr_ws_close(relays[i].client);
        }
        
        // Free stored events
        for (int j = 0; j < relays[i].events_received; j++) {
            if (relays[i].events[j]) {
                cJSON_Delete(relays[i].events[j]);
            }
        }
        
        free(relays[i].url);
        free(relays[i].events);
    }
    free(relays);
    
    return result_array;
}

// =============================================================================
// SYNCHRONOUS MULTI-RELAY PUBLISH WITH PROGRESS CALLBACKS
// =============================================================================

publish_result_t* synchronous_publish_event_with_progress(
    const char** relay_urls,
    int relay_count,
    cJSON* event,
    int* success_count,
    int relay_timeout_seconds,
    publish_progress_callback_t callback,
    void* user_data,
    int nip42_enabled,
    const unsigned char* private_key) {
    
    if (!relay_urls || relay_count <= 0 || !event || !success_count) {
        if (success_count) *success_count = 0;
        return NULL;
    }
    
    *success_count = 0;
    
    // Set default timeout if not specified
    if (relay_timeout_seconds <= 0) {
        relay_timeout_seconds = 5;  // Default 5 seconds for publishing
    }
    
    // Initialize relay connections
    relay_connection_t* relays = calloc(relay_count, sizeof(relay_connection_t));
    if (!relays) {
        return NULL;
    }
    
    // Initialize results array
    publish_result_t* results = calloc(relay_count, sizeof(publish_result_t));
    if (!results) {
        free(relays);
        return NULL;
    }
    
    // Get event ID for tracking
    cJSON* event_id_json = cJSON_GetObjectItem(event, "id");
    const char* event_id = NULL;
    if (event_id_json && cJSON_IsString(event_id_json)) {
        event_id = cJSON_GetStringValue(event_id_json);
    }
    
    // Setup connections
    int active_relays = 0;
    time_t start_time = time(NULL);
    
    for (int i = 0; i < relay_count; i++) {
        relays[i].url = strdup(relay_urls[i]);
        relays[i].state = RELAY_STATE_CONNECTING;
        relays[i].last_activity = start_time;
        results[i] = PUBLISH_ERROR; // Default to error

        // Initialize NIP-42 authentication fields
        relays[i].auth_state = NOSTR_AUTH_STATE_NONE;
        memset(relays[i].auth_challenge, 0, sizeof(relays[i].auth_challenge));
        relays[i].auth_challenge_time = 0;
        relays[i].nip42_enabled = nip42_enabled;

        if (callback) {
            callback(relays[i].url, "connecting", NULL, 0, relay_count, 0, user_data);
        }
        
        // Attempt connection
        relays[i].client = nostr_ws_connect(relays[i].url);
        if (relays[i].client) {
            relays[i].state = RELAY_STATE_SUBSCRIBED; // Reuse for "publishing" state
            relays[i].request_sent_time = time(NULL);
            relays[i].last_activity = relays[i].request_sent_time;
            
            // Send EVENT message
            if (nostr_relay_send_event(relays[i].client, event) >= 0) {
                active_relays++;
                if (callback) {
                    callback(relays[i].url, "publishing", NULL, 0, relay_count, 0, user_data);
                }
            } else {
                relays[i].state = RELAY_STATE_ERROR;
                results[i] = PUBLISH_ERROR;
                if (callback) {
                    callback(relays[i].url, "error", "Failed to send EVENT message", 
                            0, relay_count, 0, user_data);
                }
            }
        } else {
            relays[i].state = RELAY_STATE_ERROR;
            results[i] = PUBLISH_ERROR;
            if (callback) {
                callback(relays[i].url, "error", "Failed to connect", 0, relay_count, 0, user_data);
            }
        }
    }
    
    if (active_relays == 0) {
        // Cleanup and return
        for (int i = 0; i < relay_count; i++) {
            free(relays[i].url);
        }
        free(relays);
        return results; // Return error results
    }
    
    // Main polling loop
    int completed_relays = relay_count - active_relays;
    
    while (active_relays > 0) {
        time_t current_time = time(NULL);
        
        for (int i = 0; i < relay_count; i++) {
            relay_connection_t* relay = &relays[i];
            
            // Skip finished relays
            if (relay->state == RELAY_STATE_EOSE_RECEIVED || // Reuse for "completed"
                relay->state == RELAY_STATE_TIMED_OUT ||
                relay->state == RELAY_STATE_ERROR) {
                continue;
            }
            
            // Check for timeout
            time_t time_since_activity = current_time - relay->last_activity;
            if (time_since_activity > relay_timeout_seconds) {
                relay->state = RELAY_STATE_TIMED_OUT;
                results[i] = PUBLISH_TIMEOUT;
                active_relays--;
                completed_relays++;
                
                if (callback) {
                    callback(relay->url, "timeout", "No response within timeout", 
                            *success_count, relay_count, completed_relays, user_data);
                }
                
                if (relay->client) {
                    nostr_ws_close(relay->client);
                    relay->client = NULL;
                }
                continue;
            }
            
            // Try to receive message (short timeout to keep UI responsive)
            char buffer[8192];
            int len = nostr_ws_receive(relay->client, buffer, sizeof(buffer)-1, 50);
            
            if (len > 0) {
                buffer[len] = '\0';
                relay->last_activity = current_time;
                
                // Parse message
                char* msg_type = NULL;
                cJSON* parsed = NULL;
                if (nostr_parse_relay_message(buffer, &msg_type, &parsed) == 0) {

                    if (msg_type && strcmp(msg_type, "AUTH") == 0) {
                        // Handle AUTH challenge message: ["AUTH", <challenge-string>]
                        if (relay->nip42_enabled && private_key && cJSON_IsArray(parsed) && cJSON_GetArraySize(parsed) >= 2) {
                            cJSON* challenge_json = cJSON_GetArrayItem(parsed, 1);
                            if (cJSON_IsString(challenge_json)) {
                                const char* challenge = cJSON_GetStringValue(challenge_json);

                                // Store challenge and attempt authentication
                                strncpy(relay->auth_challenge, challenge, sizeof(relay->auth_challenge) - 1);
                                relay->auth_challenge[sizeof(relay->auth_challenge) - 1] = '\0';
                                relay->auth_challenge_time = time(NULL);
                                relay->auth_state = NOSTR_AUTH_STATE_CHALLENGE_RECEIVED;

                                // Create and send authentication event
                                cJSON* auth_event = nostr_nip42_create_auth_event(challenge, relay->url, private_key, 0);
                                if (auth_event) {
                                    char* auth_message = nostr_nip42_create_auth_message(auth_event);
                                    if (auth_message) {
                                        if (nostr_ws_send_text(relay->client, auth_message) >= 0) {
                                            relay->auth_state = NOSTR_AUTH_STATE_AUTHENTICATING;
                                            if (callback) {
                                                callback(relay->url, "authenticating", NULL, 0, relay_count, completed_relays, user_data);
                                            }
                                        }
                                        free(auth_message);
                                    }
                                    cJSON_Delete(auth_event);
                                }
                            }
                        }

                    } else if (msg_type && strcmp(msg_type, "OK") == 0) {
                        // Handle OK message: ["OK", <event_id>, <true|false>, <message>]
                        if (cJSON_IsArray(parsed) && cJSON_GetArraySize(parsed) >= 3) {
                            cJSON* ok_event_id = cJSON_GetArrayItem(parsed, 1);
                            cJSON* accepted = cJSON_GetArrayItem(parsed, 2);
                            cJSON* message = cJSON_GetArrayItem(parsed, 3);

                            // Verify this OK is for our event
                            if (ok_event_id && cJSON_IsString(ok_event_id) && event_id &&
                                strcmp(cJSON_GetStringValue(ok_event_id), event_id) == 0) {

                                relay->state = RELAY_STATE_EOSE_RECEIVED; // Reuse for "completed"
                                active_relays--;
                                completed_relays++;

                                const char* ok_message = "";
                                if (message && cJSON_IsString(message)) {
                                    ok_message = cJSON_GetStringValue(message);
                                }

                                if (accepted && cJSON_IsBool(accepted) && cJSON_IsTrue(accepted)) {
                                    // Event accepted
                                    results[i] = PUBLISH_SUCCESS;
                                    (*success_count)++;
                                    
                                    if (callback) {
                                        callback(relay->url, "accepted", ok_message,
                                                *success_count, relay_count, completed_relays, user_data);
                                    }
                                } else {
                                    // Event rejected
                                    results[i] = PUBLISH_REJECTED;
                                    
                                    if (callback) {
                                        callback(relay->url, "rejected", ok_message,
                                                *success_count, relay_count, completed_relays, user_data);
                                    }
                                }

                                // Close connection
                                nostr_ws_close(relay->client);
                                relay->client = NULL;
                            }
                        }
                    }
                    
                    if (msg_type) free(msg_type);
                    if (parsed) cJSON_Delete(parsed);
                }
            }
        }
        
        // Small sleep to prevent busy waiting
        usleep(10000); // 10ms
    }
    
    // Final callback with summary
    if (callback) {
        callback(NULL, "all_complete", NULL, *success_count, relay_count, relay_count, user_data);
    }
    
    // Cleanup relay connections
    for (int i = 0; i < relay_count; i++) {
        if (relays[i].client) {
            nostr_ws_close(relays[i].client);
        }
        free(relays[i].url);
    }
    free(relays);
    
    return results;
}
