/*
 * 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;
}