use crate::Day;

pub struct Day16(Packet);

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Debug)]
pub struct Packet {
    version: u32,
    type_id: PacketType,
    sub_packets: Vec<Packet>,
    value: u64,
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Clone, Debug)]
pub enum LengthType {
    Bits,
    Packets,
}


impl From<u8> for LengthType {
    fn from(a: u8) -> Self {
        use LengthType::{Bits, Packets};
        match a {
            0 => Bits,
            _ => Packets
        }
    }
}

#[derive(Ord, PartialOrd, Eq, PartialEq, Copy, Clone, Debug)]
enum PacketType {
    Sum,
    Product,
    Minimum,
    Maximum,
    Literal,
    Greater,
    Less,
    Eq,
}

impl From<u8> for PacketType {
    fn from(n: u8) -> Self {
        use PacketType::{Sum, Product, Minimum, Maximum, Literal, Greater, Less, Eq};
        match n {
            0 => Sum,
            1 => Product,
            2 => Minimum,
            3 => Maximum,
            4 => Literal,
            5 => Greater,
            6 => Less,
            7 => Eq,
            _ => Literal
        }
    }
}

fn hex_to_bytes(hex: char) -> [u8; 4] {
    match hex {
        '0' => [0, 0, 0, 0],
        '1' => [0, 0, 0, 1],
        '2' => [0, 0, 1, 0],
        '3' => [0, 0, 1, 1],
        '4' => [0, 1, 0, 0],
        '5' => [0, 1, 0, 1],
        '6' => [0, 1, 1, 0],
        '7' => [0, 1, 1, 1],
        '8' => [1, 0, 0, 0],
        '9' => [1, 0, 0, 1],
        'A' => [1, 0, 1, 0],
        'B' => [1, 0, 1, 1],
        'C' => [1, 1, 0, 0],
        'D' => [1, 1, 0, 1],
        'E' => [1, 1, 1, 0],
        'F' => [1, 1, 1, 1],
        _ => [0, 0, 0, 0],
    }
}

fn parse_str<T: AsRef<str>>(content: T) -> Vec<u8> {
    content.as_ref().chars().flat_map(hex_to_bytes).collect::<Vec<_>>()
}

impl Day for Day16 {
    fn init(content: String) -> anyhow::Result<Self> {
        let txt = parse_str(content);
        let mut iter = txt.iter().copied();
        Ok(Self(parse_packet(&mut iter).ok_or_else(|| anyhow::Error::msg("Could not parse packet!"))?))
    }

    fn part1(&self) -> anyhow::Result<String> {
        fn calc(packet: &Packet) -> u32 {
            (packet.version) + if packet.type_id == PacketType::Literal {
                0
            } else {
                packet.sub_packets.iter().map(calc).sum()
            }
        }
        Ok(format!("{}", calc(&self.0)))
    }

    fn part2(&self) -> anyhow::Result<String> {
        fn calc(packet: &Packet) -> u64 {
            match packet.type_id {
                PacketType::Sum => packet.sub_packets.iter().map(calc).sum(),
                PacketType::Product => packet.sub_packets.iter().map(calc).product(),
                PacketType::Minimum => packet.sub_packets.iter().map(calc).min().unwrap(),
                PacketType::Maximum => packet.sub_packets.iter().map(calc).max().unwrap(),
                PacketType::Literal => packet.value,
                PacketType::Greater => (calc(&packet.sub_packets[0]) > calc(&packet.sub_packets[1])) as u64,
                PacketType::Less => (calc(&packet.sub_packets[0]) < calc(&packet.sub_packets[1])) as u64,
                PacketType::Eq => (calc(&packet.sub_packets[0]) == calc(&packet.sub_packets[1])) as u64,
            }
        }
        Ok(format!("{}", calc(&self.0)))
    }
}

fn parse_packet<T: Iterator<Item=u8>>(it: &mut T) -> Option<Packet> {
    let version = it.take(3).enumerate().fold(0, |a, (idx, num)| a | (u32::from(num) << (2 - idx)));
    let type_id = it.take(3).enumerate().fold(0, |a, (idx, num)| a | (num << (2 - idx))).into();


    Some(Packet {
        version,
        type_id,
        sub_packets: if type_id == PacketType::Literal { Vec::new() } else { decode_other(it)? },
        value: if type_id == PacketType::Literal { decode_literal(it) } else { 0 },
    })
}

fn decode_other<T: Iterator<Item=u8>>(it: &mut T) -> Option<Vec<Packet>> {
    let t: LengthType = it.next()?.into();
    return match t {
        LengthType::Bits => {
            let l = it.take(15).fold(0, |a, b| (a << 1) | (b as usize));
            let to_parse = it.take(l).collect::<Vec<_>>();
            let mut new_iter = to_parse.iter().copied().peekable();
            let mut a = Vec::new();
            while new_iter.peek().is_some() {
                a.push(parse_packet(&mut new_iter)?);
            }
            Some(a)
        }
        LengthType::Packets => {
            let l = it.take(11).fold(0, |a, b| (a << 1) | (b as usize));
            let mut packets = Vec::new();
            for _ in 0..l {
                let p = parse_packet(it)?;
                packets.push(p);
            }
            Some(packets)
        }
    };
}

fn decode_literal<T: Iterator<Item=u8>>(it: &mut T) -> u64 {
    let mut z = it.take(5).collect::<Vec<_>>();
    let mut n = 0;
    loop {
        for a in 0..4 {
            n <<= 1;
            n += u64::from(z[a + 1]);
        }

        if z[0] == 0 {
            break;
        }
        z = it.take(5).collect::<Vec<_>>();
    }
    n
}

#[cfg(test)]
mod tests {
    use crate::day16::{Day16, PacketType, parse_packet, parse_str};
    use crate::day::Day;
    use crate::day_tests;

    #[test]
    fn test_literal() {
        let inp = parse_str("D2FE28");
        let mut i = inp.iter().copied();
        if let Some(r) = parse_packet(&mut i) {
            if r.type_id == PacketType::Literal {
                assert_eq!(r.value, 2021);
            } else {
                assert!(false);
            }
        } else {
            assert!(false);
        }
    }

    #[test]
    fn test_packet_bits() {
        let inp = parse_str("38006F45291200");
        let mut i = inp.iter().copied();
        if let Some(r) = parse_packet(&mut i) {
            if r.type_id != PacketType::Literal {
                let p = &r.sub_packets[0];
                assert_eq!(p.value, 10);
                let p = &r.sub_packets[1];
                assert_eq!(p.value, 20);
            } else {
                assert!(false);
            }
        } else {
            assert!(false);
        }
    }

    #[test]
    fn test_packet_packets() {
        let inp = parse_str("EE00D40C823060");
        let mut i = inp.iter().copied();
        if let Some(r) = parse_packet(&mut i) {
            if r.type_id != PacketType::Literal {
                let p = &r.sub_packets[0];
                assert_eq!(p.value, 1);
                let p = &r.sub_packets[1];
                assert_eq!(p.value, 2);
                let p = &r.sub_packets[2];
                assert_eq!(p.value, 3);
            } else {
                assert!(false);
            }
        } else {
            assert!(false);
        }
    }

    #[test]
    fn part1_tests() -> anyhow::Result<()> {
        let d = Day16::init("8A004A801A8002F478".to_string())?;
        assert_eq!(d.part1()?, "16");
        let d = Day16::init("620080001611562C8802118E34".to_string())?;
        assert_eq!(d.part1()?, "12");
        let d = Day16::init("C0015000016115A2E0802F182340".to_string())?;
        assert_eq!(d.part1()?, "23");
        let d = Day16::init("A0016C880162017C3686B18A3D4780".to_string())?;
        assert_eq!(d.part1()?, "31");
        Ok(())
    }

    #[test]
    fn part2_tests() -> anyhow::Result<()> {
        fn calc(s: &str) -> anyhow::Result<String> {
            let d = Day16::init(s.to_string())?;
            d.part2()
        }
        assert_eq!("3", calc("C200B40A82")?);
        assert_eq!("54", calc("04005AC33890")?);
        assert_eq!("7", calc("880086C3E88112")?);
        assert_eq!("9", calc("CE00C43D881120")?);
        assert_eq!("1", calc("D8005AC2A8F0")?);
        assert_eq!("0", calc("F600BC2D8F")?);
        assert_eq!("0", calc("9C005AC2F8F0")?);
        assert_eq!("1", calc("9C0141080250320F1802104A08")?);

        Ok(())
    }

    day_tests!(Day16, "16", "979", "277110354175");
}