use nostr_rs_storage_base::{Error, EventFilter, Index, SecondaryIndex, Storage};
use nostr_rs_types::types::{Event, Filter};
use std::{
    cmp::min,
    collections::{BTreeMap, VecDeque},
    sync::{
        atomic::{AtomicUsize, Ordering},
        Arc,
    },
};
use tokio::sync::RwLock;

mod cursor;

#[derive(Default)]
pub struct Indexes {
    // Vec<u8> is used instead of Id, since references can be partial keys
    author: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
    tags: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
    kind: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
    ids_by_time: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
    /// silly index to store the keys of the indexes, but it helps to use the same code
    ids: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
    local_events: RwLock<BTreeMap<Vec<u8>, Vec<u8>>>,
}

#[derive(Default)]
pub struct Memory {
    db: Arc<RwLock<BTreeMap<[u8; 32], Event>>>,
    indexes: Arc<Indexes>,
    indexers_running: Arc<AtomicUsize>,
}

#[async_trait::async_trait]
impl Storage for Memory {
    type Cursor<'a> = cursor::Cursor<'a>;

    fn is_flushing(&self) -> bool {
        self.indexers_running.load(Ordering::SeqCst) > 0
    }

    async fn store(&self, event: &Event) -> Result<bool, Error> {
        let mut db = self.db.write().await;
        if db.get(&event.id.0).is_some() {
            return Ok(false);
        }

        db.insert(event.id.0, event.clone());

        let event = event.clone();
        let indexes = self.indexes.clone();
        let indexer_running = self.indexers_running.clone();
        let _ = indexer_running.fetch_add(1, Ordering::SeqCst);
        tokio::spawn(async move {
            // build indexes asynchronously to avoid locking contention and make
            // `storing` slower, as the stream will hold a read lock while iterating
            let secondary_index = SecondaryIndex::new(&event.id, event.created_at());
            let event_id = event.id.clone().0.to_vec();

            let time_desc = u64::MAX
                .checked_sub(
                    event
                        .created_at()
                        .timestamp_micros()
                        .try_into()
                        .unwrap_or_default(),
                )
                .unwrap_or_default()
                .to_be_bytes();

            indexes
                .ids
                .write()
                .await
                .insert(event_id.clone(), event_id.clone());

            indexes
                .ids_by_time
                .write()
                .await
                .insert(secondary_index.index_by(time_desc), event_id.clone());

            indexes
                .author
                .write()
                .await
                .insert(secondary_index.index_by(event.author()), event_id.clone());

            let kind: u32 = event.kind().into();

            indexes.kind.write().await.insert(
                secondary_index.index_by(kind.to_be_bytes()),
                event_id.clone(),
            );

            for event_tags in event
                .tags()
                .iter()
                .filter_map(|tag| tag.clone().into_bytes())
            {
                indexes
                    .tags
                    .write()
                    .await
                    .insert(secondary_index.index_by(event_tags), event_id.clone());
            }
            let _ = indexer_running.fetch_sub(1, Ordering::SeqCst);
        });
        Ok(true)
    }

    async fn set_local_event(&self, event: &Event) -> Result<(), Error> {
        let mut local_events = self.indexes.local_events.write().await;
        local_events.insert(
            SecondaryIndex::new(&event.id, event.created_at()).index_by([]),
            event.id.0.to_vec(),
        );
        Ok(())
    }

    async fn get_event<T: AsRef<[u8]> + Send + Sync>(&self, id: T) -> Result<Option<Event>, Error> {
        let events = self.db.read().await;
        let mut id_filter = [0u8; 32];
        let id = id.as_ref();
        let slice = 0..min(id.len(), 32);
        id_filter[slice.clone()].copy_from_slice(&id[slice]);

        if let Some((db_id, value)) = events.range(id_filter..).next() {
            if db_id.starts_with(id) {
                return Ok(Some(value.clone()));
            }
        }

        return Ok(None);
    }

    async fn get_by_filter(&self, mut query: Filter) -> Result<Self::Cursor<'_>, Error> {
        let limit = if query.limit == 0 {
            None
        } else {
            Some(query.limit.try_into()?)
        };

        let query_index: Index = (&mut query).into();

        let (index, index_prefixes) = match query_index {
            Index::Tag(tag_name, tags) => (
                self.indexes.tags.read().await,
                tags.into_iter()
                    .map(|tag| tag.into_bytes(&tag_name))
                    .collect::<VecDeque<_>>(),
            ),
            Index::Author(authors) => (
                self.indexes.author.read().await,
                authors
                    .into_iter()
                    .map(|c| c.into_bytes())
                    .collect::<VecDeque<_>>(),
            ),
            Index::Id(ids) => (
                self.indexes.ids.read().await,
                ids.into_iter()
                    .map(|c| c.into_bytes())
                    .collect::<VecDeque<_>>(),
            ),
            Index::Kind(kind) => (
                self.indexes.kind.read().await,
                kind.into_iter()
                    .map(|kind| kind.into_bytes())
                    .collect::<VecDeque<_>>(),
            ),

            Index::TableScan => (
                self.indexes.ids_by_time.read().await,
                vec![Vec::new()].into(), // all keys
            ),
        };

        let filter: EventFilter = query.into();

        Ok(Self::Cursor {
            db: self,
            index,
            index_prefixes,
            current_index_key: Vec::new(),
            filter: filter.into_option(),
            limit,
            returned: 0,
            last_prefix: None,
            future_event: None,
        })
    }

    async fn get_local_events(&self, limit: Option<usize>) -> Result<Self::Cursor<'_>, Error> {
        Ok(Self::Cursor {
            db: self,
            index: self.indexes.local_events.read().await,
            index_prefixes: vec![Vec::new()].into(), // all keys
            current_index_key: Vec::new(),
            filter: None,
            limit,
            returned: 0,
            last_prefix: None,
            future_event: None,
        })
    }
}

#[cfg(test)]
mod test {
    use super::*;
    async fn new_instance(_path: &str) -> Memory {
        Memory::default()
    }

    async fn destroy_instance(_path: &str) {}

    nostr_rs_storage_base::storage_test!(Memory, new_instance, destroy_instance);
}
