//! Bounded parallel generation of RandomX segments, serially chained with authenticated seed wraps.
//! This is a calibrated sequential-work assumption, not a proven VDF or wall clock.
mod randomx;
#[cfg(target_os = "linux")]
pub mod process;
pub mod paste;
use anyhow::{ensure, Result};
use chacha20poly1305::{
    aead::{Aead, KeyInit, Payload},
    ChaCha20Poly1305, Nonce, XChaCha20Poly1305, XNonce,
};
pub use ed25519_dalek::SigningKey;
use ed25519_dalek::{Signature, Signer, VerifyingKey};
use hkdf::Hkdf;
use rand::{rngs::OsRng, RngCore};
pub use randomx::{Dataset, RandomXMode};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use zeroize::Zeroizing;

pub const RANDOMX_VERSION: &str = "2.0.1";
pub const RANDOMX_COMMIT: &str = "aaafe71322df6602c21a5c72937ac284724ae561";
pub const SEGMENTS: usize = 96;
pub const PREVIOUS_SEGMENTS: usize = 84;
pub const LEGACY_SEGMENTS: usize = 7;
pub const MAX_ITERATIONS: u64 = 1_000_000_000_000;
pub const MAX_MANIFEST_BYTES: usize = 16_384;
const DOMAIN: &[u8] = b"relay-timelock-v1";

/// JSON transport; signature covers fixed-width canonical binary encoding, not JSON.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Manifest {
    pub version: u32,
    pub epoch: u64,
    pub randomx: String,
    pub randomx_commit: String,
    pub randomx_algorithm: String,
    pub dataset_key: String,
    pub iterations: u64,
    pub segments: u32,
    pub seed_1: String,
    /// One 12-byte nonce + 48-byte AEAD ciphertext per segment, in order.
    pub wrapped_keys: Vec<String>,
    pub key_commitment: String,
}
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct SignedManifest {
    pub puzzle: Manifest,
    pub service_public_key: String,
    pub service_signature: String,
}
pub struct GeneratedPuzzle {
    pub manifest: SignedManifest,
    pub epoch_key: Zeroizing<[u8; 32]>,
}
pub const MAX_RECORD_PLAINTEXT_BYTES: usize = 16 * 1024 * 1024;
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct EncryptedRecord {
    pub version: u32,
    pub epoch: u64,
    pub sequence: u64,
    pub puzzle_id: String,
    pub nonce: String,
    pub ciphertext: String,
}
impl EncryptedRecord {
    pub fn parse_json(input: &[u8]) -> Result<Self> {
        ensure!(
            input.len() <= 2 * MAX_RECORD_PLAINTEXT_BYTES + 2048,
            "encrypted record too large"
        );
        Ok(serde_json::from_slice(input)?)
    }
}
fn record_context(manifest: &SignedManifest, sequence: u64) -> Result<Vec<u8>> {
    let mut context = b"relay-encrypted-record-v1\0".to_vec();
    context.extend(1u32.to_be_bytes());
    context.extend(manifest.puzzle.epoch.to_be_bytes());
    context.extend(sequence.to_be_bytes());
    context.extend(bytes::<32>(&manifest.id()?)?);
    Ok(context)
}
fn record_key(manifest: &SignedManifest, epoch_key: &[u8; 32]) -> Result<Zeroizing<[u8; 32]>> {
    ensure!(
        hex::encode(Sha256::digest(epoch_key)) == manifest.puzzle.key_commitment,
        "epoch key commitment mismatch"
    );
    let mut key = Zeroizing::new([0; 32]);
    Hkdf::<Sha256>::new(Some(b"relay-record-key-v1"), epoch_key)
        .expand(&bytes::<32>(&manifest.id()?)?, key.as_mut())
        .map_err(|_| anyhow::anyhow!("HKDF expansion"))?;
    Ok(key)
}
pub fn encrypt_record(
    manifest: &SignedManifest,
    epoch_key: &[u8; 32],
    sequence: u64,
    plaintext: &[u8],
) -> Result<EncryptedRecord> {
    ensure!(
        plaintext.len() <= MAX_RECORD_PLAINTEXT_BYTES,
        "record plaintext too large"
    );
    let key = record_key(manifest, epoch_key)?;
    let mut nonce = [0; 24];
    OsRng.fill_bytes(&mut nonce);
    let context = record_context(manifest, sequence)?;
    let ciphertext = XChaCha20Poly1305::new_from_slice(key.as_ref())
        .unwrap()
        .encrypt(
            XNonce::from_slice(&nonce),
            Payload {
                msg: plaintext,
                aad: &context,
            },
        )
        .map_err(|_| anyhow::anyhow!("record encryption failed"))?;
    Ok(EncryptedRecord {
        version: 1,
        epoch: manifest.puzzle.epoch,
        sequence,
        puzzle_id: manifest.id()?,
        nonce: hex::encode(nonce),
        ciphertext: hex::encode(ciphertext),
    })
}
pub fn decrypt_record(
    manifest: &SignedManifest,
    epoch_key: &[u8; 32],
    record: &EncryptedRecord,
) -> Result<Zeroizing<Vec<u8>>> {
    ensure!(
        record.version == 1
            && record.epoch == manifest.puzzle.epoch
            && record.puzzle_id == manifest.id()?,
        "record belongs to another puzzle or protocol"
    );
    ensure!(
        record.ciphertext.len() >= 32
            && record.ciphertext.len() <= 2 * (MAX_RECORD_PLAINTEXT_BYTES + 16),
        "invalid ciphertext length"
    );
    let key = record_key(manifest, epoch_key)?;
    let nonce = bytes::<24>(&record.nonce)?;
    let ciphertext = hex::decode(&record.ciphertext)?;
    let context = record_context(manifest, record.sequence)?;
    let plaintext = XChaCha20Poly1305::new_from_slice(key.as_ref())
        .unwrap()
        .decrypt(
            XNonce::from_slice(&nonce),
            Payload {
                msg: &ciphertext,
                aad: &context,
            },
        )
        .map_err(|_| anyhow::anyhow!("record authentication failed"))?;
    Ok(Zeroizing::new(plaintext))
}
fn bytes<const N: usize>(text: &str) -> Result<[u8; N]> {
    ensure!(
        text.len() == N * 2
            && text
                .bytes()
                .all(|c| c.is_ascii_digit() || (b'a'..=b'f').contains(&c)),
        "noncanonical hex or wrong length"
    );
    let mut out = [0; N];
    hex::decode_to_slice(text, &mut out)?;
    Ok(out)
}
fn random32() -> Zeroizing<[u8; 32]> {
    let mut x = Zeroizing::new([0; 32]);
    OsRng.fill_bytes(x.as_mut());
    x
}
impl Manifest {
    pub fn canonical_bytes(&self) -> Result<Vec<u8>> {
        ensure!(
            self.version == 1
                && self.randomx == RANDOMX_VERSION
                && self.randomx_commit == RANDOMX_COMMIT
                && self.randomx_algorithm == "v2",
            "unsupported puzzle protocol"
        );
        ensure!(
            [LEGACY_SEGMENTS, PREVIOUS_SEGMENTS, SEGMENTS].contains(&(self.segments as usize))
                && self.wrapped_keys.len() == self.segments as usize
                && (1..=MAX_ITERATIONS).contains(&self.iterations),
            "invalid work parameters"
        );
        let mut out = b"relay-timelock-signed-manifest-v1\0".to_vec();
        out.extend(self.version.to_be_bytes());
        out.extend(self.epoch.to_be_bytes());
        out.extend(RANDOMX_VERSION.as_bytes());
        out.extend(RANDOMX_COMMIT.as_bytes());
        out.extend(b"v2");
        out.extend(bytes::<32>(&self.dataset_key)?);
        out.extend(self.iterations.to_be_bytes());
        out.extend(self.segments.to_be_bytes());
        out.extend(bytes::<32>(&self.seed_1)?);
        for wrapped in &self.wrapped_keys {
            out.extend(bytes::<60>(wrapped)?);
        }
        out.extend(bytes::<32>(&self.key_commitment)?);
        Ok(out)
    }
}
impl SignedManifest {
    pub fn parse_json(input: &[u8]) -> Result<Self> {
        ensure!(
            input.len() <= MAX_MANIFEST_BYTES,
            "manifest exceeds size limit"
        );
        let manifest: Self = serde_json::from_slice(input)?;
        manifest.puzzle.canonical_bytes()?;
        bytes::<32>(&manifest.service_public_key)?;
        bytes::<64>(&manifest.service_signature)?;
        Ok(manifest)
    }
    pub fn verify(&self, expected_public_key: &[u8; 32]) -> Result<()> {
        ensure!(
            bytes::<32>(&self.service_public_key)? == *expected_public_key,
            "service signer does not match attestation pin"
        );
        VerifyingKey::from_bytes(expected_public_key)?.verify_strict(
            &self.puzzle.canonical_bytes()?,
            &Signature::from_bytes(&bytes::<64>(&self.service_signature)?),
        )?;
        Ok(())
    }
    pub fn canonical_bytes(&self) -> Result<Vec<u8>> {
        let mut out = self.puzzle.canonical_bytes()?;
        out.extend(bytes::<32>(&self.service_public_key)?);
        out.extend(bytes::<64>(&self.service_signature)?);
        Ok(out)
    }
    pub fn id(&self) -> Result<String> {
        Ok(hex::encode(Sha256::digest(self.canonical_bytes()?)))
    }
}

fn input(epoch: u64, segment: usize, iteration: u64, x: &[u8; 32]) -> Zeroizing<Vec<u8>> {
    let mut out = Zeroizing::new(DOMAIN.to_vec());
    out.extend(epoch.to_be_bytes());
    out.extend((segment as u32).to_be_bytes());
    out.extend(iteration.to_be_bytes());
    out.extend(x);
    out
}
fn wrap_context(puzzle: &Manifest, segment: usize) -> Result<Vec<u8>> {
    let mut out = b"relay-timelock-wrap-v1\0".to_vec();
    out.extend(puzzle.epoch.to_be_bytes());
    out.extend((segment as u32).to_be_bytes());
    out.extend(puzzle.iterations.to_be_bytes());
    out.extend(bytes::<32>(&puzzle.dataset_key)?);
    out.extend(bytes::<32>(&puzzle.key_commitment)?);
    Ok(out)
}
fn wrap_key(y: &[u8; 32], context: &[u8]) -> Result<Zeroizing<[u8; 32]>> {
    let mut key = Zeroizing::new([0; 32]);
    Hkdf::<Sha256>::new(Some(DOMAIN), y)
        .expand(context, key.as_mut())
        .map_err(|_| anyhow::anyhow!("HKDF expansion"))?;
    Ok(key)
}
fn unwrap(puzzle: &Manifest, segment: usize, y: &[u8; 32]) -> Result<Zeroizing<[u8; 32]>> {
    let context = wrap_context(puzzle, segment)?;
    let key = wrap_key(y, &context)?;
    let wrapped = bytes::<60>(&puzzle.wrapped_keys[segment])?;
    let plaintext = Zeroizing::new(
        ChaCha20Poly1305::new_from_slice(key.as_ref())
            .unwrap()
            .decrypt(
                Nonce::from_slice(&wrapped[..12]),
                Payload {
                    msg: &wrapped[12..],
                    aad: &context,
                },
            )
            .map_err(|_| {
                anyhow::anyhow!(
                    "segment {} authentication failed (bad checkpoint or manifest)",
                    segment + 1
                )
            })?,
    );
    Ok(Zeroizing::new(plaintext.as_slice().try_into()?))
}

/// All workers share one initialized RandomX dataset/cache. Only s1 is published.
pub fn generate(
    epoch: u64,
    iterations: u64,
    signing_key: &SigningKey,
    mode: RandomXMode,
) -> Result<GeneratedPuzzle> {
    generate_with_rng(epoch, iterations, signing_key, mode, |bytes| {
        OsRng.try_fill_bytes(bytes).map_err(|_| anyhow::anyhow!("OS randomness failed"))
    })
}

/// Generate using a caller-supplied fallible entropy source. Nitro callers
/// provide direct NSM randomness; standalone CLI callers retain OS randomness.
/// No fallback is attempted if the supplied source fails.
pub fn generate_with_rng(
    epoch: u64,
    iterations: u64,
    signing_key: &SigningKey,
    mode: RandomXMode,
    fill: impl FnMut(&mut [u8]) -> Result<()>,
) -> Result<GeneratedPuzzle> {
    let workers = std::thread::available_parallelism()?.get().min(SEGMENTS);
    generate_with_rng_workers(epoch, iterations, signing_key, mode, workers, fill)
}

/// The worker count affects generation throughput only; the published puzzle
/// always requires all 96 segments to be recovered serially.
pub fn generate_with_rng_workers(
    epoch: u64,
    iterations: u64,
    signing_key: &SigningKey,
    mode: RandomXMode,
    workers: usize,
    fill: impl FnMut(&mut [u8]) -> Result<()>,
) -> Result<GeneratedPuzzle> {
    generate_with_progress(epoch, iterations, signing_key, mode, workers, fill, |_, _| {})
}

/// Report only aggregate work increments, never a chain state or seed.
/// The callback runs once per 10,000 hashes and at each group's completion.
pub fn generate_with_progress(
    epoch: u64, iterations: u64, signing_key: &SigningKey, mode: RandomXMode,
    workers: usize, mut fill: impl FnMut(&mut [u8]) -> Result<()>,
    progress: impl Fn(u64, bool) + Sync,
) -> Result<GeneratedPuzzle> {
    generate_internal(epoch, iterations, signing_key, mode, workers, &mut fill, progress, None)
}

/// Generate using an exec-isolated single-threaded supervisor and process workers.
/// The executable must dispatch `process::dispatch()` before starting any threads.
#[cfg(target_os = "linux")]
pub fn generate_with_process_progress(
    executable: &std::path::Path,
    epoch: u64, iterations: u64, signing_key: &SigningKey, mode: RandomXMode,
    workers: usize, mut fill: impl FnMut(&mut [u8]) -> Result<()>,
    progress: impl Fn(u64, bool) + Sync,
) -> Result<GeneratedPuzzle> {
    generate_internal(epoch, iterations, signing_key, mode, workers, &mut fill, progress, Some(executable))
}

fn generate_internal(
    epoch: u64, iterations: u64, signing_key: &SigningKey, mode: RandomXMode,
    workers: usize, mut fill: impl FnMut(&mut [u8]) -> Result<()>,
    progress: impl Fn(u64, bool) + Sync, executable: Option<&std::path::Path>,
) -> Result<GeneratedPuzzle> {
    ensure!((1..=SEGMENTS).contains(&workers), "invalid generation worker count");
    ensure!(
        (1..=MAX_ITERATIONS).contains(&iterations),
        "invalid iterations"
    );
    let mut dataset_key = Zeroizing::new([0u8; 32]);
    fill(dataset_key.as_mut())?;
    let mut seeds = Vec::with_capacity(SEGMENTS);
    for _ in 0..SEGMENTS {
        let mut seed = Zeroizing::new([0u8; 32]);
        fill(seed.as_mut())?;
        seeds.push(seed);
    }
    let mut epoch_key = Zeroizing::new([0u8; 32]);
    fill(epoch_key.as_mut())?;
    // Obtain all entropy before expensive work; a failed entropy source never
    // wastes a production-sized generation or publishes a partial puzzle.
    let mut nonces = [[0u8; 12]; SEGMENTS];
    for nonce in &mut nonces { fill(nonce)?; }
    let ys = if let Some(executable) = executable {
        #[cfg(target_os = "linux")]
        { process::generate(executable, epoch, iterations, mode, workers, &dataset_key, &seeds, &progress)? }
        #[cfg(not(target_os = "linux"))]
        { let _ = executable; anyhow::bail!("process generation requires Linux"); }
    } else {
    let dataset = Dataset::new(dataset_key.as_ref(), mode)?;
    // Work stealing keeps the pool busy without creating a VM per segment.
    // Each worker owns one VM; only the initialized dataset/cache is shared.
    let next = std::sync::atomic::AtomicUsize::new(0);
    std::thread::scope(|scope| -> Result<Vec<Zeroizing<[u8; 32]>>> {
        let tasks: Vec<_> = (0..workers).map(|_| {
            let (dataset, seeds, next, progress) = (&dataset, &seeds, &next, &progress);
            scope.spawn(move || -> Result<Vec<(usize, Zeroizing<[u8; 32]>)>> {
                let mut vm = dataset.vm()?;
                let mut completed = Vec::new();
                loop {
                    let segment = next.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
                    if segment >= SEGMENTS { break; }
                    let mut x = seeds[segment].clone();
                    for iteration in 0..iterations {
                        *x = vm.hash(&input(epoch, segment, iteration, &x));
                        if (iteration + 1) % 10_000 == 0 { progress(10_000, false); }
                    }
                    progress(iterations % 10_000, true);
                    completed.push((segment, x));
                }
                Ok(completed)
            })
        }).collect();
        let mut ordered: Vec<_> = (0..SEGMENTS).map(|_| Zeroizing::new([0; 32])).collect();
        for task in tasks {
            for (segment, output) in task.join()
                .map_err(|_| anyhow::anyhow!("RandomX worker panicked"))?? {
                ordered[segment] = output;
            }
        }
        Ok(ordered)
    })?
    };
    let mut puzzle = Manifest {
        version: 1,
        epoch,
        randomx: RANDOMX_VERSION.into(),
        randomx_commit: RANDOMX_COMMIT.into(),
        randomx_algorithm: "v2".into(),
        dataset_key: hex::encode(*dataset_key),
        iterations,
        segments: SEGMENTS as u32,
        seed_1: hex::encode(*seeds[0]),
        wrapped_keys: vec![String::new(); SEGMENTS],
        key_commitment: hex::encode(Sha256::digest(*epoch_key)),
    };
    for segment in 0..SEGMENTS {
        let context = wrap_context(&puzzle, segment)?;
        let key = wrap_key(&ys[segment], &context)?;
        let plaintext = if segment + 1 == SEGMENTS {
            &epoch_key
        } else {
            &seeds[segment + 1]
        };
        let nonce = nonces[segment];
        let cipher = ChaCha20Poly1305::new_from_slice(key.as_ref())
            .unwrap()
            .encrypt(
                Nonce::from_slice(&nonce),
                Payload {
                    msg: plaintext.as_ref(),
                    aad: &context,
                },
            )
            .map_err(|_| anyhow::anyhow!("AEAD encryption"))?;
        let mut wrapped = nonce.to_vec();
        wrapped.extend(cipher);
        puzzle.wrapped_keys[segment] = hex::encode(wrapped);
    }
    let signature = signing_key.sign(&puzzle.canonical_bytes()?);
    let manifest = SignedManifest {
        puzzle,
        service_public_key: hex::encode(signing_key.verifying_key().to_bytes()),
        service_signature: hex::encode(signature.to_bytes()),
    };
    Ok(GeneratedPuzzle {
        manifest,
        epoch_key,
    })
}

/// Local progress hint. Digest detects corruption, not a malicious writer. Final
/// AEAD and epoch commitment validate resumed computation before returning a key.
#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Checkpoint {
    pub version: u32,
    pub manifest_id: String,
    pub segment: u32,
    pub iteration: u64,
    pub x: String,
    pub digest: String,
}
impl Checkpoint {
    fn checksum(&self) -> Result<String> {
        let mut out = b"relay-timelock-checkpoint-v1\0".to_vec();
        out.extend(self.version.to_be_bytes());
        out.extend(bytes::<32>(&self.manifest_id)?);
        out.extend(self.segment.to_be_bytes());
        out.extend(self.iteration.to_be_bytes());
        out.extend(bytes::<32>(&self.x)?);
        Ok(hex::encode(Sha256::digest(out)))
    }
    fn new(id: &str, segment: usize, iteration: u64, x: &[u8; 32]) -> Result<Self> {
        let mut cp = Self {
            version: 1,
            manifest_id: id.into(),
            segment: segment as u32,
            iteration,
            x: hex::encode(x),
            digest: String::new(),
        };
        cp.digest = cp.checksum()?;
        Ok(cp)
    }
    pub fn parse_json(input: &[u8]) -> Result<Self> {
        ensure!(input.len() <= 2048, "checkpoint too large");
        Ok(serde_json::from_slice(input)?)
    }
    pub fn validate(&self, manifest: &SignedManifest) -> Result<()> {
        ensure!(
            self.version == 1 && self.manifest_id == manifest.id()?,
            "checkpoint belongs to another puzzle"
        );
        ensure!(
            self.segment < manifest.puzzle.segments && self.iteration <= manifest.puzzle.iterations,
            "invalid checkpoint coordinates"
        );
        ensure!(self.digest == self.checksum()?, "checkpoint corrupted");
        if self.segment == 0 && self.iteration == 0 {
            ensure!(
                self.x == manifest.puzzle.seed_1,
                "invalid initial checkpoint"
            );
        }
        Ok(())
    }
}

pub fn solve(
    manifest: &SignedManifest,
    expected_public_key: &[u8; 32],
    mode: RandomXMode,
    checkpoint: Option<Checkpoint>,
    checkpoint_interval: u64,
    mut save: impl FnMut(&Checkpoint) -> Result<()>,
) -> Result<Zeroizing<[u8; 32]>> {
    manifest.verify(expected_public_key)?;
    ensure!(checkpoint_interval > 0, "zero checkpoint interval");
    let id = manifest.id()?;
    let puzzle = &manifest.puzzle;
    let cp = checkpoint.unwrap_or(Checkpoint::new(&id, 0, 0, &bytes::<32>(&puzzle.seed_1)?)?);
    cp.validate(manifest)?;
    save(&cp)?;
    let dataset = Dataset::new(&bytes::<32>(&puzzle.dataset_key)?, mode)?;
    let mut vm = dataset.vm()?;
    let mut x = Zeroizing::new(bytes::<32>(&cp.x)?);
    for segment in cp.segment as usize..puzzle.segments as usize {
        let start = if segment == cp.segment as usize {
            cp.iteration
        } else {
            0
        };
        for iteration in start..puzzle.iterations {
            *x = vm.hash(&input(puzzle.epoch, segment, iteration, &x));
            if (iteration + 1) % checkpoint_interval == 0 {
                save(&Checkpoint::new(&id, segment, iteration + 1, &x)?)?;
            }
        }
        x = unwrap(puzzle, segment, &x)?;
        if segment + 1 < puzzle.segments as usize {
            save(&Checkpoint::new(&id, segment + 1, 0, &x)?)?;
        }
    }
    ensure!(
        hex::encode(Sha256::digest(*x)) == puzzle.key_commitment,
        "epoch key commitment mismatch"
    );
    Ok(x)
}

#[derive(Serialize)]
pub struct Calibration {
    pub randomx: &'static str,
    pub algorithm: &'static str,
    pub samples: u64,
    pub setup_seconds: f64,
    pub hash_seconds: f64,
    pub hashes_per_second: f64,
    pub workers: usize,
    pub worker_seconds: Vec<f64>,
    pub fastest_worker_hashes_per_second: f64,
    pub slowest_worker_hashes_per_second: f64,
    pub iterations_for_one_day: u64,
    pub mode: String,
}
pub fn calibrate(samples: u64, mode: RandomXMode) -> Result<Calibration> {
    calibrate_workers(samples, 1, mode)
}
pub fn calibrate_workers(samples: u64, workers: usize, mode: RandomXMode) -> Result<Calibration> {
    ensure!((1..=MAX_ITERATIONS).contains(&samples), "invalid samples");
    ensure!(
        (1..=SEGMENTS).contains(&workers),
        "workers must be 1 through 96"
    );
    let started = std::time::Instant::now();
    let dataset = Dataset::new(random32().as_ref(), mode)?;
    let setup_seconds = started.elapsed().as_secs_f64();
    let started = std::time::Instant::now();
    let barrier = std::sync::Barrier::new(workers);
    let worker_seconds = std::thread::scope(|scope| -> Result<Vec<f64>> {
        let tasks: Vec<_> = (0..workers)
            .map(|segment| {
                let dataset = &dataset;
                let barrier = &barrier;
                scope.spawn(move || -> Result<f64> {
                    // Synchronize start before allocation so a failing allocation does
                    // not strand the other workers on an unreachable barrier.
                    barrier.wait();
                    let mut vm = dataset.vm()?;
                    let started = std::time::Instant::now();
                    let mut x = random32();
                    for iteration in 0..samples {
                        *x = vm.hash(&input(0, segment, iteration, &x));
                    }
                    Ok(started.elapsed().as_secs_f64())
                })
            })
            .collect();
        tasks
            .into_iter()
            .map(|task| {
                task.join()
                    .map_err(|_| anyhow::anyhow!("calibration worker panicked"))?
            })
            .collect()
    })?;
    let hash_seconds = started.elapsed().as_secs_f64();
    let rate = samples as f64 * workers as f64 / hash_seconds;
    let fastest = samples as f64 / worker_seconds.iter().copied().fold(f64::INFINITY, f64::min);
    let slowest = samples as f64 / worker_seconds.iter().copied().fold(0.0, f64::max);
    Ok(Calibration {
        randomx: RANDOMX_VERSION,
        algorithm: "v2",
        samples,
        setup_seconds,
        hash_seconds,
        hashes_per_second: rate,
        workers,
        worker_seconds,
        fastest_worker_hashes_per_second: fastest,
        slowest_worker_hashes_per_second: slowest,
        iterations_for_one_day: (fastest * 86400.0).ceil() as u64,
        mode: format!("{mode:?}"),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn worker_counts_preserve_the_same_signed_chain() -> Result<()> {
        let signer = SigningKey::from_bytes(&[42; 32]);
        let mut previous = None;
        for workers in [7, 12] {
            let mut calls = 0u8;
            let counted = std::sync::atomic::AtomicU64::new(0);
            let groups = std::sync::atomic::AtomicU64::new(0);
            let generated = generate_with_progress(92, 1, &signer, RandomXMode::Light,
                workers, |bytes| { calls += 1; bytes.fill(calls); Ok(()) }, |n, done| {
                    counted.fetch_add(n, std::sync::atomic::Ordering::Relaxed);
                    if done { groups.fetch_add(1, std::sync::atomic::Ordering::Relaxed); }
                })?;
            assert_eq!(counted.load(std::sync::atomic::Ordering::Relaxed), SEGMENTS as u64);
            assert_eq!(groups.load(std::sync::atomic::Ordering::Relaxed), SEGMENTS as u64);
            assert_eq!(generated.manifest.puzzle.segments, 96);
            let json = serde_json::to_vec(&generated.manifest)?;
            assert!(json.len() <= MAX_MANIFEST_BYTES);
            if let Some(expected) = previous { assert_eq!(json, expected); }
            previous = Some(json);
        }
        Ok(())
    }

    #[test]
    fn legacy_seven_segment_fixture_still_solves() -> Result<()> {
        let raw = std::fs::read(concat!(env!("CARGO_MANIFEST_DIR"),
            "/../../python/attested-relay/tests/fixtures/native-puzzle.json"))?;
        let manifest = SignedManifest::parse_json(&raw)?;
        assert_eq!(manifest.puzzle.segments, 7);
        let key = solve(&manifest, &bytes::<32>(&manifest.service_public_key)?,
            RandomXMode::Light, None, 1, |_| Ok(()))?;
        assert_eq!(hex::encode(Sha256::digest(*key)), manifest.puzzle.key_commitment);
        let mut malformed = manifest.clone();
        malformed.puzzle.segments = 84;
        assert!(malformed.puzzle.canonical_bytes().is_err());
        malformed.puzzle.segments = 8;
        assert!(malformed.puzzle.canonical_bytes().is_err());
        Ok(())
    }

    #[test]
    fn supplied_entropy_drives_native_generation_and_fails_closed() -> Result<()> {
        let signer = SigningKey::from_bytes(&[42; 32]);
        // Every entropy call, including the last wrap nonce, is required before
        // dataset initialization. There is no hidden OS-RNG recovery path.
        for failure in 0..(2 * SEGMENTS + 2) {
            let mut calls = 0;
            let result = generate_with_rng(90, 1, &signer, RandomXMode::Light, |bytes| {
                let current = calls;
                calls += 1;
                if current == failure { anyhow::bail!("entropy unavailable"); }
                bytes.fill(current as u8 + 1);
                Ok(())
            });
            assert!(result.is_err());
            assert_eq!(calls, failure + 1);
        }
        let mut calls = 0u8;
        let generated = generate_with_rng(90, 1, &signer, RandomXMode::Light, |bytes| {
            calls += 1;
            bytes.fill(calls);
            Ok(())
        })?;
        assert_eq!(calls as usize, 2 * SEGMENTS + 2);
        assert_eq!(generated.manifest.puzzle.dataset_key, hex::encode([1; 32]));
        assert_eq!(generated.manifest.puzzle.seed_1, hex::encode([2; 32]));
        assert_eq!(*generated.epoch_key, [(SEGMENTS + 2) as u8; 32]);
        let solved = solve(&generated.manifest, &signer.verifying_key().to_bytes(),
            RandomXMode::Light, None, 1, |_| Ok(()))?;
        assert_eq!(*solved, [(SEGMENTS + 2) as u8; 32]);
        Ok(())
    }

    #[test]
    fn real_randomx_roundtrip_restart_and_tampering() -> Result<()> {
        let signer = SigningKey::from_bytes(&[42; 32]);
        let public = signer.verifying_key().to_bytes();
        let generated = generate(91, 2, &signer, RandomXMode::Light)?;
        let json = serde_json::to_vec(&generated.manifest)?;
        let manifest = SignedManifest::parse_json(&json)?;
        let mut saved = Vec::new();
        let key = solve(&manifest, &public, RandomXMode::Light, None, 1, |cp| {
            saved.push(cp.clone());
            Ok(())
        })?;
        assert_eq!(*key, *generated.epoch_key);
        let record = encrypt_record(&manifest, &key, 17, b"request and response")?;
        assert_eq!(
            &**decrypt_record(&manifest, &key, &record)?,
            b"request and response"
        );
        let mut tampered_record = record.clone();
        tampered_record.sequence += 1;
        assert!(decrypt_record(&manifest, &key, &tampered_record).is_err());
        tampered_record = record.clone();
        tampered_record.epoch += 1;
        assert!(decrypt_record(&manifest, &key, &tampered_record).is_err());
        assert!(decrypt_record(&manifest, &[3; 32], &record).is_err());
        // Resume midway and on a segment boundary; every saved state is validated.
        for cp in [&saved[1], &saved[3], saved.last().unwrap()] {
            let recovered = solve(
                &manifest,
                &public,
                RandomXMode::Light,
                Some(cp.clone()),
                1,
                |_| Ok(()),
            )?;
            assert_eq!(*recovered, *key);
        }
        let mut wrong = saved[0].clone();
        wrong.x = hex::encode([(SEGMENTS + 2) as u8; 32]);
        assert!(wrong.validate(&manifest).is_err());
        wrong.digest = wrong.checksum()?;
        assert!(solve(
            &manifest,
            &public,
            RandomXMode::Light,
            Some(wrong),
            1,
            |_| Ok(())
        )
        .is_err());
        let mut tampered = manifest.clone();
        tampered.puzzle.epoch += 1;
        assert!(tampered.verify(&public).is_err());
        assert!(manifest.verify(&[1; 32]).is_err());
        let mut tampered = manifest.clone();
        tampered.puzzle.wrapped_keys.swap(0, 1);
        assert!(tampered.verify(&public).is_err());
        // Even a properly signed malformed AEAD wrapper fails in the solver.
        tampered.service_signature =
            hex::encode(signer.sign(&tampered.puzzle.canonical_bytes()?).to_bytes());
        assert!(solve(&tampered, &public, RandomXMode::Light, None, 1, |_| Ok(())).is_err());
        let mut changed = manifest.clone();
        changed.puzzle.iterations = 0;
        assert!(changed.puzzle.canonical_bytes().is_err());
        assert!(SignedManifest::parse_json(&vec![b' '; MAX_MANIFEST_BYTES + 1]).is_err());
        let mut object: serde_json::Value = serde_json::from_slice(&json)?;
        object["extra"] = true.into();
        assert!(SignedManifest::parse_json(&serde_json::to_vec(&object)?).is_err());
        Ok(())
    }
}