//! A database instance is a set of ground atoms (facts).

use std::collections::{HashMap, HashSet};
use std::error::Error;
use std::fmt;

use super::atom::Atom;

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum InstanceError {
    NonGroundFact(Atom),
}

impl fmt::Display for InstanceError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            InstanceError::NonGroundFact(atom) => {
                write!(f, "facts must be ground atoms: {}", atom)
            }
        }
    }
}

impl Error for InstanceError {}

/// A database instance containing ground atoms.
#[derive(Debug, Clone, Default)]
pub struct Instance {
    facts_by_predicate: HashMap<String, HashSet<Atom>>,
    len: usize,
}

impl Instance {
    /// Create an empty instance.
    pub fn new() -> Self {
        Instance {
            facts_by_predicate: HashMap::new(),
            len: 0,
        }
    }

    /// Try to add a fact to the instance. Returns true if the fact was new.
    pub fn try_add(&mut self, fact: Atom) -> Result<bool, InstanceError> {
        if !fact.is_ground() {
            return Err(InstanceError::NonGroundFact(fact));
        }

        let bucket = self
            .facts_by_predicate
            .entry(fact.predicate.clone())
            .or_default();
        let inserted = bucket.insert(fact);
        if inserted {
            self.len += 1;
        }

        Ok(inserted)
    }

    /// Add a fact to the instance. Returns true if the fact was new.
    ///
    /// Panics if the atom is not ground.
    pub fn add(&mut self, fact: Atom) -> bool {
        self.try_add(fact).expect("facts must be ground atoms")
    }

    /// Check if the instance contains a fact.
    pub fn contains(&self, fact: &Atom) -> bool {
        self.facts_by_predicate
            .get(&fact.predicate)
            .is_some_and(|facts| facts.contains(fact))
    }

    /// Get the number of facts.
    pub fn len(&self) -> usize {
        self.len
    }

    /// Check if the instance is empty.
    pub fn is_empty(&self) -> bool {
        self.len == 0
    }

    /// Iterate over all facts.
    pub fn iter(&self) -> impl Iterator<Item = &Atom> {
        self.facts_by_predicate
            .values()
            .flat_map(|facts| facts.iter())
    }

    /// Get all facts with a given predicate.
    pub fn facts_for_predicate(&self, predicate: &str) -> Vec<&Atom> {
        self.facts_by_predicate
            .get(predicate)
            .into_iter()
            .flat_map(|facts| facts.iter())
            .collect()
    }
}

impl fmt::Display for Instance {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "Instance {{")?;
        for fact in self.iter() {
            writeln!(f, "  {}", fact)?;
        }
        write!(f, "}}")
    }
}

impl FromIterator<Atom> for Instance {
    fn from_iter<T: IntoIterator<Item = Atom>>(iter: T) -> Self {
        let mut instance = Instance::new();
        for atom in iter {
            instance.add(atom);
        }
        instance
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::chase::term::Term;

    #[test]
    fn test_instance_operations() {
        let mut instance = Instance::new();

        let fact1 = Atom::new(
            "Parent",
            vec![Term::constant("alice"), Term::constant("bob")],
        );
        let fact2 = Atom::new(
            "Parent",
            vec![Term::constant("bob"), Term::constant("carol")],
        );

        assert!(instance.add(fact1.clone()));
        assert!(instance.add(fact2.clone()));
        assert!(!instance.add(fact1.clone())); // Duplicate

        assert_eq!(instance.len(), 2);
        assert!(instance.contains(&fact1));
    }

    #[test]
    fn test_facts_for_predicate() {
        let instance: Instance = vec![
            Atom::new(
                "Parent",
                vec![Term::constant("alice"), Term::constant("bob")],
            ),
            Atom::new("Person", vec![Term::constant("alice")]),
        ]
        .into_iter()
        .collect();

        assert_eq!(instance.facts_for_predicate("Parent").len(), 1);
        assert_eq!(instance.facts_for_predicate("Person").len(), 1);
        assert_eq!(instance.facts_for_predicate("Other").len(), 0);
    }

    #[test]
    fn test_try_add_rejects_non_ground_facts() {
        let mut instance = Instance::new();
        let fact = Atom::new("Parent", vec![Term::var("X"), Term::constant("bob")]);

        let error = instance.try_add(fact).unwrap_err();

        assert!(matches!(error, InstanceError::NonGroundFact(_)));
        assert!(instance.is_empty());
    }
}