use std::collections::HashMap;
use std::hash::{BuildHasher, Hash};

use crate::ArcInfo;

use super::ArcStorage;

/// Arc storage with hash maps.
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde1", derive(serde::Serialize, serde::Deserialize))]
pub struct HashMapStorage<I, A>
where
    I: Hash + Eq,
{
    forward: HashMap<I, HashMap<I, usize>>,
    reverse: HashMap<I, HashMap<I, usize>>,
    arc_attributes: HashMap<usize, ArcInfo<I, A>>,
    next_arc_id: usize,
}

impl<I, A> HashMapStorage<I, A>
where
    I: Hash + Eq,
{
    /// Create a new `HashMapStorage`.
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl<I, A> PartialEq for HashMapStorage<I, A>
where
    I: PartialEq + Hash + Eq + Clone,
    A: PartialEq,
{
    fn eq(&self, other: &Self) -> bool {
        self.arc_iter()
            .all(|arc| other.arc(&arc.tail, &arc.head) == Some(&arc.attributes))
    }
}

impl<I, A> Default for HashMapStorage<I, A>
where
    I: Hash + Eq,
{
    fn default() -> Self {
        Self {
            forward: HashMap::default(),
            reverse: HashMap::default(),
            arc_attributes: HashMap::default(),
            next_arc_id: Default::default(),
        }
    }
}

impl<I, A> ArcStorage<I, A> for HashMapStorage<I, A>
where
    I: Hash + Eq + Clone,
{
    fn with_capacity(n_nodes: usize, m_arcs: usize) -> Self {
        Self {
            forward: HashMap::with_capacity(n_nodes),
            reverse: HashMap::with_capacity(n_nodes),
            arc_attributes: HashMap::with_capacity(m_arcs),
            next_arc_id: 0,
        }
    }

    fn m_arcs(&self) -> usize {
        self.arc_attributes.len()
    }

    fn arc(&self, tail: &I, head: &I) -> Option<&A> {
        self.forward
            .get(tail)
            .and_then(|x| x.get(head))
            .and_then(|id| self.arc_attributes.get(id).map(|a| &a.attributes))
    }

    fn arc_mut(&mut self, tail: &I, head: &I) -> Option<&mut A> {
        self.forward
            .get(tail)
            .and_then(|x| x.get(head))
            .and_then(|id| self.arc_attributes.get_mut(id).map(|a| &mut a.attributes))
    }

    fn contains_arc(&self, tail: &I, head: &I) -> bool {
        self.forward
            .get(tail)
            .and_then(|x| x.get(head))
            .is_some_and(|id| self.arc_attributes.contains_key(id))
    }

    fn add_arc<T: Into<crate::ArcInfo<I, A>>>(&mut self, arc: T) {
        let arc = arc.into();
        let id = self.next_arc_id;
        self.next_arc_id += 1;

        self.forward
            .entry(arc.tail.clone())
            .or_default()
            .insert(arc.head.clone(), id);

        self.reverse
            .entry(arc.head.clone())
            .or_default()
            .insert(arc.tail.clone(), id);

        self.arc_attributes.insert(id, arc);
    }

    fn remove_arc(&mut self, tail: &I, head: &I) -> Option<A> {
        let forward = self.forward.get_mut(tail)?;
        let reverse = self.reverse.get_mut(head)?;

        let forward_id = forward.get(head).copied();
        let reverse_id = reverse.get(tail).copied();

        match (forward_id, reverse_id) {
            (None, None) => None,
            (Some(x), Some(y)) if x == y => {
                forward.remove(head);
                reverse.remove(tail);
                self.arc_attributes.remove(&x).map(|a| a.attributes)
            }
            (Some(_) | None, Some(_)) | (Some(_), None) => unreachable!(),
        }
    }

    fn arc_iter<'a>(&'a self) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.forward
            .iter()
            .flat_map(|(_tail, f)| f.values().map(|id| &self.arc_attributes[id]))
    }

    fn arc_iter_mut<'a>(&'a mut self) -> impl Iterator<Item = &'a mut ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.arc_attributes.values_mut()
    }

    fn forward_arcs<'a>(&'a self, node: &'a I) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.forward
            .get(node)
            .map_or_else(std::collections::hash_map::Iter::default, |f| f.iter())
            .map(|(_, id)| {
                self.arc_attributes
                    .get(id)
                    .expect("Forward arcs should always point to a valid node")
            })
    }

    fn reverse_arcs<'a>(&'a self, node: &'a I) -> impl Iterator<Item = &'a ArcInfo<I, A>> + 'a
    where
        A: 'a,
        I: 'a,
    {
        self.reverse
            .get(node)
            .map_or_else(std::collections::hash_map::Iter::default, |f| f.iter())
            .map(|(_, id)| {
                self.arc_attributes
                    .get(id)
                    .expect("Reverse arcs should always point to a valid node")
            })
    }

    fn has_forward_arcs(&self, node: &I) -> bool {
        self.forward.get(node).is_some_and(|a| !a.is_empty())
    }

    fn has_reverse_arcs(&self, node: &I) -> bool {
        self.reverse.get(node).is_some_and(|a| !a.is_empty())
    }

    fn remove_arcs_with_node(&mut self, id: &I) {
        if let Some(reverse) = self.forward.remove(id) {
            for (rid, arc_id) in reverse {
                if let Some(reverse_tails) = self.reverse.get_mut(&rid) {
                    reverse_tails.remove(id);
                }
                self.arc_attributes.remove(&arc_id);
            }
        }
    }
}

impl<I, A> FromIterator<ArcInfo<I, A>> for HashMapStorage<I, A>
where
    I: Hash + Eq + Clone,
{
    fn from_iter<T: IntoIterator<Item = ArcInfo<I, A>>>(iter: T) -> Self {
        let mut arcs = Self::new();
        for arc in iter {
            arcs.add_arc(arc);
        }

        arcs
    }
}

impl<I, A> IntoIterator for HashMapStorage<I, A>
where
    I: Hash + Eq,
{
    type Item = ArcInfo<I, A>;

    type IntoIter = std::collections::hash_map::IntoValues<usize, ArcInfo<I, A>>;

    fn into_iter(self) -> Self::IntoIter {
        self.arc_attributes.into_values()
    }
}

impl<I, A> From<HashMap<(I, I), A>> for HashMapStorage<I, A>
where
    I: Hash + Eq + Clone,
{
    fn from(value: HashMap<(I, I), A>) -> Self {
        let mut hms = Self::new();

        for ((to, from), attr) in value {
            hms.add_arc((to, from, attr));
        }

        hms
    }
}

impl<I, A, S> From<HashMapStorage<I, A>> for HashMap<(I, I), A, S>
where
    I: Hash + Eq,
    S: BuildHasher + Default,
{
    fn from(value: HashMapStorage<I, A>) -> Self {
        value
            .into_iter()
            .map(|a: ArcInfo<I, A>| ((a.tail, a.head), a.attributes))
            .collect()
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use crate::{ArcStorage, HashMapStorage};

    #[test]
    fn from_hashmap() {
        let map = HashMap::from([((0, 1), 'a'), ((1, 2), 'b')]);
        let hms: HashMapStorage<u8, char> = HashMapStorage::from(map);

        assert_eq!(
            hms.forward_arcs(&0)
                .map(|a| (a.head, a.attributes))
                .collect::<Vec<_>>(),
            vec![(1, 'a')]
        );
        assert_eq!(
            hms.forward_arcs(&1)
                .map(|a| (a.head, a.attributes))
                .collect::<Vec<_>>(),
            vec![(2, 'b')]
        );
    }
}
