use nom::{
    branch::alt,
    bytes::complete::take,
    character::complete::{char as nom_char, one_of},
    combinator::{consumed, map, map_res},
    error::FromExternalError,
    multi::many0,
    IResult,
};
use std::fs;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_16.txt")?;
    let binary = convert_to_binary(&input).unwrap().1;
    let packet = parse_packet(&binary).unwrap().1;
    dbg!(packet.version_sum());
    dbg!(packet.eval());

    Ok(())
}

#[derive(Debug)]
struct Packet {
    version: u8,
    data: PacketData,
}

impl Packet {
    fn version_sum(&self) -> u32 {
        self.version as u32
            + match &self.data {
                PacketData::Literal(_) => 0,
                PacketData::Operator(data) => {
                    data.sub_packets.iter().map(|p| p.version_sum()).sum()
                }
            }
    }

    fn eval(&self) -> u64 {
        match &self.data {
            PacketData::Literal(val) => *val,
            PacketData::Operator(data) => data.eval(),
        }
    }
}

#[derive(Debug)]
enum PacketData {
    Literal(u64),
    Operator(OperatorData),
}

#[derive(Debug)]
struct OperatorData {
    type_id: OperatorType,
    sub_packets: Vec<Packet>,
}

impl OperatorData {
    fn eval(&self) -> u64 {
        match &self.type_id {
            OperatorType::Sum => self.sub_packets.iter().map(|p| p.eval()).sum(),
            OperatorType::Product => self.sub_packets.iter().map(|p| p.eval()).product(),
            OperatorType::Minimum => self.sub_packets.iter().map(|p| p.eval()).min().unwrap_or(0),
            OperatorType::Maximum => self.sub_packets.iter().map(|p| p.eval()).max().unwrap_or(0),
            OperatorType::GreaterThan => {
                if self.sub_packets[0].eval() > self.sub_packets[1].eval() {
                    1
                } else {
                    0
                }
            }
            OperatorType::LessThan => {
                if self.sub_packets[0].eval() < self.sub_packets[1].eval() {
                    1
                } else {
                    0
                }
            }
            OperatorType::EqualTo => {
                if self.sub_packets[0].eval() == self.sub_packets[1].eval() {
                    1
                } else {
                    0
                }
            }
        }
    }
}

#[derive(Debug)]
enum OperatorType {
    Sum,
    Product,
    Minimum,
    Maximum,
    GreaterThan,
    LessThan,
    EqualTo,
}

#[derive(Debug, thiserror::Error)]
enum OperatorError {
    #[error("type was not a valid operator")]
    InvalidType,
}

impl TryFrom<u8> for OperatorType {
    type Error = OperatorError;
    fn try_from(num: u8) -> Result<Self, OperatorError> {
        match num {
            0 => Ok(Self::Sum),
            1 => Ok(Self::Product),
            2 => Ok(Self::Minimum),
            3 => Ok(Self::Maximum),
            5 => Ok(Self::GreaterThan),
            6 => Ok(Self::LessThan),
            7 => Ok(Self::EqualTo),
            _ => Err(OperatorError::InvalidType),
        }
    }
}

fn convert_to_binary(input: &str) -> IResult<&str, String> {
    map(
        many0(map(one_of("0123456789ABCDEF"), |hex_char| {
            let digit = hex_char.to_digit(16).unwrap();
            format!("{:04b}", digit)
        })),
        |bin_strings| bin_strings.join(""),
    )(input)
}

fn parse_packet(input: &str) -> IResult<&str, Packet> {
    let (input, version) = parse_bits3(input)?;
    let (input, type_id) = parse_bits3(input)?;
    if type_id == 4 {
        let (input, literal) = parse_literal(input)?;
        Ok((
            input,
            Packet {
                version,
                data: PacketData::Literal(literal),
            },
        ))
    } else {
        let (input, sub_packets) =
            alt((parse_sub_packets_bits_mode, parse_sub_packets_count_mode))(input)?;
        Ok((
            input,
            Packet {
                version,
                data: PacketData::Operator(OperatorData {
                    type_id: type_id.try_into().expect("Invalid operator"),
                    sub_packets,
                }),
            },
        ))
    }
}

fn parse_bits3(input: &str) -> IResult<&str, u8> {
    map_res(take(3usize), |bin_str| u8::from_str_radix(bin_str, 2))(input)
}

fn parse_literal(input: &str) -> IResult<&str, u64> {
    let (input, mut chunks) = many0(parse_literal_continuing_chunk)(input)?;
    let (input, last_chunk) = parse_literal_last_chunk(input)?;
    chunks.push(last_chunk);
    let binary_num = chunks.join("");
    let num = u64::from_str_radix(&binary_num, 2).map_err(|e| {
        nom::Err::Error(nom::error::Error::from_external_error(
            input,
            nom::error::ErrorKind::MapRes,
            e,
        ))
    })?;
    Ok((input, num))
}

fn parse_literal_continuing_chunk(input: &str) -> IResult<&str, &str> {
    let (input, _) = nom_char('1')(input)?;
    take(4usize)(input)
}

fn parse_literal_last_chunk(input: &str) -> IResult<&str, &str> {
    let (input, _) = nom_char('0')(input)?;
    take(4usize)(input)
}

fn parse_sub_packets_bits_mode(input: &str) -> IResult<&str, Vec<Packet>> {
    let (input, _) = nom_char('0')(input)?;
    let (mut input, length_in_bits) =
        map_res(take(15usize), |bin_str| usize::from_str_radix(bin_str, 2))(input)?;
    let mut consumed_by_subpackets = String::new();
    let mut sub_packets = Vec::new();
    while consumed_by_subpackets.len() < length_in_bits {
        let (next_input, (next_consumed, next_packet)) = consumed(parse_packet)(input)?;
        input = next_input;
        consumed_by_subpackets += next_consumed;
        sub_packets.push(next_packet);
    }
    Ok((input, sub_packets))
}

fn parse_sub_packets_count_mode(input: &str) -> IResult<&str, Vec<Packet>> {
    let (input, _) = nom_char('1')(input)?;
    let (mut input, number_of_packets) =
        map_res(take(11usize), |bin_str| u16::from_str_radix(bin_str, 2))(input)?;
    let mut sub_packets = Vec::new();
    for _ in 0..number_of_packets {
        let (next_input, next_packet) = parse_packet(input)?;
        input = next_input;
        sub_packets.push(next_packet);
    }
    Ok((input, sub_packets))
}