use std::collections::HashMap;

extern crate rusty_dsp;

use rusty_dsp::error::DspError;

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 struct Encoder<'a> {
    shift_state: ShiftState,
    buffer: Option<u8>,
    iterator: Box<dyn Iterator<Item = u8> + 'a>,
}

impl<'a> Encoder<'a> {
    pub fn from<I>(iterator: I) -> Encoder<'a>
    where
        I: Iterator<Item = u8> + 'a,
    {
        Encoder {
            shift_state: ShiftState::None,
            buffer: None,
            iterator: Box::new(iterator),
        }
    }
}

impl<'a> Iterator for Encoder<'a> {
    type Item = Result<u8, DspError>;

    fn next(&mut self) -> Option<Self::Item> {
        if let Some(c) = self.buffer {
            self.buffer = None;
            return Some(Ok(c));
        }

        let byte = match self.iterator.next() {
            Some(b) => b,
            None => return None,
        };

        let letter = code_from_letter(&byte);
        let figure = code_from_figure(&byte);

        match (letter, figure, self.shift_state) {
            (Some(l), Some(_), ShiftState::None) => {
                self.shift_state = ShiftState::Letters;
                self.buffer = Some(l);

                Some(Ok(SHIFT_LETTERS))
            }
            (Some(l), Some(_), _) => Some(Ok(l)),
            (Some(l), None, ShiftState::Letters) => Some(Ok(l)),
            (Some(l), None, _) => {
                self.shift_state = ShiftState::Letters;
                self.buffer = Some(l);

                Some(Ok(SHIFT_LETTERS))
            }
            (None, Some(f), ShiftState::Figures) => Some(Ok(f)),
            (None, Some(f), _) => {
                self.shift_state = ShiftState::Figures;
                self.buffer = Some(f);

                Some(Ok(SHIFT_FIGURES))
            }
            (None, None, _) => Some(Err(DspError::Error(format!(
                "Unable to decode symbol: {:?}",
                &byte
            )))),
        }
    }
}

pub fn encode(msg: &str) -> Result<Vec<u8>, DspError> {
    let encoder = Encoder::from(msg.to_ascii_uppercase().into_bytes().into_iter());
    encoder.collect()
}

pub struct Decoder<'a> {
    shift_state: ShiftState,
    iterator: Box<dyn Iterator<Item = u8> + 'a>,
}

impl<'a> Decoder<'a> {
    pub fn from<I>(iterator: I) -> Decoder<'a>
    where
        I: Iterator<Item = u8> + 'a,
    {
        Decoder {
            shift_state: ShiftState::Letters,
            iterator: Box::new(iterator),
        }
    }
}

impl<'a> Iterator for Decoder<'a> {
    // Option<Result<u8, u8>> for nice error handling if somebody wants to use collect
    // See https://doc.rust-lang.org/stable/rust-by-example/error/iter_result.html
    type Item = Result<u8, DspError>;

    fn next(&mut self) -> Option<Self::Item> {
        let mut byte = match self.iterator.next() {
            Some(b) => b,
            None => return None,
        };

        match byte {
            SHIFT_LETTERS => {
                self.shift_state = ShiftState::Letters;
                byte = match self.iterator.next() {
                    Some(b) => b,
                    None => return None,
                };
            }
            SHIFT_FIGURES => {
                self.shift_state = ShiftState::Figures;
                byte = match self.iterator.next() {
                    Some(b) => b,
                    None => return None,
                };
            }
            _ => {}
        };

        match (byte, self.shift_state) {
            (c, ShiftState::Letters) => match letter_from_code(c) {
                Some(l) => Some(Ok(l)),
                None => Some(Err(DspError::Error(format!(
                    "Unable to decode symbol: {:?}",
                    c
                )))),
            },
            (c, ShiftState::Figures) => match figure_from_code(c) {
                Some(f) => Some(Ok(f)),
                None => Some(Err(DspError::Error(format!(
                    "Unable to decode symbol: {:?}",
                    c
                )))),
            },
            (_, ShiftState::None) => panic!("Shift state is expectedly None"),
        }
    }
}

pub fn decode(msg: &[u8]) -> Result<String, DspError> {
    let decoder = Decoder::from(msg.iter().copied());
    match decoder.collect() {
        Ok(v) => Ok(String::from_utf8(v).unwrap()),
        Err(c) => Err(c),
    }
}

#[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
        )
    }
}