//! Base32 and base64, in the two exact dialects atproto uses.
//!
//! Both are here rather than taken from a crate because both are needed in
//! their *strict* form, and a general-purpose codec is lenient by default: it
//! accepts padding where the protocol forbids it, and it accepts a final
//! character whose unused low bits are set, which decodes fine and re-encodes
//! to a different string. Either would break the one property this whole
//! crate exists to hold — that a value has exactly one encoding — and neither
//! is something a caller can switch off after the fact.
//!
//! The two dialects:
//!
//! * base32, lowercase, no padding, RFC 4648 alphabet. This is multibase `b`,
//!   which is the only string form atproto permits for a CID.
//! * base64, standard alphabet (`+` and `/`, not the URL-safe pair), no
//!   padding. This is the `$bytes` representation in the JSON form of the
//!   data model.

/// Why a string is not valid in one of the two alphabets.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum DecodeError {
    /// A character outside the alphabet, including `=`: padding is refused
    /// rather than tolerated, because the protocol specifies unpadded and a
    /// padded string is a second encoding of the same bytes.
    #[error("invalid character {0:?}")]
    InvalidCharacter(char),
    /// A length that no whole number of bytes can produce, such as a single
    /// leftover base64 character. The value is how many bits were left over.
    #[error("truncated: {0} leftover bits encode no whole byte")]
    Truncated(u32),
    /// The final character carries bits that the decoded bytes do not use.
    /// Decoding would silently discard them and re-encoding would produce a
    /// different string, so it is refused.
    #[error("the final character has {0} non-zero unused bits")]
    UnusedBits(u32),
}

/// The lowercase RFC 4648 base32 alphabet, which is multibase `b`.
const BASE32: &[u8; 32] = b"abcdefghijklmnopqrstuvwxyz234567";

/// The standard RFC 4648 base64 alphabet: `+` and `/`, not the URL-safe pair.
const BASE64: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

/// Decodes a bit-packed string in `alphabet`, `bits` bits per character.
///
/// One routine covers both bases because they differ only in the alphabet and
/// the group size. The two strictness rules — no padding, no set bits beyond
/// the last whole byte — are the same rules in both.
fn decode(text: &str, alphabet: &[u8], bits: u32) -> Result<Vec<u8>, DecodeError> {
    let mut out = Vec::with_capacity(text.len() * bits as usize / 8);
    let mut buffer: u32 = 0;
    let mut held: u32 = 0;
    for ch in text.chars() {
        let value = alphabet
            .iter()
            .position(|c| char::from(*c) == ch)
            .ok_or(DecodeError::InvalidCharacter(ch))?;
        buffer = (buffer << bits) | value as u32;
        held += bits;
        if held >= 8 {
            held -= 8;
            out.push((buffer >> held) as u8);
        }
    }
    // What is left in `buffer` is the tail of the last character. It must be
    // zero, and there must not be enough of it to have been a whole byte.
    if held >= bits {
        return Err(DecodeError::Truncated(held));
    }
    let leftover = buffer & ((1 << held) - 1);
    if leftover != 0 {
        return Err(DecodeError::UnusedBits(leftover.count_ones()));
    }
    Ok(out)
}

/// Encodes bytes bit-packed into `alphabet`, `bits` bits per character.
fn encode(bytes: &[u8], alphabet: &[u8], bits: u32) -> String {
    let mut out = String::with_capacity(bytes.len() * 8 / bits as usize + 1);
    let mut buffer: u32 = 0;
    let mut held: u32 = 0;
    for byte in bytes {
        buffer = (buffer << 8) | u32::from(*byte);
        held += 8;
        while held >= bits {
            held -= bits;
            out.push(alphabet[((buffer >> held) & ((1 << bits) - 1)) as usize] as char);
        }
    }
    if held > 0 {
        let index = (buffer << (bits - held)) & ((1 << bits) - 1);
        out.push(alphabet[index as usize] as char);
    }
    out
}

/// Decodes lowercase unpadded base32.
pub fn base32_decode(text: &str) -> Result<Vec<u8>, DecodeError> {
    decode(text, BASE32, 5)
}

/// Encodes lowercase unpadded base32.
pub fn base32_encode(bytes: &[u8]) -> String {
    encode(bytes, BASE32, 5)
}

/// Decodes standard unpadded base64.
pub fn base64_decode(text: &str) -> Result<Vec<u8>, DecodeError> {
    decode(text, BASE64, 6)
}

/// Encodes standard unpadded base64.
pub fn base64_encode(bytes: &[u8]) -> String {
    encode(bytes, BASE64, 6)
}

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

    #[test]
    fn base64_round_trips_every_tail_length() {
        for len in 0..8usize {
            let bytes: Vec<u8> = (0..len)
                .map(|i| (i as u8).wrapping_mul(37).wrapping_add(9))
                .collect();
            let text = base64_encode(&bytes);
            assert!(!text.contains('='), "{text:?} is padded");
            assert_eq!(base64_decode(&text), Ok(bytes));
        }
    }

    #[test]
    fn base32_round_trips_every_tail_length() {
        for len in 0..8usize {
            let bytes: Vec<u8> = (0..len)
                .map(|i| (i as u8).wrapping_mul(53).wrapping_add(3))
                .collect();
            let text = base32_encode(&bytes);
            assert_eq!(base32_decode(&text), Ok(bytes));
        }
    }

    #[test]
    fn padding_is_a_character_outside_the_alphabet() {
        assert_eq!(
            base64_decode("YWJjZA=="),
            Err(DecodeError::InvalidCharacter('='))
        );
    }

    #[test]
    fn a_set_bit_beyond_the_last_byte_is_refused() {
        // "QQ" decodes to one byte and leaves four bits; "QR" sets one of
        // them. A lenient decoder returns the same byte for both, which is
        // two strings for one value.
        assert_eq!(base64_decode("QQ"), Ok(vec![0x41]));
        assert_eq!(base64_decode("QR"), Err(DecodeError::UnusedBits(1)));
    }

    #[test]
    fn a_lone_trailing_character_encodes_no_byte() {
        assert_eq!(base64_decode("QQQQQ"), Err(DecodeError::Truncated(6)));
    }

    #[test]
    fn the_url_safe_alphabet_is_not_the_one_atproto_uses() {
        assert_eq!(base64_decode("-_"), Err(DecodeError::InvalidCharacter('-')));
    }

    #[test]
    fn uppercase_base32_is_a_different_multibase() {
        assert_eq!(
            base32_decode("MFRGG"),
            Err(DecodeError::InvalidCharacter('M'))
        );
    }
}