jam-tooling 0.1.28

use anyhow::anyhow;
use jam_std_common::hash_raw;
use jam_types::{CoreIndex, Memo, ServiceId, VecSet, MEMO_LEN};
use std::time::Duration;

pub type Blob = Vec<u8>;

#[derive(Clone, Debug)]
#[allow(clippy::len_without_is_empty)]
pub enum DataIdLen {
	HashLen([u8; 32], usize),
	Data(Vec<u8>),
}

impl DataIdLen {
	pub fn hash(&self) -> [u8; 32] {
		match self {
			DataIdLen::HashLen(hash, _) => *hash,
			DataIdLen::Data(data) => hash_raw(data),
		}
	}
	pub fn len(&self) -> usize {
		match self {
			DataIdLen::HashLen(_, l) => *l,
			DataIdLen::Data(data) => data.len(),
		}
	}
	pub fn data(&self) -> Option<&[u8]> {
		match self {
			DataIdLen::HashLen(..) => None,
			DataIdLen::Data(ref data) => Some(data),
		}
	}
	pub fn into_data(self) -> Option<Vec<u8>> {
		match self {
			DataIdLen::HashLen(..) => None,
			DataIdLen::Data(d) => Some(d),
		}
	}
}

pub fn data_id_len(s: &str) -> anyhow::Result<DataIdLen> {
	match s {
		_ if s.starts_with("0x") &&
			s.len() > 67 &&
			s.bytes().skip(2).take(64).all(|x| x.is_ascii_hexdigit()) &&
			s.find(':') == Some(66) &&
			s[67..].parse::<usize>().is_ok() =>
		{
			let mut hash = [0u8; 32];
			for (i, c) in s.as_bytes()[2..66].chunks(2).enumerate() {
				hash[i] = u8::from_str_radix(std::str::from_utf8(c)?, 16)?;
			}
			Ok(DataIdLen::HashLen(hash, s[67..].parse::<usize>().expect("see above")))
		},
		_ if s.len() > 65 &&
			s.bytes().take(64).all(|x| x.is_ascii_hexdigit()) &&
			s.find(':') == Some(64) &&
			s[65..].parse::<usize>().is_ok() =>
		{
			let mut hash = [0u8; 32];
			for (i, c) in s.as_bytes()[..64].chunks(2).enumerate() {
				hash[i] = u8::from_str_radix(std::str::from_utf8(c)?, 16)?;
			}
			Ok(DataIdLen::HashLen(hash, s[65..].parse::<usize>().expect("see above")))
		},
		_ => Ok(DataIdLen::Data(blob(s)?)),
	}
}

#[allow(dead_code)]
#[derive(Clone, Debug)]
pub enum DataId {
	Hash([u8; 32]),
	Data(Vec<u8>),
}

#[allow(dead_code)]
#[allow(clippy::len_without_is_empty)]
impl DataId {
	pub fn hash(&self) -> [u8; 32] {
		match self {
			DataId::Hash(hash) => *hash,
			DataId::Data(data) => hash_raw(data),
		}
	}
	pub fn len(&self) -> Option<usize> {
		match self {
			DataId::Hash(_) => None,
			DataId::Data(data) => Some(data.len()),
		}
	}
	pub fn data(&self) -> Option<&[u8]> {
		match self {
			DataId::Hash(..) => None,
			DataId::Data(ref data) => Some(data),
		}
	}
}

pub fn data_id(s: &str) -> anyhow::Result<DataId> {
	match s {
		_ if s.len() == 64 && s.bytes().all(|x| x.is_ascii_hexdigit()) => {
			let mut hash = [0u8; 32];
			for (i, c) in s.as_bytes().chunks(2).enumerate() {
				hash[i] = u8::from_str_radix(std::str::from_utf8(c)?, 16)?;
			}
			Ok(DataId::Hash(hash))
		},
		_ if s.starts_with("0x") &&
			s.len() == 66 &&
			s.bytes().skip(2).all(|x| x.is_ascii_hexdigit()) =>
		{
			let mut hash = [0u8; 32];
			for (i, c) in s.as_bytes().chunks(2).skip(1).enumerate() {
				hash[i] = u8::from_str_radix(std::str::from_utf8(c)?, 16)?;
			}
			Ok(DataId::Hash(hash))
		},
		_ => Ok(DataId::Data(blob(s)?)),
	}
}

pub fn blob(s: &str) -> anyhow::Result<Blob> {
	match s {
		_ if std::path::Path::new(s).exists() => Ok(std::fs::read(s)?),
		_ if s.starts_with("0x") &&
			s.len().is_multiple_of(2) &&
			s.bytes().skip(2).all(|x| x.is_ascii_hexdigit()) =>
		{
			let mut inner = Vec::with_capacity((s.len() - 2) / 2);
			for c in s.as_bytes()[2..].chunks(2) {
				inner.push(u8::from_str_radix(std::str::from_utf8(c)?, 16)?);
			}
			Ok(inner)
		},
		_ => Ok(s.as_bytes().to_vec()),
	}
}

pub fn gas(s: &str) -> anyhow::Result<u64> {
	if s == "max" {
		Ok(u64::MAX)
	} else {
		Ok(s.parse::<u64>()?)
	}
}

pub fn memo(s: &str) -> anyhow::Result<Memo> {
	match s {
		_ if s.len() == MEMO_LEN * 2 && s.bytes().all(|x| x.is_ascii_hexdigit()) => {
			let mut inner = Memo::default();
			for (i, c) in s.as_bytes().chunks(2).enumerate() {
				inner[i] = u8::from_str_radix(std::str::from_utf8(c)?, 16)?;
			}
			Ok(inner)
		},
		_ if s.len() <= MEMO_LEN => {
			let mut inner = Memo::default();
			inner.0[..s.len()].copy_from_slice(s.as_bytes());
			Ok(inner)
		},
		_ => Err(anyhow!("Memo too long")),
	}
}

pub fn service_id(s: &str) -> anyhow::Result<ServiceId> {
	Ok(ServiceId::from_str_radix(s, 16)?)
}

pub fn exact_bytes(s: &str) -> anyhow::Result<u64> {
	let i = s
		.rfind(char::is_numeric)
		.ok_or_else(|| anyhow!("Failed to parse number from {s:?}"))?;
	let (number_str, prefix) = s.split_at(i + 1);
	let prefix = prefix.trim();
	if !matches!(prefix.chars().last(), Some('B' | 'b')) {
		return Err(anyhow!("Failed to parse size in bytes: invalid prefix {prefix:?}"))
	}
	let prefix = &prefix[..prefix.len() - 1];
	let scale = get_iec_scale(prefix)
		.ok_or_else(|| anyhow!("Failed to parse size in bytes: invalid prefix {prefix:?}"))?;
	let number: u64 = number_str.trim().parse()?;
	number
		.checked_mul(scale)
		.ok_or_else(|| anyhow!("The size is too big: {number} * {scale}"))
}

fn get_iec_scale(prefix: &str) -> Option<u64> {
	for (iec_prefix, scale) in IEC_TABLE.iter() {
		if iec_prefix.eq_ignore_ascii_case(prefix) {
			return Some(*scale)
		}
	}
	None
}

pub(crate) const IEC_TABLE: [(&str, u64); 7] = [
	("Ei", 1024_u64.pow(6)),
	("Pi", 1024_u64.pow(5)),
	("Ti", 1024_u64.pow(4)),
	("Gi", 1024_u64.pow(3)),
	("Mi", 1024_u64.pow(2)),
	("Ki", 1024_u64),
	("", 1_u64),
];

pub fn exact_duration(s: &str) -> anyhow::Result<Duration> {
	let (number_str, unit) = match s.find(|ch: char| !ch.is_numeric()) {
		Some(i) => (&s[..i], s[i..].trim()),
		None => (s, ""),
	};
	let scale =
		get_duration_scale(unit).ok_or_else(|| anyhow!("Invalid duration unit {unit:?}"))?;
	let number: u64 = number_str.trim().parse()?;
	let nanos = number
		.checked_mul(scale)
		.ok_or_else(|| anyhow!("The duration is too big: {number} * {scale}"))?;
	Ok(Duration::from_nanos(nanos))
}

fn get_duration_scale(their_unit: &str) -> Option<u64> {
	for (our_unit, scale) in DURATION_TABLE.iter() {
		if our_unit.eq_ignore_ascii_case(their_unit) {
			return Some(*scale)
		}
	}
	None
}

pub(crate) const DURATION_TABLE: [(&str, u64); 9] = [
	("w", 60 * 60 * 24 * 7 * 1_000_000_000),
	("d", 60 * 60 * 24 * 1_000_000_000),
	("h", 60 * 60 * 1_000_000_000),
	("m", 60 * 1_000_000_000),
	("s", 1_000_000_000),
	("", 1_000_000_000),
	("ms", 1_000_000),
	("µs", 1_000),
	("ns", 1),
];

/// Hex-encoded 32-byte Ed25519 public key.
pub fn ed25519_public_key(s: &str) -> anyhow::Result<[u8; 32]> {
	let s = s.trim();
	hex::FromHex::from_hex(s).map_err(|e| anyhow!("Invalid Ed25519 public key {s:?}: {e}"))
}

/// Parse comma-separated list of cores or core ranges: "1,2,4-6,10". Ranges are inclusive.
///
/// NOTE This function doesn't check if core indices are correct since the protocol parameters might
/// haven't been initialized at the time of command line arguments parsing.
pub fn cores(s: &str) -> anyhow::Result<VecSet<CoreIndex>> {
	let error = || {
		anyhow!(
			"Expects comma-separated list of cores or core ranges, e.g. \"1,2,4-6,10\". \
            Ranges are inclusive."
		)
	};
	let mut cores = VecSet::new();
	for core_spec in s.split(',') {
		let range = if core_spec.contains('-') {
			let mut iter = core_spec.splitn(2, '-');
			let min_core: CoreIndex =
				iter.next().ok_or_else(error)?.parse().map_err(|_| error())?;
			let max_core: CoreIndex =
				iter.next().ok_or_else(error)?.parse().map_err(|_| error())?;
			min_core..=max_core
		} else {
			let core: CoreIndex = core_spec.parse()?;
			core..=core
		};
		for core in range {
			cores.insert(core);
		}
	}
	Ok(cores)
}

pub fn import_spec(input: &str) -> anyhow::Result<jam_types::ImportSpec> {
	let mut parts = input.splitn(3, ':');
	let malformed = anyhow!("import must be of form \"target_prefix:hash:index\"");
	let Some(prefix) = parts.next() else {
		return Err(malformed);
	};
	let is_segment_root = if prefix == "wp" {
		false
	} else if prefix == "sr" {
		true
	} else {
		return Err(anyhow!(
			"invalid prefix for imports, must be `sr` (segment root) or `wp` (work package)"
		))
	};
	let Some(hash_str) = parts.next() else {
		return Err(malformed);
	};
	let hash: [u8; 32] = decode_hex(hash_str)?;
	let Some(index_str) = parts.next() else {
		return Err(malformed);
	};
	let index = index_str.parse::<u16>()?;
	Ok(if is_segment_root {
		jam_types::ImportSpec {
			root: jam_types::RootIdentifier::Direct(jam_types::SegmentTreeRoot(hash)),
			index,
		}
	} else {
		jam_types::ImportSpec {
			root: jam_types::RootIdentifier::Indirect(jam_types::WorkPackageHash(hash)),
			index,
		}
	})
}

/// Decode a hex string to a fixed-size byte array, accepting the length parameter
fn decode_hex<const N: usize>(s: &str) -> anyhow::Result<[u8; N]> {
	let s = s.strip_prefix("0x").unwrap_or(s);
	if s.len() != N * 2 {
		return Err(anyhow!("Invalid hex length: expected {} characters, got {}", N * 2, s.len()))
	}
	let mut result = [0u8; N];
	for i in 0..N {
		let byte_str = &s[i * 2..(i + 1) * 2];
		result[i] = u8::from_str_radix(byte_str, 16)
			.map_err(|_| anyhow!("Invalid hex character at position {}", i * 2))?;
	}
	Ok(result)
}

#[cfg(test)]
mod tests {
	use super::*;

	#[test]
	fn exact_bytes_works() {
		assert_eq!(1, exact_bytes("1B").unwrap());
		assert_eq!(1024, exact_bytes("1KiB").unwrap());
		assert_eq!(4 * 1024 * 1024, exact_bytes(" 4 mib ").unwrap());
	}

	#[test]
	fn exact_duration_works() {
		assert_eq!(Duration::from_secs(1), exact_duration("1 s").unwrap());
		assert_eq!(Duration::from_secs(60), exact_duration("1m").unwrap());
		assert_eq!(Duration::from_secs(60 * 60), exact_duration("1h").unwrap());
		assert_eq!(Duration::from_secs(60 * 60 * 24), exact_duration("1d").unwrap());
		assert_eq!(Duration::from_secs(60 * 60 * 24 * 7), exact_duration("1w").unwrap());
	}

	#[test]
	fn cores_work() {
		assert_eq!(&[1, 2, 4, 5, 6, 10], cores("1,2,4-6,10").unwrap().as_ref());
		assert_eq!(&[1, 2, 3, 4, 5, 6, 7, 8, 9, 10], cores("1-4,2-10").unwrap().as_ref());
	}

	#[test]
	fn ed25519_public_key_works() {
		assert_eq!(
			[0x11_u8; 32],
			ed25519_public_key("1111111111111111111111111111111111111111111111111111111111111111")
				.unwrap()
		);
	}
}