use anyhow::{anyhow, Result};
use daggy::{Dag, Walker, NodeIndex};
use petgraph::dot::{Dot, Config};
use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::env::args;
use std::fs::File;
use std::io::{BufRead, BufReader, BufWriter};
use std::io::prelude::*;
use std::path::Path;

type BagCount = usize;
type BagColor = String;
type Rules = Dag<BagColor, BagCount>;
type Index = HashMap<String, NodeIndex>;

// The example data in
// simplified form:    A       | E       |       D
//                       1 x B |   1 x G |      / \
// A contain 1B, 2C.     2 x C |   2 x H |     | A |
// D contain 3B, 4C.   D       | G       |     |/ \|
// B contain 1E.         3 x B |   3 x F |     B __C
// C contain 2E, 9F.     4 x C |   4 x I |     |/  |
// E contain 1G, 2H.   B       | H       |     E   |
// G contain 3F, 4I.     1 x E |   5 x F |    /|  /
// H contain 5F, 6I.   C       |   6 x I |   H G /
// F contain NONE.       2 x E | F       |  / \|/
// I contain NONE.       9 x F | I       | I   F

fn read_input<T: AsRef<Path>>(filename: T) -> Result<(Rules, Index)> {
    let f = File::open(filename)?;
    let reader = BufReader::new(f);

    let outer_re = Regex::new(
        r"(?x)
        ^(?P<outer>[a-z\s]+)\ bags?\ contain\ 
        (?P<inner>(?:(?:\d+\ [a-z\s]+|no\ other)\ bags?[.,]\s?)+)$"
    )?;

    let inner_re = Regex::new(
        r"(?x)
        (?:(?:(?P<count>\d+)\ (?P<color>[a-z\s]+)|(?P<empty>no\ other))\ bags?[.,]\s?)"
    )?;

    let mut rules: Rules = Dag::new();
    let mut rules_index = HashMap::new();

    for line in reader.lines() {
        let line = line?;
        let caps_outer = outer_re.captures(&line).ok_or_else(|| anyhow!("Invalid input"))?;
        let container_color = caps_outer.name("outer")
            .ok_or_else(|| anyhow!("Invalid input"))?.as_str().to_string();
        let parent_index = match rules_index.get(&String::from(&container_color)) {
            None => {
                rules.add_node(container_color.clone())
            },
            Some(existing_index) => {
                *existing_index
            },
        };
        rules_index.insert(container_color, parent_index);
        for m in inner_re.find_iter(caps_outer.name("inner")
            .ok_or_else(|| anyhow!("Invalid input"))?.as_str())
        {
            let s = m.as_str();
            let caps_inner = inner_re.captures(&s).ok_or_else(|| anyhow!("Invalid input"))?;
            match caps_inner.name("empty") {
                Some(_) => { },
                None => {
                    let content_color = caps_inner.name("color")
                        .ok_or_else(|| anyhow!("Invalid input"))?.as_str();
                    let content_count: usize = caps_inner.name("count")
                        .ok_or_else(|| anyhow!("Invalid input"))?.as_str().parse()?;
                    match rules_index.get(content_color) {
                        None => {
                            let (_, child_index) = rules.add_child(
                                parent_index, content_count, String::from(content_color)
                            );
                            rules_index.insert(String::from(content_color), child_index);
                        },
                        Some(existing_index) => {
                            let _ = rules.add_edge(parent_index, *existing_index, content_count);
                        },
                    }
                }
            }
        }
    }

     Ok((rules, rules_index))
}

fn plot_graph(rules: Rules) -> Result<()>{
    let f = File::create("graph.dot")?;
    let mut writer = BufWriter::new(f);
    let g = rules.graph();
    write!(writer, "{}", Dot::with_config(g, &[Config::EdgeNoLabel]))?;
    Ok(())
}

fn recurse_part1(start: NodeIndex, rules: &Rules, mut visited: &mut HashSet<BagColor>) -> usize {
    let mut ans = 0;
    let node = rules.node_weight(start).unwrap();

    if ! visited.contains(node) {
        ans += 1;
        visited.insert(String::from(node));
        let mut parent_walker = rules.parents(start);
        while let Some(walk) = parent_walker.walk_next(&rules) {
            let ( _edge, next_node_index ) = walk;
            ans += recurse_part1(next_node_index, rules, &mut visited);
        }
    }

    ans
}

fn recurse_part2(start: NodeIndex, rules: &Rules) -> usize {
    let mut ans=0;

    rules.node_weight(start).unwrap();
        let mut child_walker = rules.children(start);
        while let Some(walk) = child_walker.walk_next(&rules) {
            let ( edge_index, next_node_index ) = walk;
            let content_count = rules.edge_weight(edge_index).unwrap();
            ans += content_count * (1 + recurse_part2(next_node_index, rules));
        }

    ans
}

fn part1(rules: &Rules, index: &Index) -> Result<usize> {
    let start_index = index.get(&String::from("shiny gold")).ok_or(anyhow!("Bag is missing"))?;
    let mut visited: HashSet<BagColor> = HashSet::new();
    let ans = recurse_part1(*start_index, rules, &mut visited);

    Ok(ans-1)
}

fn part2(rules: &Rules, index: &Index) -> Result<usize> {
    let start_index = index.get(&String::from("shiny gold")).ok_or(anyhow!("Bag is missing"))?;
    Ok(recurse_part2(*start_index, rules))
}

fn main() {
    let ( do_part_1, do_part_2 ) = aoc::do_parts();

    let filename = match args().nth(1) {
        Some(f) => f,
        None => {
            eprintln!("Missing input filename");
            std::process::exit(1);
        },
    };
    match read_input(filename) {
        Ok((input, index)) => {
            if do_part_1 {
                let solution = part1(&input, &index);
                println!("Part1, {:?}", solution);
            }
            if do_part_2 {
                let solution = part2(&input, &index);
                println!("Part2, {:?}", solution);
            }
            match plot_graph(input) {
                Ok(_) => println!("Wrote a graphviz file. Please feel free to plot it."),
                Err(err) => eprintln!("Could not write graphviz file: {}", err),
            }
        },
        Err(err) => eprintln!("Could not read input: {}", err),
    }
}