signer/src/miner.rs

//! NIP-13 proof-of-work mining coordinator.
//!
//! Port of `miner.c`. Multi-threaded best-effort mining that:
//! - Runs N worker threads, each trying nonces with a unique stride
//! - Tracks the best event (highest difficulty) across all threads
//! - Stops when target difficulty is reached OR timeout expires
//! - Always returns the best result found

use nostr_core::types::{Event, SecretKey, Tag};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::{Duration, Instant};

/// Mine result.
#[derive(Debug, Clone)]
pub struct MineResult {
    pub best_event_json: String,
    pub achieved_difficulty: i32,
    pub target_difficulty: i32,
    pub target_reached: bool,
    pub elapsed_sec: u64,
    pub total_attempts: u64,
}

/// Shared state across worker threads.
struct MineShared {
    best_difficulty: i32,
    best_nonce: u64,
    total_attempts: u64,
    target_reached: bool,
    stop: bool,
}

/// Run the mining coordinator.
///
/// `event_json` — the unsigned event JSON (will be parsed and modified)
/// `private_key` — 32-byte secp256k1 private key (for signing the mined event)
/// `target_difficulty` — target leading zero bits (0 = no target, mine for full timeout)
/// `thread_count` — number of mining threads (1..32)
/// `timeout_sec` — time budget in seconds
pub fn miner_run(
    event_json: &str,
    private_key: &[u8; 32],
    target_difficulty: i32,
    thread_count: i32,
    timeout_sec: u64,
) -> Result<MineResult, crate::SignerError> {
    let threads = thread_count.clamp(1, 32) as usize;
    let timeout = if timeout_sec == 0 { 600 } else { timeout_sec };
    let deadline = Instant::now() + Duration::from_secs(timeout);

    // Parse the unsigned event
    let mut event: Event =
        serde_json::from_str(event_json).map_err(|_| crate::SignerError::InvalidInput)?;

    // Ensure there's a nonce tag (will be updated by workers)
    let has_nonce = event.tags.iter().any(|t| t.kind() == "nonce");
    if !has_nonce {
        let mut tag = Tag::with_value("nonce", "0");
        tag.0.push(target_difficulty.to_string());
        event.tags.push(tag);
    }

    let shared = Arc::new(Mutex::new(MineShared {
        best_difficulty: 0,
        best_nonce: 0,
        total_attempts: 0,
        target_reached: false,
        stop: false,
    }));

    let start = Instant::now();
    let mut handles = Vec::new();

    for thread_id in 0..threads {
        let shared = Arc::clone(&shared);
        let event = event.clone();
        let target = target_difficulty;
        let deadline = deadline;

        let handle = thread::spawn(move || {
            mine_worker(
                thread_id as u64,
                threads as u64,
                &shared,
                &event,
                target,
                deadline,
            );
        });
        handles.push(handle);
    }

    for handle in handles {
        let _ = handle.join();
    }

    let shared = Arc::try_unwrap(shared)
        .map_err(|_| crate::SignerError::Internal("mining thread still holds shared state".into()))?
        .into_inner()
        .map_err(|_| crate::SignerError::Internal("mining shared state poisoned".into()))?;

    let elapsed = start.elapsed().as_secs();

    // Reconstruct the best event with the best nonce and sign it
    let best_event_json = if shared.best_difficulty > 0 {
        let nonce_tag_index = event.tags.iter().position(|t| t.kind() == "nonce");
        if let Some(idx) = nonce_tag_index {
            event.tags[idx] = Tag::with_value("nonce", &shared.best_nonce.to_string());
            event.tags[idx].0.push(target_difficulty.to_string());
        }
        let sk = SecretKey::from_bytes(*private_key);
        let signed = nips::nip001::create_and_sign_event(
            event.kind,
            &event.content,
            event.tags.clone(),
            &sk,
            event.created_at,
        )
        .map_err(|_| crate::SignerError::CryptoFailed)?;
        serde_json::to_string(&signed).map_err(|_| crate::SignerError::InvalidInput)?
    } else {
        // No event mined — return the original unsigned event
        event_json.to_string()
    };

    Ok(MineResult {
        best_event_json,
        achieved_difficulty: shared.best_difficulty,
        target_difficulty,
        target_reached: shared.target_reached,
        elapsed_sec: elapsed,
        total_attempts: shared.total_attempts,
    })
}

/// Per-worker mining loop.
fn mine_worker(
    thread_id: u64,
    thread_count: u64,
    shared: &Arc<Mutex<MineShared>>,
    event: &Event,
    target_difficulty: i32,
    deadline: Instant,
) {
    let mut nonce: u64 = thread_id;
    let stride = thread_count;
    let nonce_tag_index = event.tags.iter().position(|t| t.kind() == "nonce");

    let mut local_event = event.clone();
    let mut local_best_difficulty: i32 = 0;
    let mut local_best_nonce: u64 = 0;
    let mut attempts: u64 = 0;

    while Instant::now() < deadline {
        // Check if we should stop (target reached by another thread)
        {
            let s = shared.lock().unwrap();
            if s.stop {
                // Flush our attempts
                drop(s);
                let mut s = shared.lock().unwrap();
                s.total_attempts += attempts;
                return;
            }
        }

        // Update nonce tag
        if let Some(idx) = nonce_tag_index {
            local_event.tags[idx] = Tag::with_value("nonce", &nonce.to_string());
            local_event.tags[idx].0.push(target_difficulty.to_string());
        }

        // Recompute event ID
        if let Ok(new_id) = local_event.compute_id() {
            let difficulty = nips::nip013::count_leading_zero_bits(&new_id) as i32;

            if difficulty > local_best_difficulty {
                local_best_difficulty = difficulty;
                local_best_nonce = nonce;

                // Check if target reached
                if target_difficulty > 0 && difficulty >= target_difficulty {
                    let mut s = shared.lock().unwrap();
                    if local_best_difficulty > s.best_difficulty {
                        s.best_difficulty = local_best_difficulty;
                        s.best_nonce = local_best_nonce;
                    }
                    s.total_attempts += attempts;
                    s.target_reached = true;
                    s.stop = true;
                    return;
                }
            }
        }

        attempts += 1;
        nonce += stride;

        // Periodically flush to shared
        if attempts % 10000 == 0 {
            let mut s = shared.lock().unwrap();
            s.total_attempts += 10000;
            if local_best_difficulty > s.best_difficulty {
                s.best_difficulty = local_best_difficulty;
                s.best_nonce = local_best_nonce;
            }
            attempts = 0;
        }
    }

    // Final flush
    let mut s = shared.lock().unwrap();
    s.total_attempts += attempts;
    if local_best_difficulty > s.best_difficulty {
        s.best_difficulty = local_best_difficulty;
        s.best_nonce = local_best_nonce;
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nostr_core::crypto::keys::generate_keypair;
    use nostr_core::types::Kind;

    #[test]
    fn test_mine_low_difficulty() {
        let (sk, pk) = generate_keypair();
        let event = Event::new(pk, 1234567890, Kind::Text, vec![], "hello mining");

        let event_json = serde_json::to_string(&event).unwrap();
        let priv_bytes = sk.as_bytes();

        let result = miner_run(&event_json, &priv_bytes, 4, 2, 5).unwrap();

        assert!(result.achieved_difficulty >= 4, "achieved: {}", result.achieved_difficulty);
        assert!(result.target_reached);
        assert!(result.total_attempts > 0);
    }

    #[test]
    fn test_mine_zero_target_returns_best() {
        let (sk, pk) = generate_keypair();
        let event = Event::new(pk, 1234567890, Kind::Text, vec![], "test zero target");

        let event_json = serde_json::to_string(&event).unwrap();
        let priv_bytes = sk.as_bytes();

        // target=0 means mine for the full timeout, return best found
        let result = miner_run(&event_json, &priv_bytes, 0, 2, 1).unwrap();

        // Should have found something with at least 0 difficulty
        assert!(result.achieved_difficulty >= 0);
        assert!(!result.best_event_json.is_empty());
    }

    #[test]
    fn test_mine_invalid_event_json() {
        let (sk, _) = generate_keypair();
        let priv_bytes = sk.as_bytes();

        assert!(miner_run("not valid json", &priv_bytes, 4, 1, 5).is_err());
    }
}