use sha2::{Digest, Sha256};

pub enum PubkeyDecodeError {
	BadChecksum,
	BadFormat,
}

/// Decode pubkey from base 58
/// Verify checksum if provided
/// Accepts any checksum length >=3
pub fn format_pubkey(raw: &str) -> Result<String, PubkeyDecodeError> {
	let mut iter = raw.splitn(2, ':');
	let raw_pubkey = iter.next().ok_or(PubkeyDecodeError::BadFormat)?;
	if raw_pubkey.len() < 43 || raw_pubkey.len() > 44 {
		return Err(PubkeyDecodeError::BadFormat);
	}
	let mut pubkey = [0u8; 32];
	bs58::decode(raw_pubkey)
		.into(&mut pubkey)
		.map_err(|_| PubkeyDecodeError::BadFormat)?;

	let mut hasher1 = Sha256::new();
	hasher1.update(&pubkey.as_ref()[..32]);
	let mut hasher2 = Sha256::new();
	hasher2.update(hasher1.finalize());
	let checksum = bs58::encode(hasher2.finalize()).into_string();

	if let Some(given_checksum) = iter.next() {
		if given_checksum.len() < 3 {
			return Err(PubkeyDecodeError::BadFormat);
		}
		if !checksum.starts_with(given_checksum) {
			return Err(PubkeyDecodeError::BadChecksum);
		}
	}
	Ok(format!(
		"{}:{}",
		bs58::encode(pubkey).into_string(),
		&checksum[..3]
	))
}