/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cell::RefCell;
use std::collections::vec_deque::Iter;
use std::collections::vec_deque::IterMut;
use std::collections::VecDeque;
use std::fs::read_dir;
use std::fs::DirEntry;
use std::io::ErrorKind;
use std::path::Path;
use std::path::PathBuf;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::Result;
use parking_lot::Mutex;
use types::Key;
use types::NodeInfo;

use crate::datapack::DataPack;
use crate::datapack::DataPackVersion;
use crate::datastore::Delta;
use crate::datastore::HgIdDataStore;
use crate::datastore::HgIdMutableDeltaStore;
use crate::datastore::Metadata;
use crate::datastore::StoreResult;
use crate::historypack::HistoryPack;
use crate::historypack::HistoryPackVersion;
use crate::historystore::HgIdHistoryStore;
use crate::historystore::HgIdMutableHistoryStore;
use crate::localstore::ExtStoredPolicy;
use crate::localstore::LocalStore;
use crate::localstore::StoreFromPath;
use crate::mutabledatapack::MutableDataPack;
use crate::mutablehistorypack::MutableHistoryPack;
use crate::repack::Repackable;
use crate::types::StoreKey;
use crate::uniondatastore::UnionHgIdDataStore;
use crate::unionhistorystore::UnionHgIdHistoryStore;

/// Naive implementation of a store that order its underlying stores based on how recently we found
/// data in them. This helps in reducing the number of stores that are iterated on.
///
/// The implementation uses a `VecDeque` and always moves the last used store to the front.
///
/// Open source crates were considered, but none support both having an unbounded size, and being
/// able to move one element to the front.
struct LruStore<T> {
    stores: VecDeque<T>,
}

impl<T> LruStore<T> {
    fn new() -> Self {
        Self {
            stores: VecDeque::new(),
        }
    }

    /// Move the store at `index` at the front.
    ///
    /// From a complexity standpoint, the complexity is at worst O(N). In practice, we're expecting
    /// the most recent stores to be near the beginning, which would reduce the observed cost of
    /// this. A more efficient datastructure should allow for a lower complexity.
    fn update(&mut self, index: usize) {
        if let Some(store) = self.stores.remove(index) {
            self.stores.push_front(store);
        }
    }

    /// Add the store at the front of the `LruStore`.
    fn add(&mut self, store: T) {
        self.stores.push_front(store);
    }

    /// Remove an element from the `LruStore`. The order will not be preserved.
    fn remove(&mut self, index: usize) -> T {
        self.stores.swap_remove_back(index).unwrap()
    }

    /// Iterates over all the element, the most recently used items will be returned first.
    fn iter(&self) -> Iter<T> {
        self.stores.iter()
    }

    /// Iterates over all the element, the most recently used items will be returned first.
    fn iter_mut(&mut self) -> IterMut<T> {
        self.stores.iter_mut()
    }
}

impl<'a, T> IntoIterator for &'a LruStore<T> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<T> From<Vec<T>> for LruStore<T> {
    fn from(other: Vec<T>) -> Self {
        Self {
            stores: other.into(),
        }
    }
}

#[derive(PartialEq)]
pub enum CorruptionPolicy {
    IGNORE,
    REMOVE,
}

struct PackStoreInner<T> {
    pack_dir: PathBuf,
    extension: &'static str,
    corruption_policy: CorruptionPolicy,
    extstored_policy: ExtStoredPolicy,
    scan_frequency: Duration,
    last_scanned: RefCell<Option<Instant>>,
    packs: RefCell<LruStore<T>>,
    max_bytes: Option<u64>,
    current_bytes: AtomicU64,
}

/// A `PackStore` automatically keeps track of packfiles in a given directory. New on-disk
/// packfiles will be periodically scanned and opened accordingly.
pub struct PackStore<T> {
    inner: Mutex<PackStoreInner<T>>,
}

pub type DataPackStore = PackStore<DataPack>;
pub type HistoryPackStore = PackStore<HistoryPack>;

struct PackStoreOptions {
    pack_dir: PathBuf,
    scan_frequency: Duration,
    extension: &'static str,
    corruption_policy: CorruptionPolicy,
    max_bytes: Option<u64>,
    extstored_policy: ExtStoredPolicy,
}

impl PackStoreOptions {
    fn new() -> Self {
        Self {
            pack_dir: PathBuf::new(),
            scan_frequency: Duration::from_secs(10),
            extension: "",
            corruption_policy: CorruptionPolicy::IGNORE,
            max_bytes: None,
            extstored_policy: ExtStoredPolicy::Use,
        }
    }

    fn directory<P: AsRef<Path>>(mut self, dir: P) -> Self {
        self.pack_dir = PathBuf::from(dir.as_ref());
        self
    }

    #[cfg(test)]
    fn scan_frequency(mut self, frequency: Duration) -> Self {
        self.scan_frequency = frequency;
        self
    }

    fn extension(mut self, extension: &'static str) -> Self {
        self.extension = extension;
        self
    }

    /// When a packfile is detected to be corrupted, should we automatically remove it from disk or
    /// simply ignore it?
    fn corruption_policy(mut self, corruption_policy: CorruptionPolicy) -> Self {
        self.corruption_policy = corruption_policy;
        self
    }

    fn max_bytes(mut self, max_bytes: Option<u64>) -> Self {
        self.max_bytes = max_bytes;
        self
    }

    /// Should externally stored blobs (LFS) be used or ignored?
    fn extstored_policy(mut self, extstored_policy: ExtStoredPolicy) -> Self {
        self.extstored_policy = extstored_policy;
        self
    }

    fn build<T>(self) -> PackStore<T> {
        PackStore {
            inner: Mutex::new(PackStoreInner {
                pack_dir: self.pack_dir,
                scan_frequency: self.scan_frequency,
                extension: self.extension,
                corruption_policy: self.corruption_policy,
                extstored_policy: self.extstored_policy,
                last_scanned: RefCell::new(None),
                packs: RefCell::new(LruStore::new()),
                max_bytes: self.max_bytes,
                current_bytes: AtomicU64::new(0),
            }),
        }
    }
}

impl<T: LocalStore + Repackable + StoreFromPath> PackStore<T> {
    /// Force a rescan to be performed. Since rescan are expensive, this should only be called for
    /// out-of-process created packfiles.
    pub fn force_rescan(&self) {
        let packstore = self.inner.lock();
        packstore.last_scanned.replace(None);
    }

    /// Add a packfile to this store.
    fn add_pack(&self, pack: T) -> Result<()> {
        let inner = self.inner.lock();
        let size = pack.size();
        inner.packs.borrow_mut().add(pack);
        let current_bytes = inner.current_bytes.fetch_add(size, Ordering::SeqCst) + size;

        if let Some(max_bytes) = inner.max_bytes {
            if current_bytes > max_bytes {
                if inner.delete_old_packs().is_ok() {
                    let _ = inner.rescan()?;
                } else {
                    // If the delete fails, give up and move on. We don't want to block the user
                    // operation on this maintenance work.
                }
            }
        }

        Ok(())
    }
}

impl DataPackStore {
    /// Build a new DataPackStore. The default rescan rate is 10 seconds.
    ///
    /// Only use for data that can be recoverd from the network, corrupted datapacks will be
    /// automatically removed from disk.
    pub fn new<P: AsRef<Path>>(
        pack_dir: P,
        corruption_policy: CorruptionPolicy,
        max_bytes: Option<u64>,
        extstored_policy: ExtStoredPolicy,
    ) -> Self {
        PackStoreOptions::new()
            .directory(pack_dir)
            .corruption_policy(corruption_policy)
            .max_bytes(max_bytes)
            .extstored_policy(extstored_policy)
            .extension("datapack")
            .build()
    }
}

impl HistoryPackStore {
    /// Build a new HistoryPackStore. The default rescan rate is 10 seconds.
    ///
    /// Only use for data that can be recoverd from the network, corrupted datapacks will be
    /// automatically removed from disk.
    pub fn new<P: AsRef<Path>>(
        pack_dir: P,
        corruption_policy: CorruptionPolicy,
        max_bytes: Option<u64>,
    ) -> Self {
        PackStoreOptions::new()
            .directory(pack_dir)
            .corruption_policy(corruption_policy)
            .max_bytes(max_bytes)
            .extension("histpack")
            .build()
    }
}

impl<T: LocalStore + Repackable + StoreFromPath> PackStoreInner<T> {
    /// Open new on-disk packfiles, and close removed ones.
    fn rescan(&self) -> Result<()> {
        let mut new_packs = Vec::new();

        let mut new_size = 0;
        for entry in self.get_pack_paths()?.into_iter() {
            if let Ok(pack) = T::from_path(&entry.path(), self.extstored_policy) {
                new_size += pack.size();
                new_packs.push(pack);
            }
        }

        self.packs.replace(new_packs.into());
        self.current_bytes.store(new_size, Ordering::SeqCst);
        Ok(())
    }

    fn delete_old_packs(&self) -> Result<()> {
        if let Some(max_bytes) = self.max_bytes {
            let mut entries = vec![];
            for entry in self.get_pack_paths()? {
                let metadata = match entry.metadata() {
                    Ok(m) => m,
                    Err(_) => continue,
                };
                let modified = match metadata.modified() {
                    Ok(m) => m,
                    Err(_) => continue,
                };

                entries.push((entry, modified, metadata.len()));
            }

            // Sort by reverse modified to get them in newest first order.
            entries.sort_by(|a, b| b.1.cmp(&a.1));

            let mut size = 0;
            for entry in entries.into_iter() {
                if size >= max_bytes {
                    // Delete the remaining packs
                    match T::from_path(&entry.0.path(), self.extstored_policy) {
                        Ok(pack) => pack.delete()?,
                        Err(_) => continue,
                    };
                } else {
                    size += entry.2;
                }
            }
        }
        Ok(())
    }

    fn get_pack_paths(&self) -> Result<Vec<DirEntry>> {
        let readdir = match read_dir(&self.pack_dir) {
            Ok(readdir) => readdir,
            Err(e) => {
                if e.kind() == ErrorKind::NotFound {
                    return Ok(vec![]);
                } else {
                    return Err(e.into());
                }
            }
        };

        let mut result = vec![];
        for entry in readdir {
            let entry = entry?;
            if entry.file_type()?.is_file() {
                let path = entry.path();
                if let Some(ext) = path.extension() {
                    if ext == self.extension {
                        result.push(entry);
                    }
                }
            }
        }

        Ok(result)
    }

    /// Scan the store when too much time has passed since the last scan. Returns whether the
    /// filesystem was actually scanned.
    fn try_scan(&self) -> Result<bool> {
        let now = Instant::now();

        let needs_scan = match *self.last_scanned.borrow() {
            Some(last_scanned) => now.duration_since(last_scanned) >= self.scan_frequency,
            None => true,
        };
        if needs_scan {
            self.rescan()?;
            self.last_scanned.replace(Some(now));
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Execute the `op` function. May call `rescan` when `op` fails with `KeyError`.
    fn run<R, F>(&self, op: F) -> Result<Option<R>>
    where
        F: Fn(&T) -> Result<Option<R>>,
    {
        for _ in 0..2 {
            let mut found = None;
            {
                let mut corrupted = Vec::new();

                let mut lrustore = self.packs.try_borrow_mut()?;
                for (index, store) in lrustore.iter_mut().enumerate() {
                    match op(store) {
                        Ok(None) => continue,
                        Ok(Some(result)) => {
                            found = Some((index, result));
                            break;
                        }
                        Err(_) => {
                            corrupted.push(index);
                        }
                    }
                }

                if !corrupted.is_empty() {
                    for store_index in corrupted.into_iter().rev() {
                        let store = lrustore.remove(store_index);
                        if self.corruption_policy == CorruptionPolicy::REMOVE {
                            let _ = store.delete();
                        }
                    }
                }
            }

            if let Some((index, result)) = found {
                self.packs.borrow_mut().update(index);
                return Ok(Some(result));
            }

            // We didn't find anything, let's try to probe the filesystem to discover new packfiles
            // and retry.
            if !self.try_scan()? {
                break;
            }
        }

        Ok(None)
    }
}

impl<T: LocalStore + Repackable + StoreFromPath> LocalStore for PackStore<T> {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        // Since the packfiles are loaded lazily, it's possible that `get_missing` is called before
        // any packfiles have been loaded. Let's tentatively scan the store before iterating over
        // all the known packs.
        let packstore = self.inner.lock();
        packstore.try_scan()?;

        let initial_keys = Ok(keys.to_vec());
        let packs = packstore.packs.try_borrow()?;
        packs
            .into_iter()
            .fold(initial_keys, |missing_keys, store| match missing_keys {
                Ok(missing_keys) => store.get_missing(&missing_keys),
                Err(e) => Err(e),
            })
    }
}

impl HgIdDataStore for DataPackStore {
    fn get(&self, key: StoreKey) -> Result<StoreResult<Vec<u8>>> {
        let res = self
            .inner
            .lock()
            .run(|store| match store.get(key.clone())? {
                StoreResult::Found(content) => Ok(Some(content)),
                StoreResult::NotFound(_) => Ok(None),
            })?;

        match res {
            None => Ok(StoreResult::NotFound(key)),
            Some(content) => Ok(StoreResult::Found(content)),
        }
    }

    fn get_meta(&self, key: StoreKey) -> Result<StoreResult<Metadata>> {
        let res = self
            .inner
            .lock()
            .run(|store| match store.get_meta(key.clone())? {
                StoreResult::Found(meta) => Ok(Some(meta)),
                StoreResult::NotFound(_) => Ok(None),
            })?;

        match res {
            None => Ok(StoreResult::NotFound(key)),
            Some(meta) => Ok(StoreResult::Found(meta)),
        }
    }

    fn refresh(&self) -> Result<()> {
        let inner = self.inner.lock();
        inner.last_scanned.replace(None);
        Ok(())
    }
}

impl HgIdHistoryStore for HistoryPackStore {
    fn get_node_info(&self, key: &Key) -> Result<Option<NodeInfo>> {
        self.inner.lock().run(|store| store.get_node_info(key))
    }

    fn refresh(&self) -> Result<()> {
        let inner = self.inner.lock();
        inner.last_scanned.replace(None);
        Ok(())
    }
}

struct MutableDataPackStoreInner {
    pack_store: Arc<DataPackStore>,
    mutable_pack: Arc<MutableDataPack>,
    union_store: UnionHgIdDataStore<Arc<dyn HgIdDataStore>>,
}

/// A `MutableDataPackStore` allows both reading and writing to data packfiles.
pub struct MutableDataPackStore {
    inner: MutableDataPackStoreInner,
    pending: AtomicU64,
    result_packs: Arc<Mutex<Vec<PathBuf>>>,
    max_pending_bytes: u64,
}

impl MutableDataPackStore {
    pub fn new(
        pack_dir: impl AsRef<Path>,
        corruption_policy: CorruptionPolicy,
        max_pending_bytes: u64,
        max_bytes: Option<u64>,
        extstored_policy: ExtStoredPolicy,
    ) -> Result<Self> {
        let pack_store = Arc::new(DataPackStore::new(
            pack_dir.as_ref(),
            corruption_policy,
            max_bytes,
            extstored_policy,
        ));
        let mutable_pack = Arc::new(MutableDataPack::new(pack_dir, DataPackVersion::One));
        let mut union_store: UnionHgIdDataStore<Arc<dyn HgIdDataStore>> = UnionHgIdDataStore::new();
        union_store.add(pack_store.clone());
        union_store.add(mutable_pack.clone());

        Ok(Self {
            inner: MutableDataPackStoreInner {
                pack_store,
                mutable_pack,
                union_store,
            },
            pending: AtomicU64::new(0),
            result_packs: Arc::new(Mutex::new(Vec::new())),
            max_pending_bytes,
        })
    }

    fn inner_flush(&self) -> Result<()> {
        self.pending.store(0, Ordering::SeqCst);
        if let Some(paths) = self.inner.mutable_pack.flush()? {
            let mut result_packs = self.result_packs.lock();
            for path in paths {
                let datapack = DataPack::new(
                    path.as_path(),
                    self.inner.pack_store.inner.lock().extstored_policy,
                )?;
                self.inner.pack_store.add_pack(datapack)?;
                result_packs.push(path);
            }
        }
        Ok(())
    }
}

impl HgIdDataStore for MutableDataPackStore {
    fn get(&self, key: StoreKey) -> Result<StoreResult<Vec<u8>>> {
        self.inner.union_store.get(key)
    }

    fn get_meta(&self, key: StoreKey) -> Result<StoreResult<Metadata>> {
        self.inner.union_store.get_meta(key)
    }

    fn refresh(&self) -> Result<()> {
        self.inner.union_store.refresh()
    }
}

impl LocalStore for MutableDataPackStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.inner.union_store.get_missing(keys)
    }
}

impl HgIdMutableDeltaStore for MutableDataPackStore {
    fn add(&self, delta: &Delta, metadata: &Metadata) -> Result<()> {
        self.inner.mutable_pack.add(delta, metadata)?;
        let pending = self
            .pending
            .fetch_add(delta.data.len() as u64, Ordering::SeqCst)
            + (delta.data.len() as u64);
        if pending >= self.max_pending_bytes {
            self.inner_flush()?;
        }
        Ok(())
    }

    /// Flush the current mutable datapack to disk and add it to the `PackStore`.
    fn flush(&self) -> Result<Option<Vec<PathBuf>>> {
        self.inner_flush()?;
        let mut packs = self.result_packs.lock();
        let result = std::mem::take(&mut *packs);

        Ok(if result.len() > 0 { Some(result) } else { None })
    }
}

struct MutableHistoryPackStoreInner {
    pack_store: Arc<HistoryPackStore>,
    mutable_pack: Arc<MutableHistoryPack>,
    union_store: UnionHgIdHistoryStore<Arc<dyn HgIdHistoryStore>>,
}

/// A `MutableHistoryPackStore` allows both reading and writing to history packfiles.
pub struct MutableHistoryPackStore {
    inner: MutableHistoryPackStoreInner,
    pending: AtomicU64,
    result_packs: Arc<Mutex<Vec<PathBuf>>>,
    max_pending: u64,
}

impl MutableHistoryPackStore {
    pub fn new(
        pack_dir: impl AsRef<Path>,
        corruption_policy: CorruptionPolicy,
        max_pending: u64,
        max_bytes: Option<u64>,
    ) -> Result<Self> {
        let pack_store = Arc::new(HistoryPackStore::new(
            pack_dir.as_ref(),
            corruption_policy,
            max_bytes,
        ));
        let mutable_pack = Arc::new(MutableHistoryPack::new(pack_dir, HistoryPackVersion::One));
        let mut union_store: UnionHgIdHistoryStore<Arc<dyn HgIdHistoryStore>> =
            UnionHgIdHistoryStore::new();
        union_store.add(pack_store.clone());
        union_store.add(mutable_pack.clone());

        Ok(Self {
            inner: MutableHistoryPackStoreInner {
                pack_store,
                mutable_pack,
                union_store,
            },
            pending: AtomicU64::new(0),
            result_packs: Arc::new(Mutex::new(Vec::new())),
            max_pending,
        })
    }

    fn inner_flush(&self) -> Result<()> {
        self.pending.store(0, Ordering::SeqCst);
        if let Some(paths) = self.inner.mutable_pack.flush()? {
            let mut result_packs = self.result_packs.lock();
            for path in paths {
                let histpack = HistoryPack::new(path.as_path())?;
                self.inner.pack_store.add_pack(histpack)?;
                result_packs.push(path);
            }
        }
        Ok(())
    }
}

impl HgIdHistoryStore for MutableHistoryPackStore {
    fn get_node_info(&self, key: &Key) -> Result<Option<NodeInfo>> {
        self.inner.union_store.get_node_info(key)
    }

    fn refresh(&self) -> Result<()> {
        self.inner.union_store.refresh()
    }
}

impl LocalStore for MutableHistoryPackStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.inner.union_store.get_missing(keys)
    }
}

impl HgIdMutableHistoryStore for MutableHistoryPackStore {
    fn add(&self, key: &Key, info: &NodeInfo) -> Result<()> {
        self.inner.mutable_pack.add(key, info)?;
        let pending = self.pending.fetch_add(1, Ordering::SeqCst) + 1;
        if pending >= self.max_pending {
            self.inner_flush()?;
        }
        Ok(())
    }

    /// Flush the current mutable historypack to disk and add it to the `PackStore`.
    fn flush(&self) -> Result<Option<Vec<PathBuf>>> {
        self.inner_flush()?;
        let mut packs = self.result_packs.lock();
        let result = std::mem::take(&mut *packs);

        Ok(if result.len() > 0 { Some(result) } else { None })
    }
}

#[cfg(test)]
mod tests {
    use std::fs;
    use std::fs::OpenOptions;

    use minibytes::Bytes;
    use rand::SeedableRng;
    use rand_chacha::ChaChaRng;
    use tempfile::TempDir;
    use types::testutil::*;

    use super::*;
    use crate::datapack::tests::make_datapack;
    use crate::historypack::tests::get_nodes;
    use crate::historypack::tests::make_historypack;

    #[test]
    fn test_datapack_created_before() -> Result<()> {
        let tempdir = TempDir::new()?;

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        let store = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            None,
            ExtStoredPolicy::Use,
        );
        let stored = store.get(StoreKey::hgid(k))?;
        assert_eq!(
            stored,
            StoreResult::Found(revision.0.data.as_ref().to_vec())
        );
        Ok(())
    }

    #[test]
    fn test_datapack_get_missing() -> Result<()> {
        let tempdir = TempDir::new()?;

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: Some(key("a", "1")),
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        let store = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            None,
            ExtStoredPolicy::Use,
        );
        let missing = store.get_missing(&vec![StoreKey::from(k)])?;
        assert_eq!(missing.len(), 0);
        Ok(())
    }

    #[test]
    fn test_datapack_created_after() -> Result<()> {
        let tempdir = TempDir::new()?;
        let store = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            None,
            ExtStoredPolicy::Use,
        );

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        let stored = store.get(StoreKey::hgid(k))?;
        assert_eq!(
            stored,
            StoreResult::Found(revision.0.data.as_ref().to_vec())
        );
        Ok(())
    }

    #[test]
    fn test_slow_rescan() {
        let tempdir = TempDir::new().unwrap();
        let store = PackStoreOptions::new()
            .directory(&tempdir)
            .extension("datapack")
            .scan_frequency(Duration::from_secs(1000))
            .build();

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        store.get(StoreKey::hgid(k)).unwrap();

        let k = key("a", "3");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        let k = StoreKey::hgid(k);
        assert_eq!(store.get(k.clone()).unwrap(), StoreResult::NotFound(k));
    }

    #[test]
    fn test_force_rescan() -> Result<()> {
        let tempdir = TempDir::new()?;
        let store = PackStoreOptions::new()
            .directory(&tempdir)
            .extension("datapack")
            .scan_frequency(Duration::from_secs(1000))
            .build();

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        store.get(StoreKey::hgid(k))?;

        let k = key("a", "3");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision.clone()]);

        store.force_rescan();
        assert_eq!(
            store.get(StoreKey::hgid(k))?,
            StoreResult::Found(vec![1, 2, 3, 4])
        );
        Ok(())
    }

    #[test]
    fn test_refresh() -> Result<()> {
        let tempdir = TempDir::new()?;
        let store = PackStoreOptions::new()
            .directory(&tempdir)
            .extension("datapack")
            .scan_frequency(Duration::from_secs(10000))
            .build();

        let k = key("a", "2");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision]);

        store.get(StoreKey::hgid(k))?;

        let k = key("a", "3");
        let revision = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision]);

        assert_eq!(
            store.get(StoreKey::hgid(k.clone()))?,
            StoreResult::NotFound(StoreKey::hgid(k.clone())),
        );
        store.refresh().unwrap();
        assert_eq!(
            store.get(StoreKey::hgid(k))?,
            StoreResult::Found(vec![1, 2, 3, 4])
        );
        Ok(())
    }

    #[test]
    fn test_histpack() -> Result<()> {
        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let tempdir = TempDir::new()?;
        let store = HistoryPackStore::new(&tempdir, CorruptionPolicy::REMOVE, None);

        let nodes = get_nodes(&mut rng);
        make_historypack(&tempdir, &nodes);
        for (key, info) in nodes.iter() {
            let response: NodeInfo = store.get_node_info(key)?.unwrap();
            assert_eq!(response, *info);
        }

        Ok(())
    }

    #[test]
    fn test_lrustore_order() -> Result<()> {
        let tempdir = TempDir::new()?;

        let k1 = key("a", "2");
        let revision1 = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k1.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision1.clone()]);

        let k2 = key("b", "3");
        let revision2 = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k2.clone(),
            },
            Default::default(),
        );
        make_datapack(&tempdir, &vec![revision2.clone()]);

        let packstore = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            None,
            ExtStoredPolicy::Use,
        );

        let k2 = StoreKey::hgid(k2);
        let _ = packstore.get(k2.clone())?;
        assert!(
            packstore.inner.lock().packs.borrow().stores[0]
                .get(k2)
                .is_ok()
        );

        let k1 = StoreKey::hgid(k1);
        let _ = packstore.get(k1.clone())?;
        assert!(
            packstore.inner.lock().packs.borrow().stores[0]
                .get(k1)
                .is_ok()
        );

        Ok(())
    }

    #[test]
    fn test_rescan_no_dir() -> Result<()> {
        let tempdir = TempDir::new()?;
        let mut non_present_tempdir = tempdir.into_path();
        non_present_tempdir.push("non_present");
        let store = HistoryPackStore::new(&non_present_tempdir, CorruptionPolicy::REMOVE, None);

        let store = store.inner.lock();
        store.rescan()
    }

    #[test]
    fn test_corrupted() {
        let tempdir = TempDir::new().unwrap();

        let k1 = key("a", "2");
        let revision1 = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k1.clone(),
            },
            Default::default(),
        );
        let path = make_datapack(&tempdir, &vec![revision1.clone()])
            .pack_path()
            .to_path_buf();

        let metadata = fs::metadata(&path).unwrap();
        let mut permissions = metadata.permissions();
        permissions.set_readonly(false);
        fs::set_permissions(&path, permissions).unwrap();

        let datapack = OpenOptions::new().write(true).open(path).unwrap();
        datapack
            .set_len(datapack.metadata().unwrap().len() / 2)
            .unwrap();

        let packstore = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            None,
            ExtStoredPolicy::Use,
        );
        let k1 = StoreKey::hgid(k1);
        assert_eq!(
            packstore.get(k1.clone()).unwrap(),
            StoreResult::NotFound(k1)
        );
    }

    #[test]
    fn test_ignore_corrupted() -> Result<()> {
        let tempdir = TempDir::new()?;

        let k1 = key("a", "2");
        let revision1 = (
            Delta {
                data: Bytes::from(&[1, 2, 3, 4][..]),
                base: None,
                key: k1.clone(),
            },
            Default::default(),
        );
        let path = make_datapack(&tempdir, &vec![revision1.clone()])
            .pack_path()
            .to_path_buf();

        let metadata = fs::metadata(&path).unwrap();
        let mut permissions = metadata.permissions();
        permissions.set_readonly(false);
        fs::set_permissions(&path, permissions).unwrap();

        let datapack = OpenOptions::new().write(true).open(path)?;
        datapack.set_len(datapack.metadata()?.len() / 2)?;

        assert_eq!(read_dir(&tempdir)?.count(), 2);

        let packstore = DataPackStore::new(
            &tempdir,
            CorruptionPolicy::IGNORE,
            None,
            ExtStoredPolicy::Use,
        );
        let k1 = StoreKey::hgid(k1);
        assert_eq!(packstore.get(k1.clone())?, StoreResult::NotFound(k1));

        assert_eq!(read_dir(&tempdir)?.count(), 2);
        Ok(())
    }

    #[test]
    fn test_add_flush() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableDataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            1000,
            None,
            ExtStoredPolicy::Use,
        )?;

        let k1 = key("a", "2");
        let delta = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: Some(key("a", "1")),
            key: k1.clone(),
        };

        packstore.add(&delta, &Default::default())?;
        packstore.flush()?;
        Ok(())
    }

    #[test]
    fn test_add_get_delta() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableDataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            1000,
            None,
            ExtStoredPolicy::Use,
        )?;

        let k1 = key("a", "2");
        let delta = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: None,
            key: k1.clone(),
        };

        packstore.add(&delta, &Default::default())?;
        let stored = packstore.get(StoreKey::hgid(k1))?;
        assert_eq!(stored, StoreResult::Found(delta.data.as_ref().to_vec()));
        Ok(())
    }

    #[test]
    fn test_add_flush_get_delta() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableDataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            1000,
            None,
            ExtStoredPolicy::Use,
        )?;

        let k1 = key("a", "2");
        let delta = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: None,
            key: k1.clone(),
        };

        packstore.add(&delta, &Default::default())?;
        packstore.flush()?;
        let stored = packstore.get(StoreKey::hgid(k1))?;
        assert_eq!(stored, StoreResult::Found(delta.data.as_ref().to_vec()));
        Ok(())
    }

    #[test]
    fn test_histpack_add_get() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore =
            MutableHistoryPackStore::new(&tempdir, CorruptionPolicy::REMOVE, 1000, None)?;

        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let nodes = get_nodes(&mut rng);
        for (key, info) in &nodes {
            packstore.add(key, info)?;
        }

        for (key, info) in nodes {
            let nodeinfo = packstore.get_node_info(&key)?.unwrap();
            assert_eq!(nodeinfo, info);
        }
        Ok(())
    }

    #[test]
    fn test_histpack_add_flush_get() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore =
            MutableHistoryPackStore::new(&tempdir, CorruptionPolicy::REMOVE, 1000, None)?;

        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let nodes = get_nodes(&mut rng);
        for (key, info) in &nodes {
            packstore.add(key, info)?;
        }

        packstore.flush()?;

        for (key, info) in nodes {
            let nodeinfo = packstore.get_node_info(&key)?.unwrap();
            assert_eq!(nodeinfo, info);
        }
        Ok(())
    }

    #[test]
    fn test_histpack_auto_flush() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableHistoryPackStore::new(&tempdir, CorruptionPolicy::REMOVE, 0, None)?;

        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let nodes = get_nodes(&mut rng);
        for (key, info) in &nodes {
            packstore.add(key, info)?;
        }

        let packs = packstore.flush().unwrap().unwrap();
        assert_eq!(packs.len(), 3);

        for (key, info) in nodes {
            let nodeinfo = packstore.get_node_info(&key)?.unwrap();
            assert_eq!(nodeinfo, info);
        }
        Ok(())
    }

    #[test]
    fn test_datapack_auto_flush() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableDataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            0,
            None,
            ExtStoredPolicy::Ignore,
        )?;

        let k1 = key("a", "1");
        let delta1 = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: None,
            key: k1.clone(),
        };
        let k2 = key("a", "2");
        let delta2 = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: None,
            key: k2.clone(),
        };
        let k3 = key("a", "3");
        let delta3 = Delta {
            data: Bytes::from(&[1, 2, 3, 4][..]),
            base: None,
            key: k3.clone(),
        };

        packstore.add(&delta1, &Default::default())?;
        packstore.add(&delta2, &Default::default())?;
        packstore.add(&delta3, &Default::default())?;

        let packs = packstore.flush().unwrap().unwrap();
        assert_eq!(packs.len(), 3);
        Ok(())
    }

    #[test]
    fn test_datapack_flush_empty() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore = MutableDataPackStore::new(
            &tempdir,
            CorruptionPolicy::REMOVE,
            1000,
            None,
            ExtStoredPolicy::Use,
        )?;
        packstore.flush()?;
        Ok(())
    }

    #[test]
    fn test_histpack_flush_empty() -> Result<()> {
        let tempdir = TempDir::new()?;
        let packstore =
            MutableHistoryPackStore::new(&tempdir, CorruptionPolicy::REMOVE, 1000, None)?;
        packstore.flush()?;
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::cell::RefCell;
use std::cell::RefMut;
use std::fmt::Debug;
use std::io;
use std::io::BufWriter;
use std::io::Write;

use anyhow::Result;

/// A `PackWriter` will buffers all the writes to `T` and count the total number of bytes written.
pub struct PackWriter<T: Write> {
    data: RefCell<BufWriter<T>>,
    bytes_written: u64,
}

impl<T: 'static + Write + Debug + Send + Sync> PackWriter<T> {
    pub fn new(value: T) -> PackWriter<T> {
        PackWriter {
            data: RefCell::new(BufWriter::new(value)),
            bytes_written: 0,
        }
    }

    /// Flush the buffered data to the underlying writer.
    pub fn flush_inner(&self) -> Result<()> {
        let ret = self.data.try_borrow_mut()?.flush()?;
        Ok(ret)
    }

    /// Return the number of bytes written. Note that due to the buffering nature of a
    /// `PackWriter`, not all the data may have reached the underlying writer.
    pub fn bytes_written(&self) -> u64 {
        self.bytes_written
    }

    /// Return a mutable reference on the underlying writer. It's not recommended to write to it.
    pub fn get_mut(&self) -> RefMut<T> {
        let cell = self.data.borrow_mut();
        RefMut::map(cell, |w| w.get_mut())
    }

    /// Flush the buffered data and return the underlying writer.
    pub fn into_inner(self) -> Result<T> {
        let ret = self.data.into_inner().into_inner()?;
        Ok(ret)
    }
}

impl<T: Write> Write for PackWriter<T> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        let ret = self.data.get_mut().write(buf)?;
        self.bytes_written += ret as u64;
        Ok(ret)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.data.get_mut().flush()
    }
}

#[cfg(test)]
mod tests {
    use std::io::Seek;
    use std::io::SeekFrom;

    use byteorder::ReadBytesExt;
    use byteorder::WriteBytesExt;
    use tempfile::tempfile;

    use super::*;

    #[test]
    fn test_bytes_written() {
        let mut file = PackWriter::new(tempfile().unwrap());
        file.write_u8(10).unwrap();

        assert_eq!(file.bytes_written(), 1);
    }

    #[test]
    fn test_write() {
        let mut file = PackWriter::new(tempfile().unwrap());
        file.write_u8(10).unwrap();

        // into_inner() flushes its internal buffer.
        let mut inner = file.into_inner().unwrap();
        inner.seek(SeekFrom::Start(0)).unwrap();
        let data = inner.read_u8().unwrap();
        assert_eq!(data, 10);
    }

    #[test]
    fn test_read_without_drain() {
        let mut file = PackWriter::new(tempfile().unwrap());
        file.write_u8(10).unwrap();

        let mut inner = file.get_mut();
        inner.seek(SeekFrom::Start(0)).unwrap();
        assert!(inner.read_u8().is_err());
    }

    #[test]
    fn test_flush_inner() {
        let mut file = PackWriter::new(tempfile().unwrap());
        file.write_u8(10).unwrap();
        file.flush_inner().unwrap();

        let mut inner = file.get_mut();
        inner.seek(SeekFrom::Start(0)).unwrap();
        let data = inner.read_u8().unwrap();
        assert_eq!(data, 10);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use minibytes::Bytes;

/// TODO(T48685378): Handle redacted content in a less hacky way.
pub static REDACTED_CONTENT: &[u8] = b"PoUOK1GkdH6Xtx5j9WKYew3dZXspyfkahcNkhV6MJ4rhyNICTvX0nxmbCImFoT0oHAF9ivWGaC6ByswQZUgf1nlyxcDcahHknJS15Vl9Lvc4NokYhMg0mV1rapq1a4bhNoUI9EWTBiAkYmkadkO3YQXV0TAjyhUQWxxLVskjOwiiFPdL1l1pdYYCLTE3CpgOoxQV3EPVxGUPh1FGfk7F9Myv22qN1sUPSNN4h3IFfm2NNPRFgWPDsqAcaQ7BUSKa\n";
static REDACTED_MESSAGE: &[u8] = b"This version of the file is redacted and you are not allowed to access it. Update or rebase to a newer commit.\n";

pub fn is_redacted(data: &Bytes) -> bool {
    data.as_ref() == REDACTED_CONTENT
}

pub fn redact_if_needed(data: Bytes) -> Bytes {
    if is_redacted(&data) {
        Bytes::from_static(REDACTED_MESSAGE)
    } else {
        data
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use crate::datastore::HgIdMutableDeltaStore;
use crate::datastore::RemoteDataStore;
use crate::historystore::HgIdMutableHistoryStore;
use crate::historystore::RemoteHistoryStore;

pub trait HgIdRemoteStore: Send + Sync {
    fn datastore(
        self: Arc<Self>,
        store: Arc<dyn HgIdMutableDeltaStore>,
    ) -> Arc<dyn RemoteDataStore>;
    fn historystore(
        self: Arc<Self>,
        store: Arc<dyn HgIdMutableHistoryStore>,
    ) -> Arc<dyn RemoteHistoryStore>;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashSet;
use std::fs;
use std::io::Error as IoError;
use std::io::ErrorKind as IoErrorKind;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::format_err;
use anyhow::Error;
use anyhow::Result;
use configmodel::convert::ByteCount;
use configmodel::Config;
use configmodel::ConfigExt;
use minibytes::Bytes;
use thiserror::Error;
use types::Key;

use crate::datapack::DataPack;
use crate::datapack::DataPackVersion;
use crate::datastore::HgIdDataStore;
use crate::datastore::HgIdMutableDeltaStore;
use crate::datastore::StoreResult;
use crate::historypack::HistoryPack;
use crate::historypack::HistoryPackVersion;
use crate::historystore::HgIdHistoryStore;
use crate::historystore::HgIdMutableHistoryStore;
use crate::localstore::ExtStoredPolicy;
use crate::localstore::LocalStore;
use crate::localstore::StoreFromPath;
use crate::metadatastore::MetadataStore;
use crate::mutabledatapack::MutableDataPack;
use crate::mutablehistorypack::MutableHistoryPack;
use crate::mutablepack::MutablePack;
use crate::types::StoreKey;
use crate::LegacyStore;

#[derive(Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
pub enum RepackLocation {
    Local,
    Shared,
}

#[derive(Copy, Clone, PartialEq, Eq, Ord, PartialOrd)]
pub enum RepackKind {
    Incremental,
    Full,
}

pub trait ToKeys {
    fn to_keys(&self) -> Vec<Result<Key>>;
}

pub trait Repackable {
    fn delete(self) -> Result<()>;
    fn size(&self) -> u64;
}

fn repack_datapack(data_pack: &DataPack, mut_pack: &mut MutableDataPack) -> Result<()> {
    for k in data_pack.to_keys() {
        let key = k?;

        if let Some(chain) = data_pack.get_delta_chain(&key)? {
            for delta in chain.iter() {
                if mut_pack.contains(&StoreKey::hgid(delta.key.clone()))? {
                    break;
                }

                // If we managed to get a delta, the metadata must be present.
                match data_pack.get_meta(StoreKey::hgid(delta.key.clone()))? {
                    StoreResult::Found(meta) => mut_pack.add(&delta, &meta)?,
                    _ => {}
                }
            }
        }
    }

    Ok(())
}

#[derive(Debug, Error)]
enum RepackFailure {
    #[error("Repack failure: {0:?}")]
    Total(Vec<(PathBuf, Error)>),

    #[error("Repack successful but with errors: {0:?}")]
    Partial(Vec<(PathBuf, Error)>),
}

/// Repack all pack files in the paths iterator. Once repacked, the repacked packs will be removed
/// from the filesystem.
fn repack_packs<T: MutablePack, U: LocalStore + Repackable + ToKeys + StoreFromPath>(
    paths: impl IntoIterator<Item = PathBuf> + Clone,
    mut mut_pack: T,
    repack_pack: impl Fn(&U, &mut T) -> Result<()>,
) -> Result<Option<PathBuf>> {
    if paths.clone().into_iter().count() <= 1 {
        if let Some(path) = paths.into_iter().next() {
            return Ok(Some(path));
        } else {
            return Ok(None);
        }
    }

    let mut repacked = vec![];
    let mut errors = vec![];

    for path in paths {
        match U::from_path(&path, ExtStoredPolicy::Use) {
            Ok(pack) => {
                if let Err(e) = repack_pack(&pack, &mut mut_pack) {
                    errors.push((path.clone(), e));
                } else {
                    repacked.push(path);
                }
            }
            Err(e) => {
                if let Some(e) = e.downcast_ref::<IoError>() {
                    if e.kind() == IoErrorKind::NotFound {
                        continue;
                    }
                }
                errors.push((path.clone(), e));
            }
        }
    }

    if repacked.is_empty() {
        if !errors.is_empty() {
            return Err(RepackFailure::Total(errors).into());
        } else {
            // Nothing to repack
            return Ok(None);
        }
    }

    let new_pack_path = mut_pack.close_pack()?.unwrap();
    let new_pack = U::from_path(&new_pack_path, ExtStoredPolicy::Use)?;

    let mut successfully_repacked = 0;
    for path in repacked {
        if *path != new_pack_path {
            let pack = match U::from_path(&path, ExtStoredPolicy::Use) {
                Ok(pack) => pack,
                Err(_e) => {
                    // We were about to remove this file, let's just ignore the failures to open
                    // it.
                    successfully_repacked += 1;
                    continue;
                }
            };

            let keys = pack
                .to_keys()
                .into_iter()
                .filter_map(|res| res.ok())
                .map(StoreKey::hgid)
                .collect::<Vec<_>>();
            let missing = new_pack.get_missing(&keys)?;

            if missing.is_empty() {
                let _ = pack.delete();
                successfully_repacked += 1;
            } else {
                errors.push((path.clone(), format_err!("{:?}", missing)));
            }
        } else {
            successfully_repacked += 1;
        }
    }

    if successfully_repacked == 0 {
        Err(RepackFailure::Total(errors).into())
    } else if !errors.is_empty() {
        Err(RepackFailure::Partial(errors).into())
    } else {
        Ok(Some(new_pack_path))
    }
}

fn repack_datapacks(
    paths: impl IntoIterator<Item = PathBuf> + Clone,
    outdir: &Path,
) -> Result<Option<PathBuf>> {
    let mut_pack = MutableDataPack::new(outdir, DataPackVersion::One);

    repack_packs(paths, mut_pack, repack_datapack)
}

fn repack_historypack(history_pack: &HistoryPack, mut_pack: &mut MutableHistoryPack) -> Result<()> {
    for k in history_pack.to_keys() {
        let key = k?;
        if let Some(hgid) = history_pack.get_node_info(&key)? {
            mut_pack.add(&key, &hgid)?;
        }
    }

    Ok(())
}

fn repack_historypacks(
    paths: impl IntoIterator<Item = PathBuf> + Clone,
    outdir: &Path,
) -> Result<Option<PathBuf>> {
    let mut_pack = MutableHistoryPack::new(outdir, HistoryPackVersion::One);

    repack_packs(paths, mut_pack, repack_historypack)
}

/// List all the pack files in the directory `dir` that ends with `extension`.
fn list_packs(dir: &Path, extension: &str) -> Result<Vec<PathBuf>> {
    let mut dirents = fs::read_dir(dir)?
        .filter_map(|e| match e {
            Err(_) => None,
            Ok(entry) => {
                let entrypath = entry.path();
                if entrypath.extension() == Some(extension.as_ref()) {
                    Some(entrypath.with_extension(""))
                } else {
                    None
                }
            }
        })
        .collect::<Vec<PathBuf>>();
    dirents.sort_unstable();
    Ok(dirents)
}

/// Select all the packs from `packs` that needs to be repacked during an incremental repack.
///
/// The filtering is fairly basic and is intended to reduce the fragmentation of pack files.
fn filter_incrementalpacks(
    packs: Vec<PathBuf>,
    extension: &str,
    config: &dyn Config,
) -> Result<Vec<PathBuf>> {
    // The overall maximum pack size.
    let max_pack_size: u64 = {
        if extension == "histpack" {
            config.get_or("repack", "maxhistpacksize", || {
                // Per 100MB of histpack size, the memory consumption is over 1GB,
                // thus repacking 4GB would need over 40GB of RAM.
                ByteCount::from(400 * 1024 * 1024)
            })?
        } else {
            config.get_or("repack", "maxdatapacksize", || {
                ByteCount::from(4 * 1024 * 1024 * 1024)
            })?
        }
    }
    .value();
    // The size limit for any individual pack.
    let size_limit: u64 = config
        .get_or("repack", "sizelimit", || ByteCount::from(100 * 1024 * 1024))?
        .value();
    // The maximum number of packs we want to have after repack (overrides `size_limit`).
    let max_packs: usize = config.get_or("repack", "maxpacks", || 50)?;

    let mut packssizes = packs
        .into_iter()
        .map(|p| {
            let size = p
                .with_extension(extension)
                .metadata()
                .and_then(|m| Ok(m.len()))
                .unwrap_or(u64::max_value());
            (p, size)
        })
        .collect::<Vec<(PathBuf, u64)>>();

    // Sort by file size in increasing order
    packssizes.sort_unstable_by(|a, b| a.1.cmp(&b.1));

    let mut num_packs = packssizes.len();
    let mut accumulated_sizes = 0;
    Ok(packssizes
        .into_iter()
        .take_while(|e| {
            if e.1 + accumulated_sizes > max_pack_size {
                return false;
            }

            if e.1 > size_limit && num_packs < max_packs {
                false
            } else {
                accumulated_sizes += e.1;
                num_packs -= 1;

                true
            }
        })
        .map(|e| e.0)
        .collect())
}

/// Fallback for `repack` for when no `ContentStore`/`MetadataStore` were passed in. Will simply
/// use the legacy code path to write the content of the packfiles to a packfile.
fn repack_no_store(path: PathBuf, kind: RepackKind, config: &dyn Config) -> Result<()> {
    let mut datapacks = list_packs(&path, "datapack")?;
    let mut histpacks = list_packs(&path, "histpack")?;

    if kind == RepackKind::Incremental {
        datapacks = filter_incrementalpacks(datapacks, "datapack", config)?;
        histpacks = filter_incrementalpacks(histpacks, "histpack", config)?;
    }

    let datapack_res = repack_datapacks(datapacks, &path).map(|_| ());
    let histpack_res = repack_historypacks(histpacks, &path).map(|_| ());

    datapack_res.and(histpack_res)
}

fn repack_datapack_to_contentstore(
    paths: Vec<PathBuf>,
    store: &Arc<dyn LegacyStore>,
    location: RepackLocation,
) -> Result<()> {
    let mut repacked = Vec::with_capacity(paths.len());
    let mut errors = vec![];

    let mut seen = HashSet::new();
    for path in paths {
        let pack = match DataPack::new(&path, ExtStoredPolicy::Use) {
            Ok(pack) => pack,
            Err(_) => continue,
        };

        let res = (|| -> Result<()> {
            for key in pack.to_keys() {
                let key = key?;
                if !seen.contains(&key) {
                    if let StoreResult::Found(content) = store.get(StoreKey::hgid(key.clone()))? {
                        match store.get_meta(StoreKey::hgid(key.clone()))? {
                            StoreResult::Found(meta) => {
                                store.add_pending(&key, Bytes::from(content), meta, location)?;
                                seen.insert(key);
                            }
                            _ => {}
                        }
                    }
                }
            }
            Ok(())
        })();

        match res {
            Ok(_) => repacked.push(path),
            Err(e) => errors.push((path, e)),
        }
    }

    if repacked.is_empty() {
        return Err(RepackFailure::Total(errors).into());
    }

    let new_packs = store
        .commit_pending(location)?
        .unwrap_or_else(|| vec![])
        .into_iter()
        .collect::<HashSet<PathBuf>>();

    for path in repacked {
        // TODO: This is a bit fragile as a bug in commit_pending not returning a path could lead
        // to data loss. A better return type would avoid this.
        if !new_packs.contains(&path) {
            match DataPack::new(&path, ExtStoredPolicy::Use) {
                Ok(pack) => pack.delete()?,
                Err(_) => continue,
            }
        }
    }

    if !errors.is_empty() {
        Err(RepackFailure::Partial(errors).into())
    } else {
        Ok(())
    }
}

fn repack_histpack_to_metadatastore(
    paths: Vec<PathBuf>,
    store: &MetadataStore,
    location: RepackLocation,
) -> Result<()> {
    let mut repacked = Vec::with_capacity(paths.len());
    let mut errors = vec![];

    for path in paths {
        let pack = match HistoryPack::new(&path) {
            Ok(pack) => pack,
            Err(_) => continue,
        };

        let res = (|| -> Result<()> {
            for key in pack.to_keys() {
                let key = key?;

                if let Some(nodeinfo) = store.get_node_info(&key)? {
                    store.add_pending(&key, nodeinfo, location)?;
                }
            }

            Ok(())
        })();

        match res {
            Ok(_) => repacked.push(path),
            Err(e) => errors.push((path, e)),
        }
    }

    if repacked.is_empty() {
        return Err(RepackFailure::Total(errors).into());
    }

    let new_packs = store
        .commit_pending(location)?
        .unwrap_or_else(|| vec![])
        .into_iter()
        .collect::<HashSet<PathBuf>>();

    for path in repacked {
        if !new_packs.contains(&path) {
            match HistoryPack::new(&path) {
                Ok(pack) => pack.delete()?,
                Err(_) => continue,
            }
        }
    }

    if !errors.is_empty() {
        Err(RepackFailure::Partial(errors).into())
    } else {
        Ok(())
    }
}

/// Read blobs and metadata contained in the packfiles from `path` and write them back to the
/// `stores`.
///
/// The primary goal of `repack` is to reduce the performance effect of having many packfiles on
/// disk. This is done by writing all the data (and metadata) of the several packfiles onto one.
///
/// The secondary goal of `repack` is for file format changes, packfile are for instance holding
/// LFS pointers, and by virtue of writing these pointers to a `ContentStore`, these will be
/// written to the `LfsStore` instead of to a packfile.
///
/// When `RepackKind::Incremental` is passed in, only a subset of the packfiles will be repacked in
/// order to minimize CPU cost.
///
/// When `stores` is None, a much dumber repack operation is performed, where only the primary goal
/// is fullfilled.
pub fn repack(
    path: PathBuf,
    stores: Option<(Arc<dyn LegacyStore>, Arc<MetadataStore>)>,
    kind: RepackKind,
    location: RepackLocation,
    config: &dyn Config,
) -> Result<()> {
    let (content, metadata) = match stores {
        Some((content, metadata)) => (content, metadata),
        None => return repack_no_store(path, kind, config),
    };

    let mut datapacks = list_packs(&path, "datapack")?;
    let mut histpacks = list_packs(&path, "histpack")?;

    if kind == RepackKind::Incremental {
        // We may be filtering out packfiles that contain LFS pointers, reducing the effectiveness
        // of the secondary goal of repack. To fully perform this secondary goal, a full repack
        // will be necessary, to keep incremental repacks simple.
        datapacks = filter_incrementalpacks(datapacks, "datapack", config)?;
        histpacks = filter_incrementalpacks(histpacks, "histpack", config)?;
    }

    if !datapacks.is_empty() {
        repack_datapack_to_contentstore(datapacks, &content, location)?;
    }

    if !histpacks.is_empty() {
        repack_histpack_to_metadatastore(histpacks, &metadata, location)?;
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::fs::set_permissions;
    use std::fs::File;
    use std::fs::OpenOptions;
    use std::io::Write;

    use minibytes::Bytes;
    use rand::SeedableRng;
    use rand_chacha::ChaChaRng;
    use tempfile::TempDir;
    use types::testutil::*;

    use super::*;
    use crate::datapack::tests::make_datapack;
    use crate::datastore::Delta;
    use crate::historypack::tests::get_nodes;
    use crate::historypack::tests::make_historypack;
    use crate::testutil::empty_config;

    #[test]
    fn test_repack_filter_incremental() -> Result<()> {
        let tempdir = TempDir::new()?;
        let get_packfile_path = |size: usize| tempdir.path().join(format!("{}.datapack", size));
        let get_packfile_paths = |sizes: &[usize]| {
            sizes
                .iter()
                .map(|size| get_packfile_path(*size))
                .collect::<Vec<PathBuf>>()
        };

        let packs: Vec<PathBuf> = [100, 200, 300, 400, 500]
            .iter()
            .map(|size| {
                let path = get_packfile_path(*size);
                let mut file = File::create(&path).unwrap();
                let bytes = vec![0; *size];
                file.write_all(&bytes).unwrap();
                path
            })
            .collect();

        let config = empty_config();
        assert_eq!(
            filter_incrementalpacks(packs.clone(), "datapack", &config)?,
            get_packfile_paths(&[100, 200, 300, 400, 500])
        );

        let config = {
            let mut config = empty_config();
            config.insert("repack.sizelimit".to_string(), "300".to_string());
            config
        };
        assert_eq!(
            filter_incrementalpacks(packs.clone(), "datapack", &config)?,
            get_packfile_paths(&[100, 200, 300])
        );

        let config = {
            let mut config = empty_config();
            config.insert("repack.maxdatapacksize".to_string(), "1k".to_string());
            config
        };
        assert_eq!(
            filter_incrementalpacks(packs.clone(), "datapack", &config)?,
            get_packfile_paths(&[100, 200, 300, 400])
        );

        let config = {
            let mut config = empty_config();
            config.insert("repack.maxhistpacksize".to_string(), "1k".to_string());
            config
        };
        assert!(filter_incrementalpacks(packs.clone(), "histpack", &config)?.is_empty());

        // We have 5 packs pre-repack and we want to make sure we have no more
        // than 2 packs post-repack.
        let config = {
            let mut config = empty_config();
            config.insert("repack.sizelimit".to_string(), "300".to_string());
            config.insert("repack.maxpacks".to_string(), "2".to_string());
            config
        };
        assert_eq!(
            filter_incrementalpacks(packs, "datapack", &config)?,
            get_packfile_paths(&[100, 200, 300, 400])
        );

        Ok(())
    }

    #[test]
    fn test_repack_no_datapack() {
        let tempdir = TempDir::new().unwrap();

        let newpath = repack_datapacks(vec![].into_iter(), tempdir.path());
        assert!(newpath.is_ok());
        let newpath = newpath.unwrap();
        assert_eq!(newpath, None);
    }

    #[test]
    fn test_repack_one_datapack() {
        let tempdir = TempDir::new().unwrap();

        let revisions = vec![(
            Delta {
                data: Bytes::from(&[1u8, 2, 3, 4][..]),
                base: None,
                key: key("a", "1"),
            },
            Default::default(),
        )];

        let pack = make_datapack(&tempdir, &revisions);
        let newpath = repack_datapacks(
            vec![pack.base_path().to_path_buf()].into_iter(),
            tempdir.path(),
        );
        assert!(newpath.is_ok());
        let newpath2 = newpath.unwrap().unwrap();
        assert_eq!(newpath2.with_extension("datapack"), pack.pack_path());
        let datapack = DataPack::new(&newpath2, ExtStoredPolicy::Use);
        assert!(datapack.is_ok());
        let newpack = datapack.unwrap();
        assert_eq!(
            newpack
                .to_keys()
                .into_iter()
                .collect::<Result<Vec<Key>>>()
                .unwrap(),
            revisions
                .iter()
                .map(|d| d.0.key.clone())
                .collect::<Vec<Key>>()
        );
    }

    #[test]
    fn test_repack_multiple_datapacks() {
        let tempdir = TempDir::new().unwrap();
        let mut revisions = Vec::new();
        let mut paths = Vec::new();

        for i in 1..=2 {
            let base = key("a", &i.to_string());
            let rev = vec![
                (
                    Delta {
                        data: Bytes::from(&[1, 2, 3, 4][..]),
                        base: None,
                        key: base.clone(),
                    },
                    Default::default(),
                ),
                (
                    Delta {
                        data: Bytes::from(&[1, 2, 3, 4][..]),
                        base: Some(base),
                        key: key("a", &(100 + i).to_string()),
                    },
                    Default::default(),
                ),
            ];
            let pack = make_datapack(&tempdir, &rev);
            let path = pack.base_path().to_path_buf();
            revisions.push(rev);
            paths.push(path);
        }

        let newpath = repack_datapacks(paths.into_iter(), tempdir.path());
        assert!(newpath.is_ok());
        let newpack = DataPack::new(&newpath.unwrap().unwrap(), ExtStoredPolicy::Use).unwrap();
        assert_eq!(
            newpack
                .to_keys()
                .into_iter()
                .collect::<Result<Vec<Key>>>()
                .unwrap(),
            revisions
                .iter()
                .flatten()
                .map(|d| d.0.key.clone())
                .collect::<Vec<Key>>()
        );
    }

    #[test]
    fn test_repack_missing_files() {
        let tempdir = TempDir::new().unwrap();

        let paths = vec![PathBuf::from("foo.datapack"), PathBuf::from("bar.datapack")];
        let res = repack_datapacks(paths.clone().into_iter(), tempdir.path());

        assert!(res.unwrap().is_none());
    }

    #[test]
    fn test_repack_corrupted() {
        let tempdir = TempDir::new().unwrap();
        let mut revisions = Vec::new();
        let mut paths = Vec::new();

        for i in 1..=2 {
            let base = key("a", &i.to_string());
            let rev = vec![
                (
                    Delta {
                        data: Bytes::from(&[1, 2, 3, 4][..]),
                        base: None,
                        key: base.clone(),
                    },
                    Default::default(),
                ),
                (
                    Delta {
                        data: Bytes::from(&[1, 2, 3, 4][..]),
                        base: Some(base),
                        key: key("a", &(100 + i).to_string()),
                    },
                    Default::default(),
                ),
            ];
            let pack = make_datapack(&tempdir, &rev);
            let path = pack.base_path().to_path_buf();
            revisions.push(rev);
            paths.push(path);
        }

        let mut to_corrupt = paths.get(0).unwrap().clone();
        to_corrupt.set_extension("datapack");
        let mut perms = to_corrupt.metadata().unwrap().permissions();
        perms.set_readonly(false);
        set_permissions(to_corrupt.clone(), perms).unwrap();
        let mut file = OpenOptions::new()
            .write(true)
            .open(to_corrupt.clone())
            .unwrap();
        file.write_all(b"FOOBARBAZ").unwrap();
        drop(file);

        let res = repack_datapacks(paths.into_iter(), tempdir.path())
            .err()
            .unwrap();

        if let Some(RepackFailure::Partial(errors)) = res.downcast_ref() {
            assert_eq!(errors.iter().count(), 1);
            to_corrupt.set_extension("");
            assert!(errors.iter().find(|(p, _)| p == &to_corrupt).is_some());
        } else {
            assert!(false);
        }
    }

    #[test]
    fn test_repack_one_historypack() {
        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let tempdir = TempDir::new().unwrap();

        let nodes = get_nodes(&mut rng);

        let pack = make_historypack(&tempdir, &nodes);
        let newpath = repack_historypacks(
            vec![pack.base_path().to_path_buf()].into_iter(),
            tempdir.path(),
        );
        assert!(newpath.is_ok());
        let newpack = HistoryPack::new(&newpath.unwrap().unwrap()).unwrap();

        for (ref key, _) in nodes.iter() {
            let response = newpack.get_node_info(key).unwrap().unwrap();
            assert_eq!(&response, nodes.get(key).unwrap());
        }
    }

    #[test]
    fn test_repack_multiple_historypack() {
        let mut rng = ChaChaRng::from_seed([0u8; 32]);
        let tempdir = TempDir::new().unwrap();
        let mut nodes = HashMap::new();
        let mut paths = Vec::new();

        for _ in 0..2 {
            let hgid = get_nodes(&mut rng);
            let pack = make_historypack(&tempdir, &hgid);
            let path = pack.base_path().to_path_buf();

            nodes.extend(hgid.into_iter());
            paths.push(path);
        }

        let newpath = repack_historypacks(paths.into_iter(), tempdir.path());
        assert!(newpath.is_ok());
        let newpack = HistoryPack::new(&newpath.unwrap().unwrap()).unwrap();

        for (key, _) in nodes.iter() {
            let response = newpack.get_node_info(&key).unwrap().unwrap();
            assert_eq!(&response, nodes.get(key).unwrap());
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::fmt::Debug;
use std::slice::SliceIndex;

use anyhow::Result;
use thiserror::Error;

#[derive(Debug, Error)]
#[error("SliceOutOfBounds Error: {0:?}")]
struct SliceOutOfBoundsError(String);

pub trait SliceExt<'a, T> {
    fn get_err<I>(self, range: I) -> Result<&'a I::Output>
    where
        I: SliceIndex<[T]> + Clone + Debug;
}

impl<'a, T> SliceExt<'a, T> for &'a [T] {
    fn get_err<I>(self, range: I) -> Result<&'a I::Output>
    where
        I: SliceIndex<[T]> + Clone + Debug,
    {
        self.get(range.clone()).ok_or_else(|| {
            SliceOutOfBoundsError(format!(
                "slice (len {:?}) too short to read range {:?}",
                self.len(),
                range
            ))
            .into()
        })
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::path::Path;
use std::sync::Arc;

use anyhow::Error;
use anyhow::Result;
use async_trait::async_trait;
use edenapi::EdenApi;
use edenapi::EdenApiError;
use edenapi::Response;
use edenapi::ResponseMeta;
use edenapi::Stats;
use edenapi_types::EdenApiServerError;
use edenapi_types::FileAttributes;
use edenapi_types::FileContent;
use edenapi_types::FileEntry;
use edenapi_types::FileResponse;
use edenapi_types::FileSpec;
use edenapi_types::HistoryEntry;
use edenapi_types::TreeAttributes;
use edenapi_types::TreeEntry;
use futures::prelude::*;
#[cfg(test)]
pub use lfs_mocks::*;
use minibytes::Bytes;
use types::Key;
use types::NodeInfo;
use types::Parents;

use crate::datastore::Delta;
use crate::datastore::HgIdDataStore;
use crate::datastore::HgIdMutableDeltaStore;
use crate::datastore::Metadata;
use crate::datastore::RemoteDataStore;
use crate::datastore::StoreResult;
use crate::historystore::HgIdHistoryStore;
use crate::historystore::HgIdMutableHistoryStore;
use crate::historystore::RemoteHistoryStore;
use crate::localstore::LocalStore;
use crate::remotestore::HgIdRemoteStore;
use crate::scmstore::file::LazyFile;
use crate::types::StoreKey;

pub fn delta(data: &str, base: Option<Key>, key: Key) -> Delta {
    Delta {
        data: Bytes::copy_from_slice(data.as_bytes()),
        base,
        key,
    }
}

pub struct FakeHgIdRemoteStore {
    data: Option<HashMap<Key, (Bytes, Option<u64>)>>,
    hist: Option<HashMap<Key, NodeInfo>>,
}

impl FakeHgIdRemoteStore {
    pub fn new() -> FakeHgIdRemoteStore {
        Self {
            data: None,
            hist: None,
        }
    }

    pub fn data(&mut self, map: HashMap<Key, (Bytes, Option<u64>)>) {
        self.data = Some(map)
    }

    pub fn hist(&mut self, map: HashMap<Key, NodeInfo>) {
        self.hist = Some(map)
    }
}

impl HgIdRemoteStore for FakeHgIdRemoteStore {
    fn datastore(
        self: Arc<Self>,
        store: Arc<dyn HgIdMutableDeltaStore>,
    ) -> Arc<dyn RemoteDataStore> {
        assert!(self.data.is_some());

        Arc::new(FakeRemoteDataStore {
            store,
            map: self.data.as_ref().unwrap().clone(),
        })
    }

    fn historystore(
        self: Arc<Self>,
        store: Arc<dyn HgIdMutableHistoryStore>,
    ) -> Arc<dyn RemoteHistoryStore> {
        assert!(self.hist.is_some());

        Arc::new(FakeRemoteHistoryStore {
            store,
            map: self.hist.as_ref().unwrap().clone(),
        })
    }
}

struct FakeRemoteDataStore {
    store: Arc<dyn HgIdMutableDeltaStore>,
    map: HashMap<Key, (Bytes, Option<u64>)>,
}

impl RemoteDataStore for FakeRemoteDataStore {
    fn prefetch(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        for k in keys {
            match k {
                StoreKey::HgId(k) => {
                    let (data, flags) = self.map.get(&k).ok_or_else(|| Error::msg("Not found"))?;
                    let delta = Delta {
                        data: data.clone(),
                        base: None,
                        key: k.clone(),
                    };
                    self.store.add(
                        &delta,
                        &Metadata {
                            size: Some(data.len() as u64),
                            flags: *flags,
                        },
                    )?;
                }
                StoreKey::Content(_, _) => continue,
            }
        }

        self.store.get_missing(keys)
    }

    fn upload(&self, _keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        unimplemented!()
    }
}

impl HgIdDataStore for FakeRemoteDataStore {
    fn get(&self, key: StoreKey) -> Result<StoreResult<Vec<u8>>> {
        match self.prefetch(&[key.clone()]) {
            Err(_) => Ok(StoreResult::NotFound(key)),
            Ok(_) => self.store.get(key),
        }
    }

    fn get_meta(&self, key: StoreKey) -> Result<StoreResult<Metadata>> {
        match self.prefetch(&[key.clone()]) {
            Err(_) => Ok(StoreResult::NotFound(key)),
            Ok(_) => self.store.get_meta(key),
        }
    }

    fn refresh(&self) -> Result<()> {
        Ok(())
    }
}

impl LocalStore for FakeRemoteDataStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.store.get_missing(keys)
    }
}

struct FakeRemoteHistoryStore {
    store: Arc<dyn HgIdMutableHistoryStore>,
    map: HashMap<Key, NodeInfo>,
}

impl RemoteHistoryStore for FakeRemoteHistoryStore {
    fn prefetch(&self, keys: &[StoreKey]) -> Result<()> {
        for k in keys {
            match k {
                StoreKey::HgId(k) => self
                    .store
                    .add(&k, self.map.get(&k).ok_or_else(|| Error::msg("Not found"))?)?,
                StoreKey::Content(_, _) => continue,
            }
        }

        Ok(())
    }
}

impl HgIdHistoryStore for FakeRemoteHistoryStore {
    fn get_node_info(&self, key: &Key) -> Result<Option<NodeInfo>> {
        match self.prefetch(&[StoreKey::hgid(key.clone())]) {
            Err(_) => Ok(None),
            Ok(()) => self.store.get_node_info(key),
        }
    }

    fn refresh(&self) -> Result<()> {
        Ok(())
    }
}

impl LocalStore for FakeRemoteHistoryStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.store.get_missing(keys)
    }
}

#[derive(Default)]
pub struct FakeEdenApi {
    files: HashMap<Key, (Bytes, Option<u64>)>,
    trees: HashMap<Key, Bytes>,
    history: HashMap<Key, NodeInfo>,
}

impl FakeEdenApi {
    pub fn new() -> Self {
        Default::default()
    }

    pub fn files(self, files: impl IntoIterator<Item = (Key, Bytes)>) -> Self {
        Self {
            files: files
                .into_iter()
                .map(|(key, bytes)| (key, (bytes, None)))
                .collect(),
            ..self
        }
    }

    /// See revisionstore::types::datastore::Metadata for how to construct these flags.
    ///
    /// Hint: None, or Some(Metadata::LFS_FLAG) are all you'll ever need.
    pub fn files_with_flags(self, files: HashMap<Key, (Bytes, Option<u64>)>) -> Self {
        Self { files, ..self }
    }

    pub fn trees(self, trees: HashMap<Key, Bytes>) -> Self {
        Self { trees, ..self }
    }

    pub fn history(self, history: HashMap<Key, NodeInfo>) -> Self {
        Self { history, ..self }
    }

    pub fn into_arc(self) -> Arc<Self> {
        Arc::new(self)
    }

    fn get_files(
        map: &HashMap<Key, (Bytes, Option<u64>)>,
        reqs: impl Iterator<Item = FileSpec>,
    ) -> Result<Response<FileResponse>, EdenApiError> {
        let entries = reqs
            .filter_map(|spec| {
                let parents = Parents::default();
                let mut entry = FileEntry::new(spec.key.clone(), parents);

                let (data, flags) = map.get(&spec.key)?.clone();
                let metadata = Metadata {
                    flags,
                    size: Some(data.len() as u64),
                };
                let data = data.to_vec().into();
                let content = FileContent {
                    hg_file_blob: data,
                    metadata,
                };

                if spec.attrs.aux_data {
                    // TODO(meyer): Compute aux data directly.
                    let mut file = LazyFile::EdenApi(entry.clone().with_content(content.clone()));
                    let aux = file.aux_data().ok()?;
                    entry = entry.with_aux_data(aux.into());
                }

                if spec.attrs.content {
                    entry = entry.with_content(content);
                }

                Some(Ok(FileResponse {
                    key: spec.key,
                    result: Ok(entry),
                }))
            })
            .collect::<Vec<_>>();

        Ok(Response {
            entries: Box::pin(stream::iter(entries)),
            stats: Box::pin(future::ok(Stats::default())),
        })
    }

    fn get_trees(
        map: &HashMap<Key, Bytes>,
        keys: Vec<Key>,
    ) -> Result<Response<Result<TreeEntry, EdenApiServerError>>, EdenApiError> {
        let entries = keys
            .into_iter()
            .filter_map(|key| {
                let data = map.get(&key)?.clone();
                let parents = Parents::default();
                let data = data.to_vec().into();
                Some(Ok(Ok(TreeEntry::new(key, data, parents))))
            })
            .collect::<Vec<_>>();

        Ok(Response {
            entries: Box::pin(stream::iter(entries)),
            stats: Box::pin(future::ok(Stats::default())),
        })
    }
}

#[async_trait]
impl EdenApi for FakeEdenApi {
    async fn health(&self) -> Result<ResponseMeta, EdenApiError> {
        Ok(ResponseMeta::default())
    }

    async fn files(&self, keys: Vec<Key>) -> Result<Response<FileResponse>, EdenApiError> {
        Self::get_files(
            &self.files,
            keys.into_iter().map(|key| FileSpec {
                key,
                attrs: FileAttributes {
                    content: true,
                    aux_data: false,
                },
            }),
        )
    }

    async fn files_attrs(
        &self,
        reqs: Vec<FileSpec>,
    ) -> Result<Response<FileResponse>, EdenApiError> {
        Self::get_files(&self.files, reqs.into_iter())
    }

    async fn history(
        &self,
        keys: Vec<Key>,
        _length: Option<u32>,
    ) -> Result<Response<HistoryEntry>, EdenApiError> {
        let entries = keys
            .into_iter()
            .filter_map(|key| {
                let nodeinfo = self.history.get(&key)?.clone();
                Some(Ok(HistoryEntry { key, nodeinfo }))
            })
            .collect::<Vec<_>>();

        Ok(Response {
            entries: Box::pin(stream::iter(entries)),
            stats: Box::pin(future::ok(Stats::default())),
        })
    }

    async fn trees(
        &self,
        keys: Vec<Key>,
        _attrs: Option<TreeAttributes>,
    ) -> Result<Response<Result<TreeEntry, EdenApiServerError>>, EdenApiError> {
        Self::get_trees(&self.trees, keys)
    }
}

pub fn make_config(dir: impl AsRef<Path>) -> BTreeMap<String, String> {
    [
        ("remotefilelog.reponame", "test".to_string()),
        (
            "remotefilelog.cachepath",
            dir.as_ref().display().to_string(),
        ),
        (
            "remotefilelog.cachekey",
            "cca::hg:rust_unittest".to_string(),
        ),
    ]
    .iter()
    .map(|(k, v)| (k.to_string(), v.clone()))
    .collect()
}

#[cfg(test)]
pub(crate) fn empty_config() -> BTreeMap<String, String> {
    BTreeMap::new()
}

#[cfg(test)]
mod lfs_mocks {
    use lfs_protocol::ObjectAction;
    use lfs_protocol::ObjectError;
    use lfs_protocol::ObjectStatus;
    use lfs_protocol::Operation;
    use lfs_protocol::RequestObject;
    use lfs_protocol::ResponseBatch;
    use lfs_protocol::ResponseObject;
    use lfs_protocol::Sha256 as LfsSha256;
    use lfs_protocol::Transfer;
    use mockito::mock;
    use mockito::Mock;
    use types::Sha256;

    use super::*;

    pub struct TestBlob {
        pub oid: &'static str,
        pub size: usize,
        pub content: Bytes,
        pub sha: Sha256,
        pub response: Vec<&'static [u8]>,
        pub error: bool,
        pub chunk_size: Option<usize>,
    }

    pub fn example_blob() -> TestBlob {
        use std::str::FromStr;

        let blob_oid = "fc613b4dfd6736a7bd268c8a0e74ed0d1c04a959f59dd74ef2874983fd443fc9";
        let content = b"master";

        TestBlob {
            oid: blob_oid,
            size: 6,
            content: Bytes::from(&content[..]),
            sha: Sha256::from_str(blob_oid).unwrap(),
            response: vec![content],
            error: false,
            chunk_size: None,
        }
    }

    pub fn example_blob2() -> TestBlob {
        use std::str::FromStr;
        let blob2_oid = "ca3e228a1d8d845064112c4e92781f6b8fc2501f0aa0e415d4a1dcc941485b24";
        let content = b"1.44.0";
        TestBlob {
            oid: blob2_oid,
            size: 6,
            content: Bytes::from(&content[..]),
            sha: Sha256::from_str(blob2_oid).unwrap(),
            response: vec![content],
            error: false,
            chunk_size: None,
        }
    }

    pub fn nonexistent_blob() -> TestBlob {
        use std::str::FromStr;
        let blob3_oid = "0000000000000000000000000000000000000000000000000000000000000000";
        TestBlob {
            oid: blob3_oid,
            size: 0,
            content: Bytes::from(&b""[..]),
            sha: Sha256::from_str(blob3_oid).unwrap(),
            response: vec![b"not_reached"],
            error: true,
            chunk_size: None,
        }
    }

    pub fn get_lfs_batch_mock(status: usize, blobs: &[&TestBlob]) -> Mock {
        let objects = blobs
            .iter()
            .map(|tb| {
                let object = RequestObject {
                    oid: LfsSha256(tb.sha.into_inner()),
                    size: tb.size as u64,
                };

                let status = if tb.error {
                    ObjectStatus::Err {
                        error: ObjectError {
                            code: 404,
                            message: "".into(),
                        },
                    }
                } else {
                    ObjectStatus::Ok {
                        authenticated: false,
                        actions: vec![(
                            Operation::Download,
                            ObjectAction {
                                href: format!("{}/repo/download/{}", mockito::server_url(), tb.oid)
                                    .as_str()
                                    .try_into()
                                    .unwrap(),
                                expires_at: None,
                                expires_in: None,
                                header: None,
                            },
                        )]
                        .into_iter()
                        .collect(),
                    }
                };

                ResponseObject { object, status }
            })
            .collect();

        let r = ResponseBatch {
            transfer: Transfer::Basic,
            objects,
        };

        let json_response = serde_json::to_string(&r).unwrap();

        mock("POST", "/repo/objects/batch")
            .with_status(status)
            .with_body(json_response)
            .create()
    }

    pub fn get_lfs_download_mock(status: usize, blob: &TestBlob) -> Vec<Mock> {
        let mut mocks = vec![];
        for response in blob.response.iter() {
            let m = mock("GET", format!("/repo/download/{}", blob.oid).as_str())
                .with_status(status)
                .with_body(response)
                .with_header("content-type", "application/octet-stream");

            mocks.push(m);
        }

        let mocks = if let Some(chunk_size) = blob.chunk_size {
            let mut i = 0;
            let mut chunked_mocks: Vec<Mock> = vec![];

            for mock in mocks.into_iter() {
                let m = mock.with_status(206).match_header(
                    "Range",
                    format!("bytes={}-{}", i, i + chunk_size - 1).as_str(),
                );
                chunked_mocks.push(m);
                i += chunk_size;
            }

            chunked_mocks
        } else {
            mocks
        };

        mocks.into_iter().map(|m| m.create()).collect()
    }

    pub fn make_lfs_config(dir: impl AsRef<Path>, agent_sufix: &str) -> BTreeMap<String, String> {
        let mut config = make_config(dir);
        let mut set = |key: &str, value: &str| {
            config.insert(key.to_string(), value.to_string());
        };
        set(
            "lfs.url",
            &[mockito::server_url(), "/repo".to_string()].concat(),
        );
        set("lfs.use-client-certs", "false");
        set(
            "experimental.lfs.user-agent",
            &format!("mercurial/revisionstore/unittests/{}", agent_sufix),
        );
        set("lfs.threshold", "4");
        set("remotefilelog.lfs", "true");
        set("lfs.moveafterupload", "true");

        config
    }
}

#[cfg(test)]
pub(crate) fn setconfig(
    config: &mut BTreeMap<String, String>,
    section: &str,
    name: &str,
    value: &str,
) {
    config.insert(format!("{}.{}", section, name), value.to_string());
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! Implement traits defined by other crates.

use std::sync::Arc;

use anyhow::anyhow;
use anyhow::format_err;
use anyhow::Result;
use async_trait::async_trait;
use futures::stream;
use futures::stream::BoxStream;
use futures::Stream;
use futures::StreamExt;
use hgstore::strip_metadata;
use minibytes::Bytes;
use storemodel::ReadFileContents;
use storemodel::RefreshableReadFileContents;
use tokio::runtime::Handle;
use types::Key;

use crate::scmstore::fetch::FetchMode;
use crate::scmstore::FileAttributes;
use crate::scmstore::FileStore;
use crate::RemoteDataStore;
use crate::StoreKey;
use crate::StoreResult;

// Wrapper types to workaround Rust's orphan rule.
#[derive(Clone)]
pub struct ArcFileStore(pub Arc<FileStore>);

pub struct ArcRemoteDataStore<T: ?Sized>(pub Arc<T>);

#[async_trait]
impl<T> ReadFileContents for ArcRemoteDataStore<T>
where
    T: RemoteDataStore + 'static + ?Sized,
{
    type Error = anyhow::Error;

    async fn read_file_contents(&self, keys: Vec<Key>) -> BoxStream<Result<(Bytes, Key)>> {
        stream_data_from_remote_data_store(self.0.clone(), keys)
            .map(|result| match result {
                Ok((data, key, _copy_from)) => Ok((data, key)),
                Err(err) => Err(err),
            })
            .boxed()
    }

    async fn read_rename_metadata(
        &self,
        keys: Vec<Key>,
    ) -> BoxStream<Result<(Key, Option<Key>), Self::Error>> {
        stream_data_from_remote_data_store(self.0.clone(), keys)
            .map(|result| match result {
                Ok((_data, key, copy_from)) => Ok((key, copy_from)),
                Err(err) => Err(err),
            })
            .boxed()
    }
}

#[async_trait]
impl ReadFileContents for ArcFileStore {
    type Error = anyhow::Error;

    async fn read_file_contents(&self, keys: Vec<Key>) -> BoxStream<Result<(Bytes, Key)>> {
        stream_data_from_scmstore(self.0.clone(), keys)
            .map(|result| match result {
                Ok((data, key, _copy_from)) => Ok((data, key)),
                Err(err) => Err(err),
            })
            .boxed()
    }

    async fn read_rename_metadata(
        &self,
        keys: Vec<Key>,
    ) -> BoxStream<Result<(Key, Option<Key>), Self::Error>> {
        stream_data_from_scmstore(self.0.clone(), keys)
            .map(|result| match result {
                Ok((_data, key, copy_from)) => Ok((key, copy_from)),
                Err(err) => Err(err),
            })
            .boxed()
    }
}

impl RefreshableReadFileContents for ArcFileStore {
    fn refresh(&self) -> Result<()> {
        FileStore::refresh(&self.0)
    }
}

const PREFETCH_CHUNK_SIZE: usize = 1000;
const FETCH_PARALLELISM: usize = 20;

fn stream_data_from_remote_data_store<DS: RemoteDataStore + Clone + 'static>(
    store: DS,
    keys: Vec<Key>,
) -> impl Stream<Item = Result<(Bytes, Key, Option<Key>)>> {
    stream::iter(keys.into_iter().map(StoreKey::HgId))
        .chunks(PREFETCH_CHUNK_SIZE)
        .map(move |chunk| {
            let store = store.clone();
            Handle::current().spawn_blocking(move || {
                let mut data = vec![];
                match store.prefetch(&chunk) {
                    Err(e) => {
                        data.push(Err(e));
                    }
                    Ok(_) => {
                        for store_key in chunk.iter() {
                            let key = match store_key {
                                StoreKey::HgId(key) => key,
                                _ => unreachable!(),
                            };
                            let store_result = store.get(store_key.clone());
                            let result = match store_result {
                                Err(err) => Err(err),
                                Ok(StoreResult::Found(data)) => strip_metadata(&data.into())
                                    .map(|(d, copy_from)| (d, key.clone(), copy_from)),
                                Ok(StoreResult::NotFound(k)) => {
                                    Err(format_err!("{:?} not found in store", k))
                                }
                            };
                            let is_err = result.is_err();
                            data.push(result);
                            if is_err {
                                break;
                            }
                        }
                    }
                };
                stream::iter(data.into_iter())
            })
        })
        .buffer_unordered(FETCH_PARALLELISM)
        .map(|r| {
            r.unwrap_or_else(|_| {
                stream::iter(vec![Err(anyhow!("background fetch join error"))].into_iter())
            })
        })
        .flatten()
}

fn stream_data_from_scmstore(
    store: Arc<FileStore>,
    keys: Vec<Key>,
) -> impl Stream<Item = Result<(Bytes, Key, Option<Key>)>> {
    stream::iter(keys.into_iter())
        .chunks(PREFETCH_CHUNK_SIZE)
        .map(move |chunk| {
            let store = store.clone();
            Handle::current().spawn_blocking(move || {
                let mut data = vec![];
                for result in store.fetch(
                    chunk.iter().cloned(),
                    FileAttributes::CONTENT,
                    FetchMode::AllowRemote,
                ) {
                    let result = match result {
                        Err(err) => Err(err.into()),
                        Ok((key, mut file)) => file
                            .file_content_with_copy_info()
                            .map(|(content, copy_from)| (content, key, copy_from)),
                    };
                    let is_err = result.is_err();
                    data.push(result);
                    if is_err {
                        break;
                    }
                }
                stream::iter(data.into_iter())
            })
        })
        .buffer_unordered(FETCH_PARALLELISM)
        .map(|r| {
            r.unwrap_or_else(|_| {
                stream::iter(vec![Err(anyhow!("background fetch join error"))].into_iter())
            })
        })
        .flatten()
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::io::Write;

use edenapi_types::Blake3;
use edenapi_types::ContentId;
use edenapi_types::Sha1;
use minibytes::Bytes;
#[cfg(any(test, feature = "for-tests"))]
use quickcheck_arbitrary_derive::Arbitrary;
use serde_derive::Deserialize;
use serde_derive::Serialize;
use types::Key;
use types::Sha256;

/// Kind of content hash stored in the LFS pointer. Adding new types is acceptable, re-ordering or
/// removal is forbidden.
#[derive(
    Serialize,
    Deserialize,
    PartialEq,
    Eq,
    Debug,
    Hash,
    Ord,
    PartialOrd,
    Clone
)]
#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
pub enum ContentHash {
    Sha256(#[serde(with = "types::serde_with::sha256::tuple")] Sha256),
}

#[derive(Clone, PartialEq, Eq, Debug, Hash, Ord, PartialOrd)]
#[cfg_attr(any(test, feature = "for-tests"), derive(Arbitrary))]
pub enum StoreKey {
    HgId(Key),
    /// The Key is a temporary workaround to being able to fallback from the LFS protocol to the
    /// non-LFS protocol. Do not depend on it as it will be removed.
    Content(ContentHash, Option<Key>),
}

impl ContentHash {
    pub fn sha256(data: &Bytes) -> Self {
        use sha2::Digest;

        let mut hash = sha2::Sha256::new();
        hash.update(data);

        let bytes: [u8; Sha256::len()] = hash.finalize().into();
        ContentHash::Sha256(Sha256::from(bytes))
    }

    pub(crate) fn content_id(data: &Bytes) -> ContentId {
        use blake2::digest::Update;
        use blake2::digest::VariableOutput;
        use blake2::VarBlake2b;

        // cribbed from pyedenapi
        let mut hash = VarBlake2b::new_keyed(b"content", ContentId::len());
        hash.update(data);
        let mut ret = [0u8; ContentId::len()];
        hash.finalize_variable(|res| {
            if let Err(e) = ret.as_mut().write_all(res) {
                panic!(
                    "{}-byte array must work with {}-byte blake2b: {:?}",
                    ContentId::len(),
                    ContentId::len(),
                    e
                );
            }
        });
        ContentId::from(ret)
    }

    pub(crate) fn sha1(data: &Bytes) -> Sha1 {
        use sha1::Digest;

        let mut hash = sha1::Sha1::new();
        hash.update(data);

        let bytes: [u8; Sha1::len()] = hash.finalize().into();
        Sha1::from(bytes)
    }

    pub(crate) fn seeded_blake3(data: &Bytes) -> Blake3 {
        #[cfg(not(fbcode_build))]
        {
            use blake3::Hasher;
            let key = "20220728-2357111317192329313741#".as_bytes();
            let mut ret = [0; Blake3::len()];
            ret.copy_from_slice(key);
            let mut hasher = Hasher::new_keyed(&ret);
            hasher.update(data.as_ref());
            let hashed_bytes: [u8; Blake3::len()] = hasher.finalize().into();
            Blake3::from(hashed_bytes)
        }
        #[cfg(fbcode_build)]
        {
            use blake3_c_ffi::Hasher;
            let key = blake3_constant::BLAKE3_HASH_KEY.as_bytes();
            let mut ret = [0; Blake3::len()];
            ret.copy_from_slice(key);
            let mut hasher = Hasher::new_keyed(&ret);
            hasher.update(data.as_ref());
            let mut hashed_bytes = [0; Blake3::len()];
            hasher.finalize(&mut hashed_bytes);
            Blake3::from(hashed_bytes)
        }
    }

    pub fn unwrap_sha256(self) -> Sha256 {
        match self {
            ContentHash::Sha256(hash) => hash,
        }
    }
}

impl StoreKey {
    pub fn hgid(key: Key) -> Self {
        StoreKey::HgId(key)
    }

    pub fn content(hash: ContentHash) -> Self {
        StoreKey::Content(hash, None)
    }

    pub fn maybe_into_key(self) -> Option<Key> {
        match self {
            StoreKey::HgId(key) => Some(key),
            StoreKey::Content(_, k) => k,
        }
    }

    pub fn maybe_as_key(&self) -> Option<&Key> {
        match self {
            StoreKey::HgId(ref key) => Some(key),
            StoreKey::Content(_, ref k) => k.as_ref(),
        }
    }
}

impl From<Key> for StoreKey {
    fn from(k: Key) -> Self {
        StoreKey::HgId(k)
    }
}

impl<'a> From<&'a Key> for StoreKey {
    fn from(k: &'a Key) -> Self {
        StoreKey::HgId(k.clone())
    }
}

impl From<ContentHash> for StoreKey {
    fn from(hash: ContentHash) -> Self {
        StoreKey::Content(hash, None)
    }
}

impl<'a> From<&'a ContentHash> for StoreKey {
    fn from(hash: &'a ContentHash) -> Self {
        StoreKey::Content(hash.clone(), None)
    }
}