//! Hash helpers: SHA-256 and hash-to-prime.
//!
//! `hash_prime` follows the shape of chiavdf's `HashPrime`
//! (vendor/chiavdf/src/proof_common.h): expand the seed with SHA-256 under a counter until
//! `bits` bits are filled, force the requested bits on, force odd, and retry until the
//! candidate passes a probable-prime test. The counter is a 64-bit big-endian suffix here,
//! where chiavdf increments the seed bytes in place. Same idea, not byte-compatible.

use rug::integer::{IsPrime, Order};
use rug::Integer;
use sha2::{Digest, Sha256};

pub const PRIME_REPS: u32 = 30;

pub fn sha256(parts: &[&[u8]]) -> [u8; 32] {
    let mut h = Sha256::new();
    for p in parts {
        h.update(p);
    }
    h.finalize().into()
}

/// Deterministic prime of exactly `bits` bits derived from `seed`.
/// Bits listed in `set_bits` are forced to 1 (chiavdf passes {0, 1, 2, bits-1} for a
/// discriminant, so that p = 7 mod 8 and the top bit is set, and {bits-1} for the
/// Fiat-Shamir prime).
pub fn hash_prime(seed: &[u8], bits: u32, set_bits: &[u32]) -> Integer {
    assert!(bits % 8 == 0 && bits >= 64);
    let nbytes = (bits / 8) as usize;
    let mut ctr: u64 = 0;
    loop {
        let mut blob: Vec<u8> = Vec::with_capacity(nbytes + 32);
        while blob.len() < nbytes {
            let h = sha256(&[seed, &ctr.to_be_bytes()]);
            ctr += 1;
            blob.extend_from_slice(&h);
        }
        blob.truncate(nbytes);
        let mut p = Integer::from_digits(&blob, Order::MsfBe);
        for &b in set_bits {
            p.set_bit(b, true);
        }
        p.set_bit(0, true);
        if p.is_probably_prime(PRIME_REPS) != IsPrime::No {
            return p;
        }
    }
}

/// Fixed-width big-endian encoding of a non-negative integer.
pub fn int_to_bytes(x: &Integer, width: usize) -> Vec<u8> {
    assert!(x.cmp0() != std::cmp::Ordering::Less);
    let d = x.to_digits::<u8>(Order::MsfBe);
    assert!(d.len() <= width, "integer wider than {} bytes", width);
    let mut out = vec![0u8; width - d.len()];
    out.extend_from_slice(&d);
    out
}

/// Signed fixed-width encoding: one sign byte (0 or 1) then the magnitude.
pub fn signed_to_bytes(x: &Integer, width: usize) -> Vec<u8> {
    let mut out = Vec::with_capacity(width + 1);
    out.push(if x.cmp0() == std::cmp::Ordering::Less { 1 } else { 0 });
    let mag = Integer::from(x.abs_ref());
    out.extend_from_slice(&int_to_bytes(&mag, width));
    out
}

pub fn bytes_to_int(b: &[u8]) -> Integer {
    Integer::from_digits(b, Order::MsfBe)
}

pub fn hex(b: &[u8]) -> String {
    b.iter().map(|x| format!("{:02x}", x)).collect()
}

pub fn from_hex(s: &str) -> Option<Vec<u8>> {
    if s.len() % 2 != 0 {
        return None;
    }
    (0..s.len() / 2)
        .map(|i| u8::from_str_radix(&s[2 * i..2 * i + 2], 16).ok())
        .collect()
}