hpl-toolkit 0.0.5

use anchor_lang::prelude::*;
use std::collections::HashMap;

#[cfg(feature = "compression")]
use crate::compression::*;

#[cfg(feature = "schema")]
use crate::schema::*;

use super::{short_vec::ShortVecIter, Inner, ShortVec};

#[cfg_attr(feature = "debug", derive(Debug))]
#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
pub struct KeyValue<K: Inner, V: Inner> {
    pub key: K,
    pub value: V,
}

impl<K: Inner, V: Inner> Into<(K, V)> for KeyValue<K, V> {
    fn into(self) -> (K, V) {
        (self.key, self.value)
    }
}
impl<'i, K: Inner, V: Inner> Into<(&'i K, &'i V)> for &'i KeyValue<K, V> {
    fn into(self) -> (&'i K, &'i V) {
        (&self.key, &self.value)
    }
}
impl<'i, K: Inner, V: Inner> Into<(&'i K, &'i mut V)> for &'i mut KeyValue<K, V> {
    fn into(self) -> (&'i K, &'i mut V) {
        (&self.key, &mut self.value)
    }
}
impl<K: Inner, V: Inner> From<(&K, &V)> for KeyValue<K, V> {
    fn from((key, value): (&K, &V)) -> Self {
        KeyValue {
            key: key.to_owned(),
            value: value.to_owned(),
        }
    }
}
impl<K: Inner, V: Inner> From<(K, V)> for KeyValue<K, V> {
    fn from((key, value): (K, V)) -> Self {
        Self { key, value }
    }
}

#[cfg_attr(feature = "debug", derive(Debug))]
#[derive(AnchorSerialize, AnchorDeserialize, Clone, PartialEq)]
pub struct VecMap<K: Inner + PartialEq, V: Inner + PartialEq> {
    inner: ShortVec<KeyValue<K, V>>,
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> VecMap<K, V> {
    pub fn new() -> Self {
        Self {
            inner: ShortVec::new(),
        }
    }

    pub fn intoinner(self) -> ShortVec<KeyValue<K, V>> {
        self.inner
    }

    pub fn insert(&mut self, key: K, value: V) {
        for KeyValue { key: k, value: v } in self.inner.inner_mut() {
            if k == &key {
                *v = value;
                return;
            }
        }
        self.inner.push(KeyValue { key, value });
    }

    pub fn get(&self, key: &K) -> Option<&V> {
        self.inner
            .inner()
            .iter()
            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
    }

    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        self.inner
            .inner_mut()
            .iter_mut()
            .find_map(|KeyValue { key: k, value: v }| if k == key { Some(v) } else { None })
    }

    pub fn iter<'i>(&'i self) -> impl Iterator<Item = (&'i K, &'i V)> {
        let a = self.inner.inner().iter().map(|i| -> (&K, &V) { i.into() });
        a
    }

    pub fn iter_mut<'i>(&'i mut self) -> impl Iterator<Item = (&'i K, &'i mut V)> {
        let a = self
            .inner
            .inner_mut()
            .iter_mut()
            .map(|i| -> (&K, &mut V) { i.into() });

        a
    }

    pub fn remove(&mut self, key: &K) {
        self.inner
            .inner_mut()
            .retain(|KeyValue { key: k, .. }| k != key);
    }
}

// Custom iterator for VecMap
pub struct VecMapIter<K: Inner + PartialEq, V: Inner + PartialEq> {
    inner: ShortVecIter<KeyValue<K, V>>,
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> Iterator for VecMapIter<K, V>
where
    K: Clone,
    V: Clone,
{
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.inner.next().map(|kv| (kv.key, kv.value))
    }
}
impl<K: Inner + PartialEq, V: Inner + PartialEq> FromIterator<(K, V)> for VecMap<K, V> {
    fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> Self {
        let mut vec_map = VecMap::new();
        for (key, value) in iter {
            vec_map.inner.push(KeyValue { key, value });
        }
        vec_map
    }
}
impl<K: Inner + PartialEq, V: Inner + PartialEq> IntoIterator for VecMap<K, V>
where
    K: Inner + PartialEq,
    V: Inner + PartialEq,
{
    type Item = (K, V);
    type IntoIter = VecMapIter<K, V>;

    fn into_iter(self) -> Self::IntoIter {
        VecMapIter {
            inner: self.inner.into_iter(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> FromIterator<KeyValue<K, V>> for VecMap<K, V> {
    fn from_iter<T: IntoIterator<Item = KeyValue<K, V>>>(iter: T) -> Self {
        Self {
            inner: iter.into_iter().collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> From<&HashMap<K, V>> for VecMap<K, V> {
    fn from(map: &HashMap<K, V>) -> Self {
        Self {
            inner: map
                .iter()
                .map(|(k, v)| KeyValue {
                    key: k.clone(),
                    value: v.clone(),
                })
                .collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> From<HashMap<K, V>> for VecMap<K, V> {
    fn from(map: HashMap<K, V>) -> Self {
        Self {
            inner: map
                .into_iter()
                .map(|(k, v)| KeyValue { key: k, value: v })
                .collect(),
        }
    }
}

impl<K: Inner + PartialEq, V: Inner + PartialEq> From<Vec<KeyValue<K, V>>> for VecMap<K, V> {
    fn from(map: Vec<KeyValue<K, V>>) -> Self {
        Self { inner: map.into() }
    }
}

#[cfg(feature = "compression")]
impl<K: Inner + PartialEq + ToNode, V: Inner + PartialEq + ToNode> ToNode for VecMap<K, V> {
    fn to_node(&self) -> [u8; 32] {
        let mut seeds: Vec<[u8; 32]> = vec![];
        for (key, value) in self.iter() {
            seeds.push(key.to_node());
            seeds.push(value.to_node());
        }
        let seeds_refs: Vec<&[u8]> = seeds.iter().map(|node| &node[..]).collect();
        anchor_lang::solana_program::keccak::hashv(&seeds_refs[..]).to_bytes()
    }
}

#[cfg(feature = "schema")]
impl<K: Inner + PartialEq + ToSchema, V: Inner + PartialEq + ToSchema> ToSchema for VecMap<K, V> {
    fn schema() -> Schema {
        Schema::VecMap(Box::new(K::schema()), Box::new(V::schema()))
    }

    fn schema_value(&self) -> SchemaValue {
        let mut schema = VecMap::<SchemaValue, SchemaValue>::new();
        self.inner
            .inner()
            .iter()
            .for_each(|KeyValue { key, value }| {
                schema.insert(key.schema_value(), value.schema_value());
            });
        SchemaValue::VecMap(schema)
    }
}