use std::collections::hash_map::Entry;
use std::collections::HashMap;
use std::error::Error;
use std::fmt::Debug;
use std::hash::Hash;
use std::sync::{Arc, Mutex};

use tokio::sync::Mutex as AsyncMutex;
use tokio::time::{Duration, Instant};

use crate::io::Stream;

pub trait CacheKey: Eq + Hash + Debug {}
impl<T: Eq + Hash + Debug> CacheKey for T {}

pub trait StreamCreator: Fn() -> Result<Stream, Box<dyn Error>> + Send + Sync + 'static {}
impl<T: Fn() -> Result<Stream, Box<dyn Error>> + Send + Sync + 'static> StreamCreator for T {}

struct CacheEntry {
    stream: Arc<AsyncMutex<Stream>>,
    last_activity: Instant,
}

const DEFAULT_CLEANUP_INTERVAL: Duration = Duration::from_secs(60);

pub struct StreamsCache<F: StreamCreator, K: CacheKey> {
    new_stream_creator: F,
    entries: Mutex<HashMap<K, CacheEntry>>,
    idle_entry_timeout: Duration,
    cleanup_interval: Duration,
    last_cleanup: Mutex<Instant>,
}

impl<F: StreamCreator, K: CacheKey> StreamsCache<F, K> {
    pub fn new(
        new_stream_creator: F,
        idle_entry_timeout: Duration,
        cleanup_interval: Duration,
    ) -> Self {
        Self {
            new_stream_creator,
            entries: Mutex::new(HashMap::new()),
            idle_entry_timeout,
            cleanup_interval,
            last_cleanup: Mutex::new(Instant::now()),
        }
    }

    pub fn with_default_cleanup_duration(
        new_stream_creator: F,
        idle_entry_timeout: Duration,
    ) -> Self {
        StreamsCache::new(
            new_stream_creator,
            idle_entry_timeout,
            DEFAULT_CLEANUP_INTERVAL,
        )
    }

    pub fn get(&self, key: K, now: Instant) -> Result<Arc<AsyncMutex<Stream>>, Box<dyn Error>> {
        let res = self.get_or_create_stream(key, now);
        let mut last_cleanup = self.last_cleanup.lock().unwrap();
        if now - *last_cleanup > self.cleanup_interval {
            self.cleanup_old_idle_streams(now);
            *last_cleanup = now;
        };
        res
    }

    fn get_or_create_stream(
        &self,
        key: K,
        now: Instant,
    ) -> Result<Arc<AsyncMutex<Stream>>, Box<dyn Error>> {
        let new_stream_creator = &self.new_stream_creator;
        let mut entries = self.entries.lock().unwrap();
        log::debug!("got key {:?}, entries size is {}", key, entries.len());
        let entry = match entries.entry(key) {
            Entry::Occupied(o) => o.into_mut(),
            Entry::Vacant(v) => v.insert(CacheEntry {
                stream: Arc::new(AsyncMutex::new(new_stream_creator()?)),
                last_activity: now,
            }),
        };
        entry.last_activity = now;
        Ok(entry.stream.clone())
    }

    fn cleanup_old_idle_streams(&self, now: Instant) {
        let mut entries = self.entries.lock().unwrap();
        entries.retain(|_, v| now - v.last_activity < self.idle_entry_timeout);
    }
}

#[cfg(test)]
mod tests {
    use tokio::time::Duration;
    use tokio_test::io::Builder;

    use super::*;

    #[test]
    fn stream_cache_get_new_stream_creator_failed() -> Result<(), Box<dyn Error>> {
        let cache = StreamsCache::new(
            || Err(String::from("bla").into()),
            Duration::from_secs(3 * 60),
            Duration::from_secs(60),
        );
        let res = cache.get("bla", Instant::now());
        assert!(res.is_err());
        Ok(())
    }

    #[test]
    fn stream_cache_get_new_stream_success() -> Result<(), Box<dyn Error>> {
        let cache = StreamsCache::new(
            || Ok(Stream::new(Builder::new().build(), Builder::new().build())),
            Duration::from_secs(3 * 60),
            Duration::from_secs(60),
        );

        let now = Instant::now();
        cache.get("bla", now)?;

        let entries = cache.entries.lock().unwrap();
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key("bla"));
        assert_eq!(entries.get("bla").unwrap().last_activity, now);
        let last_cleanup = cache.last_cleanup.lock().unwrap();
        assert_ne!(*last_cleanup, now);
        Ok(())
    }

    #[test]
    fn stream_cache_get_new_stream_another_exists() -> Result<(), Box<dyn Error>> {
        let t1 = Instant::now();
        let cache = StreamsCache::new(
            || Ok(Stream::new(Builder::new().build(), Builder::new().build())),
            Duration::from_secs(3 * 60),
            Duration::from_secs(60),
        );
        {
            let mut entries = cache.entries.lock().unwrap();
            let entry = CacheEntry {
                stream: Arc::new(AsyncMutex::new(Stream::new(
                    Builder::new().build(),
                    Builder::new().build(),
                ))),
                last_activity: t1,
            };
            entries.insert("bli", entry);
        }

        let mut t2 = t1.clone();
        t2 += Duration::from_secs(1);
        cache.get("bla", t2)?;

        let entries = cache.entries.lock().unwrap();
        assert_eq!(entries.len(), 2);
        assert!(entries.contains_key("bla"));
        assert_eq!(entries.get("bla").unwrap().last_activity, t2);
        assert!(entries.contains_key("bli"));
        assert_eq!(entries.get("bli").unwrap().last_activity, t1);
        let last_cleanup = cache.last_cleanup.lock().unwrap();
        assert_ne!(*last_cleanup, t2);
        Ok(())
    }

    #[test]
    fn stream_cache_get_existing_stream_success() -> Result<(), Box<dyn Error>> {
        let mut now = Instant::now();
        let cache = StreamsCache::new(
            || Err(String::from("bla").into()),
            Duration::from_secs(3 * 60),
            Duration::from_secs(60),
        );
        {
            let mut entries = cache.entries.lock().unwrap();
            let entry = CacheEntry {
                stream: Arc::new(AsyncMutex::new(Stream::new(
                    Builder::new().build(),
                    Builder::new().build(),
                ))),
                last_activity: now,
            };
            entries.insert("bla", entry);
        }

        now += Duration::from_secs(1);
        cache.get("bla", now)?;

        let entries = cache.entries.lock().unwrap();
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key("bla"));
        assert_eq!(entries.get("bla").unwrap().last_activity, now);
        let last_cleanup = cache.last_cleanup.lock().unwrap();
        assert_ne!(*last_cleanup, now);
        Ok(())
    }

    #[test]
    fn stream_cache_get_auto_cleanup() -> Result<(), Box<dyn Error>> {
        let mut now = Instant::now();
        let cache = StreamsCache::new(
            || Ok(Stream::new(Builder::new().build(), Builder::new().build())),
            Duration::from_secs(1),
            Duration::from_secs(1),
        );
        {
            let mut entries = cache.entries.lock().unwrap();
            let entry = CacheEntry {
                stream: Arc::new(AsyncMutex::new(Stream::new(
                    Builder::new().build(),
                    Builder::new().build(),
                ))),
                last_activity: now,
            };
            entries.insert("bli", entry);
        }

        now += Duration::from_secs(5);
        cache.get("bla", now)?;

        let entries = cache.entries.lock().unwrap();
        assert_eq!(entries.len(), 1);
        assert!(entries.contains_key("bla"));
        assert_eq!(entries.get("bla").unwrap().last_activity, now);
        let last_cleanup = cache.last_cleanup.lock().unwrap();
        assert_eq!(*last_cleanup, now);
        Ok(())
    }
}