use std::{cell::RefCell, cmp::min, collections::BTreeMap, rc::Rc, str::FromStr};

use aoc_2022::*;

type Input = Filesystem;
type Output = usize;

type FileHandle = Rc<RefCell<AocFile>>;

#[derive(Debug)]
enum AocFile {
    File(usize),
    Directory(BTreeMap<String, FileHandle>),
}

impl Default for AocFile {
    fn default() -> Self {
        AocFile::Directory(BTreeMap::new())
    }
}

impl AocFile {
    fn dir() -> FileHandle {
        Rc::new(RefCell::new(AocFile::Directory(BTreeMap::new())))
    }

    fn file(size: usize) -> FileHandle {
        Rc::new(RefCell::new(AocFile::File(size)))
    }

    fn add_file(&mut self, name: &str, size: usize) {
        if let AocFile::Directory(files) = self {
            files.insert(name.to_string(), AocFile::file(size));
            return;
        }

        panic!("cannot cd in file")
    }

    fn cd(&mut self, dir: &str) -> FileHandle {
        if let AocFile::Directory(files) = self {
            return files.entry(dir.to_string()).or_default().clone();
        }

        panic!("cannot cd in file")
    }

    fn size_under(&self, max_size: usize) -> (bool, usize, usize) {
        match self {
            AocFile::File(s) => (false, 0, *s),
            AocFile::Directory(files) => {
                let (running_total, size) = files
                    .values()
                    .map(|f| f.borrow().size_under(max_size))
                    .fold((0_usize, 0_usize), |(mut running, size), (dir, r, s)| {
                        if dir && s <= max_size {
                            running += s;
                        }

                        (r + running, size + s)
                    });

                (true, running_total, size)
            }
        }
    }

    fn smallest_bigger(&self, required: usize) -> (bool, usize, usize) {
        match self {
            AocFile::File(s) => (false, usize::MAX, *s),
            AocFile::Directory(files) => {
                let (mut dir_min, size) = files
                    .values()
                    .map(|f| f.borrow().smallest_bigger(required))
                    .fold(
                        (usize::MAX, 0_usize),
                        |(mut current_min, size), (dir, r_min, s)| {
                            if dir && s >= required {
                                // we got back from a subdirectory
                                current_min = min(current_min, s);
                            }

                            // also check local solutions in subdirectories
                            (min(current_min, r_min), size + s)
                        },
                    );

                if size >= required {
                    dir_min = min(dir_min, size);
                }

                (true, dir_min, size)
            }
        }
    }
}

struct Filesystem {
    root: FileHandle,
    cwd: Vec<FileHandle>,
}

impl Filesystem {
    fn new() -> Filesystem {
        let root = AocFile::dir();

        Filesystem {
            root: root.clone(),
            cwd: vec![root],
        }
    }

    // [MARK] Helper functions for ‹FromStr› trait

    fn cd(&mut self, dir: &str) {
        match dir {
            ".." => {
                self.cwd.pop();
            }
            "/" => {
                self.cwd.clear();
                self.cwd.push(self.root.clone());
            }
            _ => {
                let idx = self.cwd.len() - 1;
                let subdir = self.cwd[idx].borrow_mut().cd(dir);
                self.cwd.push(subdir);
            }
        }
    }

    fn touch(&mut self, name: &str, size: usize) {
        let idx = self.cwd.len() - 1;
        self.cwd[idx].borrow_mut().add_file(name, size);
    }

    fn run_command(&mut self, command: &str) {
        if command.starts_with("cd ") {
            let parts = command.split_ascii_whitespace().collect_vec();
            self.cd(parts[1]);
            return;
        }

        for file in command.lines().skip(1) {
            let parts = file.split_ascii_whitespace().collect_vec();
            if parts[0] != "dir" {
                let name = parts[1];
                let size: usize = parts[0].parse().unwrap();
                self.touch(name, size);
            }
        }
    }

    // [MARK] Helper functions for ‹FromStr› trait

    fn size_under(&self, max_size: usize) -> usize {
        self.root.borrow().size_under(max_size).1
    }

    fn purge(&self, total: usize, needed: usize) -> usize {
        let used = self.root.borrow().size_under(0).2;

        // to_be_freed >= needed - (total - used)
        let to_be_freed = needed - (total - used);

        self.root.borrow().smallest_bigger(to_be_freed).1
    }
}

impl FromStr for Filesystem {
    type Err = Report;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut fs = Filesystem::new();

        for command in s.trim_start_matches("$ ").split("\n$ ") {
            fs.run_command(command);
        }

        Ok(fs)
    }
}

struct Day07;
impl Solution<Input, Output> for Day07 {
    fn parse_input<P: AsRef<Path>>(pathname: P) -> Input {
        file_to_string(pathname).parse().unwrap()
    }

    fn part_1(input: &Input) -> Output {
        input.size_under(100000)
    }

    fn part_2(input: &Input) -> Output {
        input.purge(70000000, 30000000)
    }
}

fn main() -> Result<()> {
    Day07::main()
}

test_sample!(day_07, Day07, 95437, 24933642);