use std::collections as col;
use std::error;
use std::fmt;
use std::rc;
use std::result;

use stringcache::StringCache;

#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum Error {
    StartingRuleNotDefined(rc::Rc<str>),
    RuleMentionsUnknownSymbol {
        rule: rc::Rc<str>,
        symbol: rc::Rc<str>,
    },
}

impl error::Error for Error {}

impl<'a> fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Error::StartingRuleNotDefined(name) => write!(
                f,
                "Grammar starts with symbol {},
                    but that symbol was never defined",
                name
            ),
            Error::RuleMentionsUnknownSymbol { rule, symbol } => write!(
                f,
                "A rule for symbol {rule} refers to symbol {symbol},
                    but that symbol was never defined",
                rule = rule,
                symbol = symbol,
            ),
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub enum Symbol {
    Terminal(char),
    NonTerminal(rc::Rc<str>),
}

impl From<char> for Symbol {
    fn from(c: char) -> Symbol {
        Symbol::Terminal(c)
    }
}

impl<'a> From<&'a rc::Rc<str>> for Symbol {
    fn from(s: &rc::Rc<str>) -> Symbol {
        Symbol::NonTerminal(s.clone())
    }
}

impl From<rc::Rc<str>> for Symbol {
    fn from(s: rc::Rc<str>) -> Symbol {
        Symbol::NonTerminal(s)
    }
}

impl fmt::Display for Symbol {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Symbol::Terminal(c) => write!(f, "{:?}", c),
            Symbol::NonTerminal(name) => write!(f, "{}", name),
        }
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Rule {
    name: rc::Rc<str>,
    symbols: Vec<Symbol>,
}

impl Rule {
    pub fn name(&self) -> rc::Rc<str> {
        self.name.clone()
    }
    pub fn symbols(&self) -> &[Symbol] {
        &self.symbols
    }
    pub fn symbol_at(&self, index: usize) -> Option<&Symbol> {
        self.symbols.get(index)
    }
}

impl fmt::Display for Rule {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{} ->", self.name)?;
        for each in self.symbols.iter() {
            write!(f, " {}", each)?;
        }

        Ok(())
    }
}

#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct RuleBuilder<'a> {
    grammar: &'a mut GrammarBuilder,
    name: rc::Rc<str>,
    symbols: Vec<Symbol>,
}

impl<'a> RuleBuilder<'a> {
    fn new(
        grammar: &'a mut GrammarBuilder,
        name: rc::Rc<str>,
    ) -> RuleBuilder<'a> {
        RuleBuilder {
            grammar: grammar,
            name: name,
            symbols: vec![],
        }
    }

    pub fn nonterminal<S: Into<String>>(mut self, name: S) -> Self {
        let name = self.grammar.cache.get_or_intern(name);
        self.symbols.push(Symbol::NonTerminal(name));
        self
    }

    pub fn terminal(mut self, c: char) -> Self {
        self.symbols.push(Symbol::Terminal(c));
        self
    }

    pub fn add(self) {
        let name = self.name;
        let symbols = self.symbols;

        if symbols.is_empty() {
            self.grammar.nullables.insert(name.clone());
        }

        self.grammar
            .rules
            .entry(name.clone())
            .or_default()
            .insert(Rule { name, symbols });
    }
}

#[derive(Clone, Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Grammar {
    start: rc::Rc<str>,
    rules: col::BTreeMap<rc::Rc<str>, col::BTreeSet<Rule>>,
    nullables: col::BTreeSet<rc::Rc<str>>,
}

impl Grammar {
    pub fn start(&self) -> rc::Rc<str> {
        self.start.clone()
    }

    pub fn rules(
        &self,
        name: &rc::Rc<str>,
    ) -> Option<impl Iterator<Item = &Rule>> {
        self.rules.get(name).map(|r| r.iter())
    }

    pub fn nullable(&self, name: &rc::Rc<str>) -> Option<bool> {
        if !self.rules.contains_key(name) {
            return None;
        }

        Some(self.nullables.contains(name))
    }
}

#[derive(Debug, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct GrammarBuilder {
    cache: StringCache,
    start: rc::Rc<str>,
    rules: col::BTreeMap<rc::Rc<str>, col::BTreeSet<Rule>>,
    nullables: col::BTreeSet<rc::Rc<str>>,
}

impl GrammarBuilder {
    pub fn new<S: Into<String>>(start: S) -> GrammarBuilder {
        let mut cache = StringCache::new();
        let start = cache.get_or_intern(start);
        let rules = col::BTreeMap::new();
        let nullables = col::BTreeSet::new();

        GrammarBuilder {
            cache,
            start,
            rules,
            nullables,
        }
    }

    pub fn rule<S: Into<String>>(&mut self, name: S) -> RuleBuilder {
        let name = self.cache.get_or_intern(name);
        RuleBuilder::new(self, name)
    }

    fn propagate_nullability(&mut self) -> Result<()> {
        // Build a map of which non-terminal symbols are called by which rules.
        let mut rules_by_callee: col::BTreeMap<_, col::BTreeSet<&Rule>> =
            col::BTreeMap::new();
        for alternatives in self.rules.values() {
            for rule in alternatives {
                for symbol in rule.symbols.iter() {
                    if let Symbol::NonTerminal(name) = symbol {
                        if !self.rules.contains_key(name) {
                            return Err(Error::RuleMentionsUnknownSymbol {
                                rule: rule.name.clone(),
                                symbol: name.clone(),
                            });
                        }
                        rules_by_callee
                            .entry(name.clone())
                            .or_default()
                            .insert(rule);
                    }
                }
            }
        }

        // current_nullables is the set of symbols whose nullability needs
        // to be propagated.
        let mut current_nullables = col::BTreeSet::new();
        current_nullables.append(&mut self.nullables);

        // next_nullables is the set of symbols discovered to be nullable
        // by exploring the callers of symbols in current_nullables.
        let mut next_nullables = col::BTreeSet::new();

        // A handy fallback for symbols that have no callers.
        let no_callers = col::BTreeSet::new();

        loop {
            // For each symbol whose nullability we need to propagate...
            for nullable in current_nullables.iter() {
                // For each rule that calls that symbol...
                for rule in
                    rules_by_callee.get(nullable).unwrap_or(&no_callers).iter()
                {
                    // The rule is nullable if it has no terminal symbols,
                    // and all the non-terminal symbols are known to be
                    // nullable.
                    let rule_nullable = rule.symbols.iter().all(|s| match s {
                        Symbol::Terminal(_) => false,
                        Symbol::NonTerminal(name) => {
                            self.nullables.contains(name)
                                || current_nullables.contains(name)
                        }
                    });

                    let rule_known = self.nullables.contains(&rule.name)
                        || current_nullables.contains(&rule.name);

                    if rule_nullable && !rule_known {
                        next_nullables.insert(rule.name.clone());
                    }
                }
            }

            self.nullables.append(&mut current_nullables);

            if next_nullables.is_empty() {
                break;
            }

            current_nullables.append(&mut next_nullables);
        }

        Ok(())
    }

    pub fn build(mut self) -> Result<Grammar> {
        if !self.rules.contains_key(&self.start) {
            return Err(Error::StartingRuleNotDefined(self.start.clone()));
        }

        self.propagate_nullability()?;

        Ok(Grammar {
            start: self.start,
            rules: self.rules,
            nullables: self.nullables,
        })
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn grammar_needs_start_symbol_defined() {
        let actual = GrammarBuilder::new("Start")
            .build()
            .expect_err("Built a grammar without rules?");
        let expected = Error::StartingRuleNotDefined("Start".into());

        assert_eq!(actual, expected);
    }

    #[test]
    fn nonterminal_symbols_must_have_rules() {
        let mut builder = GrammarBuilder::new("Start");
        builder
            .rule("Start")
            .terminal('(')
            .nonterminal("Inner")
            .terminal(')')
            .add();

        let actual = builder
            .build()
            .expect_err("Built a grammar with missing rules?");
        let expected = Error::RuleMentionsUnknownSymbol {
            rule: "Start".into(),
            symbol: "Inner".into(),
        };

        assert_eq!(actual, expected);
    }

    #[test]
    fn get_info_for_unknown_symbol_returns_none() {
        let mut builder = GrammarBuilder::new("Start");
        builder.rule("Start").terminal('A').add();

        let grammar = builder.build().expect("Could not build grammar?");
        let name: rc::Rc<str> = "Bogus".into();

        assert!(grammar.rules(&name).is_none());
        assert!(grammar.nullable(&name).is_none());
    }

    #[test]
    fn symbol_with_empty_rule_is_nullable() {
        let mut builder = GrammarBuilder::new("Start");
        builder.rule("Start").terminal('A').add();
        builder.rule("Start").add();

        let grammar = builder.build().expect("Could not build grammar?");
        let name: rc::Rc<str> = "Start".into();

        assert_eq!(grammar.nullable(&name), Some(true));
    }

    #[test]
    fn symbol_without_empty_rule_is_not_nullable() {
        let mut builder = GrammarBuilder::new("Start");
        builder.rule("Start").terminal('A').add();
        builder
            .rule("Start")
            .nonterminal("Start")
            .terminal('A')
            .add();

        let grammar = builder.build().expect("Could not build grammar?");
        let name: rc::Rc<str> = "Start".into();

        assert_eq!(grammar.nullable(&name), Some(false));
    }

    #[test]
    fn symbol_nullable_iff_all_nullable_rule() {
        let mut builder = GrammarBuilder::new("Start");
        // An unused, nullable symbol. It should be OK if it has no callers.
        builder.rule("Unused").add();
        builder.rule("AlwaysEmpty").add();
        builder.rule("NeverEmpty").terminal('N').add();

        // This symbol refers to a nullable symbol, but no rule *only* mentions
        // nullable symbols, so it is not nullable.
        builder
            .rule("NotNullable")
            .nonterminal("AlwaysEmpty")
            .nonterminal("NeverEmpty")
            .add();
        builder
            .rule("NotNullable")
            .nonterminal("NeverEmpty")
            .nonterminal("AlwaysEmpty")
            .add();
        builder
            .rule("NotNullable")
            .terminal('x')
            .nonterminal("AlwaysEmpty")
            .add();

        // This symbol has a rule that contains only nullable symbols,
        // so it is nullable.
        builder
            .rule("Nullable")
            .nonterminal("AlwaysEmpty")
            .nonterminal("AlwaysEmpty")
            .add();
        builder
            .rule("Nullable")
            .nonterminal("NeverEmpty")
            .nonterminal("NeverEmpty")
            .add();

        // A starting symbol for the grammar.
        builder
            .rule("Start")
            .nonterminal("Nullable")
            .nonterminal("NotNullable")
            .add();

        let grammar = builder.build().expect("Could not build grammar?");

        let nullable_name: rc::Rc<str> = "Nullable".into();
        assert_eq!(grammar.nullable(&nullable_name), Some(true));

        let not_nullable_name: rc::Rc<str> = "NotNullable".into();
        assert_eq!(grammar.nullable(&not_nullable_name), Some(false));
    }
}