#![feature(iter_array_chunks)]

/// --- Part Two ---
///
/// Now, you just need to put all of the packets in the right order.  Disregard the blank lines in
/// your list of received packets.
///
/// The distress signal protocol also requires that you include two additional divider packets:
///
/// [[2]]
/// [[6]]
///
/// Using the same rules as before, organize all packets - the ones in your list of received
/// packets as well as the two divider packets - into the correct order.
///
/// For the example above, the result of putting the packets in the correct order is:
///
/// ```
/// []
/// [[]]
/// [[[]]]
/// [1,1,3,1,1]
/// [1,1,5,1,1]
/// [[1],[2,3,4]]
/// [1,[2,[3,[4,[5,6,0]]]],8,9]
/// [1,[2,[3,[4,[5,6,7]]]],8,9]
/// [[1],4]
/// [[2]]
/// [3]
/// [[4,4],4,4]
/// [[4,4],4,4,4]
/// [[6]]
/// [7,7,7]
/// [7,7,7,7]
/// [[8,7,6]]
/// [9]
/// ```
///
/// Afterward, locate the divider packets.  To find the decoder key for this distress signal, you
/// need to determine the indices of the two divider packets and multiply them together.  (The
/// first packet is at index 1, the second packet is at index 2, and so on.) In this example, the
/// divider packets are 10th and 14th, and so the decoder key is 140.
///
/// Organize all of the packets into the correct order.  What is the decoder key for the distress
/// signal?
use clap::Parser;
use itertools::Itertools;
use nom::branch::alt;
use nom::bytes::complete::tag;
use nom::character::complete::i8;
use nom::combinator::{map, opt};
use nom::error::{ContextError, ErrorKind as NomErrorKind, ParseError};
use nom::multi::many1;
use nom::sequence::{delimited, terminated};
use nom::IResult;

use std::cmp::Ordering;
use std::fs::File;
use std::io::prelude::*;
use std::io::BufReader;
use std::path::PathBuf;

const FILEPATH: &'static str = "examples/input.txt";

pub type Input<'a> = &'a str;
pub type Result<'a, T> = IResult<Input<'a>, T, Error<Input<'a>>>;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ErrorKind {
    Nom(NomErrorKind),
    Context(&'static str),
    Custom(String),
}

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Cli {
    #[clap(short, long, default_value = FILEPATH)]
    file: PathBuf,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Error<I> {
    pub errors: Vec<(I, ErrorKind)>,
}

#[derive(Clone, Debug, PartialEq, Eq)]
enum ListEntry {
    Value(Option<i8>),
    List(Vec<ListEntry>),
}

impl Ord for ListEntry {
    fn cmp(&self, other: &Self) -> Ordering {
        let res = compare(&self, other);
        match res {
            Some(x) => match x {
                true => Ordering::Less,
                false => Ordering::Greater,
            },
            None => Ordering::Equal,
        }
    }
}

impl PartialOrd for ListEntry {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<I> ParseError<I> for Error<I> {
    fn from_error_kind(input: I, kind: NomErrorKind) -> Self {
        let errors = vec![(input, ErrorKind::Nom(kind))];
        Self { errors }
    }

    fn append(input: I, kind: NomErrorKind, mut other: Self) -> Self {
        other.errors.push((input, ErrorKind::Nom(kind)));
        other
    }
}

impl<I> ContextError<I> for Error<I> {
    fn add_context(input: I, ctx: &'static str, mut other: Self) -> Self {
        other.errors.push((input, ErrorKind::Context(ctx)));
        other
    }
}

fn parse_list(input: &str) -> Result<ListEntry> {
    alt((
        map(i8, |v| ListEntry::Value(Some(v))),
        map(tag("[]"), |_| ListEntry::Value(None)),
        delimited(
            tag("["),
            map(many1(terminated(parse_list, opt(tag(",")))), |list| {
                ListEntry::List(list)
            }),
            tag("]"),
        ),
    ))(input)
}

fn compare(lhs: &ListEntry, rhs: &ListEntry) -> Option<bool> {
    use ListEntry::*;
    match (lhs, rhs) {
        (Value(lv), Value(rv)) => {
            if lv == rv {
                None
            } else {
                Some(lv < rv)
            }
        }
        (Value(None), List(_)) => Some(true),
        (Value(_), List(_)) => compare(&ListEntry::List(vec![lhs.clone()]), rhs),
        (List(_), Value(None)) => Some(false),
        (List(_), Value(_)) => compare(lhs, &ListEntry::List(vec![rhs.clone()])),
        (List(lv), List(rv)) => lv
            .iter()
            .zip_longest(rv.iter())
            .find_map(|pair| match pair {
                itertools::EitherOrBoth::Both(l, r) => compare(l, r),
                itertools::EitherOrBoth::Right(_) => Some(true),
                itertools::EitherOrBoth::Left(_) => Some(false),
            }),
    }
}

fn main() {
    let divider_packets: [ListEntry; 2] = [
        ListEntry::List(vec![ListEntry::Value(Some(2))]),
        ListEntry::List(vec![ListEntry::Value(Some(6))]),
    ];

    let args = Cli::parse();

    let file = File::open(&args.file).unwrap();
    let reader = BufReader::new(file);

    let res: usize = reader
        .lines()
        .filter_map(|line| {
            let s = line.unwrap();
            if s.is_empty() {
                return None;
            }
            Some(parse_list(s.as_str()).unwrap().1)
        })
        .chain(divider_packets.clone().into_iter())
        .sorted()
        .zip(1..)
        .filter(|(x, _)| x == &divider_packets[0] || x == &divider_packets[1])
        .map(|(_, idx)| idx)
        .take(2)
        .product();

    println!("{res}");
}