use nom::{
    character::complete::{alpha1, char as nom_char, line_ending},
    combinator::map,
    multi::separated_list1,
    sequence::tuple,
    IResult,
};
use std::{collections::BTreeMap, fs};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let input = fs::read_to_string("inputs/day_12.txt")?;
    let cave_system = parse_cave_system(&input).unwrap().1;
    dbg!(cave_system.find_all_paths(0).len());
    dbg!(cave_system.find_all_paths(1).len());

    Ok(())
}

#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
struct CaveId(String);

impl CaveId {
    fn big(&self) -> bool {
        self.0.chars().next().map_or(false, char::is_uppercase)
    }
    fn terminal(&self) -> bool {
        self.0 == "start" || self.0 == "end"
    }
}

#[derive(Debug)]
struct Cave {
    id: CaveId,
    exits: Vec<CaveId>,
}

impl Cave {
    fn new(id: CaveId) -> Cave {
        Cave {
            id,
            exits: Vec::new(),
        }
    }
}

#[derive(Debug, Clone)]
struct CavePath {
    path: Vec<CaveId>,
    remaining_small_cave_visits: usize,
}

impl CavePath {
    fn new(max_small_cave_visits: usize) -> CavePath {
        CavePath {
            path: vec![CaveId("start".into())],
            remaining_small_cave_visits: max_small_cave_visits,
        }
    }

    fn push(&mut self, next: CaveId) {
        if !next.big() && self.path.iter().any(|visited| visited == &next) {
            self.remaining_small_cave_visits -= 1;
        }
        self.path.push(next);
    }

    fn can_go_to(&self, dest: &CaveId) -> bool {
        dest.big()
            || !self.path.iter().any(|visited| visited == dest)
            || (!dest.terminal() && self.remaining_small_cave_visits > 0)
    }

    fn is_complete(&self) -> bool {
        self.tail() == &CaveId("end".into())
    }

    fn tail(&self) -> &CaveId {
        self.path
            .get(self.path.len() - 1)
            .expect("There should not be empty paths")
    }
}

#[derive(Debug)]
struct CaveSystem {
    caves: BTreeMap<CaveId, Cave>,
}

impl CaveSystem {
    fn new(tunnels: Vec<(CaveId, CaveId)>) -> CaveSystem {
        let mut caves = BTreeMap::new();
        for (tunnel_start, tunnel_end) in tunnels {
            let start = caves
                .entry(tunnel_start.clone())
                .or_insert(Cave::new(tunnel_start.clone()));
            start.exits.push(tunnel_end.clone());

            let end = caves
                .entry(tunnel_end.clone())
                .or_insert(Cave::new(tunnel_end.clone()));
            end.exits.push(tunnel_start);
        }
        CaveSystem { caves }
    }

    fn find_all_paths(&self, max_small_cave_visits: usize) -> Vec<CavePath> {
        self.find_paths(CavePath::new(max_small_cave_visits))
    }

    fn find_paths(&self, active_path: CavePath) -> Vec<CavePath> {
        if active_path.is_complete() {
            vec![active_path]
        } else {
            let current = self.caves.get(active_path.tail()).expect("Unknown path");
            current
                .exits
                .iter()
                .filter(|next| active_path.can_go_to(next))
                .flat_map(|next| {
                    let mut active_path = active_path.clone();
                    active_path.push(next.clone());
                    self.find_paths(active_path)
                })
                .collect()
        }
    }
}

fn parse_cave_system(input: &str) -> IResult<&str, CaveSystem> {
    map(parse_tunnels, CaveSystem::new)(input)
}

fn parse_tunnels(input: &str) -> IResult<&str, Vec<(CaveId, CaveId)>> {
    separated_list1(line_ending, parse_tunnel)(input)
}

fn parse_tunnel(input: &str) -> IResult<&str, (CaveId, CaveId)> {
    map(
        tuple((parse_cave_id, nom_char('-'), parse_cave_id)),
        |(a, _, b)| (a, b),
    )(input)
}

fn parse_cave_id(input: &str) -> IResult<&str, CaveId> {
    map(alpha1, |s: &str| CaveId(s.to_owned()))(input)
}