use std::{
    borrow::Borrow,
    collections::{HashMap, HashSet},
    fmt,
    hash::Hash,
};

pub trait Prefix: Sized + Eq + Hash + Clone {
    fn prefix(&self) -> Option<Self>;
}

#[derive(PartialEq)]
pub enum PrefixNode<K: Prefix, V> {
    Prefix(HashSet<K>),
    Leaf(V),
}

#[derive(Default)]
pub struct PrefixMap<K: Prefix, V> {
    map: HashMap<K, PrefixNode<K, V>>,
}

impl<K: Prefix, V> PrefixNode<K, V> {
    pub fn empty_prefix() -> Self {
        Self::Prefix(HashSet::new())
    }

    pub fn add_suffix(&mut self, suffix: K) {
        if let Self::Prefix(set) = self {
            set.insert(suffix);
        } else {
            *self = Self::Prefix(HashSet::from_iter([suffix]));
        }
    }
}

impl<K: Prefix + fmt::Debug, V: fmt::Debug> fmt::Debug for PrefixNode<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Leaf(leaf) => f.debug_struct("Leaf").field("value", leaf).finish(),
            Self::Prefix(suffixes) => f
                .debug_struct("Prefix")
                .field("suffixes", suffixes)
                .finish(),
        }
    }
}

impl<K: Prefix, V> PrefixMap<K, V> {
    pub fn new() -> Self {
        Self {
            map: HashMap::new(),
        }
    }

    pub fn insert(&mut self, key: K, value: V) {
        if let Some(prefix) = key.prefix() {
            self.map
                .entry(prefix)
                .or_insert_with(PrefixNode::empty_prefix)
                .add_suffix(key.clone());
        }

        // TODO remove all suffixes of `key`, if those exist
        self.map.insert(key, PrefixNode::Leaf(value));
    }

    pub fn get<N>(&self, key: &N) -> Option<&PrefixNode<K, V>>
    where
        K: Borrow<N>,
        N: Eq + Hash + ?Sized,
    {
        self.map.get(key)
    }
}

impl<K: Prefix, V> FromIterator<(K, V)> for PrefixMap<K, V> {
    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
        let mut this = Self::new();
        for (k, v) in iter {
            this.insert(k, v);
        }
        this
    }
}

impl<K: Prefix + fmt::Debug, V: fmt::Debug> fmt::Debug for PrefixMap<K, V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut dm = f.debug_map();

        for (k, v) in self.map.iter().filter_map(|(k, v)| {
            if let PrefixNode::Leaf(leaf) = v {
                Some((k, leaf))
            } else {
                None
            }
        }) {
            dm.entry(k, v);
        }

        dm.finish()
    }
}
