use std::collections::HashMap;

const LETTERS: &[u8] = b"\0E\nA SIU\rDRJNFCKTZLWHYPQOBG\0MXV\0";
const FIGURES: &[u8] = b"\03\n- \x0787\r\x054',!:(5\")2#6019?&\0./;\0";

lazy_static! {
    static ref INV_LETTERS: HashMap<u8, u8> = LETTERS
        .iter()
        .enumerate()
        .map(|(c, f)| (*f, c as u8))
        .collect();
    static ref INV_FIGURES: HashMap<u8, u8> = FIGURES
        .iter()
        .enumerate()
        .map(|(c, f)| (*f, c as u8))
        .collect();
}

fn letter_from_code(code: u8) -> Option<u8> {
    if code >= (LETTERS.len() as u8) {
        None
    } else {
        Some(LETTERS[code as usize])
    }
}

fn figure_from_code(code: u8) -> Option<u8> {
    if code >= (FIGURES.len() as u8) {
        None
    } else {
        Some(FIGURES[code as usize])
    }
}

fn code_from_letter(letter: &u8) -> Option<u8> {
    INV_LETTERS.get(letter).cloned()
}

fn code_from_figure(figure: &u8) -> Option<u8> {
    INV_FIGURES.get(figure).cloned()
}

const SHIFT_FIGURES: u8 = 0b11011u8;
const SHIFT_LETTERS: u8 = 0b11111u8;

#[derive(Copy, Clone)]
enum ShiftState {
    None,
    Letters,
    Figures,
}

pub fn encode(msg: &str) -> Result<Vec<u8>, u8> {
    let mut res = Vec::<u8>::new();

    let mut shift = ShiftState::None;

    for &byte in msg.to_ascii_uppercase().as_bytes().iter() {
        let letter = code_from_letter(&byte);
        let figure = code_from_figure(&byte);
        match (letter, figure, shift) {
            (Some(l), Some(_), ShiftState::None) => {
                res.push(SHIFT_LETTERS);
                shift = ShiftState::Letters;
                res.push(l);
            }
            (Some(l), Some(_), _) => res.push(l),
            (Some(l), None, ShiftState::Letters) => res.push(l),
            (Some(l), None, _) => {
                res.push(SHIFT_LETTERS);
                shift = ShiftState::Letters;
                res.push(l);
            }
            (None, Some(f), ShiftState::Figures) => res.push(f),
            (None, Some(f), _) => {
                res.push(SHIFT_FIGURES);
                shift = ShiftState::Figures;
                res.push(f);
            }
            (None, None, _) => return Err(byte),
        }
    }

    Ok(res)
}

pub fn decode(msg: &Vec<u8>) -> Result<String, u8> {
    let mut res = String::new();

    // Lets assume this is a sane starting point.
    let mut shift = ShiftState::Letters;

    for &code in msg.iter() {
        match (code, shift) {
            (SHIFT_LETTERS, _) => shift = ShiftState::Letters,
            (SHIFT_FIGURES, _) => shift = ShiftState::Figures,
            (c, ShiftState::Letters) => match letter_from_code(c) {
                Some(l) => res.push(l as char),
                None => return Err(c),
            },
            (c, ShiftState::Figures) => match figure_from_code(c) {
                Some(l) => res.push(l as char),
                None => return Err(c),
            },
            (_, ShiftState::None) => panic!("Shift state is expectedly None"),
        }
    }

    Ok(res)
}

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

    #[test]
    fn test_encode() {
        let encoded = encode("kaufen sie jede woche vier gute bequeme pelze xy 1234567890");
        assert_eq!(
            encoded,
            Ok(vec![
                31, 15, 3, 7, 13, 1, 12, 4, 5, 6, 1, 4, 11, 1, 9, 1, 4, 19, 24, 14, 20, 1, 4, 30,
                6, 1, 10, 4, 26, 7, 16, 1, 4, 25, 1, 23, 7, 1, 28, 1, 4, 22, 1, 18, 17, 1, 4, 29,
                21, 4, 27, 23, 19, 1, 10, 16, 21, 7, 6, 24, 22
            ])
        );
    }

    #[test]
    fn test_decode() {
        let kaufen_schleife: Vec<u8> = vec![
            31, 15, 3, 7, 13, 1, 12, 4, 5, 6, 1, 4, 11, 1, 9, 1, 4, 19, 24, 14, 20, 1, 4, 30, 6, 1,
            10, 4, 26, 7, 16, 1, 4, 25, 1, 23, 7, 1, 28, 1, 4, 22, 1, 18, 17, 1, 4, 29, 21, 4, 27,
            23, 19, 1, 10, 16, 21, 7, 6, 24, 22,
        ];

        let decoded = decode(&kaufen_schleife);
        assert_eq!(
            Ok(String::from(
                "KAUFEN SIE JEDE WOCHE VIER GUTE BEQUEME PELZE XY 1234567890"
            )),
            decoded
        )
    }
}