use crate::storage::SyncStorage;
use crate::utils::now_secs;
use dashmap::DashMap;
use pyo3::prelude::*;
use rand::seq::IteratorRandom;
use std::collections::BTreeSet;
use std::sync::{Arc, RwLock};

use super::{CacheEntry, Storage};

pub(crate) struct InMemoryStorage {
    map: DashMap<String, (Arc<CacheEntry>, Option<u64>, u64)>,
    keys_index: RwLock<BTreeSet<String>>,
}

impl InMemoryStorage {
    pub fn new() -> Self {
        Self {
            map: DashMap::new(),
            keys_index: RwLock::new(BTreeSet::new()),
        }
    }

    fn set_internal(&self, key: String, entry: Arc<CacheEntry>, ttl: Option<u64>) -> PyResult<()> {
        let expires_at = ttl.map(|t| now_secs().saturating_add(t));
        let last_accessed = now_secs();
        self.map
            .insert(key.clone(), (entry, expires_at, last_accessed));

        let mut index = self.keys_index.write().unwrap();
        index.insert(key);
        Ok(())
    }
}

impl SyncStorage for InMemoryStorage {
    #[inline]
    fn get(&self, _py: Python, key: &str) -> super::StorageResult {
        let mut entry = match self.map.get_mut(key) {
            Some(e) => e,
            None => return super::StorageResult::NotFound,
        };
        let (val, expires_at, last_accessed) = entry.value_mut();

        if expires_at.is_some_and(|expires| now_secs() > expires) {
            drop(entry);
            self.map.remove(key);
            let mut index = self.keys_index.write().unwrap();
            index.remove(key);
            return super::StorageResult::Expired;
        }

        *last_accessed = now_secs();
        super::StorageResult::Hit(Arc::clone(val), *expires_at, None)
    }

    fn set(&self, key: String, entry: Arc<CacheEntry>, ttl: Option<u64>) -> PyResult<()> {
        self.set_internal(key, entry, ttl)
    }

    #[inline]
    fn touch_batch(&self, updates: Vec<(String, Option<u64>)>) -> PyResult<()> {
        for (key, ttl) in updates {
            if let Some(mut entry) = self.map.get_mut(&key) {
                if let Some(t) = ttl {
                    entry.1 = Some(now_secs().saturating_add(t));
                }
                entry.2 = now_secs();
            }
        }
        Ok(())
    }

    #[inline]
    fn remove(&self, key: &str) -> PyResult<()> {
        if self.map.remove(key).is_some() {
            let mut index = self.keys_index.write().unwrap();
            index.remove(key);
        }
        Ok(())
    }

    fn clear(&self) -> PyResult<()> {
        self.map.clear();
        let mut index = self.keys_index.write().unwrap();
        index.clear();
        Ok(())
    }

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

    fn evict_lru(&self, count: usize) -> PyResult<Vec<String>> {
        if self.map.is_empty() || count == 0 {
            return Ok(Vec::new());
        }

        let sample_size = count.saturating_mul(5).min(self.map.len());

        let mut oldest_items: Vec<(String, u64)> = Vec::with_capacity(count);
        let mut rng = rand::rng();

        let sample_keys: Vec<String> = {
            let index = self.keys_index.read().unwrap();
            index
                .iter()
                .sample(&mut rng, sample_size)
                .into_iter()
                .cloned()
                .collect()
        };

        for key in sample_keys {
            let ts = if let Some(entry) = self.map.get(&key) {
                entry.value().2
            } else {
                continue;
            };

            if oldest_items.len() < count {
                oldest_items.push((key, ts));
                oldest_items.sort_unstable_by_key(|&(_, t)| t);
            } else if ts < oldest_items.last().unwrap().1 {
                oldest_items.pop();
                oldest_items.push((key, ts));
                oldest_items.sort_unstable_by_key(|&(_, t)| t);
            }
        }

        let to_evict: Vec<String> = oldest_items.into_iter().map(|(k, _)| k).collect();

        let mut index = self.keys_index.write().unwrap();
        for key in &to_evict {
            self.map.remove(key);
            index.remove(key);
        }

        Ok(to_evict)
    }

    fn scan_keys(&self, prefix: &str) -> Vec<(String, Option<u64>)> {
        let mut results = Vec::new();
        let index = self.keys_index.read().unwrap();

        let it = index.range::<str, _>((
            std::ops::Bound::Included(prefix),
            std::ops::Bound::Unbounded,
        ));
        for key in it {
            if !key.starts_with(prefix) {
                break;
            }
            if let Some(entry) = self.map.get(key) {
                let (_, expires_at, _) = entry.value();
                let is_expired = expires_at.is_some_and(|expires| now_secs() > expires);
                if !is_expired {
                    results.push((key.clone(), *expires_at));
                }
            }
        }
        results
    }
}

use async_trait::async_trait;

#[async_trait]
impl Storage for InMemoryStorage {
    #[inline]
    async fn get(&self, key: &str) -> super::StorageResult {
        Python::attach(|py| SyncStorage::get(self, py, key))
    }

    fn try_get_sync(&self, py: Python, key: &str) -> Option<super::StorageResult> {
        Some(SyncStorage::get(self, py, key))
    }

    async fn set(&self, key: String, entry: Arc<CacheEntry>, ttl: Option<u64>) -> PyResult<()> {
        SyncStorage::set(self, key, entry, ttl)
    }

    #[inline]
    async fn touch_batch(&self, updates: Vec<(String, Option<u64>)>) -> PyResult<()> {
        SyncStorage::touch_batch(self, updates)
    }

    #[inline]
    async fn remove(&self, key: &str) -> PyResult<()> {
        SyncStorage::remove(self, key)
    }

    async fn clear(&self) -> PyResult<()> {
        SyncStorage::clear(self)
    }

    async fn len(&self) -> usize {
        SyncStorage::len(self)
    }

    async fn evict_lru(&self, count: usize) -> PyResult<Vec<String>> {
        SyncStorage::evict_lru(self, count)
    }

    async fn scan_keys(&self, prefix: &str) -> Vec<(String, Option<u64>)> {
        SyncStorage::scan_keys(self, prefix)
    }

    fn try_set_sync(
        &self,
        _py: Python,
        key: String,
        entry: Arc<CacheEntry>,
        ttl: Option<u64>,
    ) -> PyResult<()> {
        SyncStorage::set(self, key, entry, ttl)
    }

    fn try_remove_sync(&self, key: &str) -> PyResult<()> {
        SyncStorage::remove(self, key)
    }

    fn try_clear_sync(&self) -> PyResult<()> {
        SyncStorage::clear(self)
    }

    fn try_len_sync(&self) -> Option<usize> {
        Some(SyncStorage::len(self))
    }

    fn try_evict_lru_sync(&self, count: usize) -> Option<PyResult<Vec<String>>> {
        Some(SyncStorage::evict_lru(self, count))
    }

    fn try_scan_keys_sync(&self, prefix: &str) -> Option<Vec<(String, Option<u64>)>> {
        Some(SyncStorage::scan_keys(self, prefix))
    }

    fn is_sync_storage(&self) -> bool {
        true
    }
}
