/*
 * NOSTR Core Library - NIP-013: Proof of Work
 */

#include "nip013.h"
#include "nip001.h"
#include "utils.h"
#include "../cjson/cJSON.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <stdint.h>
#include "../nostr_core/nostr_common.h"

/**
 * Count leading zero bits in a hash (NIP-13 reference implementation)
 */
static int zero_bits(unsigned char b) {
    int n = 0;
    
    if (b == 0)
        return 8;
    
    while (b >>= 1)
        n++;
    
    return 7-n;
}

/**
 * Find the number of leading zero bits in a hash (NIP-13 reference implementation)
 */
static int count_leading_zero_bits(unsigned char *hash) {
    int bits, total, i;
    for (i = 0, total = 0; i < 32; i++) {
        bits = zero_bits(hash[i]);
        total += bits;
        if (bits != 8)
            break;
    }
    return total;
}

/**
 * Add or update nonce tag with target difficulty
 */
static int update_nonce_tag_with_difficulty(cJSON* tags, uint64_t nonce, int target_difficulty) {
    if (!tags) return -1;
    
    // Look for existing nonce tag and remove it
    cJSON* tag = NULL;
    int index = 0;
    cJSON_ArrayForEach(tag, tags) {
        if (cJSON_IsArray(tag) && cJSON_GetArraySize(tag) >= 2) {
            cJSON* tag_type = cJSON_GetArrayItem(tag, 0);
            if (tag_type && cJSON_IsString(tag_type) && 
                strcmp(cJSON_GetStringValue(tag_type), "nonce") == 0) {
                // Remove existing nonce tag
                cJSON_DetachItemFromArray(tags, index);
                cJSON_Delete(tag);
                break;
            }
        }
        index++;
    }
    
    // Add new nonce tag with format: ["nonce", "<nonce>", "<target_difficulty>"]
    cJSON* nonce_tag = cJSON_CreateArray();
    if (!nonce_tag) return -1;
    
    char nonce_str[32];
    char difficulty_str[16];
    snprintf(nonce_str, sizeof(nonce_str), "%llu", (unsigned long long)nonce);
    snprintf(difficulty_str, sizeof(difficulty_str), "%d", target_difficulty);
    
    cJSON_AddItemToArray(nonce_tag, cJSON_CreateString("nonce"));
    cJSON_AddItemToArray(nonce_tag, cJSON_CreateString(nonce_str));
    cJSON_AddItemToArray(nonce_tag, cJSON_CreateString(difficulty_str));
    
    cJSON_AddItemToArray(tags, nonce_tag);
    return 0;
}

/**
 * Helper function to replace event content with successful PoW result
 */
static void replace_event_content(cJSON* target_event, cJSON* source_event) {
    // Remove old fields
    cJSON_DeleteItemFromObject(target_event, "id");
    cJSON_DeleteItemFromObject(target_event, "sig");
    cJSON_DeleteItemFromObject(target_event, "tags");
    cJSON_DeleteItemFromObject(target_event, "created_at");
    
    // Copy new fields from successful event
    cJSON* id = cJSON_GetObjectItem(source_event, "id");
    cJSON* sig = cJSON_GetObjectItem(source_event, "sig");
    cJSON* tags = cJSON_GetObjectItem(source_event, "tags");
    cJSON* created_at = cJSON_GetObjectItem(source_event, "created_at");
    
    if (id) cJSON_AddItemToObject(target_event, "id", cJSON_Duplicate(id, 1));
    if (sig) cJSON_AddItemToObject(target_event, "sig", cJSON_Duplicate(sig, 1));
    if (tags) cJSON_AddItemToObject(target_event, "tags", cJSON_Duplicate(tags, 1));
    if (created_at) cJSON_AddItemToObject(target_event, "created_at", cJSON_Duplicate(created_at, 1));
}

/**
 * Add NIP-13 Proof of Work to an event
 * 
 * @param event The event to add proof of work to
 * @param private_key The private key for re-signing the event
 * @param target_difficulty Target number of leading zero bits (default: 4 if 0)
 * @param max_attempts Maximum number of mining attempts (default: 10,000,000 if <= 0)
 * @param progress_report_interval How often to call progress callback (default: 10,000 if <= 0)
 * @param timestamp_update_interval How often to update timestamp (default: 10,000 if <= 0)
 * @param progress_callback Optional callback for mining progress
 * @param user_data User data for progress callback
 * @return NOSTR_SUCCESS on success, error code on failure
 */
int nostr_add_proof_of_work(cJSON* event, const unsigned char* private_key, 
                           int target_difficulty, int max_attempts,
                           int progress_report_interval, int timestamp_update_interval,
                           void (*progress_callback)(int current_difficulty, uint64_t nonce, void* user_data),
                           void* user_data) {
    if (!event || !private_key) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    // Set default difficulty if not specified (but allow 0 to disable PoW)
    if (target_difficulty < 0) {
        target_difficulty = 4;
    }
    
    // If target_difficulty is 0, skip proof of work entirely
    if (target_difficulty == 0) {
        return NOSTR_SUCCESS;
    }
    
    // Set default values for parameters
    if (max_attempts <= 0) {
        max_attempts = 10000000;  // 10 million default
    }
    if (progress_report_interval <= 0) {
        progress_report_interval = 10000;  // Every 10,000 attempts
    }
    if (timestamp_update_interval <= 0) {
        timestamp_update_interval = 10000;  // Every 10,000 attempts
    }
    
    // Extract event data for reconstruction
    cJSON* kind_item = cJSON_GetObjectItem(event, "kind");
    cJSON* content_item = cJSON_GetObjectItem(event, "content");
    cJSON* created_at_item = cJSON_GetObjectItem(event, "created_at");
    cJSON* tags_item = cJSON_GetObjectItem(event, "tags");
    
    if (!kind_item || !content_item || !created_at_item || !tags_item) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    int kind = (int)cJSON_GetNumberValue(kind_item);
    const char* content = cJSON_GetStringValue(content_item);
    time_t original_timestamp = (time_t)cJSON_GetNumberValue(created_at_item);
    
    uint64_t nonce = 0;
    int attempts = 0;
    time_t current_timestamp = original_timestamp;
    
    // PoW difficulty tracking variables
    int best_difficulty_this_round = 0;
    int best_difficulty_overall = 0;
    
    // Mining loop
    while (attempts < max_attempts) {
        // Update timestamp based on timestamp_update_interval
        if (attempts % timestamp_update_interval == 0) {
            current_timestamp = time(NULL);
#ifdef ENABLE_DEBUG_LOGGING
            FILE* f = fopen("debug.log", "a"); 
            if (f) { 
                fprintf(f, "PoW mining: %d attempts, best this round: %d, overall best: %d, goal: %d\n", 
                        attempts, best_difficulty_this_round, best_difficulty_overall, target_difficulty); 
                fclose(f); 
            }
#endif
            // Reset best difficulty for the new round
            best_difficulty_this_round = 0;
        }
        
        // Call progress callback at specified intervals
        if (progress_callback && (attempts % progress_report_interval == 0)) {
            progress_callback(best_difficulty_overall, nonce, user_data);
        }
        
        // Create working copy of tags and add nonce
        cJSON* working_tags = cJSON_Duplicate(tags_item, 1);
        if (!working_tags) {
            return NOSTR_ERROR_MEMORY_FAILED;
        }
        
        if (update_nonce_tag_with_difficulty(working_tags, nonce, target_difficulty) != 0) {
            cJSON_Delete(working_tags);
            return NOSTR_ERROR_CRYPTO_FAILED;
        }
        
        // Create and sign new event with current nonce
        cJSON* test_event = nostr_create_and_sign_event(kind, content, working_tags, 
                                                       private_key, current_timestamp);
        cJSON_Delete(working_tags);
        
        if (!test_event) {
            return NOSTR_ERROR_CRYPTO_FAILED;
        }
        
        // Check PoW difficulty
        cJSON* id_item = cJSON_GetObjectItem(test_event, "id");
        if (!id_item || !cJSON_IsString(id_item)) {
            cJSON_Delete(test_event);
            return NOSTR_ERROR_CRYPTO_FAILED;
        }
        
        const char* event_id = cJSON_GetStringValue(id_item);
        unsigned char hash[32];
        if (nostr_hex_to_bytes(event_id, hash, 32) != NOSTR_SUCCESS) {
            cJSON_Delete(test_event);
            return NOSTR_ERROR_CRYPTO_FAILED;
        }
        
        // Count leading zero bits using NIP-13 method
        int current_difficulty = count_leading_zero_bits(hash);
        
        // Update difficulty tracking
        if (current_difficulty > best_difficulty_this_round) {
            best_difficulty_this_round = current_difficulty;
        }
        if (current_difficulty > best_difficulty_overall) {
            best_difficulty_overall = current_difficulty;
        }
        
        // Check if we've reached the target
        if (current_difficulty >= target_difficulty) {
#ifdef ENABLE_DEBUG_LOGGING
            FILE* f = fopen("debug.log", "a"); 
            if (f) { 
                fprintf(f, "PoW SUCCESS: Found difficulty %d (target %d) at nonce %llu after %d attempts\n", 
                        current_difficulty, target_difficulty, (unsigned long long)nonce, attempts + 1);
                
                // Print the final event JSON
                char* event_json = cJSON_Print(test_event);
                if (event_json) {
                    fprintf(f, "Final event: %s\n", event_json);
                    free(event_json);
                }
                fclose(f); 
            }
#endif
            
            // Copy successful result back to input event
            replace_event_content(event, test_event);
            cJSON_Delete(test_event);
            return NOSTR_SUCCESS;
        }
        
        cJSON_Delete(test_event);
        nonce++;
        attempts++;
    }
    
#ifdef ENABLE_DEBUG_LOGGING
    // Debug logging - failure
    FILE* f = fopen("debug.log", "a"); 
    if (f) { 
        fprintf(f, "PoW FAILED: Mining failed after %d attempts\n", max_attempts); 
        fclose(f); 
    }
#endif
    
    // If we reach here, we've exceeded max attempts
    return NOSTR_ERROR_CRYPTO_FAILED;
}

/**
 * Calculate PoW difficulty (leading zero bits) for an event ID
 *
 * @param event_id_hex Hexadecimal event ID string (64 characters)
 * @return Number of leading zero bits, or negative error code on failure
 */
int nostr_calculate_pow_difficulty(const char* event_id_hex) {
    if (!event_id_hex) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    // Validate hex string length (should be 64 characters for SHA-256)
    size_t hex_len = strlen(event_id_hex);
    if (hex_len != 64) {
        return NOSTR_ERROR_NIP13_CALCULATION;
    }
    
    // Convert hex to bytes
    unsigned char hash[32];
    if (nostr_hex_to_bytes(event_id_hex, hash, 32) != NOSTR_SUCCESS) {
        return NOSTR_ERROR_NIP13_CALCULATION;
    }
    
    // Use existing NIP-13 reference implementation
    return count_leading_zero_bits(hash);
}

/**
 * Extract nonce information from event tags
 *
 * @param event Complete event JSON object
 * @param nonce_out Output pointer for nonce value (can be NULL)
 * @param target_difficulty_out Output pointer for target difficulty (can be NULL)
 * @return NOSTR_SUCCESS if nonce tag found, NOSTR_ERROR_NIP13_NO_NONCE_TAG if not found, other error codes on failure
 */
int nostr_extract_nonce_info(cJSON* event, uint64_t* nonce_out, int* target_difficulty_out) {
    if (!event) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    // Initialize output values
    if (nonce_out) *nonce_out = 0;
    if (target_difficulty_out) *target_difficulty_out = -1;
    
    // Get tags array
    cJSON* tags = cJSON_GetObjectItem(event, "tags");
    if (!tags || !cJSON_IsArray(tags)) {
        return NOSTR_ERROR_NIP13_NO_NONCE_TAG;
    }
    
    // Search for nonce tag
    cJSON* tag = NULL;
    cJSON_ArrayForEach(tag, tags) {
        if (!cJSON_IsArray(tag) || cJSON_GetArraySize(tag) < 2) {
            continue;
        }
        
        // Check if this is a nonce tag
        cJSON* tag_type = cJSON_GetArrayItem(tag, 0);
        if (!tag_type || !cJSON_IsString(tag_type)) {
            continue;
        }
        
        const char* tag_name = cJSON_GetStringValue(tag_type);
        if (!tag_name || strcmp(tag_name, "nonce") != 0) {
            continue;
        }
        
        // Extract nonce value (second element)
        cJSON* nonce_item = cJSON_GetArrayItem(tag, 1);
        if (!nonce_item || !cJSON_IsString(nonce_item)) {
            return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
        }
        
        const char* nonce_str = cJSON_GetStringValue(nonce_item);
        if (!nonce_str) {
            return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
        }
        
        // Parse nonce value
        char* endptr;
        uint64_t nonce_val = strtoull(nonce_str, &endptr, 10);
        if (*endptr != '\0') {
            return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
        }
        
        if (nonce_out) *nonce_out = nonce_val;
        
        // Extract target difficulty (third element, optional)
        if (cJSON_GetArraySize(tag) >= 3) {
            cJSON* target_item = cJSON_GetArrayItem(tag, 2);
            if (target_item && cJSON_IsString(target_item)) {
                const char* target_str = cJSON_GetStringValue(target_item);
                if (target_str) {
                    char* endptr2;
                    long target_val = strtol(target_str, &endptr2, 10);
                    if (*endptr2 == '\0' && target_val >= 0) {
                        if (target_difficulty_out) *target_difficulty_out = (int)target_val;
                    }
                }
            }
        }
        
        return NOSTR_SUCCESS;
    }
    
    // No nonce tag found
    return NOSTR_ERROR_NIP13_NO_NONCE_TAG;
}

/**
 * Validate just the nonce tag format (without PoW calculation)
 *
 * @param nonce_tag_array JSON array representing a nonce tag
 * @return NOSTR_SUCCESS if valid format, error code otherwise
 */
int nostr_validate_nonce_tag(cJSON* nonce_tag_array) {
    if (!nonce_tag_array || !cJSON_IsArray(nonce_tag_array)) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    int array_size = cJSON_GetArraySize(nonce_tag_array);
    if (array_size < 2) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    // First element should be "nonce"
    cJSON* tag_type = cJSON_GetArrayItem(nonce_tag_array, 0);
    if (!tag_type || !cJSON_IsString(tag_type)) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    const char* tag_name = cJSON_GetStringValue(tag_type);
    if (!tag_name || strcmp(tag_name, "nonce") != 0) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    // Second element should be a valid nonce (numeric string)
    cJSON* nonce_item = cJSON_GetArrayItem(nonce_tag_array, 1);
    if (!nonce_item || !cJSON_IsString(nonce_item)) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    const char* nonce_str = cJSON_GetStringValue(nonce_item);
    if (!nonce_str) {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    // Validate nonce is a valid number
    char* endptr;
    strtoull(nonce_str, &endptr, 10);
    if (*endptr != '\0') {
        return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
    }
    
    // Third element (target difficulty) is optional, but if present should be valid
    if (array_size >= 3) {
        cJSON* target_item = cJSON_GetArrayItem(nonce_tag_array, 2);
        if (target_item && cJSON_IsString(target_item)) {
            const char* target_str = cJSON_GetStringValue(target_item);
            if (target_str) {
                char* endptr2;
                strtol(target_str, &endptr2, 10);
                if (*endptr2 != '\0') {
                    return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
                }
            }
        }
    }
    
    return NOSTR_SUCCESS;
}

/**
 * Validate Proof of Work for an event according to NIP-13
 *
 * @param event Complete event JSON object to validate
 * @param min_difficulty Minimum difficulty required by the relay (0 = no requirement)
 * @param validation_flags Bitflags for validation options
 * @param result_info Optional output struct for detailed results (can be NULL)
 * @return NOSTR_SUCCESS if PoW is valid and meets requirements, error code otherwise
 */
int nostr_validate_pow(cJSON* event, int min_difficulty, int validation_flags,
                      nostr_pow_result_t* result_info) {
    if (!event) {
        return NOSTR_ERROR_INVALID_INPUT;
    }
    
    // Initialize result info if provided
    if (result_info) {
        result_info->actual_difficulty = 0;
        result_info->committed_target = -1;
        result_info->nonce_value = 0;
        result_info->has_nonce_tag = 0;
        result_info->error_detail[0] = '\0';
    }
    
    // Get event ID for PoW calculation
    cJSON* id_item = cJSON_GetObjectItem(event, "id");
    if (!id_item || !cJSON_IsString(id_item)) {
        if (result_info) {
            snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                    "Missing or invalid event ID");
        }
        return NOSTR_ERROR_EVENT_INVALID_ID;
    }
    
    const char* event_id = cJSON_GetStringValue(id_item);
    if (!event_id) {
        if (result_info) {
            snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                    "Event ID is not a string");
        }
        return NOSTR_ERROR_EVENT_INVALID_ID;
    }
    
    // Calculate actual PoW difficulty
    int actual_difficulty = nostr_calculate_pow_difficulty(event_id);
    if (actual_difficulty < 0) {
        if (result_info) {
            snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                    "Failed to calculate PoW difficulty");
        }
        return actual_difficulty; // Return the specific error from calculation
    }
    
    if (result_info) {
        result_info->actual_difficulty = actual_difficulty;
    }
    
    // Check if minimum difficulty requirement is met
    if (min_difficulty > 0 && actual_difficulty < min_difficulty) {
        if (result_info) {
            snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                    "Insufficient difficulty: %d < %d", actual_difficulty, min_difficulty);
        }
        return NOSTR_ERROR_NIP13_INSUFFICIENT;
    }
    
    // Extract nonce information if validation flags require it
    uint64_t nonce_value = 0;
    int committed_target = -1;
    int nonce_extract_result = nostr_extract_nonce_info(event, &nonce_value, &committed_target);
    
    if (result_info) {
        result_info->nonce_value = nonce_value;
        result_info->committed_target = committed_target;
        result_info->has_nonce_tag = (nonce_extract_result == NOSTR_SUCCESS) ? 1 : 0;
    }
    
    // Validate nonce tag presence if required
    if (validation_flags & NOSTR_POW_VALIDATE_NONCE_TAG) {
        if (nonce_extract_result == NOSTR_ERROR_NIP13_NO_NONCE_TAG) {
            if (result_info) {
                snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                        "Missing required nonce tag");
            }
            return NOSTR_ERROR_NIP13_NO_NONCE_TAG;
        } else if (nonce_extract_result != NOSTR_SUCCESS) {
            if (result_info) {
                snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                        "Invalid nonce tag format");
            }
            return nonce_extract_result;
        }
        
        // If strict format validation is enabled, validate the nonce tag structure
        if (validation_flags & NOSTR_POW_STRICT_FORMAT) {
            cJSON* tags = cJSON_GetObjectItem(event, "tags");
            if (tags && cJSON_IsArray(tags)) {
                cJSON* tag = NULL;
                cJSON_ArrayForEach(tag, tags) {
                    if (cJSON_IsArray(tag) && cJSON_GetArraySize(tag) >= 2) {
                        cJSON* tag_type = cJSON_GetArrayItem(tag, 0);
                        if (tag_type && cJSON_IsString(tag_type)) {
                            const char* tag_name = cJSON_GetStringValue(tag_type);
                            if (tag_name && strcmp(tag_name, "nonce") == 0) {
                                int validation_result = nostr_validate_nonce_tag(tag);
                                if (validation_result != NOSTR_SUCCESS) {
                                    if (result_info) {
                                        snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                                                "Nonce tag failed strict format validation");
                                    }
                                    return validation_result;
                                }
                                break;
                            }
                        }
                    }
                }
            }
        }
    }
    
    // Validate committed target difficulty if required
    if (validation_flags & NOSTR_POW_VALIDATE_TARGET_COMMIT) {
        if (nonce_extract_result == NOSTR_SUCCESS && committed_target == -1) {
            if (result_info) {
                snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                        "Missing committed target difficulty in nonce tag");
            }
            return NOSTR_ERROR_NIP13_INVALID_NONCE_TAG;
        }
        
        // Check for target/difficulty mismatch if rejecting lower targets
        if (validation_flags & NOSTR_POW_REJECT_LOWER_TARGET) {
            // According to NIP-13: "if you require 40 bits to reply to your thread and see
            // a committed target of 30, you can safely reject it even if the note has 40 bits difficulty"
            // This means we reject if committed_target < min_difficulty, not actual_difficulty
            if (committed_target != -1 && min_difficulty > 0 && committed_target < min_difficulty) {
                if (result_info) {
                    snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                            "Committed target (%d) is lower than required minimum (%d)",
                            committed_target, min_difficulty);
                }
                return NOSTR_ERROR_NIP13_TARGET_MISMATCH;
            }
        }
    }
    
    // All validations passed
    if (result_info) {
        snprintf(result_info->error_detail, sizeof(result_info->error_detail),
                "PoW validation successful");
    }
    
    return NOSTR_SUCCESS;
}