use std::collections::HashMap;

/// Given a Vec<u8> give a base64 encoded string.
pub fn encode(v: &Vec<u8>) -> String {
    if v.len() == 0 {
        return String::from("");
    }

    let lookup_table = get_encode_lookup(&String::from("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"));
    let pads: usize = get_number_of_pads(&v);

    let mut buffer: Vec<u8> = vec![0; v.len() + pads];
    let mut output: String = String::from("");

    for i in 0..v.len() {
        buffer[i] = v[i];
    }

    let mut i = 0;
    while i < buffer.len() {
        let mut number: usize = 0;

        number += buffer[i] as usize;
        number = number << 8;
        number += buffer[i+1] as usize;
        number = number << 8;
        number += buffer[i+2] as usize;

        output.push_str(&get_segments(
            &lookup_table,
            &buffer.len(),
            number,
            &i,
            &pads));

        i = i + 3;
    }

    return output;
}

/// Given a base64 encoded string, give a Vec<u8>.
pub fn decode(base64_string: &String) -> Vec<u8> {
    if base64_string == "" {
        let empty_vec: Vec<u8> = vec![];
        return empty_vec;
    }

    let decode_lookup: HashMap<char, usize> = get_decode_lookup(&String::from("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"));
    let base64_vector: Vec<char> = base64_string.chars().collect();
    let mut output_vector: Vec<u8> = vec![];
    let pad_count: usize = get_pad_count(&base64_vector);

    for c in &base64_vector {
        print!("{}, ", &decode_lookup[c]);
    }
    println!("");


    let mut i: usize = 0;
    while i < base64_vector.len() {
        let mut number: usize = 0;

        number += decode_lookup[&base64_vector[i]];
        number = number << 6;
        number += decode_lookup[&base64_vector[i+1]];
        number = number << 6;
        number += decode_lookup[&base64_vector[i+2]];
        number = number << 6;
        number += decode_lookup[&base64_vector[i+3]];

        if i != base64_vector.len() - 4 {
          output_vector.push(((number & 0xff0000) >> 16) as u8);
          output_vector.push(((number & 0x00ff00) >> 8) as u8);
          output_vector.push((number & 0x0000ff) as u8);
        } else {
          if pad_count == 0 {
            output_vector.push(((number & 0xff0000) >> 16) as u8);
            output_vector.push(((number & 0x00ff00) >> 8) as u8);
            output_vector.push((number & 0x0000ff) as u8);
          } else if pad_count == 1 {
            output_vector.push(((number & 0xff0000) >> 16) as u8);
            output_vector.push(((number & 0x00ff00) >> 8) as u8);
          } else if pad_count == 2 {
            output_vector.push(((number & 0xff0000) >> 16) as u8);
          } else {
            panic!("Invalid pad_count");
          }
        }

        i = i + 4;
    }

    return output_vector;
}

fn get_encode_lookup(s: &String) -> Vec<char> {
    s.chars().collect()
}

fn get_decode_lookup(s: &String) -> HashMap<char, usize> {
    let v: Vec<char> = s.chars().collect();
    let mut i: usize = 0;
    let mut reverse_lookup = HashMap::new();

    for c in v {
        reverse_lookup.insert(c, i);
        i = i + 1;
    }

    reverse_lookup.insert('=', 0);

    return reverse_lookup;
}

fn get_number_of_pads(bytes: &Vec<u8>) -> usize {

    if bytes.len() % 3 != 0 {
        return 3 - (bytes.len() % 3);
    }

    return 0;
}

fn get_pad_count(bytes: &Vec<char>) -> usize {
    let mut pad_count = 0;
    let mut i: usize = 0;

    while i < bytes.len() {
        if bytes[i] == '=' {
            pad_count += 1;
        }
        i = i + 1;
    }

    return pad_count;
}

fn get_segments(lookup_table: &Vec<char>, data_length: &usize, number: usize, i: &usize, pads: &usize) -> String {

    let mut output: String = String::from("");
    let mut segment: usize;
    let mask: usize = 0x3f;

    if *i == *data_length - 3 && *pads != 0 {
        if *pads == 1 {
            segment = (number >> 18) & mask;
            output.push(lookup_table[segment]);
            segment = (number >> 12) & mask;
            output.push(lookup_table[segment]);
            segment = (number >> 6) & mask;
            output.push(lookup_table[segment]);
            output.push('=');
        } else if *pads == 2 {
            segment = (number >> 18) & mask;
            output.push(lookup_table[segment]);
            segment = (number >> 12) & mask;
            output.push(lookup_table[segment]);
            output.push('=');
            output.push('=');
        } else {
            panic!("pads had invalid value???");
        }
    } else {
        segment = (number >> 18) & mask;
        output.push(lookup_table[segment]);
        segment = (number >> 12) & mask;
        output.push(lookup_table[segment]);
        segment = (number >> 6) & mask;
        output.push(lookup_table[segment]);
        segment = number & mask;
        output.push(lookup_table[segment]);
    }

    return output;
}

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

    fn get_plain_strings() -> Vec<String> {
        let plain_strings: Vec<String> = vec![
            String::from("Foo Bar"),
            String::from("This string\nhas newlines and\ttabs"),
            String::from("mañana"),
            String::from("Iñtërnâtiônàlizætiøn☃💩"),
            String::from("🇺🇸🇺🇸")
        ];

        return plain_strings;
    }

    fn get_base64_strings() -> Vec<String> {
        let base64_strings: Vec<String> = vec![
            String::from("Rm9vIEJhcg=="),
            String::from("VGhpcyBzdHJpbmcKaGFzIG5ld2xpbmVzIGFuZAl0YWJz"),
            String::from("bWHDsWFuYQ=="),
            String::from("ScOxdMOrcm7DonRpw7Ruw6BsaXrDpnRpw7hu4piD8J+SqQ=="),
            String::from("8J+HuvCfh7jwn4e68J+HuA==")
        ];

        return base64_strings;
    }

    // encode() tests.
    #[test]
    fn should_encode_empty_vector_correctly() {
        let actual = encode(&vec![]);
        assert!(actual == String::from(""));
    }

    #[test]
    fn should_encode_plain_strings_correctly() {
        let plain_strings = get_plain_strings();
        let base64_strings = get_base64_strings();

        let mut i: usize = 0;
        while i < plain_strings.len() {
            let test_vector = &plain_strings[i].as_bytes().to_vec();
            let actual = encode(&test_vector);
            assert_eq!(actual, base64_strings[i]);
            i = i + 1;
        }
    }

    // decode() tests.
    #[test]
    fn decode_empty_string_should_return_empty_vector() {
        let actual: Vec<u8> = decode(&String::from(""));
        assert_eq!(actual, []);
    }

    #[test]
    fn decode_should_decode_base64_strings_correctly() {
        let base64_strings = get_base64_strings();
        let plain_strings = get_plain_strings();

        let mut i: usize = 0;
        while i < base64_strings.len() {
            let test_string = &base64_strings[i];
            let actual: String = match String::from_utf8(decode(&test_string)) {
                Ok(s) => s,
                Err(e) => panic!("Invalid utf-8 sequence: {:?}", e),
            };

            assert_eq!(actual, plain_strings[i]);
            i = i + 1;
        }
    }

    // Helper function tests.
    #[test]
    fn get_number_of_pads_should_return_zero() {
        let test_vec = String::from("Foo").into_bytes();
        assert!(get_number_of_pads(&test_vec) == 0);
    }

    #[test]
    fn get_number_of_pads_should_return_one() {
        let test_vec = String::from("Fooox").into_bytes();
        assert!(get_number_of_pads(&test_vec) == 1);
    }

    #[test]
    fn get_number_of_pads_should_return_two() {
        let test_vec = String::from("Foo Bar").into_bytes();
        assert!(get_number_of_pads(&test_vec) == 2);
    }

    #[test]
    fn get_pad_count_should_return_zero() {
        let input: Vec<char> = String::from("FooBar").chars().collect();
        let actual: usize = get_pad_count(&input);
        assert_eq!(actual, 0);
    }

    #[test]
    fn get_pad_count_should_return_one() {
        let input: Vec<char> = String::from("FooBar=").chars().collect();
        let actual: usize = get_pad_count(&input);
        assert_eq!(actual, 1);
    }

    #[test]
    fn get_pad_count_should_return_two() {
        let input: Vec<char> = String::from("FooBar==").chars().collect();
        let actual: usize = get_pad_count(&input);
        assert_eq!(actual, 2);
    }
}