use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;

/// Test whether `n` is prime using trial division up to √n.
/// Deliberately not optimized — we *want* to burn CPU.
fn is_prime(n: u64) -> bool {
    if n < 2 {
        return false;
    }
    if n < 4 {
        return true;
    }
    if n % 2 == 0 || n % 3 == 0 {
        return false;
    }
    let mut i = 5u64;
    while i.saturating_mul(i) <= n {
        if n % i == 0 || n % (i + 2) == 0 {
            return false;
        }
        i += 6;
    }
    true
}

/// CPU-burning worker that counts primes by brute-force trial division.
///
/// Each thread tests a disjoint slice of odd numbers (interleaved by stride)
/// and atomically increments shared counters for primes found and numbers tested.
/// Uses constant memory regardless of how long it runs.
pub fn prime_worker(
    running: Arc<AtomicBool>,
    primes_found: Arc<AtomicU64>,
    numbers_tested: Arc<AtomicU64>,
    thread_id: u64,
    num_threads: u64,
) {
    // Each thread works on odd numbers: 3 + 2*thread_id, stepping by 2*num_threads.
    // Thread 0 also accounts for the prime 2 by adding 1 to its first batch.
    let stride = 2 * num_threads;
    let mut n = 3 + 2 * thread_id;
    let batch_size = 500u64;

    if thread_id == 0 {
        // Count prime 2 once
        primes_found.fetch_add(1, Ordering::Relaxed);
    }

    while running.load(Ordering::Relaxed) {
        let mut local_primes = 0u64;

        for _ in 0..batch_size {
            if is_prime(n) {
                local_primes += 1;
            }
            n += stride;
        }

        primes_found.fetch_add(local_primes, Ordering::Relaxed);
        numbers_tested.fetch_add(batch_size, Ordering::Relaxed);
    }
}