/// RecordStore
/// Keeps an LRU cache of dht keys and their associated subkey valuedata.
/// Instances of this store are used for 'local' (persistent) and 'remote' (ephemeral) dht key storage.
/// This store does not perform any validation on the schema, and all ValueRecordData passed in must have been previously validated.
/// Uses an in-memory store for the records, backed by the TableStore. Subkey data is LRU cached and rotated out by a limits policy,
/// and backed to the TableStore for persistence.
use super::*;
use hashlink::LruCache;

pub struct RecordStore<D>
where
    D: fmt::Debug + Clone + RkyvArchive + RkyvSerialize<DefaultVeilidRkyvSerializer>,
    for<'t> <D as RkyvArchive>::Archived: CheckBytes<RkyvDefaultValidator<'t>>,
    <D as RkyvArchive>::Archived: RkyvDeserialize<D, VeilidSharedDeserializeMap>,
{
    table_store: TableStore,
    name: String,
    limits: RecordStoreLimits,

    record_table: Option<TableDB>,
    subkey_table: Option<TableDB>,
    record_index: LruCache<RecordTableKey, Record<D>>,
    subkey_cache: LruCache<SubkeyTableKey, RecordData>,
    subkey_cache_total_size: LimitedSize<usize>,
    total_storage_space: LimitedSize<u64>,

    dead_records: Vec<(RecordTableKey, Record<D>)>,
    changed_records: HashSet<RecordTableKey>,

    purge_dead_records_mutex: Arc<AsyncMutex<()>>,
}

/// The result of the do_get_value_operation
#[derive(Default, Debug)]
pub struct SubkeyResult {
    /// The subkey value if we got one
    pub value: Option<SignedValueData>,
    /// The descriptor if we got a fresh one or empty if no descriptor was needed
    pub descriptor: Option<SignedValueDescriptor>,
}

impl<D> RecordStore<D>
where
    D: fmt::Debug + Clone + RkyvArchive + RkyvSerialize<DefaultVeilidRkyvSerializer>,
    for<'t> <D as RkyvArchive>::Archived: CheckBytes<RkyvDefaultValidator<'t>>,
    <D as RkyvArchive>::Archived: RkyvDeserialize<D, VeilidSharedDeserializeMap>,
{
    pub fn new(table_store: TableStore, name: &str, limits: RecordStoreLimits) -> Self {
        let subkey_cache_size = limits.subkey_cache_size as usize;
        let limit_subkey_cache_total_size = limits
            .max_subkey_cache_memory_mb
            .map(|mb| mb * 1_048_576usize);
        let limit_max_storage_space = limits
            .max_storage_space_mb
            .map(|mb| mb as u64 * 1_048_576u64);

        Self {
            table_store,
            name: name.to_owned(),
            limits,
            record_table: None,
            subkey_table: None,
            record_index: LruCache::new(limits.max_records.unwrap_or(usize::MAX)),
            subkey_cache: LruCache::new(subkey_cache_size),
            subkey_cache_total_size: LimitedSize::new(
                "subkey_cache_total_size",
                0,
                limit_subkey_cache_total_size,
            ),
            total_storage_space: LimitedSize::new(
                "total_storage_space",
                0,
                limit_max_storage_space,
            ),
            dead_records: Vec::new(),
            changed_records: HashSet::new(),
            purge_dead_records_mutex: Arc::new(AsyncMutex::new(())),
        }
    }

    pub async fn init(&mut self) -> EyreResult<()> {
        let record_table = self
            .table_store
            .open(&format!("{}_records", self.name), 1)
            .await?;
        let subkey_table = self
            .table_store
            .open(&&format!("{}_subkeys", self.name), 1)
            .await?;

        // Pull record index from table into a vector to ensure we sort them
        let record_table_keys = record_table.get_keys(0).await?;
        let mut record_index_saved: Vec<(RecordTableKey, Record<D>)> =
            Vec::with_capacity(record_table_keys.len());
        for rtk in record_table_keys {
            if let Some(vr) = record_table.load_rkyv::<Record<D>>(0, &rtk).await? {
                let rik = RecordTableKey::try_from(rtk.as_ref())?;
                record_index_saved.push((rik, vr));
            }
        }

        // Sort the record index by last touched time and insert in sorted order
        record_index_saved.sort_by(|a, b| a.1.last_touched().cmp(&b.1.last_touched()));
        let mut dead_records = Vec::new();
        for ri in record_index_saved {
            // total the storage space
            self.total_storage_space
                .add(mem::size_of::<RecordTableKey>() as u64)
                .unwrap();
            self.total_storage_space
                .add(ri.1.total_size() as u64)
                .unwrap();
            if let Err(_) = self.total_storage_space.commit() {
                // If we overflow the limit, kill off the record
                dead_records.push((ri.0, ri.1));
                continue;
            }

            // add to index and ensure we deduplicate in the case of an error
            if let Some(v) = self.record_index.insert(ri.0, ri.1, |k, v| {
                // If the configuration change, we only want to keep the 'limits.max_records' records
                dead_records.push((k, v));
            }) {
                // This shouldn't happen, but deduplicate anyway
                log_stor!(warn "duplicate record in table: {:?}", ri.0);
                dead_records.push((ri.0, v));
            }
        }
        for (k, v) in dead_records {
            self.add_dead_record(k, v);
        }

        self.record_table = Some(record_table);
        self.subkey_table = Some(subkey_table);
        Ok(())
    }

    fn add_dead_record(&mut self, key: RecordTableKey, record: Record<D>) {
        self.dead_records.push((key, record));
    }

    fn mark_record_changed(&mut self, key: RecordTableKey) {
        self.changed_records.insert(key);
    }

    fn add_to_subkey_cache(&mut self, key: SubkeyTableKey, record_data: RecordData) {
        let record_data_total_size = record_data.total_size();
        // Write to subkey cache
        let mut dead_size = 0usize;
        if let Some(old_record_data) = self.subkey_cache.insert(key, record_data, |_, v| {
            // LRU out
            dead_size += v.total_size();
        }) {
            // Old data
            dead_size += old_record_data.total_size();
        }
        self.subkey_cache_total_size.sub(dead_size).unwrap();
        self.subkey_cache_total_size
            .add(record_data_total_size)
            .unwrap();

        // Purge over size limit
        while self.subkey_cache_total_size.commit().is_err() {
            if let Some((_, v)) = self.subkey_cache.remove_lru() {
                self.subkey_cache_total_size.saturating_sub(v.total_size());
            } else {
                self.subkey_cache_total_size.rollback();

                log_stor!(error "subkey cache should not be empty, has {} bytes unaccounted for",self.subkey_cache_total_size.get());

                self.subkey_cache_total_size.set(0);
                self.subkey_cache_total_size.commit().unwrap();
                break;
            }
        }
    }

    fn remove_from_subkey_cache(&mut self, key: SubkeyTableKey) {
        if let Some(dead_record_data) = self.subkey_cache.remove(&key) {
            self.subkey_cache_total_size
                .saturating_sub(dead_record_data.total_size());
            self.subkey_cache_total_size.commit().unwrap();
        }
    }

    async fn purge_dead_records(&mut self, lazy: bool) {
        let purge_dead_records_mutex = self.purge_dead_records_mutex.clone();
        let _lock = if lazy {
            match asyncmutex_try_lock!(purge_dead_records_mutex) {
                Some(v) => v,
                None => {
                    // If not ready now, just skip it if we're lazy
                    return;
                }
            }
        } else {
            // Not lazy, must wait
            purge_dead_records_mutex.lock().await
        };

        // Delete dead keys
        if self.dead_records.is_empty() {
            return;
        }

        let record_table = self.record_table.clone().unwrap();
        let subkey_table = self.subkey_table.clone().unwrap();

        let rt_xact = record_table.transact();
        let st_xact = subkey_table.transact();
        let dead_records = mem::take(&mut self.dead_records);
        for (k, v) in dead_records {
            // Record should already be gone from index
            if self.record_index.contains_key(&k) {
                log_stor!(error "dead record found in index: {:?}", k);
            }

            // Delete record
            if let Err(e) = rt_xact.delete(0, &k.bytes()) {
                log_stor!(error "record could not be deleted: {}", e);
            }

            // Delete subkeys
            let stored_subkeys = v.stored_subkeys();
            for sk in stored_subkeys.iter() {
                // From table
                let stk = SubkeyTableKey {
                    key: k.key,
                    subkey: sk,
                };
                let stkb = stk.bytes();
                if let Err(e) = st_xact.delete(0, &stkb) {
                    log_stor!(error "subkey could not be deleted: {}", e);
                }

                // From cache
                self.remove_from_subkey_cache(stk);
            }

            // Remove from total size
            self.total_storage_space
                .saturating_sub(mem::size_of::<RecordTableKey>() as u64);
            self.total_storage_space
                .saturating_sub(v.total_size() as u64);
            self.total_storage_space.commit().unwrap();
        }
        if let Err(e) = rt_xact.commit().await {
            log_stor!(error "failed to commit record table transaction: {}", e);
        }
        if let Err(e) = st_xact.commit().await {
            log_stor!(error "failed to commit subkey table transaction: {}", e);
        }
    }

    async fn flush_changed_records(&mut self) {
        // touch records
        if self.changed_records.is_empty() {
            return;
        }

        let record_table = self.record_table.clone().unwrap();

        let rt_xact = record_table.transact();
        let changed_records = mem::take(&mut self.changed_records);
        for rtk in changed_records {
            // Get the changed record and save it to the table
            if let Some(r) = self.record_index.peek(&rtk) {
                if let Err(e) = rt_xact.store_rkyv(0, &rtk.bytes(), r) {
                    log_stor!(error "failed to save record: {}", e);
                }
            }
        }
        if let Err(e) = rt_xact.commit().await {
            log_stor!(error "failed to commit record table transaction: {}", e);
        }
    }

    pub async fn tick(&mut self) -> EyreResult<()> {
        self.flush_changed_records().await;
        self.purge_dead_records(true).await;
        Ok(())
    }

    pub async fn new_record(&mut self, key: TypedKey, record: Record<D>) -> VeilidAPIResult<()> {
        let rtk = RecordTableKey { key };
        if self.record_index.contains_key(&rtk) {
            apibail_internal!("record already exists");
        }

        // Get record table
        let Some(record_table) = self.record_table.clone() else {
            apibail_internal!("record store not initialized");
        };

        // If over size limit, dont create record
        self.total_storage_space
            .add((mem::size_of::<RecordTableKey>() + record.total_size()) as u64)
            .unwrap();
        if !self.total_storage_space.check_limit() {
            self.total_storage_space.rollback();
            apibail_try_again!();
        }

        // Save to record table
        record_table
            .store_rkyv(0, &rtk.bytes(), &record)
            .await
            .map_err(VeilidAPIError::internal)?;

        // Save to record index
        let mut dead_records = Vec::new();
        if let Some(v) = self.record_index.insert(rtk, record, |k, v| {
            dead_records.push((k, v));
        }) {
            // Shouldn't happen but log it
            log_stor!(warn "new duplicate record in table: {:?}", rtk);
            self.add_dead_record(rtk, v);
        }
        for dr in dead_records {
            self.add_dead_record(dr.0, dr.1);
        }

        // Update storage space
        self.total_storage_space.commit().unwrap();

        Ok(())
    }

    pub async fn delete_record(&mut self, key: TypedKey) -> VeilidAPIResult<()> {
        // Get the record table key
        let rtk = RecordTableKey { key };

        // Remove record from the index
        let Some(record) = self.record_index.remove(&rtk) else {
            apibail_key_not_found!(key);
        };

        self.add_dead_record(rtk, record);

        self.purge_dead_records(false).await;

        Ok(())
    }

    pub(super) fn with_record<R, F>(&mut self, key: TypedKey, f: F) -> Option<R>
    where
        F: FnOnce(&Record<D>) -> R,
    {
        // Get record from index
        let mut out = None;
        let rtk = RecordTableKey { key };
        if let Some(record) = self.record_index.get_mut(&rtk) {
            // Callback
            out = Some(f(record));

            // Touch
            record.touch(get_aligned_timestamp());
        }
        if out.is_some() {
            self.mark_record_changed(rtk);
        }

        out
    }

    pub(super) fn peek_record<R, F>(&self, key: TypedKey, f: F) -> Option<R>
    where
        F: FnOnce(&Record<D>) -> R,
    {
        // Get record from index
        let mut out = None;
        let rtk = RecordTableKey { key };
        if let Some(record) = self.record_index.peek(&rtk) {
            // Callback
            out = Some(f(record));
        }
        out
    }

    pub(super) fn with_record_mut<R, F>(&mut self, key: TypedKey, f: F) -> Option<R>
    where
        F: FnOnce(&mut Record<D>) -> R,
    {
        // Get record from index
        let mut out = None;
        let rtk = RecordTableKey { key };
        if let Some(record) = self.record_index.get_mut(&rtk) {
            // Callback
            out = Some(f(record));

            // Touch
            record.touch(get_aligned_timestamp());
        }
        if out.is_some() {
            self.mark_record_changed(rtk);
        }

        out
    }

    // pub fn get_descriptor(&mut self, key: TypedKey) -> Option<SignedValueDescriptor> {
    //     self.with_record(key, |record| record.descriptor().clone())
    // }

    pub async fn get_subkey(
        &mut self,
        key: TypedKey,
        subkey: ValueSubkey,
        want_descriptor: bool,
    ) -> VeilidAPIResult<Option<SubkeyResult>> {
        // record from index
        let Some((subkey_count, has_subkey, opt_descriptor)) = self.with_record(key, |record| {
            (record.subkey_count(), record.stored_subkeys().contains(subkey), if want_descriptor {
                Some(record.descriptor().clone())
            } else {
                None
            })
        }) else {
            // Record not available
            return Ok(None);
        };

        // Check if the subkey is in range
        if subkey as usize >= subkey_count {
            apibail_invalid_argument!("subkey out of range", "subkey", subkey);
        }

        // See if we have this subkey stored
        if !has_subkey {
            // If not, return no value but maybe with descriptor
            return Ok(Some(SubkeyResult {
                value: None,
                descriptor: opt_descriptor,
            }));
        }

        // Get subkey table
        let Some(subkey_table) = self.subkey_table.clone() else {
            apibail_internal!("record store not initialized");
        };

        // If subkey exists in subkey cache, use that
        let stk = SubkeyTableKey { key, subkey };
        if let Some(record_data) = self.subkey_cache.get_mut(&stk) {
            let out = record_data.signed_value_data().clone();

            return Ok(Some(SubkeyResult {
                value: Some(out),
                descriptor: opt_descriptor,
            }));
        }
        // If not in cache, try to pull from table store if it is in our stored subkey set
        let Some(record_data) = subkey_table
            .load_rkyv::<RecordData>(0, &stk.bytes())
            .await
            .map_err(VeilidAPIError::internal)? else {
                apibail_internal!("failed to get subkey that was stored");
            };

        let out = record_data.signed_value_data().clone();

        // Add to cache, do nothing with lru out
        self.add_to_subkey_cache(stk, record_data);

        return Ok(Some(SubkeyResult {
            value: Some(out),
            descriptor: opt_descriptor,
        }));
    }

    pub(crate) async fn peek_subkey(
        &self,
        key: TypedKey,
        subkey: ValueSubkey,
        want_descriptor: bool,
    ) -> VeilidAPIResult<Option<SubkeyResult>> {
        // record from index
        let Some((subkey_count, has_subkey, opt_descriptor)) = self.peek_record(key, |record| {
            (record.subkey_count(), record.stored_subkeys().contains(subkey), if want_descriptor {
                Some(record.descriptor().clone())
            } else {
                None
            })
        }) else {
            // Record not available
            return Ok(None);
        };

        // Check if the subkey is in range
        if subkey as usize >= subkey_count {
            apibail_invalid_argument!("subkey out of range", "subkey", subkey);
        }

        // See if we have this subkey stored
        if !has_subkey {
            // If not, return no value but maybe with descriptor
            return Ok(Some(SubkeyResult {
                value: None,
                descriptor: opt_descriptor,
            }));
        }

        // Get subkey table
        let Some(subkey_table) = self.subkey_table.clone() else {
            apibail_internal!("record store not initialized");
        };

        // If subkey exists in subkey cache, use that
        let stk = SubkeyTableKey { key, subkey };
        if let Some(record_data) = self.subkey_cache.peek(&stk) {
            let out = record_data.signed_value_data().clone();

            return Ok(Some(SubkeyResult {
                value: Some(out),
                descriptor: opt_descriptor,
            }));
        }
        // If not in cache, try to pull from table store if it is in our stored subkey set
        let Some(record_data) = subkey_table
            .load_rkyv::<RecordData>(0, &stk.bytes())
            .await
            .map_err(VeilidAPIError::internal)? else {
                apibail_internal!("failed to peek subkey that was stored");
            };

        let out = record_data.signed_value_data().clone();

        return Ok(Some(SubkeyResult {
            value: Some(out),
            descriptor: opt_descriptor,
        }));
    }

    pub async fn set_subkey(
        &mut self,
        key: TypedKey,
        subkey: ValueSubkey,
        signed_value_data: SignedValueData,
    ) -> VeilidAPIResult<()> {
        // Check size limit for data
        if signed_value_data.value_data().data().len() > self.limits.max_subkey_size {
            apibail_invalid_argument!(
                "record subkey too large",
                "signed_value_data.value_data.data.len",
                signed_value_data.value_data().data().len()
            );
        }

        // Get record from index
        let Some((subkey_count, total_size)) = self.with_record(key, |record| {
            (record.subkey_count(), record.total_size())
        }) else {
            apibail_invalid_argument!("no record at this key", "key", key);
        };

        // Check if the subkey is in range
        if subkey as usize >= subkey_count {
            apibail_invalid_argument!("subkey out of range", "subkey", subkey);
        }

        // Get subkey table
        let Some(subkey_table) = self.subkey_table.clone() else {
            apibail_internal!("record store not initialized");
        };

        // Get the previous subkey and ensure we aren't going over the record size limit
        let mut prior_record_data_size = 0usize;

        // If subkey exists in subkey cache, use that
        let stk = SubkeyTableKey { key, subkey };
        let stk_bytes = stk.bytes();

        if let Some(record_data) = self.subkey_cache.peek(&stk) {
            prior_record_data_size = record_data.total_size();
        } else {
            // If not in cache, try to pull from table store
            if let Some(record_data) = subkey_table
                .load_rkyv::<RecordData>(0, &stk_bytes)
                .await
                .map_err(VeilidAPIError::internal)?
            {
                prior_record_data_size = record_data.total_size();
            }
        }

        // Make new record data
        let record_data = RecordData::new(signed_value_data);

        // Check new total record size
        let new_record_data_size = record_data.total_size();
        let new_total_size = total_size + new_record_data_size - prior_record_data_size;
        if new_total_size > self.limits.max_record_total_size {
            apibail_generic!("dht record too large");
        }

        // Check new total storage space
        self.total_storage_space
            .sub(prior_record_data_size as u64)
            .unwrap();
        self.total_storage_space
            .add(new_record_data_size as u64)
            .unwrap();
        if !self.total_storage_space.check_limit() {
            apibail_try_again!();
        }

        // Write subkey
        subkey_table
            .store_rkyv(0, &stk_bytes, &record_data)
            .await
            .map_err(VeilidAPIError::internal)?;

        // Write to subkey cache
        self.add_to_subkey_cache(stk, record_data);

        // Update record
        self.with_record_mut(key, |record| {
            record.store_subkey(subkey);
            record.set_record_data_size(new_record_data_size);
        })
        .expect("record should still be here");

        // Update storage space
        self.total_storage_space.commit().unwrap();

        Ok(())
    }

    /// LRU out some records until we reclaim the amount of space requested
    /// This will force a garbage collection of the space immediately
    /// If zero is passed in here, a garbage collection will be performed of dead records
    /// without removing any live records
    pub async fn reclaim_space(&mut self, space: usize) -> usize {
        let mut reclaimed = 0usize;
        while reclaimed < space {
            if let Some((k, v)) = self.record_index.remove_lru() {
                reclaimed += mem::size_of::<RecordTableKey>();
                reclaimed += v.total_size();
                self.add_dead_record(k, v);
            } else {
                break;
            }
        }
        self.purge_dead_records(false).await;
        reclaimed
    }

    pub(super) fn debug_records(&self) -> String {
        // Dump fields in an abbreviated way
        let mut out = String::new();

        out += "Record Index:\n";
        for (rik, rec) in &self.record_index {
            out += &format!(
                "  {} age={} len={} subkeys={}\n",
                rik.key.to_string(),
                debug_duration(get_timestamp() - rec.last_touched().as_u64()),
                rec.record_data_size(),
                rec.stored_subkeys(),
            );
        }
        out += &format!("Subkey Cache Count: {}\n", self.subkey_cache.len());
        out += &format!(
            "Subkey Cache Total Size: {}\n",
            self.subkey_cache_total_size.get()
        );
        out += &format!("Total Storage Space: {}\n", self.total_storage_space.get());
        out += &format!("Dead Records: {}\n", self.dead_records.len());
        for dr in &self.dead_records {
            out += &format!("  {}\n", dr.0.key.to_string());
        }
        out += &format!("Changed Records: {}\n", self.changed_records.len());
        for cr in &self.changed_records {
            out += &format!("  {}\n", cr.key.to_string());
        }

        out
    }

    pub(super) fn debug_record_info(&self, key: TypedKey) -> String {
        self.peek_record(key, |r| format!("{:#?}", r))
            .unwrap_or("Not found".to_owned())
    }

    pub(super) async fn debug_record_subkey_info(
        &self,
        key: TypedKey,
        subkey: ValueSubkey,
    ) -> String {
        match self.peek_subkey(key, subkey, true).await {
            Ok(Some(v)) => {
                format!("{:#?}", v)
            }
            Ok(None) => "Subkey not available".to_owned(),
            Err(e) => format!("{}", e),
        }
    }
}
