use hmac::{Hmac, Mac};

pub enum Algorithm {
    Sha1,
    Sha256,
    Sha512,
    Steam,
}

const STEAM_CHARS: &[u8] = b"23456789BCDFGHJKMNPQRTVWXY";

pub fn decode_base32(input: &str) -> Option<Vec<u8>> {
    let cleaned: Vec<u8> = input
        .trim()
        .bytes()
        .filter(|b| !b.is_ascii_whitespace())
        .collect();
    let trimmed: &[u8] = {
        let mut end = cleaned.len();
        while end > 0 && cleaned[end - 1] == b'=' {
            end -= 1;
        }
        &cleaned[..end]
    };
    let mut out = Vec::with_capacity(trimmed.len() * 5 / 8);
    let mut buffer: u16 = 0;
    let mut bits: u8 = 0;
    for &b in trimmed {
        let v: u16 = match b {
            b'A'..=b'Z' => u16::from(b - b'A'),
            b'a'..=b'z' => u16::from(b - b'a'),
            b'2'..=b'7' => u16::from(b - b'2') + 26,
            _ => return None,
        };
        buffer = (buffer << 5) | v;
        bits += 5;
        if bits >= 8 {
            bits -= 8;
            let byte = u8::try_from((buffer >> bits) & 0xff).ok()?;
            out.push(byte);
            buffer &= (1_u16 << bits).wrapping_sub(1);
        }
    }
    Some(out)
}

pub fn generate(
    secret: &[u8],
    unix_secs: u64,
    step: u64,
    digits: u32,
    algorithm: &Algorithm,
) -> crate::error::Result<String> {
    let totp_err = |msg: &str| crate::error::Error::Totp {
        msg: msg.to_string(),
    };
    let counter = (unix_secs / step).to_be_bytes();
    let mac: Vec<u8> =
        match algorithm {
            Algorithm::Sha1 | Algorithm::Steam => {
                let mut m = Hmac::<sha1::Sha1>::new_from_slice(secret)
                    .map_err(|e| crate::error::Error::Totp {
                        msg: format!("invalid hmac key: {e}"),
                    })?;
                m.update(&counter);
                m.finalize().into_bytes().to_vec()
            }
            Algorithm::Sha256 => {
                let mut m = Hmac::<sha2::Sha256>::new_from_slice(secret)
                    .map_err(|e| crate::error::Error::Totp {
                        msg: format!("invalid hmac key: {e}"),
                    })?;
                m.update(&counter);
                m.finalize().into_bytes().to_vec()
            }
            Algorithm::Sha512 => {
                let mut m = Hmac::<sha2::Sha512>::new_from_slice(secret)
                    .map_err(|e| crate::error::Error::Totp {
                        msg: format!("invalid hmac key: {e}"),
                    })?;
                m.update(&counter);
                m.finalize().into_bytes().to_vec()
            }
        };

    let offset = usize::from(
        *mac.last().ok_or_else(|| totp_err("empty hmac output"))? & 0x0f,
    );
    let mut truncated = u32::from_be_bytes(
        mac[offset..offset + 4]
            .try_into()
            .map_err(|_| totp_err("totp truncation failed"))?,
    ) & 0x7fff_ffff;

    let digits_usize = usize::try_from(digits)
        .map_err(|_| totp_err("digits out of range"))?;

    match algorithm {
        Algorithm::Sha1 | Algorithm::Sha256 | Algorithm::Sha512 => {
            let modulus = 10_u32
                .checked_pow(digits)
                .ok_or_else(|| totp_err("digits too large"))?;
            Ok(format!(
                "{:0width$}",
                truncated % modulus,
                width = digits_usize
            ))
        }
        Algorithm::Steam => {
            let len = u32::try_from(STEAM_CHARS.len())
                .map_err(|_| totp_err("steam alphabet too large"))?;
            let mut s = String::with_capacity(digits_usize);
            for _ in 0..digits {
                let idx = usize::try_from(truncated % len)
                    .map_err(|_| totp_err("steam index error"))?;
                s.push(char::from(STEAM_CHARS[idx]));
                truncated /= len;
            }
            Ok(s)
        }
    }
}

#[cfg(test)]
#[allow(clippy::as_conversions, clippy::cast_possible_truncation)]
mod test {
    use super::{decode_base32, generate, Algorithm, STEAM_CHARS};

    #[test]
    fn test_decode_base32_basic() {
        assert_eq!(decode_base32("MZXW6===").unwrap(), b"foo");
        assert_eq!(decode_base32("MZXW6").unwrap(), b"foo");
        assert_eq!(decode_base32("mzxw6").unwrap(), b"foo");
        assert_eq!(decode_base32("MZ XW 6").unwrap(), b"foo");
        assert_eq!(decode_base32("MZXW6YQ=").unwrap(), b"foob");
        assert_eq!(decode_base32("MZXW6YTB").unwrap(), b"fooba");
        assert_eq!(decode_base32("MZXW6YTBOI======").unwrap(), b"foobar");
        assert!(decode_base32("!!!").is_none());
    }

    #[test]
    fn test_rfc6238_sha1() {
        let secret = b"12345678901234567890";
        let cases = [
            (59_u64, "94287082"),
            (1_111_111_109, "07081804"),
            (1_111_111_111, "14050471"),
            (1_234_567_890, "89005924"),
            (2_000_000_000, "69279037"),
            (20_000_000_000, "65353130"),
        ];
        for (t, expected) in cases {
            let code = generate(secret, t, 30, 8, &Algorithm::Sha1).unwrap();
            assert_eq!(code, expected, "time={t}");
        }
    }

    #[test]
    fn test_rfc6238_sha256() {
        let secret = b"12345678901234567890123456789012";
        let code = generate(secret, 59, 30, 8, &Algorithm::Sha256).unwrap();
        assert_eq!(code, "46119246");
    }

    #[test]
    fn test_rfc6238_sha512() {
        let secret =
            b"1234567890123456789012345678901234567890123456789012345678901234";
        let code = generate(secret, 59, 30, 8, &Algorithm::Sha512).unwrap();
        assert_eq!(code, "90693936");
    }

    #[test]
    fn test_digits_six() {
        let secret = b"12345678901234567890";
        let code = generate(secret, 59, 30, 6, &Algorithm::Sha1).unwrap();
        assert_eq!(code, "287082");
    }

    #[test]
    fn test_steam() {
        let secret = decode_base32("STEAMKEY234567").unwrap();
        let code = generate(&secret, 1_000_000_000, 30, 5, &Algorithm::Steam)
            .unwrap();
        assert_eq!(code.len(), 5);
        for c in code.chars() {
            assert!(STEAM_CHARS.contains(&(c as u8)));
        }
        let pinned =
            generate(b"12345678901234567890", 59, 30, 5, &Algorithm::Steam)
                .unwrap();
        assert_eq!(pinned, "PV9M4");
    }
}
