#![allow(clippy::type_complexity)]
use core::{
    hash::{BuildHasher, Hash},
    ops::AsyncFnOnce,
};

use alloc::{collections::LinkedList, vec::Vec};

use crate::hash_table::DefaultHashBuilder;

struct LruCacheBucket<K, V> {
    data: LinkedList<(K, V)>,
    entry_count: usize,
    capacity: usize,
}

enum EntryMut<'a, K, V> {
    Vacant(&'a mut LruCacheBucket<K, V>),
    Occupied(&'a mut V),
}

pub struct LruCache<K, V, H: BuildHasher = DefaultHashBuilder> {
    buckets: Vec<LruCacheBucket<K, V>>,
    hasher: H,
}

pub struct FlushIter<'a, K, V, H: BuildHasher> {
    cache: &'a mut LruCache<K, V, H>,
}

impl<K, V, H: BuildHasher> Iterator for FlushIter<'_, K, V, H> {
    type Item = (K, V);

    fn next(&mut self) -> Option<Self::Item> {
        self.cache.pop_entry()
    }
}

impl<K, V> LruCacheBucket<K, V> {
    pub const fn empty(capacity: usize) -> Self {
        Self {
            data: LinkedList::new(),
            capacity,
            entry_count: 0,
        }
    }

    fn pop_front(&mut self) -> Option<(K, V)> {
        self.data.pop_front()
    }
}

impl<K: Eq, V> LruCacheBucket<K, V> {
    pub fn entry_mut(&mut self, key: &K) -> EntryMut<'_, K, V> {
        let mut cursor = self.data.cursor_front_mut();

        loop {
            let Some(node) = cursor.current() else {
                break;
            };

            // If entry found corresponding to the key, remove it and reinsert at the front
            if &node.0 == key {
                // Safety: just checked it's not None above, so safe
                let (key, value) = unsafe { cursor.remove_current().unwrap_unchecked() };

                self.data.push_front((key, value));

                // Safety: just pushed here, so safe
                return EntryMut::Occupied(
                    &mut unsafe { self.data.front_mut().unwrap_unchecked() }.1,
                );
            }

            cursor.move_next();
        }

        EntryMut::Vacant(self)
    }

    fn insert(&mut self, key: K, value: V) -> (&mut V, Option<(K, V)>) {
        let evicted = if self.entry_count + 1 > self.capacity {
            self.data.pop_back()
        } else {
            self.entry_count += 1;
            None
        };

        self.data.push_front((key, value));

        (
            &mut unsafe { self.data.front_mut().unwrap_unchecked() }.1,
            evicted,
        )
    }

    fn pop(&mut self, search_key: &K) -> Option<V> {
        let mut cursor = self.data.cursor_front_mut();
        while let Some((key, _)) = cursor.current() {
            if key == search_key {
                return cursor.remove_current().map(|v| v.1);
            }
        }
        None
    }
}

impl<K, V> LruCache<K, V, DefaultHashBuilder> {
    pub fn with_capacity(bucket_capacity: usize, bucket_count: usize) -> Self {
        let mut buckets = Vec::new();
        for _ in 0..bucket_count {
            buckets.push(LruCacheBucket::empty(bucket_capacity));
        }
        Self {
            buckets,
            hasher: DefaultHashBuilder::new(),
        }
    }
}

impl<K, V, H: BuildHasher> LruCache<K, V, H> {
    pub fn pop_entry(&mut self) -> Option<(K, V)> {
        self.buckets
            .iter_mut()
            .find_map(|bucket| bucket.pop_front())
    }

    pub fn flush(&mut self) -> FlushIter<'_, K, V, H> {
        FlushIter { cache: self }
    }
}

impl<K: Hash + Eq, V, H: BuildHasher> LruCache<K, V, H> {
    fn bucket_mut(&mut self, key: &K) -> &mut LruCacheBucket<K, V> {
        let h = self.hasher.hash_one(key);
        let bucket_count = self.buckets.len();
        &mut self.buckets[h as usize % bucket_count]
    }

    pub async fn try_get_or_insert_with_async<F, E>(
        &mut self,
        key: K,
        supplier: F,
    ) -> Result<(&mut V, Option<(K, V)>), E>
    where
        F: AsyncFnOnce() -> Result<V, E>,
    {
        let h = self.hasher.hash_one(&key);
        let bucket_count = self.buckets.len();
        let bucket = &mut self.buckets[h as usize % bucket_count];

        match bucket.entry_mut(&key) {
            EntryMut::Vacant(bucket) => {
                let value = supplier.async_call_once(()).await?;
                let (value, evicted) = bucket.insert(key, value);
                Ok((value, evicted))
            }
            EntryMut::Occupied(value) => Ok((value, None)),
        }
    }

    pub fn try_get_or_insert_with<F, E>(
        &mut self,
        key: K,
        supplier: F,
    ) -> Result<(&mut V, Option<(K, V)>), E>
    where
        F: FnOnce() -> Result<V, E>,
    {
        let h = self.hasher.hash_one(&key);
        let bucket_count = self.buckets.len();
        let bucket = &mut self.buckets[h as usize % bucket_count];

        match bucket.entry_mut(&key) {
            EntryMut::Vacant(bucket) => {
                let value = supplier()?;
                let (value, evicted) = bucket.insert(key, value);
                Ok((value, evicted))
            }
            EntryMut::Occupied(value) => Ok((value, None)),
        }
    }

    pub fn pop(&mut self, key: &K) -> Option<V> {
        self.bucket_mut(key).pop(key)
    }

    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
        match self.bucket_mut(key).entry_mut(key) {
            EntryMut::Occupied(value) => Some(value),
            _ => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::LruCache;

    #[test]
    fn lru_test_fill() {
        let mut cache = LruCache::with_capacity(4, 1);
        let (_, e) = cache
            .try_get_or_insert_with(0, || Ok::<i32, i32>(1))
            .unwrap();
        assert_eq!(e, None);

        let (_, e) = cache
            .try_get_or_insert_with(1, || Ok::<i32, i32>(2))
            .unwrap();
        assert_eq!(e, None);

        let (_, e) = cache
            .try_get_or_insert_with(2, || Ok::<i32, i32>(3))
            .unwrap();
        assert_eq!(e, None);

        let (_, e) = cache
            .try_get_or_insert_with(3, || Ok::<i32, i32>(4))
            .unwrap();
        assert_eq!(e, None);

        // Refresh the 0, 1 keys
        cache.get_mut(&0);
        cache.get_mut(&1);

        let (_, e) = cache
            .try_get_or_insert_with(4, || Ok::<i32, i32>(5))
            .unwrap();
        assert_eq!(e, Some((2, 3)));
    }
}
