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

const INPUT: &str = include_str!("../../data/day19");

#[derive(Clone, Debug)]
enum Rule {
    Char(char),
    Chain(Vec<u32>),
}

impl Rule {
    fn parse_part(alt: &str) -> Self {
        let alt = alt.trim();
        if let Some(alt) = alt.strip_prefix('"') {
            let alt = alt.strip_suffix('"').unwrap();
            assert_eq!(alt.len(), 1);
            Self::Char(alt.chars().next().unwrap())
        } else {
            Self::Chain(alt.split_whitespace().map(|s| s.parse::<u32>().unwrap()).collect())
        }
    }

    fn parse_line(line: &str) -> (u32, CachedRule) {
        let colon = line.find(": ").unwrap();
        let id = line[..colon].parse::<u32>().unwrap();
        let alts = line[colon+2..].split(" | ").map(Self::parse_part).collect();
        (id, CachedRule::Alternatives(alts))
    }
}

fn _make_combs(target: &mut HashSet<String>, current: &mut String, list: &[Arc<HashSet<String>>]) {
    if list.is_empty() {
        target.insert(current.clone());
    } else {
        let old_len = current.len();
        for s in &*list[0] {
            current.push_str(s);
            _make_combs(target, current, &list[1..]);
            current.truncate(old_len);
        }
    }
}

#[derive(Clone, Debug)]
enum CachedRule {
    Alternatives(Vec<Rule>),
    CachedStrings(Arc<HashSet<String>>, HashSet<usize>),
}

#[derive(Clone, Debug)]
struct Grammar {
    expansions: HashMap<u32, CachedRule>,
}

impl Grammar {
    fn parse(rule_lines: &str) -> Self {
        let expansions = rule_lines.lines().map(Rule::parse_line).collect();
        Self { expansions }
    }

    fn _build_sets(&mut self, recursive_check: &mut HashSet<u32>, ndx: u32) -> Option<Arc<HashSet<String>>> {
        let rules = self.expansions.get(&ndx).unwrap();
        let alts = match rules {
            CachedRule::Alternatives(alts) => alts,
            CachedRule::CachedStrings(cs, _) => return Some(cs.clone()),
        };

        if recursive_check.contains(&ndx) {
            return None;
        }
        recursive_check.insert(ndx);

        let mut result = HashSet::<String>::new();

        let mut found_recursion = false;
        for rule in alts.clone() {
            match rule {
                Rule::Char(c) => {
                    result.insert(c.to_string());
                },
                Rule::Chain(parts) => {
                    let mut comb = Vec::new();
                    for part_ndx in parts {
                        // abort if nested lookup loops
                        if let Some(cs) = self._build_sets(recursive_check, part_ndx) {
                            comb.push(cs);
                        } else {
                            found_recursion = true;
                        }
                    }
                    if !found_recursion {
                        _make_combs(&mut result, &mut String::new(), &comb);
                    }
                },
            }
        }
        if found_recursion {
            return None;
        }
        let lengths: HashSet<usize> = result.iter().map(|s| s.len()).collect();
        let result = Arc::new(result);
        recursive_check.remove(&ndx);
        self.expansions.insert(ndx, CachedRule::CachedStrings(result.clone(), lengths));
        Some(result)
    }

    fn optimize(&mut self) {
        let mut recursive_check = HashSet::new();
        self._build_sets(&mut recursive_check, 0);
        recursive_check.insert(0);
        for id in recursive_check.clone() {
            if let Some(CachedRule::Alternatives(alts)) = self.expansions.get(&id) {
                for alt in alts {
                    if let Rule::Chain(chain) = alt {
                        recursive_check.extend(chain);
                    }
                }
            }
        }
        let rule_ids: HashSet<u32> = self.expansions.keys().cloned().collect();
        for id in rule_ids.difference(&recursive_check) {
            self.expansions.remove(id);
        }
    }

    fn simple_match(&self, data: &str) -> bool {
        match self.expansions.get(&0) {
            Some(CachedRule::CachedStrings(cs, _)) => cs.contains(data),
            _ => panic!("not simple enough"),
        }
    }

    fn _complex_match_chain<'data>(&self, chain: &[u32], data: &'data str) -> Vec<&'data str> {
        if chain.is_empty() {
            return vec![data];
        }
        self._complex_match(chain[0], data).into_iter().map(|rem| {
            self._complex_match_chain(&chain[1..], rem)
        }).flatten().collect()
    }

    fn _complex_match<'data>(&self, start: u32, data: &'data str) -> Vec<&'data str> {
        let cr = self.expansions.get(&start).unwrap();
        match cr {
            CachedRule::Alternatives(alts) => {
                alts.iter().map(|rule| {
                    let chain = match rule {
                        Rule::Chain(chain) => chain,
                        _ => panic!("only chains here"),
                    };
                    self._complex_match_chain(&chain, data)
                }).flatten().collect()
            },
            CachedRule::CachedStrings(cs, lens) => {
                lens.iter().filter_map(|&len| {
                    if data.len() >= len && cs.contains(&data[..len]) {
                        Some(&data[len..])
                    } else {
                        None
                    }
                }).collect()
            },
        }
    }

    fn complex_match(&self, data: &str) -> bool {
        self._complex_match(0, data).into_iter().any(str::is_empty)
    }
}

fn main() {
    let input_grammar;
    let data;
    {
        let split_pos = INPUT.find("\n\n").unwrap();
        input_grammar = Grammar::parse(&INPUT[..split_pos]);
        data = INPUT[split_pos+2..].lines().map(str::to_string).collect::<Vec<String>>();
    }
    {
        let mut grammar = input_grammar.clone();
        grammar.optimize();
        println!("Lines matching simple grammar: {}", data.iter().filter(|line| grammar.simple_match(line)).count());
    }
    {
        let mut grammar = input_grammar.clone();
        grammar.expansions.extend(vec![
            Rule::parse_line("8: 42 | 42 8"),
            Rule::parse_line("11: 42 31 | 42 11 31"),
        ]);
        grammar.optimize();
/*
        println!("{:?}", grammar.expansions.keys());
        for (rule_id, rule) in grammar.expansions {
            match rule {
                CachedRule::Alternatives(alt) => println!("{} -> {:?}", rule_id, alt),
                CachedRule::CachedStrings(_, lens) => println!("{} -> string with lengths: {:?}", rule_id, lens),
            }
        }
*/
        println!("Lines matching complex grammar: {}", data.iter().filter(|line| grammar.complex_match(line)).count());
    }
}