use std::collections::{HashMap, HashSet};

pub struct CaveSystem {
    pub adjacencies: HashMap<Cave, HashSet<Cave>>,
}

#[derive(Debug, Hash, PartialEq, Eq, Clone)]
pub enum Cave {
    Big(String),
    Small(String),
    Start,
    End,
}

impl ToString for Cave {
    fn to_string(&self) -> String {
        match self {
            Cave::Big(cave) => cave.clone(),
            Cave::Small(cave) => cave.clone(),
            Cave::Start => String::from("start"),
            Cave::End => String::from("end"),
        }
    }
}

impl CaveSystem {
    pub fn list_paths(&self, may_visit_twice: bool) -> Vec<Vec<Cave>> {
        self.recurse_path(Cave::Start, vec![Cave::Start], may_visit_twice)
    }

    fn recurse_path(
        &self,
        current: Cave,
        path: Vec<Cave>,
        may_visit_twice: bool,
    ) -> Vec<Vec<Cave>> {
        let mut ret_val = Vec::new();
        for cave in self.adjacencies.get(&current).unwrap_or(&HashSet::new()) {
            match cave {
                Cave::Big(_) => {
                    let mut path = path.clone();
                    path.push(cave.clone());
                    ret_val.append(&mut self.recurse_path(cave.clone(), path, may_visit_twice));
                }
                Cave::Small(_) => {
                    let mut may_visit_twice = may_visit_twice;
                    if path.contains(cave) {
                        if may_visit_twice {
                            may_visit_twice = false;
                        } else {
                            continue;
                        }
                    };
                    let mut path = path.clone();
                    path.push(cave.clone());
                    ret_val.append(&mut self.recurse_path(cave.clone(), path, may_visit_twice));
                }
                Cave::Start => continue,
                Cave::End => {
                    let mut path = path.clone();
                    path.push(cave.clone());
                    ret_val.push(path);
                }
            }
        }
        ret_val
    }

    pub fn new(edges: Vec<[Cave; 2]>) -> Result<CaveSystem, &'static str> {
        let mut adjacencies: HashMap<Cave, HashSet<Cave>> = Default::default();
        for edge in edges {
            adjacencies
                .entry(edge[0].clone())
                .or_default()
                .insert(edge[1].clone());
            adjacencies
                .entry(edge[1].clone())
                .or_default()
                .insert(edge[0].clone());
        }
        if !adjacencies.contains_key(&Cave::Start) {
            Err("Cave system does not contain start cave")
        } else if !adjacencies.contains_key(&Cave::End) {
            Err("Cave system does not contain end cave")
        } else {
            Ok(CaveSystem { adjacencies })
        }
    }
}

#[cfg(test)]
mod test {
    use std::collections::HashSet;

    use itertools::Itertools;

    #[test]
    fn small_example_part_2() {
        let input: super::CaveSystem = super::super::parse_input(
            "start-A
start-b
A-c
A-b
b-d
A-end
b-end",
        );
        let paths: HashSet<String> = input
            .list_paths(true)
            .iter()
            .map(|path| path.iter().map(|cave| cave.to_string()).join(","))
            .collect();

        let expected_paths: HashSet<String> = "start,A,b,A,b,A,c,A,end
start,A,b,A,b,A,end
start,A,b,A,b,end
start,A,b,A,c,A,b,A,end
start,A,b,A,c,A,b,end
start,A,b,A,c,A,c,A,end
start,A,b,A,c,A,end
start,A,b,A,end
start,A,b,d,b,A,c,A,end
start,A,b,d,b,A,end
start,A,b,d,b,end
start,A,b,end
start,A,c,A,b,A,b,A,end
start,A,c,A,b,A,b,end
start,A,c,A,b,A,c,A,end
start,A,c,A,b,A,end
start,A,c,A,b,d,b,A,end
start,A,c,A,b,d,b,end
start,A,c,A,b,end
start,A,c,A,c,A,b,A,end
start,A,c,A,c,A,b,end
start,A,c,A,c,A,end
start,A,c,A,end
start,A,end
start,b,A,b,A,c,A,end
start,b,A,b,A,end
start,b,A,b,end
start,b,A,c,A,b,A,end
start,b,A,c,A,b,end
start,b,A,c,A,c,A,end
start,b,A,c,A,end
start,b,A,end
start,b,d,b,A,c,A,end
start,b,d,b,A,end
start,b,d,b,end
start,b,end"
            .lines()
            .map(String::from)
            .collect();
        dbg!(paths.symmetric_difference(&expected_paths));
        assert_eq!(paths, expected_paths);
    }
}