use std::iter::Peekable;
use std::str::Chars;

pub type Error = &'static str;

pub fn assemble(asm: &str) -> Result<Vec<u8>, Error> {
    Assembler::new(asm).assemble()
}

struct Assembler<'a> {
    iter: Peekable<Chars<'a>>,
    output: Vec<u8>,
}

enum IntPadding {
    NoPadding,
    PaddedTo(usize),
}

impl<'a> Assembler<'a> {
    fn new(buf: &'a str) -> Assembler<'a> {
        Assembler {
            iter: buf.chars().peekable(),
            output: Vec::new(),
        }
    }

    fn expect(&mut self, c: char) -> Option<()> {
        match self.iter.next() {
            Some(d) => {
                if c == d {
                    Some(())
                } else {
                    None
                }
            }
            None => None,
        }
    }

    fn emit(&mut self, s: &str) {
        self.output.extend_from_slice(s.as_bytes());
    }

    fn read_to(&mut self, any_of: &str) -> Option<String> {
        let mut s = String::new();
        loop {
            if let Some(c) = self.iter.peek()
                && any_of.contains(*c)
            {
                return Some(s);
            }

            match self.iter.next() {
                Some(c) => s.push(c),
                None => return None,
            }
        }
    }

    fn parse_int(s: &str) -> Option<usize> {
        let is_hex = s.starts_with("0x");
        let base = if is_hex { 16 } else { 10 };
        let mut ret = 0;
        let mut iter = s.chars();
        if is_hex {
            iter.nth(1);
        }

        for c in iter {
            ret *= base;
            ret += c.to_digit(base)?;
        }

        return Some(ret as usize);
    }

    fn encode_int(mut n: usize, p: IntPadding) -> String {
        let mut s = String::new();
        while n > 0 {
            if n % 2 == 1 {
                s += "C";
            } else {
                s += "I";
            }

            n /= 2;
        }

        if let IntPadding::PaddedTo(p) = p {
            for _ in s.len()..(p - 1) {
                s += "I";
            }
        }

        s += "P";
        return s;
    }

    fn emit_consts(&mut self, s: &str) {
        for c in s.chars() {
            match c {
                'I' => self.emit("C"),
                'C' => self.emit("F"),
                'F' => self.emit("P"),
                'P' => self.emit("IC"),
                _ => {}
            }
        }
    }

    fn assemble(mut self) -> Result<Vec<u8>, Error> {
        loop {
            match self.iter.next() {
                Some('I') => self.emit("C"),
                Some('C') => self.emit("F"),
                Some('F') => self.emit("P"),
                Some('P') => self.emit("IC"),
                Some('!') => {
                    self.expect('{').ok_or("Skip followed by non-{")?;
                    let n = self.read_to("}").ok_or("Skip failed to terminate")?;
                    let n = Self::parse_int(&n).ok_or("Invalid int inside skip")?;
                    self.emit("IP");
                    self.emit(&Self::encode_int(n, IntPadding::NoPadding));
                    self.expect('}').ok_or("Skip failed to terminate")?;
                }
                Some('?') => {
                    self.expect('{').ok_or("Search followed by non-?")?;
                    let s = self.read_to("}").ok_or("Search failed to terminate")?;
                    self.emit("IFF");
                    self.emit_consts(&s);
                    self.expect('}').ok_or("Search failed to terminate")?;
                }
                Some('(') => self.emit("IIP"),
                Some(')') => self.emit("IIC"),
                Some(';') => self.emit("IIC"),
                Some('{') => {
                    let n = self.read_to("},").ok_or("Reference failed terminate")?;
                    let n = Self::parse_int(&n).ok_or("Reference group index invalid")?;
                    let l = match self.iter.next() {
                        Some('}') => String::from("0"),
                        Some(',') => {
                            let l = self.read_to("}").ok_or("Reference failed to terminate")?;
                            self.expect('}');
                            l
                        }
                        _ => return Err("Reference failed to terminate"),
                    };
                    let l = Self::parse_int(&l).ok_or("Reference level invalid")?;
                    self.emit("IP");
                    self.emit(&Self::encode_int(l, IntPadding::NoPadding));
                    self.emit(&Self::encode_int(n, IntPadding::NoPadding));
                }
                Some('|') => {
                    let n = self.read_to("|").ok_or("Length failed to terminate")?;
                    let n = Self::parse_int(&n).ok_or("Length group index invalid")?;
                    self.emit("IIP");
                    self.emit(&Self::encode_int(n, IntPadding::NoPadding));
                    self.expect('|').ok_or("Length failed to terminate")?;
                }
                Some('n') => {
                    self.expect('{');
                    let n = self
                        .read_to("},")
                        .ok_or("Int constant failed to terminate")?;
                    let n = Self::parse_int(&n).ok_or("Int constant invalid")?;
                    match self.iter.next() {
                        Some('}') => {
                            self.emit(&Self::encode_int(n, IntPadding::NoPadding));
                        }
                        Some(',') => {
                            let pad = self
                                .read_to("}")
                                .ok_or("Int constant failed to terminate")?;
                            let pad =
                                Self::parse_int(&pad).ok_or("Int constant padding invalid")?;
                            self.emit_consts(&Self::encode_int(n, IntPadding::PaddedTo(pad)));
                            self.expect('}').ok_or("Int constant failed to terminate")?;
                        }
                        _ => return Err("Int constant failed to terminate"),
                    };
                }
                Some(' ') | Some('\n') => {}
                None => return Ok(self.output),
                Some(c) => {
                    println!("{}", c);
                    return Err("Invalid character");
                }
            }
        }
    }
}

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

    #[test]
    fn test_assemble_bases() {
        let asm = assemble(&"ICFP").unwrap();
        assert_eq!(b"CFPIC", &asm[..]);
    }

    #[test]
    fn test_assemble_skip() {
        let asm = assemble(&"!{1} !{0x10}").unwrap();
        assert_eq!(b"IPCPIPIIIICP", &asm[..]);
    }

    #[test]
    fn test_assemble_search() {
        let asm = assemble(&"?{ICFP}").unwrap();
        assert_eq!(b"IFFCFPIC", &asm[..]);
    }

    #[test]
    fn test_assemble_groups() {
        let asm = assemble(&"()()").unwrap();
        assert_eq!(b"IIPIICIIPIIC", &asm[..]);
    }

    #[test]
    fn test_assemble_pattern_template() {
        let asm = assemble(&"!{2} ; II ;").unwrap();
        assert_eq!(b"IPICPIICCCIIC", &asm[..]);
    }

    #[test]
    fn test_assemble_pattern_quote() {
        let asm = assemble(&"{2}").unwrap();
        assert_eq!(b"IPPICP", &asm[..]);

        let asm = assemble(&"{2,3}").unwrap();
        assert_eq!(b"IPCCPICP", &asm[..]);
    }

    #[test]
    fn test_assemble_pattern_length() {
        let asm = assemble(&"|3|").unwrap();
        assert_eq!(b"IIPCCP", &asm[..]);
    }

    #[test]
    fn test_assemble_pattern_constant() {
        let asm = assemble(&"n{4,5}").unwrap();
        assert_eq!(b"CCFCIC", &asm[..]);

        let asm = assemble(&"n{0x123,24}").unwrap();
        assert_eq!(b"FFCCCFCCFCCCCCCCCCCCCCCIC", &asm[..]);
    }
}