use rayon::prelude::*;

// Z85 alphabet (ZeroMQ RFC 32):
// "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ.-:+=^!/*?&<>()[]{}@%$#"
// 255 = invalid sentinel
const DECODE: [u8; 256] = {
    let mut t = [255u8; 256];
    let alpha =
        b"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ.-:+=^!/*?&<>()[]{}@%$#";
    let mut i = 0usize;
    while i < 85 {
        t[alpha[i] as usize] = i as u8;
        i += 1;
    }
    t
};

/// Decode one 5-char z85 group into 4 bytes.
#[inline]
fn decode_group(src: &[u8]) -> Result<[u8; 4], String> {
    let d = [
        DECODE[src[0] as usize],
        DECODE[src[1] as usize],
        DECODE[src[2] as usize],
        DECODE[src[3] as usize],
        DECODE[src[4] as usize],
    ];
    if d.iter().any(|&v| v == 255) {
        return Err(format!("invalid z85 character in group: {:?}", src));
    }
    // u64 accumulation — intermediate products can exceed u32::MAX for invalid input
    let v = d[0] as u64 * 52_200_625 // 85^4
        + d[1] as u64 * 614_125      // 85^3
        + d[2] as u64 * 7_225        // 85^2
        + d[3] as u64 * 85
        + d[4] as u64;
    Ok((v as u32).to_be_bytes())
}

/// Single-threaded z85 decode. LLVM will auto-vectorize the inner loop.
pub fn decode_scalar(input: &[u8]) -> Result<Vec<u8>, String> {
    if input.len() % 5 != 0 {
        return Err(format!(
            "input length {} is not a multiple of 5",
            input.len()
        ));
    }
    let mut out = vec![0u8; input.len() / 5 * 4];
    for (src, dst) in input.chunks_exact(5).zip(out.chunks_exact_mut(4)) {
        dst.copy_from_slice(&decode_group(src)?);
    }
    Ok(out)
}

/// Rayon-parallel z85 decode. Writes directly into a pre-allocated buffer —
/// no intermediate Vec<Vec<u8>> or flatten. Order is guaranteed by
/// IndexedParallelIterator zip.
pub fn decode_parallel(input: &[u8]) -> Result<Vec<u8>, String> {
    if input.len() % 5 != 0 {
        return Err(format!(
            "input length {} is not a multiple of 5",
            input.len()
        ));
    }
    let mut out = vec![0u8; input.len() / 5 * 4];
    input
        .par_chunks_exact(5)
        .zip(out.par_chunks_exact_mut(4))
        .try_for_each(|(src, dst)| decode_group(src).map(|g| dst.copy_from_slice(&g)))?;
    Ok(out)
}
