use std::collections::BTreeMap;

use nom::branch::alt;
use nom::bytes::complete::{tag, take_till1};
use nom::character::complete::{newline, space1};
use nom::character::{complete, is_newline};
use nom::multi::separated_list1;
use nom::IResult;

pub fn process_part_1(input: &str) -> u32 {
    let (_, commands) = parse_commands(input).unwrap();
    let directories = build_directory_tree_from_commands(commands);

    directories.values().filter(|&size| *size <= 100_000).sum()
}

pub fn process_part_2(input: &str) -> u32 {
    let (_, commands) = parse_commands(input).unwrap();
    let directories = build_directory_tree_from_commands(commands);

    let total_space = 70_000_000;
    let required_space = 30_000_000;
    let current_space = total_space - *directories.get("").unwrap();
    let wanted_space = required_space - current_space;

    directories
        .values()
        .filter(|&size| *size >= wanted_space)
        .map(|&size| size)
        .min()
        .unwrap_or(0)
}

fn build_directory_tree_from_commands(commands: Vec<Command>) -> BTreeMap<String, u32> {
    commands
        .iter()
        .fold((vec![], BTreeMap::new()), |(mut cwd, mut dirs), command| {
            match command {
                Command::Cd("/") => {
                    cwd = vec![];
                }
                Command::Cd("..") => {
                    cwd.pop();
                }
                Command::Cd(dir) => {
                    cwd.push(*dir);
                }
                Command::Ls(files) => {
                    // We don't actually care about individual files, just sum the files together
                    // so we aren't looking up the BTreeMap as often
                    let total_size = files
                        .iter()
                        .filter_map(|e| match e {
                            DirEntry::File { size, .. } => Some(size),
                            _ => None,
                        })
                        .sum();

                    for i in 0..=cwd.len() {
                        dirs.entry(cwd[..i].join("/"))
                            .and_modify(|f| *f += total_size)
                            .or_insert(total_size);
                    }
                }
            }

            (cwd, dirs)
        })
        .1
}

#[derive(Debug)]
enum Command<'a> {
    Cd(&'a str),
    Ls(Vec<DirEntry<'a>>),
}

#[derive(Debug)]
enum DirEntry<'a> {
    Dir(&'a str),
    #[allow(dead_code)]
    File {
        name: &'a str,
        size: u32,
    },
}

fn parse_commands(input: &str) -> IResult<&str, Vec<Command>> {
    separated_list1(newline, parse_command)(input)
}

fn parse_command(input: &str) -> IResult<&str, Command> {
    let (input, _) = tag("$ ")(input)?;
    alt((parse_cd, parse_ls))(input)
}

fn parse_cd(input: &str) -> IResult<&str, Command> {
    let (input, _) = tag("cd")(input)?;
    let (input, _) = space1(input)?;
    let (input, arg) = parse_file_name(input)?;
    Ok((input, Command::Cd(arg)))
}

fn parse_ls(input: &str) -> IResult<&str, Command> {
    let (input, _) = tag("ls")(input)?;
    let (input, _) = newline(input)?;
    let (input, entries) =
        separated_list1(tag("\n"), alt((parse_dir_entry_dir, parse_dir_entry_file)))(input)?;
    Ok((input, Command::Ls(entries)))
}

fn parse_dir_entry_dir(input: &str) -> IResult<&str, DirEntry> {
    let (input, _) = tag("dir ")(input)?;
    let (input, name) = parse_file_name(input)?;
    Ok((input, DirEntry::Dir(name)))
}

fn parse_dir_entry_file(input: &str) -> IResult<&str, DirEntry> {
    let (input, size) = complete::u32(input)?;
    let (input, _) = tag(" ")(input)?;
    let (input, name) = parse_file_name(input)?;
    Ok((input, DirEntry::File { name, size }))
}

fn parse_file_name(input: &str) -> IResult<&str, &str> {
    take_till1(|c| is_newline(c as u8))(input)
}

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

    const INPUT: &str = "$ cd /
$ ls
dir a
14848514 b.txt
8504156 c.dat
dir d
$ cd a
$ ls
dir e
29116 f
2557 g
62596 h.lst
$ cd e
$ ls
584 i
$ cd ..
$ cd ..
$ cd d
$ ls
4060174 j
8033020 d.log
5626152 d.ext
7214296 k";

    #[test]
    fn day1() {
        assert_eq!(process_part_1(INPUT), 95437);
    }

    #[test]
    fn day2() {
        assert_eq!(process_part_2(INPUT), 24933642);
    }
}