pub mod relation;
pub use relation::*;

use crate::time::*;
use crate::item::*;
use crate::character::*;

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

use enum_map::EnumMap;
use strum::IntoDiscriminant;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Var { index: u64 }

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Value {
    Var(Var),
    Item(Item),
    Character(Character),
    Time(Time),
}
impl From<Var> for Value { fn from(x: Var) -> Self {Value::Var(x)} }
impl From<Item> for Value { fn from(x: Item) -> Self {Value::Item(x)} }
impl From<Character> for Value { fn from(x: Character) -> Self {Value::Character(x)} }
impl From<Time> for Value { fn from(x: Time) -> Self {Value::Time(x)} }
impl TryFrom<Value> for Var {
    type Error = ();
    fn try_from(v: Value) -> Result<Self, Self::Error> {
        if let Value::Var(x) = v { Ok(x) } else { Err(()) }
    }
}
impl TryFrom<Value> for Item {
    type Error = ();
    fn try_from(v: Value) -> Result<Self, Self::Error> {
        if let Value::Item(x) = v { Ok(x) } else { Err(()) }
    }
}
impl TryFrom<Value> for Character {
    type Error = ();
    fn try_from(v: Value) -> Result<Self, Self::Error> {
        if let Value::Character(x) = v { Ok(x) } else { Err(()) }
    }
}
impl TryFrom<Value> for Time {
    type Error = ();
    fn try_from(v: Value) -> Result<Self, Self::Error> {
        if let Value::Time(x) = v { Ok(x) } else { Err(()) }
    }
}

#[derive(Debug, Default)]
pub struct Unification {
    assignment: HashMap<Var, Value>,
}
impl Unification {
    pub fn lookup(&self, v: &Var) -> Option<Value> {
        self.assignment.get(v).cloned()
    }
    pub fn add(&mut self, v: Var, x: impl Into<Value>) -> bool {
        let y = x.into();
        if let Some(old) = self.assignment.get(&v) {
            if let Value::Var(v2) = y {
                self.add(v2, *old)
            } else if let Value::Var(v2) = old {
                self.add(*v2, y)
            } else {
                *old == y
            }
        } else {
            self.assignment.insert(v, y);
            true
        }
    }
    pub fn unify_hole<A, B>(&mut self, h1: &Hole<A>, h2: &Hole<B>) -> bool
    where Value: From<A>, Value: From<B>, A: Copy, B: Copy {
        match (h1, h2) {
            (Hole::Var(v1), Hole::Var(v2)) =>
                self.add(*v1, *v2) && self.add(*v2, *v1),
            (Hole::Val(x1), Hole::Val(x2)) => Value::from(*x1) == Value::from(*x2),
            (Hole::Var(v1), Hole::Val(x2)) => self.add(*v1, *x2),
            (Hole::Val(x1), Hole::Var(v2)) => self.add(*v2, *x1),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Hole<T> {
    Var(Var),
    Val(T),
}
impl<T: Clone + TryFrom<Value>> Hole<T> {
    pub fn substitute(&self, u: &Unification) -> Self {
        match self {
            x@Self::Var(v) => if let Some(y) = u.lookup(v) && let Ok(z) = T::try_from(y) {
                Hole::Val(z)
            } else {
                x.clone()
            }
            x => x.clone(),
        }
    }
}

#[derive(Debug, Clone)]
pub struct Rule {
    pub head: Relation,
    pub body: Vec<Relation>,
}
impl Rule {
    pub fn unify_body(&self, facts: &[&Relation]) -> Option<Unification> {
        if self.body.len() != facts.len() { return None; }
        let mut u = Unification::default();
        for (r, f) in self.body.iter().zip(facts) {
            u = r.unify(f, u)?;
        }
        Some(u)
    }
}

#[derive(Debug, Default, Clone)]
pub struct Facts {
    facts: EnumMap<RelationDiscriminants, HashSet<Relation>>,
}
impl Facts {
    pub fn add(&mut self, f: Relation) -> bool {
        self.facts[f.discriminant()].insert(f)
    }
    pub fn with_sort(&self, s: RelationDiscriminants) -> &HashSet<Relation> {
        &self.facts[s]
    }
    pub fn with_sort_mut(&mut self, s: RelationDiscriminants) -> &mut HashSet<Relation> {
        &mut self.facts[s]
    }
    pub fn contains(&self, f: &Relation) -> bool {
        self.facts[f.discriminant()].contains(f)
    }
    pub fn absorb(&mut self, other: &mut Facts) {
        for (rd, fs) in self.facts.iter_mut() {
            fs.extend(other.with_sort_mut(rd).drain());
        }
    }
}

#[derive(Debug, Default, Clone)]
pub struct Database {
    facts: Facts,
    rules: Vec<Rule>,
}
impl Database {
    pub fn is_known(&self, r: &Relation) -> bool { self.facts.contains(r) }
    pub fn query<'a>(&'a mut self) -> Query<'a> {
        Query {
            db: self,
            facts: Facts::default(),
            rules: Vec::new(),
        }
    }
}

#[must_use]
#[derive(Debug)]
pub struct Query<'a> {
    db: &'a mut Database,
    facts: Facts,
    rules: Vec<Rule>,
}
#[derive(Debug, Clone, Copy)]
enum QueryFactsToConsider<'a> {
    LastIteration(&'a Facts),
    Query,
    DB,
}
impl<'a> Query<'a> {
    fn apply_rule(&self, r: &Rule, new_facts: &mut Facts, required_facts: QueryFactsToConsider) -> bool {
        // we don't need to consider every possible fact against each body term (naive evaluation)
        // instead, we can limit ourselves to the subset where at least one of the facts is new (semi-naive)
        // let found: Vec<Unification> = r.body.iter()
        //     .map(|r| {
        //         let d = r.discriminant();
        //         new_facts.with_sort(d).iter()
        //             .chain(self.facts.with_sort(d).iter())
        //             .chain(self.db.facts.with_sort(d).iter())
        //     })
        //     .multi_cartesian_product()
        //     .filter_map(|fs| r.unify_body(&fs))
        //     .collect();
        let mut found = Vec::new();
        // identify which slots of this rule body possibly have new facts to consider
        let changed_slots = r.body.iter()
            .enumerate()
            .filter_map(|(i, r)| {
                match required_facts {
                    // if we are only considering facts from the provided set
                    QueryFactsToConsider::LastIteration(rf) =>
                        if rf.with_sort(r.discriminant()).is_empty() { None } else { Some(i) },
                    // if we are only considering facts added by this query
                    QueryFactsToConsider::Query =>
                        if self.facts.with_sort(r.discriminant()).is_empty() { None } else { Some(i) },
                    // if we are considering all facts, simply consider every index
                    QueryFactsToConsider::DB => Some(i),
                }
            });
        let empty = HashSet::new();
        for cs in changed_slots {
            let attempts = r.body.iter()
                .enumerate()
                .map(|(i, r)| {
                    let d = r.discriminant();
                    match required_facts {
                        // if this is an index that changed, limit it to only the provided set of facts
                        QueryFactsToConsider::LastIteration(rf) if i == cs =>
                             rf.with_sort(d).iter().chain(empty.iter()),
                        // (or the set of facts from this query)
                        QueryFactsToConsider::Query if i == cs =>
                             self.facts.with_sort(d).iter().chain(empty.iter()),
                        // otherwise, consider all facts - a new fact might be used alongside an unchanged fact!
                        _ =>
                            self.facts.with_sort(d).iter().chain(self.db.facts.with_sort(d).iter()),
                    }
                })
                .multi_cartesian_product();
            for fs in attempts {
                if let Some(u) = r.unify_body(&fs) {
                    found.push(u);
                }
            }
        }
        let mut changed = false;
        for u in found {
            let new = r.head.substitute(&u);
            if new.is_concrete() && !self.facts.contains(&new) && !self.db.facts.contains(&new) {
                if new_facts.add(new) { changed = true; }
            }
        }
        changed
    }
    fn apply_new_once(&mut self, new_facts: &mut Facts, required_facts: QueryFactsToConsider) -> bool {
        let mut changed = false;
        match &required_facts {
            QueryFactsToConsider::DB => {
                // if we are considering only DB facts, we only need to consider new rules
                for r in self.rules.iter() {
                    if self.apply_rule(r, new_facts, required_facts) { changed = true }
                }
            },
            _ => { // otherwise, we should consider all rules
                for r in self.rules.iter().chain(self.db.rules.iter()) {
                    if self.apply_rule(r, new_facts, required_facts) { changed = true }
                }
            }
        }
        changed
    }
    pub fn fact(mut self, r: Relation) -> Self {
        if !self.db.is_known(&r) {
            self.facts.add(r);
        }
        self
    }
    pub fn rule(mut self, r: Rule) -> Self {
        self.rules.push(r);
        self
    }
    fn run(&mut self) {
        let mut new_facts = Facts::default();
        // compare the new rules with the old facts
        self.apply_new_once(&mut new_facts, QueryFactsToConsider::DB);
        // compare the new+old rules with the new facts
        self.apply_new_once(&mut new_facts, QueryFactsToConsider::Query);
        let mut prev_facts = new_facts.clone();
        self.facts.absorb(&mut new_facts);
        while self.apply_new_once(&mut new_facts, QueryFactsToConsider::LastIteration(&prev_facts)) {
            prev_facts = new_facts.clone();
            self.facts.absorb(&mut new_facts);
        }
    }
    pub fn ask(mut self) -> Facts {
        self.run();
        self.facts
    }
    pub fn learn(mut self) {
        self.run();
        self.db.facts.absorb(&mut self.facts);
        self.db.rules.extend_from_slice(&mut self.rules);
    }
}

#[derive(Debug, Default)]
pub struct Context {
    next_index: u64,
    var_names: HashMap<Var, String>,
}
impl Context {
    pub fn fresh(&mut self) -> Var {
        let ret = Var { index: self.next_index };
        self.next_index += 1;
        ret
    }
    pub fn var(&mut self, nm: &str) -> Var {
        let ret = self.fresh();
        self.var_names.insert(ret, nm.to_owned());
        ret
    }
}