use std::collections::HashSet;

use atuin_client::record::sqlite_store::SqliteStore;
use atuin_common::encryption::paseto_v4;
use atuin_domain::record::{
    Host, HostId, Record, RecordId, RecordIdx, RecordSeriesKey, RecordTag, RecordVersion,
};
use entry::KvEntry;
use eyre::{Result, eyre};
use record::KvRecord;

use crate::database::Database;

pub mod entry;
pub mod record;

#[derive(Debug, Clone)]
pub struct KvStore {
    pub record_store: SqliteStore,
    pub kv_db: Database,
    pub host_id: HostId,
    pub encryption_key: paseto_v4::Key,
}

impl KvStore {
    pub fn new(
        record_store: SqliteStore,
        kv_db: Database,
        host_id: HostId,
        encryption_key: paseto_v4::Key,
    ) -> Self {
        Self {
            record_store,
            kv_db,
            host_id,
            encryption_key,
        }
    }

    pub async fn set(&self, namespace: &str, key: &str, value: &str) -> Result<()> {
        let kv_record = KvRecord::builder()
            .namespace(namespace.to_string())
            .key(key.to_string())
            .value(Some(value.to_string()))
            .build();

        self.push_record(kv_record).await?;

        let kv = KvEntry {
            namespace: namespace.to_string(),
            key: key.to_string(),
            value: value.to_string(),
        };

        self.kv_db.save(&kv).await?;

        Ok(())
    }

    pub async fn get(&self, namespace: &str, key: &str) -> Result<Option<String>> {
        let kv = self.kv_db.load(namespace, key).await?;
        Ok(kv.map(|kv| kv.value))
    }

    pub async fn delete(&self, namespace: &str, keys: &[String]) -> Result<()> {
        for key in keys {
            let record = KvRecord::builder()
                .namespace(namespace.to_string())
                .key(key.to_string())
                .value(None)
                .build();

            self.push_record(record).await?;
            self.kv_db.delete(namespace, key).await?;
        }

        Ok(())
    }

    pub async fn list(&self, namespace: Option<&str>) -> Result<Vec<KvEntry>> {
        let entries = self.kv_db.list(namespace).await?;

        Ok(entries)
    }

    async fn push_record(&self, record: KvRecord) -> Result<(RecordId, RecordIdx)> {
        let bytes = record.serialize()?;
        let idx = self
            .record_store
            .last(&RecordSeriesKey::new(self.host_id, RecordTag::Kv))
            .await?
            .map_or(0, |p| p.idx + 1);

        let record = Record::builder()
            .host(Host::new(self.host_id))
            .version(RecordVersion::V1)
            .tag(RecordTag::Kv)
            .idx(idx)
            .data(bytes)
            .build();

        let id = record.id;

        self.record_store.push(&record.encrypt(&self.encryption_key)).await?;

        Ok((id, idx))
    }

    pub async fn build(&self) -> Result<()> {
        let mut tagged = self.record_store.all_tagged(&RecordTag::Kv).await?;
        tagged.reverse();

        let cached = self.kv_db.list(None).await?;

        let mut visited = HashSet::new();
        let mut skipped = 0;

        // Iterate through all KV records from newest to oldest;
        // only visit each KV once, inserting or deleting based on the first time we see it
        for record in tagged {
            // Skip records we can't decrypt or decode, rather than failing the entire build.
            let kv = match record.version {
                RecordVersion::V0 | RecordVersion::V1 => {
                    record.decrypt(&self.encryption_key).and_then(|decrypted| {
                        KvRecord::deserialize(&decrypted.data, &decrypted.version)
                    })
                }
                ref version => Err(eyre!("unknown kv version {version:?}")),
            };

            let kv = match kv {
                Ok(kv) => kv,
                Err(e) => {
                    tracing::warn!("failed to decode kv record, skipping: {e}");
                    skipped += 1;
                    continue;
                }
            };

            let uniq_id = format!("{}.{}", kv.namespace, kv.key);

            if visited.insert(uniq_id) {
                match kv.value {
                    Some(value) => {
                        self.kv_db
                            .save(&KvEntry {
                                namespace: kv.namespace.clone(),
                                key: kv.key.clone(),
                                value,
                            })
                            .await?;
                    }
                    None => {
                        self.kv_db.delete(kv.namespace.as_str(), kv.key.as_str()).await?;
                    }
                }
            }
        }

        // Any KVs that were in the cache but not in the tagged list should be deleted;
        // this should never happen in practice since the cache is always built from the tagged list,
        // but just in case because ** S O F T W A R E **
        for kv in cached {
            if !visited.contains(&format!("{}.{}", kv.namespace, kv.key)) {
                self.kv_db.delete(kv.namespace.as_str(), kv.key.as_str()).await?;
            }
        }

        if skipped > 0 {
            // library code that may run under the TUI or shell hooks, so no stderr here
            tracing::warn!("skipped {skipped} kv records that could not be decrypted or decoded");
        }

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use rstest::{fixture, rstest};

    use super::*;

    #[fixture]
    async fn store() -> KvStore {
        let record_store =
            SqliteStore::new(":memory:", std::time::Duration::from_secs(1)).await.unwrap();
        let kv_db = Database::new(":memory:", std::time::Duration::from_secs(1)).await.unwrap();
        let host_id = atuin_domain::record::HostId(atuin_common::utils::uuid_v7());
        let encryption_key = paseto_v4::Key::from([0; 32]);
        KvStore::new(record_store, kv_db, host_id, encryption_key)
    }

    #[rstest]
    #[tokio::test]
    async fn test_kv_store(#[future(awt)] store: KvStore) -> Result<()> {
        store.set("test", "key", "value").await.unwrap();
        let value = store.get("test", "key").await.unwrap();
        assert_eq!(value, Some("value".to_string()));

        let records = store.record_store.all_tagged(&RecordTag::Kv).await?;
        assert_eq!(records.len(), 1);

        let list = store.list(Some("test")).await.unwrap();
        let expected = vec![KvEntry {
            namespace: "test".to_string(),
            key: "key".to_string(),
            value: "value".to_string(),
        }];
        assert_eq!(list, expected);

        let ns_list = store.list(None).await.unwrap();
        assert_eq!(ns_list, expected);

        store.delete("test", &["key".to_string()]).await.unwrap();
        let value = store.get("test", "key").await.unwrap();
        assert_eq!(value, None);

        let records = store.record_store.all_tagged(&RecordTag::Kv).await?;
        assert_eq!(records.len(), 2);

        Ok(())
    }
}
