/*
 * 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::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Result;
use async_runtime::block_on;
use async_runtime::spawn_blocking;
use async_runtime::stream_to_iter;
use crossbeam::channel::Sender;
use edenapi_types::FileResponse;
use edenapi_types::FileSpec;
use futures::StreamExt;
use progress_model::AggregatingProgressBar;
use tracing::debug;
use tracing::field;
use types::errors::NetworkError;
use types::Key;
use types::Sha256;

use crate::datastore::HgIdDataStore;
use crate::datastore::RemoteDataStore;
use crate::error::ClonableError;
use crate::fetch_logger::FetchLogger;
use crate::indexedlogauxstore::AuxStore;
use crate::indexedlogauxstore::Entry as AuxDataEntry;
use crate::indexedlogdatastore::Entry;
use crate::indexedlogdatastore::IndexedLogHgIdDataStore;
use crate::indexedlogutil::StoreType;
use crate::lfs::LfsPointersEntry;
use crate::lfs::LfsRemoteInner;
use crate::lfs::LfsStore;
use crate::lfs::LfsStoreEntry;
use crate::memcache::McData;
use crate::scmstore::attrs::StoreAttrs;
use crate::scmstore::fetch::CommonFetchState;
use crate::scmstore::fetch::FetchErrors;
use crate::scmstore::fetch::KeyFetchError;
use crate::scmstore::file::metrics::FileStoreFetchMetrics;
use crate::scmstore::file::LazyFile;
use crate::scmstore::value::StoreValue;
use crate::scmstore::FileAttributes;
use crate::scmstore::FileAuxData;
use crate::scmstore::FileStore;
use crate::scmstore::StoreFile;
use crate::util;
use crate::ContentHash;
use crate::ContentStore;
use crate::EdenApiFileStore;
use crate::ExtStoredPolicy;
use crate::MemcacheStore;
use crate::Metadata;
use crate::StoreKey;

pub struct FetchState {
    common: CommonFetchState<StoreFile>,

    /// Errors encountered during fetching.
    errors: FetchErrors,

    /// LFS pointers we've discovered corresponding to a request Key.
    lfs_pointers: HashMap<Key, (LfsPointersEntry, bool)>,

    /// A table tracking if discovered LFS pointers were found in the local-only or cache / shared store.
    pointer_origin: HashMap<Sha256, StoreType>,

    /// A table tracking if each key is local-only or cache/shared so that computed aux data can be written to the appropriate store
    key_origin: HashMap<Key, StoreType>,

    /// Tracks remote fetches which match a specific regex
    fetch_logger: Option<Arc<FetchLogger>>,

    lfs_progress: Arc<AggregatingProgressBar>,

    /// Track fetch metrics,
    metrics: FileStoreFetchMetrics,

    // Config
    extstored_policy: ExtStoredPolicy,
    compute_aux_data: bool,

    lfs_enabled: bool,
}

impl FetchState {
    pub(crate) fn new(
        keys: impl Iterator<Item = Key>,
        attrs: FileAttributes,
        file_store: &FileStore,
        found_tx: Sender<Result<(Key, StoreFile), KeyFetchError>>,
        lfs_enabled: bool,
    ) -> Self {
        FetchState {
            common: CommonFetchState::new(keys, attrs, found_tx),
            errors: FetchErrors::new(),
            metrics: FileStoreFetchMetrics::default(),

            lfs_pointers: HashMap::new(),
            key_origin: HashMap::new(),
            pointer_origin: HashMap::new(),

            fetch_logger: file_store.fetch_logger.clone(),
            extstored_policy: file_store.extstored_policy,
            compute_aux_data: true,
            lfs_progress: file_store.lfs_progress.clone(),
            lfs_enabled,
        }
    }

    pub(crate) fn pending_len(&self) -> usize {
        self.common.pending_len()
    }

    pub(crate) fn pending(&self) -> Vec<Key> {
        self.common.pending.iter().cloned().collect()
    }

    pub(crate) fn metrics(&self) -> &FileStoreFetchMetrics {
        &self.metrics
    }

    /// Return all incomplete requested Keys for which additional attributes may be gathered by querying a store which provides the specified attributes.
    fn pending_all(&self, fetchable: FileAttributes) -> Vec<Key> {
        if fetchable.none() {
            return vec![];
        }
        self.common
            .pending(fetchable, self.compute_aux_data)
            .map(|(key, _attrs)| key.clone())
            .collect()
    }

    /// Returns all incomplete requested Keys for which we haven't discovered an LFS pointer, and for which additional attributes may be gathered by querying a store which provides the specified attributes.
    fn pending_nonlfs(&self, fetchable: FileAttributes) -> Vec<Key> {
        if fetchable.none() {
            return vec![];
        }
        self.common
            .pending(fetchable, self.compute_aux_data)
            .map(|(key, _attrs)| key.clone())
            .filter(|k| !self.lfs_pointers.contains_key(k))
            .collect()
    }

    /// Returns all incomplete requested Keys as Store, with content Sha256 from the LFS pointer if available, for which additional attributes may be gathered by querying a store which provides the specified attributes
    fn pending_storekey(&self, fetchable: FileAttributes) -> Vec<StoreKey> {
        if fetchable.none() {
            return vec![];
        }
        self.common
            .pending(fetchable, self.compute_aux_data)
            .map(|(key, _attrs)| key.clone())
            .map(|k| self.storekey(k))
            .collect()
    }

    /// Returns the Key as a StoreKey, as a StoreKey::Content with Sha256 from the LFS Pointer, if available, otherwise as a StoreKey::HgId.
    /// Every StoreKey returned from this function is guaranteed to have an associated Key, so unwrapping is fine.
    fn storekey(&self, key: Key) -> StoreKey {
        if let Some((ptr, _)) = self.lfs_pointers.get(&key) {
            StoreKey::Content(ContentHash::Sha256(ptr.sha256()), Some(key))
        } else {
            StoreKey::HgId(key)
        }
    }

    fn mark_complete(&mut self, key: &Key) {
        if let Some((ptr, _)) = self.lfs_pointers.remove(key) {
            self.pointer_origin.remove(&ptr.sha256());
        }
    }

    fn found_pointer(&mut self, key: Key, ptr: LfsPointersEntry, typ: StoreType, write: bool) {
        let sha256 = ptr.sha256();
        // Overwrite StoreType::Local with StoreType::Shared, but not vice versa
        match typ {
            StoreType::Shared => {
                self.pointer_origin.insert(sha256, typ);
            }
            StoreType::Local => {
                self.pointer_origin.entry(sha256).or_insert(typ);
            }
        }
        self.lfs_pointers.insert(key, (ptr, write));
    }

    fn found_attributes(&mut self, key: Key, sf: StoreFile, typ: Option<StoreType>) {
        self.key_origin
            .insert(key.clone(), typ.unwrap_or(StoreType::Shared));

        if self.common.found(key.clone(), sf) {
            self.mark_complete(&key);
        }
    }

    fn evict_to_cache(
        key: Key,
        file: LazyFile,
        indexedlog_cache: &IndexedLogHgIdDataStore,
        memcache: Option<Arc<MemcacheStore>>,
    ) -> Result<LazyFile> {
        let cache_entry = file.indexedlog_cache_entry(key.clone())?.ok_or_else(|| {
            anyhow!("expected LazyFile::EdenApi or LazyFile::Memcache, other LazyFile variants should not be written to cache")
        })?;
        if let Some(memcache) = memcache.as_ref() {
            memcache.add_mcdata(cache_entry.clone().try_into()?);
        }
        indexedlog_cache.put_entry(cache_entry)?;
        let mmap_entry = indexedlog_cache
            .get_entry(key)?
            .ok_or_else(|| anyhow!("failed to read entry back from indexedlog after writing"))?;
        Ok(LazyFile::IndexedLog(mmap_entry))
    }

    fn found_indexedlog(
        &mut self,
        key: Key,
        entry: Entry,
        typ: StoreType,
        lfs_store: Option<&LfsStore>,
    ) {
        if entry.metadata().is_lfs() && self.lfs_enabled {
            if self.extstored_policy == ExtStoredPolicy::Use {
                match entry.try_into() {
                    Ok(ptr) => {
                        if let Some(lfs_store) = lfs_store {
                            // Promote this indexedlog LFS pointer to the
                            // pointer store if it isn't already present. This
                            // should only happen when the Python LFS extension
                            // is in play.
                            if let Ok(None) = lfs_store.fetch_available(&key.clone().into()) {
                                if let Err(err) = lfs_store.add_pointer(ptr) {
                                    self.errors.keyed_error(key, err);
                                }
                            }
                        } else {
                            // If we don't have somewhere to upgrade pointer,
                            // track as a "found" pointer so it will be fetched
                            // from the remote store subsequently.
                            self.found_pointer(key, ptr, typ, true)
                        }
                    }
                    Err(err) => self.errors.keyed_error(key, err),
                }
            }
        } else {
            self.found_attributes(key, LazyFile::IndexedLog(entry).into(), Some(typ))
        }
    }

    pub(crate) fn fetch_indexedlog(
        &mut self,
        store: &IndexedLogHgIdDataStore,
        lfs_store: Option<&LfsStore>,
        typ: StoreType,
    ) {
        let pending = self.pending_nonlfs(FileAttributes::CONTENT);
        if pending.is_empty() {
            return;
        }

        debug!(
            "Checking store Indexedlog ({cache}) for {key}{more}",
            cache = match typ {
                StoreType::Shared => "cache",
                StoreType::Local => "local",
            },
            key = pending[0],
            more = if pending.len() > 1 {
                format!(" and {} more", pending.len() - 1)
            } else {
                "".into()
            },
        );

        let mut found = 0;
        let mut errors = 0;
        let mut error: Option<String> = None;

        self.metrics.indexedlog.store(typ).fetch(pending.len());
        for key in pending.into_iter() {
            let res = store.get_raw_entry(&key);
            match res {
                Ok(Some(entry)) => {
                    self.metrics.indexedlog.store(typ).hit(1);
                    found += 1;
                    self.found_indexedlog(key, entry, typ, lfs_store)
                }
                Ok(None) => {
                    self.metrics.indexedlog.store(typ).miss(1);
                }
                Err(err) => {
                    self.metrics.indexedlog.store(typ).err(1);
                    errors += 1;
                    if error.is_none() {
                        error.replace(format!("{}: {}", key, err));
                    }
                    self.errors.keyed_error(key, err)
                }
            }
        }

        if found != 0 {
            debug!(
                "    Found {found} {result}",
                found = found,
                result = if found == 1 { "result" } else { "results" }
            );
        }
        if errors != 0 {
            debug!(
                "    Errors = {errors}, Error = {error:?}",
                errors = errors,
                error = error
            );
        }
    }

    fn found_aux_indexedlog(&mut self, key: Key, entry: AuxDataEntry, typ: StoreType) {
        let aux_data: FileAuxData = entry.into();
        self.found_attributes(key, aux_data.into(), Some(typ));
    }

    pub(crate) fn fetch_aux_indexedlog(&mut self, store: &AuxStore, typ: StoreType) {
        let pending = self.pending_all(FileAttributes::AUX);
        if pending.is_empty() {
            return;
        }

        debug!(
            "Checking store AUX ({cache}) - Count = {count}",
            cache = match typ {
                StoreType::Shared => "cache",
                StoreType::Local => "local",
            },
            count = pending.len()
        );

        let mut found = 0;
        let mut errors = 0;
        let mut error: Option<String> = None;

        self.metrics.aux.store(typ).fetch(pending.len());

        for key in pending.into_iter() {
            let res = store.get(key.hgid);
            match res {
                Ok(Some(aux)) => {
                    self.metrics.aux.store(typ).hit(1);
                    found += 1;
                    self.found_aux_indexedlog(key, aux, typ)
                }
                Ok(None) => {
                    self.metrics.aux.store(typ).miss(1);
                }
                Err(err) => {
                    self.metrics.aux.store(typ).err(1);
                    errors += 1;
                    if error.is_none() {
                        error.replace(format!("{}: {}", key, err));
                    }
                    self.errors.keyed_error(key, err)
                }
            }
        }

        if found != 0 {
            debug!("    Found = {found}", found = found);
        }
        if errors != 0 {
            debug!(
                "    Errors = {errors}, Error = {error:?}",
                errors = errors,
                error = error
            );
        }
    }

    fn found_lfs(&mut self, key: Key, entry: LfsStoreEntry, typ: StoreType) {
        match entry {
            LfsStoreEntry::PointerAndBlob(ptr, blob) => {
                self.found_attributes(key, LazyFile::Lfs(blob, ptr).into(), Some(typ))
            }
            LfsStoreEntry::PointerOnly(ptr) => self.found_pointer(key, ptr, typ, false),
        }
    }

    pub(crate) fn fetch_lfs(&mut self, store: &LfsStore, typ: StoreType) {
        let pending = self.pending_storekey(FileAttributes::CONTENT);
        if pending.is_empty() {
            return;
        }

        debug!(
            "Checking store LFS ({cache}) - Count = {count}",
            cache = match typ {
                StoreType::Shared => "cache",
                StoreType::Local => "local",
            },
            count = pending.len()
        );

        let mut found = 0;
        let mut found_pointers = 0;
        let mut errors = 0;
        let mut error: Option<String> = None;

        self.metrics.lfs.store(typ).fetch(pending.len());
        for store_key in pending.into_iter() {
            let key = store_key.clone().maybe_into_key().expect(
                "no Key present in StoreKey, even though this should be guaranteed by pending_all",
            );
            match store.fetch_available(&store_key) {
                Ok(Some(entry)) => {
                    // TODO(meyer): Make found behavior w/r/t LFS pointers and content consistent
                    self.metrics.lfs.store(typ).hit(1);
                    if let LfsStoreEntry::PointerOnly(_) = &entry {
                        found_pointers += 1;
                    } else {
                        found += 1;
                    }
                    self.found_lfs(key, entry, typ)
                }
                Ok(None) => {
                    self.metrics.lfs.store(typ).miss(1);
                }
                Err(err) => {
                    self.metrics.lfs.store(typ).err(1);
                    errors += 1;
                    if error.is_none() {
                        error.replace(format!("{}: {}", key, err));
                    }
                    self.errors.keyed_error(key, err)
                }
            }
        }

        if found != 0 {
            debug!("    Found = {found}", found = found);
        }
        if found_pointers != 0 {
            debug!(
                "    Found Pointers-Only = {found_pointers}",
                found_pointers = found_pointers
            );
        }
        if errors != 0 {
            debug!(
                "    Errors = {errors}, Error = {error:?}",
                errors = errors,
                error = error
            );
        }
    }

    fn found_memcache(
        &mut self,
        entry: McData,
        indexedlog_cache: Option<&IndexedLogHgIdDataStore>,
    ) {
        let key = entry.key.clone();
        if entry.metadata.is_lfs() {
            match entry.try_into() {
                Ok(ptr) => self.found_pointer(key, ptr, StoreType::Shared, true),
                Err(err) => self.errors.keyed_error(key, err),
            }
        } else if let Some(indexedlog_cache) = indexedlog_cache {
            match Self::evict_to_cache(
                key.clone(),
                LazyFile::Memcache(entry),
                indexedlog_cache,
                None,
            ) {
                Ok(cached) => {
                    self.found_attributes(key, cached.into(), None);
                }
                Err(err) => self.errors.keyed_error(key, err),
            }
        } else {
            self.found_attributes(key, LazyFile::Memcache(entry).into(), None);
        }
    }

    fn fetch_memcache_inner(
        &mut self,
        store: &MemcacheStore,
        indexedlog_cache: Option<&IndexedLogHgIdDataStore>,
    ) -> Result<()> {
        let pending = self.pending_nonlfs(FileAttributes::CONTENT);
        if pending.is_empty() {
            return Ok(());
        }

        debug!("Fetching Memcache - Count = {count}", count = pending.len());

        self.fetch_logger
            .as_ref()
            .map(|fl| fl.report_keys(pending.iter()));

        for res in store.get_data_iter(&pending)?.into_iter() {
            match res {
                Ok(mcdata) => self.found_memcache(mcdata, indexedlog_cache),
                Err(err) => self.errors.other_error(err),
            }
        }
        Ok(())
    }

    pub(crate) fn fetch_memcache(
        &mut self,
        store: &MemcacheStore,
        indexedlog_cache: Option<&IndexedLogHgIdDataStore>,
    ) {
        if let Err(err) = self.fetch_memcache_inner(store, indexedlog_cache) {
            self.errors.other_error(err);
        }
    }

    fn found_edenapi(
        entry: FileResponse,
        indexedlog_cache: Option<Arc<IndexedLogHgIdDataStore>>,
        lfs_cache: Option<Arc<LfsStore>>,
        aux_cache: Option<Arc<AuxStore>>,
        memcache: Option<Arc<MemcacheStore>>,
    ) -> Result<(StoreFile, Option<LfsPointersEntry>)> {
        let entry = entry.result?;

        let key = entry.key.clone();
        let mut file = StoreFile::default();
        let mut lfsptr = None;

        if let Some(aux_data) = entry.aux_data() {
            let aux_data: FileAuxData = aux_data.clone().into();
            if let Some(aux_cache) = aux_cache.as_ref() {
                aux_cache.put(key.hgid, &aux_data.into())?;
            }
            file.aux_data = Some(aux_data);
        }

        if let Some(content) = entry.content() {
            if content.metadata().is_lfs() {
                let ptr: LfsPointersEntry = entry.try_into()?;
                if let Some(lfs_cache) = lfs_cache.as_ref() {
                    lfs_cache.add_pointer(ptr.clone())?;
                }
                lfsptr = Some(ptr);
            } else if let Some(indexedlog_cache) = indexedlog_cache.as_ref() {
                file.content = Some(Self::evict_to_cache(
                    key,
                    LazyFile::EdenApi(entry),
                    indexedlog_cache,
                    memcache,
                )?);
            } else {
                file.content = Some(LazyFile::EdenApi(entry));
            }
        }

        Ok((file, lfsptr))
    }

    pub(crate) fn fetch_edenapi(
        &mut self,
        store: &EdenApiFileStore,
        indexedlog_cache: Option<Arc<IndexedLogHgIdDataStore>>,
        lfs_cache: Option<Arc<LfsStore>>,
        aux_cache: Option<Arc<AuxStore>>,
        memcache: Option<Arc<MemcacheStore>>,
    ) {
        let fetchable = FileAttributes::CONTENT | FileAttributes::AUX;

        let pending = self.pending_nonlfs(fetchable);
        if pending.is_empty() {
            return;
        }

        let mut fetching_keys: HashSet<Key> = pending.iter().cloned().collect();

        debug!("Fetching EdenAPI - Count = {count}", count = pending.len());

        let mut found = 0;
        let mut found_pointers = 0;
        let mut errors = 0;
        let mut error: Option<String> = None;

        self.fetch_logger
            .as_ref()
            .map(|fl| fl.report_keys(pending.iter()));

        // TODO(meyer): Iterators or otherwise clean this up
        let pending_attrs: Vec<_> = pending
            .into_iter()
            .map(|k| {
                let actionable = self.common.actionable(&k, fetchable, self.compute_aux_data);
                FileSpec {
                    key: k,
                    attrs: actionable.into(),
                }
            })
            .collect();

        let response = match block_on(store.files_attrs(pending_attrs)).map_err(|e| e.tag_network())
        {
            Ok(r) => r,
            Err(err) => {
                let err = ClonableError::new(err);
                for key in fetching_keys.into_iter() {
                    self.errors.keyed_error(key, err.clone().into());
                }
                return;
            }
        };

        let entries = response
            .entries
            .map(move |res_entry| {
                let lfs_cache = lfs_cache.clone();
                let indexedlog_cache = indexedlog_cache.clone();
                let aux_cache = aux_cache.clone();
                let memcache = memcache.clone();
                spawn_blocking(move || {
                    res_entry.map(move |entry| {
                        (
                            entry.key.clone(),
                            Self::found_edenapi(
                                entry,
                                indexedlog_cache,
                                lfs_cache,
                                aux_cache,
                                memcache,
                            ),
                        )
                    })
                })

                // Processing a response may involve compressing the response, which
                // can be expensive. If we don't process entries fast enough, edenapi
                // can start queueing up responses which causes forever increasing
                // memory usage. So let's process responses in parallel to stay ahead
                // of download speeds.
            })
            .buffer_unordered(4);

        // Record found entries
        let mut unknown_error: Option<ClonableError> = None;
        for res in stream_to_iter(entries) {
            // TODO(meyer): This outer EdenApi error with no key sucks
            let (key, res) = match res {
                Ok(result) => match result.map_err(|e| e.tag_network()) {
                    Ok(result) => result,
                    Err(err) => {
                        if unknown_error.is_none() {
                            unknown_error.replace(ClonableError::new(err));
                        }
                        continue;
                    }
                },
                // JoinError
                Err(err) => {
                    if unknown_error.is_none() {
                        unknown_error.replace(ClonableError::new(err.into()));
                    }
                    continue;
                }
            };

            fetching_keys.remove(&key);
            match res {
                Ok((file, maybe_lfsptr)) => {
                    if let Some(lfsptr) = maybe_lfsptr {
                        found_pointers += 1;
                        self.found_pointer(key.clone(), lfsptr, StoreType::Shared, false);
                    } else {
                        found += 1;
                    }
                    self.found_attributes(key, file, Some(StoreType::Shared));
                }
                Err(err) => {
                    errors += 1;
                    if error.is_none() {
                        error.replace(format!("{}: {}", key, err));
                    }
                    self.errors.keyed_error(key, NetworkError::wrap(err))
                }
            }
        }

        for missing_key in fetching_keys.into_iter() {
            match &unknown_error {
                Some(error) => self.errors.keyed_error(missing_key, error.clone().into()),
                None => {
                    // This should never happen.
                    self.errors.keyed_error(
                        missing_key,
                        anyhow!("key not returned from files_attr request"),
                    )
                }
            };
        }

        if found != 0 {
            debug!("    Found = {found}", found = found);
        }
        if found_pointers != 0 {
            debug!(
                "    Found Pointers = {found_pointers}",
                found_pointers = found_pointers
            );
        }
        if errors != 0 {
            debug!(
                "    Errors = {errors}, Error = {error:?}",
                errors = errors,
                error = error
            );
        }

        let span = tracing::info_span!(
            "fetch_edenapi",
            downloaded = field::Empty,
            uploaded = field::Empty,
            requests = field::Empty,
            time = field::Empty,
            latency = field::Empty,
            download_speed = field::Empty,
            scmstore = true,
        );
        let _enter = span.enter();

        if let Ok(stats) = block_on(response.stats) {
            util::record_edenapi_stats(&span, &stats);
        }
    }

    pub(crate) fn fetch_lfs_remote(
        &mut self,
        store: &LfsRemoteInner,
        local: Option<Arc<LfsStore>>,
        cache: Option<Arc<LfsStore>>,
    ) {
        let errors = &mut self.errors;
        let pending: HashSet<_> = self
            .lfs_pointers
            .iter()
            .map(|(key, (ptr, write))| {
                if *write {
                    if let Some(lfs_cache) = cache.as_ref() {
                        if let Err(err) = lfs_cache.add_pointer(ptr.clone()) {
                            errors.keyed_error(key.clone(), err);
                        }
                    }
                }
                (ptr.sha256(), ptr.size() as usize)
            })
            .collect();

        let mut key_map: HashMap<Sha256, Vec<(Key, LfsPointersEntry)>> = HashMap::new();
        for (key, (ptr, _)) in self.lfs_pointers.iter() {
            let keys = key_map.entry(ptr.sha256()).or_default();
            keys.push((key.clone(), ptr.clone()));
        }

        if pending.is_empty() {
            return;
        }

        debug!("Fetching LFS - Count = {count}", count = pending.len());

        self.fetch_logger
            .as_ref()
            .map(|fl| fl.report_keys(self.lfs_pointers.keys()));

        let prog = self.lfs_progress.create_or_extend(pending.len() as u64);

        let mut keyed_errors = Vec::<(Key, anyhow::Error)>::new();
        let mut other_errors = vec![];

        // Fetch & write to local LFS stores
        let top_level_error = store.batch_fetch(
            &pending,
            |sha256, data| -> Result<()> {
                prog.increase_position(1);

                match self.pointer_origin.get(&sha256).ok_or_else(|| {
                    anyhow!(
                        "no source found for Sha256; received unexpected Sha256 from LFS server"
                    )
                })? {
                    StoreType::Local => local
                        .as_ref()
                        .expect("no lfs_local present when handling local LFS pointer")
                        .add_blob(&sha256, data.clone())?,
                    StoreType::Shared => cache
                        .as_ref()
                        .expect("no lfs_cache present when handling cache LFS pointer")
                        .add_blob(&sha256, data.clone())?,
                };

                // Unwrap is safe because the only place sha256 could come from is
                // `pending` and all of its entries were put in `key_map`.
                for (key, ptr) in key_map.get(&sha256).unwrap().iter() {
                    let mut file = StoreFile::default();
                    file.content = Some(LazyFile::Lfs(data.clone(), ptr.clone()));

                    // It's important to do this after the self.pointer_origin.get() above, since
                    // found_attributes removes the key from pointer_origin.
                    self.found_attributes(key.clone(), file, Some(StoreType::Shared));
                }

                Ok(())
            },
            |sha256, error| {
                if let Some(keys) = key_map.get(&sha256) {
                    let error = ClonableError::new(NetworkError::wrap(error));
                    for (key, _) in keys.iter() {
                        keyed_errors.push(((*key).clone(), error.clone().into()));
                    }
                } else {
                    other_errors.push(anyhow!("invalid other lfs error: {:?}", error));
                }
            },
        );

        if let Err(err) = top_level_error {
            let err = ClonableError::new(err);
            for (key, (_ptr, _write)) in self.lfs_pointers.iter() {
                self.errors.keyed_error(key.clone(), err.clone().into());
            }
        }

        for (key, error) in keyed_errors.into_iter() {
            self.errors.keyed_error(key, error);
        }
        for error in other_errors.into_iter() {
            self.errors.other_error(error);
        }
    }

    fn found_contentstore(&mut self, key: Key, bytes: Vec<u8>, meta: Metadata) {
        if meta.is_lfs() {
            self.metrics.contentstore.hit_lfsptr(1);
            // Do nothing. We're trying to avoid exposing LFS pointers to the consumer of this API.
            // We very well may need to expose LFS Pointers to the caller in the end (to match ContentStore's
            // ExtStoredPolicy behavior), but hopefully not, and if so we'll need to make it type safe.
            tracing::warn!("contentstore fallback returned serialized lfs pointer");
        } else {
            tracing::warn!(
                "contentstore fetched a file scmstore couldn't, \
                this indicates a bug or unsupported configuration: \
                fetched key '{}', found {} bytes of content with metadata {:?}.",
                key,
                bytes.len(),
                meta,
            );
            self.metrics.contentstore.hit(1);
            self.found_attributes(key, LazyFile::ContentStore(bytes.into(), meta).into(), None)
        }
    }

    fn fetch_contentstore_inner(
        &mut self,
        store: &ContentStore,
        pending: &mut Vec<StoreKey>,
    ) -> Result<()> {
        debug!(
            "ContentStore Fallback  - Count = {count}",
            count = pending.len()
        );
        let mut found = 0;
        let mut errors = 0;
        let mut error: Option<String> = None;

        store.prefetch(&pending)?;

        for store_key in pending.drain(..) {
            let key = store_key.clone().maybe_into_key().expect(
                "no Key present in StoreKey, even though this should be guaranteed by pending_storekey",
            );
            // Using the ContentStore API, fetch the hg file blob, then, if it's found, also fetch the file metadata.
            // Returns the requested file as Result<(Option<Vec<u8>>, Option<Metadata>)>
            // Produces a Result::Err if either the blob or metadata get returned an error
            let res = store
                .get(store_key.clone())
                .map(|store_result| store_result.into())
                .and_then({
                    let store_key = store_key.clone();
                    |maybe_blob| {
                        Ok((
                            maybe_blob,
                            store
                                .get_meta(store_key)
                                .map(|store_result| store_result.into())?,
                        ))
                    }
                });

            match res {
                Ok((Some(blob), Some(meta))) => {
                    found += 1;
                    self.found_contentstore(key, blob, meta)
                }
                Err(err) => {
                    self.metrics.contentstore.err(1);
                    errors += 1;
                    if error.is_none() {
                        error.replace(format!("{}: {}", key, err));
                    }
                    self.errors.keyed_error(key, err)
                }
                _ => {
                    self.metrics.contentstore.miss(1);
                }
            }
        }

        if found != 0 {
            debug!("    Found = {found}", found = found);
        }
        if errors != 0 {
            debug!(
                "    Errors = {errors}, Error = {error:?}",
                errors = errors,
                error = error
            );
        }

        Ok(())
    }

    pub(crate) fn fetch_contentstore(&mut self, store: &ContentStore) {
        let mut pending = self.pending_storekey(FileAttributes::CONTENT);
        if pending.is_empty() {
            return;
        }
        self.metrics.contentstore.fetch(pending.len());
        if let Err(err) = self.fetch_contentstore_inner(store, &mut pending) {
            debug!("ContentStore upper error - Error = {err:?}", err = err);
            self.errors.other_error(err);
            self.metrics.contentstore.err(pending.len());
        }
    }

    // TODO(meyer): Improve how local caching works. At the very least do this in the background.
    // TODO(meyer): Log errors here instead of just ignoring.
    pub(crate) fn derive_computable(
        &mut self,
        aux_cache: Option<&AuxStore>,
        aux_local: Option<&AuxStore>,
    ) {
        if !self.compute_aux_data {
            return;
        }

        for key in self.common.pending.iter().cloned().collect::<Vec<_>>() {
            if let Some(value) = self.common.found.get_mut(&key) {
                let span = tracing::debug_span!("checking derivations", %key);
                let _guard = span.enter();

                let existing_attrs = value.attrs();
                let missing = self.common.request_attrs - existing_attrs;
                let actionable = existing_attrs.with_computable() & missing;

                if actionable.aux_data {
                    let mut new = std::mem::take(value);

                    tracing::debug!("computing aux data");
                    if let Err(err) = new.compute_aux_data() {
                        self.errors.keyed_error(key.clone(), err);
                    } else {
                        tracing::debug!("computed aux data");

                        // mark complete if applicable
                        if new.attrs().has(self.common.request_attrs) {
                            tracing::debug!("marking complete");

                            match self.key_origin.get(&key).unwrap_or(&StoreType::Shared) {
                                StoreType::Shared => {
                                    if let Some(ref aux_cache) = aux_cache {
                                        if let Some(aux_data) = new.aux_data {
                                            let _ = aux_cache.put(key.hgid, &aux_data.into());
                                        }
                                    }
                                }
                                StoreType::Local => {
                                    if let Some(ref aux_local) = aux_local {
                                        if let Some(aux_data) = new.aux_data {
                                            let _ = aux_local.put(key.hgid, &aux_data.into());
                                        }
                                    }
                                }
                            }

                            // TODO(meyer): Extract out a "FetchPending" object like FetchErrors, or otherwise make it possible
                            // to share a "mark complete" implementation while holding a mutable reference to self.found.
                            self.common.pending.remove(&key);
                            self.common.found.remove(&key);
                            let new = new.mask(self.common.request_attrs);
                            let _ = self.common.found_tx.send(Ok((key.clone(), new)));
                            if let Some((ptr, _)) = self.lfs_pointers.remove(&key) {
                                self.pointer_origin.remove(&ptr.sha256());
                            }
                        } else {
                            *value = new;
                        }
                    }
                }
            }
        }
    }

    pub(crate) fn finish(self) {
        self.common.results(self.errors);
    }
}


/*
 * 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::ops::AddAssign;
use std::sync::Arc;

use parking_lot::RwLock;

use crate::scmstore::metrics::namespaced;
use crate::scmstore::metrics::ApiMetrics;
use crate::scmstore::metrics::FetchMetrics;
use crate::scmstore::metrics::LocalAndCacheFetchMetrics;
use crate::scmstore::metrics::WriteMetrics;

#[derive(Clone, Debug, Default)]
pub struct ContentStoreFetchMetrics {
    /// Only content hits are counted by common.hits, LFS pointer hits are only counted below.
    common: FetchMetrics,

    /// ContentStore returned a serialized LFS pointer instead of file content.
    lfsptr_hits: usize,
}

impl ContentStoreFetchMetrics {
    pub(crate) fn fetch(&mut self, keys: usize) {
        self.common.fetch(keys)
    }

    pub(crate) fn hit(&mut self, keys: usize) {
        self.common.hit(keys)
    }

    pub(crate) fn miss(&mut self, keys: usize) {
        self.common.miss(keys)
    }

    pub(crate) fn err(&mut self, keys: usize) {
        self.common.err(keys)
    }

    pub(crate) fn hit_lfsptr(&mut self, keys: usize) {
        self.lfsptr_hits += keys;
    }

    fn metrics(&self) -> impl Iterator<Item = (&'static str, usize)> {
        std::iter::once(("lfsptrhits", self.lfsptr_hits))
            .filter(|&(_, v)| v != 0)
            .chain(self.common.metrics())
    }
}

impl AddAssign for ContentStoreFetchMetrics {
    fn add_assign(&mut self, rhs: Self) {
        self.common += rhs.common;
        self.lfsptr_hits += rhs.lfsptr_hits;
    }
}

#[derive(Clone, Debug, Default)]
pub struct FileStoreFetchMetrics {
    pub(crate) indexedlog: LocalAndCacheFetchMetrics,
    pub(crate) lfs: LocalAndCacheFetchMetrics,
    pub(crate) aux: LocalAndCacheFetchMetrics,
    pub(crate) contentstore: ContentStoreFetchMetrics,
}

impl AddAssign for FileStoreFetchMetrics {
    fn add_assign(&mut self, rhs: Self) {
        self.indexedlog += rhs.indexedlog;
        self.lfs += rhs.lfs;
        self.aux += rhs.aux;
        self.contentstore += rhs.contentstore;
    }
}

impl FileStoreFetchMetrics {
    fn metrics(&self) -> impl Iterator<Item = (String, usize)> {
        namespaced("indexedlog", self.indexedlog.metrics())
            .chain(namespaced("lfs", self.lfs.metrics()))
            .chain(namespaced("aux", self.aux.metrics()))
            .chain(namespaced("contentstore", self.contentstore.metrics()))
    }
}

#[derive(Clone, Debug, Default)]
pub struct FileStoreWriteMetrics {
    /// Writes to the local LFS backend
    pub(crate) lfs: WriteMetrics,

    /// Writes to the local non-LFS backend
    pub(crate) nonlfs: WriteMetrics,

    /// LFS Pointer-only writes (supported only through fallback)
    pub(crate) lfsptr: WriteMetrics,
}

impl AddAssign for FileStoreWriteMetrics {
    fn add_assign(&mut self, rhs: Self) {
        self.lfs += rhs.lfs;
        self.nonlfs += rhs.nonlfs;
        self.lfsptr += rhs.lfsptr;
    }
}

impl FileStoreWriteMetrics {
    fn metrics(&self) -> impl Iterator<Item = (String, usize)> {
        namespaced("lfs", self.lfs.metrics())
            .chain(namespaced("nonlfs", self.nonlfs.metrics()))
            .chain(namespaced("lfsptr", self.lfsptr.metrics()))
    }
}

#[derive(Clone, Debug, Default)]
pub struct FileStoreApiMetrics {
    pub(crate) hg_getfilecontent: ApiMetrics,
    pub(crate) hg_addpending: ApiMetrics,
    pub(crate) hg_commitpending: ApiMetrics,
    pub(crate) hg_get: ApiMetrics,
    pub(crate) hg_getmeta: ApiMetrics,
    pub(crate) hg_refresh: ApiMetrics,
    pub(crate) hg_prefetch: ApiMetrics,
    pub(crate) hg_upload: ApiMetrics,
    pub(crate) hg_getmissing: ApiMetrics,
    pub(crate) hg_add: ApiMetrics,
    pub(crate) hg_flush: ApiMetrics,
    pub(crate) contentdatastore_blob: ApiMetrics,
    pub(crate) contentdatastore_metadata: ApiMetrics,
}

impl AddAssign for FileStoreApiMetrics {
    fn add_assign(&mut self, rhs: Self) {
        self.hg_getfilecontent += rhs.hg_getfilecontent;
        self.hg_addpending += rhs.hg_addpending;
        self.hg_commitpending += rhs.hg_commitpending;
        self.hg_get += rhs.hg_get;
        self.hg_getmeta += rhs.hg_getmeta;
        self.hg_refresh += rhs.hg_refresh;
        self.hg_prefetch += rhs.hg_prefetch;
        self.hg_upload += rhs.hg_upload;
        self.hg_getmissing += rhs.hg_getmissing;
        self.hg_add += rhs.hg_add;
        self.hg_flush += rhs.hg_flush;
        self.contentdatastore_blob += rhs.contentdatastore_blob;
        self.contentdatastore_metadata += rhs.contentdatastore_metadata;
    }
}

impl FileStoreApiMetrics {
    fn metrics(&self) -> impl Iterator<Item = (String, usize)> {
        namespaced("hg_getfilecontent", self.hg_getfilecontent.metrics())
            .chain(namespaced("hg_addpending", self.hg_addpending.metrics()))
            .chain(namespaced(
                "hg_commitpending",
                self.hg_commitpending.metrics(),
            ))
            .chain(namespaced("hg_get", self.hg_get.metrics()))
            .chain(namespaced("hg_getmeta", self.hg_getmeta.metrics()))
            .chain(namespaced("hg_refresh", self.hg_refresh.metrics()))
            .chain(namespaced("hg_prefetch", self.hg_prefetch.metrics()))
            .chain(namespaced("hg_upload", self.hg_upload.metrics()))
            .chain(namespaced("hg_getmissing", self.hg_getmissing.metrics()))
            .chain(namespaced("hg_add", self.hg_add.metrics()))
            .chain(namespaced(
                "contentdatastore_blob",
                self.contentdatastore_blob.metrics(),
            ))
            .chain(namespaced(
                "contentdatastore_metadata",
                self.contentdatastore_metadata.metrics(),
            ))
    }
}

#[derive(Debug, Default, Clone)]
pub struct FileStoreMetrics {
    pub(crate) fetch: FileStoreFetchMetrics,
    pub(crate) write: FileStoreWriteMetrics,
    pub(crate) api: FileStoreApiMetrics,
}

impl FileStoreMetrics {
    pub fn new() -> Arc<RwLock<Self>> {
        Arc::new(RwLock::new(FileStoreMetrics::default()))
    }

    pub fn metrics(&self) -> impl Iterator<Item = (String, usize)> {
        namespaced(
            "scmstore.file",
            namespaced("fetch", self.fetch.metrics())
                .chain(namespaced("write", self.write.metrics()))
                .chain(namespaced("api", self.api.metrics())),
        )
    }
}


/*
 * 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.
 */

mod fetch;
mod metrics;
mod types;

use std::collections::HashSet;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Instant;

use ::types::Key;
use ::types::RepoPathBuf;
use anyhow::anyhow;
use anyhow::bail;
use anyhow::ensure;
use anyhow::Result;
use crossbeam::channel::unbounded;
use minibytes::Bytes;
use parking_lot::Mutex;
use parking_lot::RwLock;
use progress_model::AggregatingProgressBar;
use rand::Rng;

pub(crate) use self::fetch::FetchState;
pub use self::metrics::FileStoreFetchMetrics;
pub use self::metrics::FileStoreMetrics;
pub use self::metrics::FileStoreWriteMetrics;
pub use self::types::FileAttributes;
pub use self::types::FileAuxData;
pub(crate) use self::types::LazyFile;
pub use self::types::StoreFile;
use crate::datastore::HgIdDataStore;
use crate::datastore::HgIdMutableDeltaStore;
use crate::datastore::RemoteDataStore;
use crate::fetch_logger::FetchLogger;
use crate::indexedlogauxstore::AuxStore;
use crate::indexedlogdatastore::Entry;
use crate::indexedlogdatastore::IndexedLogHgIdDataStore;
use crate::indexedlogutil::StoreType;
use crate::lfs::lfs_from_hg_file_blob;
use crate::lfs::LfsPointersEntry;
use crate::lfs::LfsRemote;
use crate::lfs::LfsStore;
use crate::memcache::MEMCACHE_DELAY;
use crate::remotestore::HgIdRemoteStore;
use crate::scmstore::activitylogger::ActivityLogger;
use crate::scmstore::fetch::FetchMode;
use crate::scmstore::fetch::FetchResults;
use crate::ContentDataStore;
use crate::ContentMetadata;
use crate::ContentStore;
use crate::Delta;
use crate::EdenApiFileStore;
use crate::ExtStoredPolicy;
use crate::LegacyStore;
use crate::LocalStore;
use crate::MemcacheStore;
use crate::Metadata;
use crate::MultiplexDeltaStore;
use crate::RepackLocation;
use crate::StoreKey;
use crate::StoreResult;

#[derive(Clone)]
pub struct FileStore {
    // Config
    // TODO(meyer): Move these to a separate config struct with default impl, etc.
    pub(crate) extstored_policy: ExtStoredPolicy,
    pub(crate) lfs_threshold_bytes: Option<u64>,
    pub(crate) cache_to_memcache: bool,
    pub(crate) edenapi_retries: i32,
    /// Allow explicitly writing serialized LFS pointers outside of tests
    pub(crate) allow_write_lfs_ptrs: bool,
    pub(crate) prefer_computing_aux_data: bool,

    // Record remote fetches
    pub(crate) fetch_logger: Option<Arc<FetchLogger>>,

    // Local-only stores
    pub(crate) indexedlog_local: Option<Arc<IndexedLogHgIdDataStore>>,
    pub(crate) lfs_local: Option<Arc<LfsStore>>,

    // Local non-lfs cache aka shared store
    pub(crate) indexedlog_cache: Option<Arc<IndexedLogHgIdDataStore>>,

    // Local LFS cache aka shared store
    pub(crate) lfs_cache: Option<Arc<LfsStore>>,

    // Memcache
    pub(crate) memcache: Option<Arc<MemcacheStore>>,

    // Remote stores
    pub(crate) lfs_remote: Option<Arc<LfsRemote>>,
    pub(crate) edenapi: Option<Arc<EdenApiFileStore>>,

    // Legacy ContentStore fallback
    pub(crate) contentstore: Option<Arc<ContentStore>>,

    // Aux Data Stores
    pub(crate) aux_local: Option<Arc<AuxStore>>,
    pub(crate) aux_cache: Option<Arc<AuxStore>>,

    // Metrics, statistics, debugging
    pub(crate) activity_logger: Option<Arc<Mutex<ActivityLogger>>>,
    pub(crate) metrics: Arc<RwLock<FileStoreMetrics>>,

    // Records the store creation time, so we can only use memcache for long running commands.
    pub(crate) creation_time: Instant,

    pub(crate) lfs_progress: Arc<AggregatingProgressBar>,

    // Don't flush on drop when we're using FileStore in a "disposable" context, like backingstore
    pub flush_on_drop: bool,
}

impl Drop for FileStore {
    fn drop(&mut self) {
        if self.flush_on_drop {
            let _ = self.flush();
        }
    }
}

impl FileStore {
    pub fn fetch(
        &self,
        keys: impl Iterator<Item = Key>,
        attrs: FileAttributes,
        fetch_mode: FetchMode,
    ) -> FetchResults<StoreFile> {
        let (found_tx, found_rx) = unbounded();
        let mut state = FetchState::new(
            keys,
            attrs,
            self,
            found_tx,
            self.lfs_threshold_bytes.is_some(),
        );

        let keys_len = state.pending_len();

        let aux_cache = self.aux_cache.clone();
        let aux_local = self.aux_local.clone();
        let indexedlog_cache = self.indexedlog_cache.clone();
        let indexedlog_local = self.indexedlog_local.clone();
        let lfs_cache = self.lfs_cache.clone();
        let lfs_local = self.lfs_local.clone();
        let memcache = self.memcache.clone();
        let edenapi = self.edenapi.clone();
        let lfs_remote = self.lfs_remote.clone();
        let contentstore = self.contentstore.clone();
        let creation_time = self.creation_time;
        let prefer_computing_aux_data = self.prefer_computing_aux_data;
        let cache_to_memcache = self.cache_to_memcache;
        let metrics = self.metrics.clone();
        let activity_logger = self.activity_logger.clone();

        let process_func = move || {
            let start_instant = Instant::now();

            let all_keys: Vec<Key> = state.pending();
            let span = tracing::span!(
                tracing::Level::DEBUG,
                "file fetch",
                id = rand::thread_rng().gen::<u16>()
            );
            let _enter = span.enter();

            if let Some(ref aux_cache) = aux_cache {
                state.fetch_aux_indexedlog(aux_cache, StoreType::Shared);
            }

            if let Some(ref aux_local) = aux_local {
                state.fetch_aux_indexedlog(aux_local, StoreType::Local);
            }

            if let Some(ref indexedlog_cache) = indexedlog_cache {
                state.fetch_indexedlog(
                    indexedlog_cache,
                    lfs_cache.as_ref().map(|v| v.as_ref()),
                    StoreType::Shared,
                );
            }

            if let Some(ref indexedlog_local) = indexedlog_local {
                state.fetch_indexedlog(
                    indexedlog_local,
                    lfs_local.as_ref().map(|v| v.as_ref()),
                    StoreType::Local,
                );
            }

            if let Some(ref lfs_cache) = lfs_cache {
                state.fetch_lfs(lfs_cache, StoreType::Shared);
            }

            if let Some(ref lfs_local) = lfs_local {
                state.fetch_lfs(lfs_local, StoreType::Local);
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if use_memcache(creation_time) {
                    if let Some(ref memcache) = memcache {
                        state.fetch_memcache(
                            memcache,
                            indexedlog_cache.as_ref().map(|s| s.as_ref()),
                        );
                    }
                }
            }

            if prefer_computing_aux_data {
                state.derive_computable(
                    aux_cache.as_ref().map(|s| s.as_ref()),
                    aux_local.as_ref().map(|s| s.as_ref()),
                );
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if let Some(ref edenapi) = edenapi {
                    state.fetch_edenapi(
                        edenapi,
                        indexedlog_cache.clone(),
                        lfs_cache.clone(),
                        aux_cache.clone(),
                        if cache_to_memcache && use_memcache(creation_time) {
                            memcache.clone()
                        } else {
                            None
                        },
                    );
                }
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if let Some(ref lfs_remote) = lfs_remote {
                    state.fetch_lfs_remote(
                        &lfs_remote.remote,
                        lfs_local.clone(),
                        lfs_cache.clone(),
                    );
                }
            }

            if let FetchMode::AllowRemote = fetch_mode {
                if let Some(ref contentstore) = contentstore {
                    state.fetch_contentstore(contentstore);
                }
            }

            state.derive_computable(
                aux_cache.as_ref().map(|s| s.as_ref()),
                aux_local.as_ref().map(|s| s.as_ref()),
            );

            metrics.write().fetch += state.metrics().clone();
            state.finish();

            if let Some(activity_logger) = activity_logger {
                if let Err(err) =
                    activity_logger
                        .lock()
                        .log_file_fetch(all_keys, attrs, start_instant.elapsed())
                {
                    tracing::error!("Error writing activity log: {}", err);
                }
            }
        };

        // Only kick off a thread if there's a substantial amount of work.
        if keys_len > 1000 {
            std::thread::spawn(process_func);
        } else {
            process_func();
        }

        FetchResults::new(Box::new(found_rx.into_iter()))
    }

    fn write_lfsptr(&self, key: Key, bytes: Bytes) -> Result<()> {
        if !self.allow_write_lfs_ptrs {
            ensure!(
                std::env::var("TESTTMP").is_ok(),
                "writing LFS pointers directly is not allowed outside of tests"
            );
        }
        let lfs_local = self.lfs_local.as_ref().ok_or_else(|| {
            anyhow!("trying to write LFS pointer but no local LfsStore is available")
        })?;

        let lfs_pointer = LfsPointersEntry::from_bytes(bytes, key.hgid)?;
        lfs_local.add_pointer(lfs_pointer)
    }

    fn write_lfs(&self, key: Key, bytes: Bytes) -> Result<()> {
        let lfs_local = self.lfs_local.as_ref().ok_or_else(|| {
            anyhow!("trying to write LFS file but no local LfsStore is available")
        })?;
        let (lfs_pointer, lfs_blob) = lfs_from_hg_file_blob(key.hgid, &bytes)?;
        let sha256 = lfs_pointer.sha256();

        lfs_local.add_blob(&sha256, lfs_blob)?;
        lfs_local.add_pointer(lfs_pointer)?;

        Ok(())
    }

    fn write_nonlfs(&self, key: Key, bytes: Bytes, meta: Metadata) -> Result<()> {
        let indexedlog_local = self.indexedlog_local.as_ref().ok_or_else(|| {
            anyhow!("trying to write non-LFS file but no local non-LFS IndexedLog is available")
        })?;
        indexedlog_local.put_entry(Entry::new(key, bytes, meta))?;

        Ok(())
    }

    pub fn write_batch(&self, entries: impl Iterator<Item = (Key, Bytes, Metadata)>) -> Result<()> {
        // TODO(meyer): Don't fail the whole batch for a single write error.
        let mut metrics = FileStoreWriteMetrics::default();
        for (key, bytes, meta) in entries {
            if meta.is_lfs() && self.lfs_threshold_bytes.is_some() {
                metrics.lfsptr.item(1);
                if let Err(e) = self.write_lfsptr(key, bytes) {
                    metrics.lfsptr.err(1);
                    return Err(e);
                }
                metrics.lfsptr.ok(1);
                continue;
            }
            let hg_blob_len = bytes.len() as u64;
            // Default to non-LFS if no LFS threshold is set
            if self
                .lfs_threshold_bytes
                .map_or(false, |threshold| hg_blob_len > threshold)
            {
                metrics.lfs.item(1);
                if let Err(e) = self.write_lfs(key, bytes) {
                    metrics.lfs.err(1);
                    return Err(e);
                }
                metrics.lfs.ok(1);
            } else {
                metrics.nonlfs.item(1);
                if let Err(e) = self.write_nonlfs(key, bytes, meta) {
                    metrics.nonlfs.err(1);
                    return Err(e);
                }
                metrics.nonlfs.ok(1);
            }
        }
        self.metrics.write().write += metrics;
        Ok(())
    }

    #[allow(unused_must_use)]
    pub fn flush(&self) -> Result<()> {
        let mut result = Ok(());
        let mut handle_error = |error| {
            tracing::error!(%error);
            result = Err(error);
        };

        if let Some(ref indexedlog_local) = self.indexedlog_local {
            indexedlog_local.flush_log().map_err(&mut handle_error);
        }

        if let Some(ref indexedlog_cache) = self.indexedlog_cache {
            indexedlog_cache.flush_log().map_err(&mut handle_error);
        }

        if let Some(ref lfs_local) = self.lfs_local {
            lfs_local.flush().map_err(&mut handle_error);
        }

        if let Some(ref lfs_cache) = self.lfs_cache {
            lfs_cache.flush().map_err(&mut handle_error);
        }

        if let Some(ref aux_local) = self.aux_local {
            aux_local.flush().map_err(&mut handle_error);
        }

        if let Some(ref aux_cache) = self.aux_cache {
            aux_cache.flush().map_err(&mut handle_error);
        }

        result
    }

    pub fn refresh(&self) -> Result<()> {
        self.metrics.write().api.hg_refresh.call(0);
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.refresh()?;
        }
        self.flush()
    }

    pub fn metrics(&self) -> Vec<(String, usize)> {
        self.metrics.read().metrics().collect()
    }

    pub fn empty() -> Self {
        FileStore {
            extstored_policy: ExtStoredPolicy::Ignore,
            lfs_threshold_bytes: None,
            edenapi_retries: 0,
            allow_write_lfs_ptrs: false,
            prefer_computing_aux_data: false,

            indexedlog_local: None,
            lfs_local: None,

            indexedlog_cache: None,
            lfs_cache: None,

            memcache: None,
            cache_to_memcache: true,

            edenapi: None,
            lfs_remote: None,

            contentstore: None,
            fetch_logger: None,
            metrics: FileStoreMetrics::new(),
            activity_logger: None,

            aux_local: None,
            aux_cache: None,

            creation_time: Instant::now(),
            lfs_progress: AggregatingProgressBar::new("fetching", "LFS"),
            flush_on_drop: true,
        }
    }

    pub fn indexedlog_local(&self) -> Option<Arc<IndexedLogHgIdDataStore>> {
        self.indexedlog_local.clone()
    }

    pub fn indexedlog_cache(&self) -> Option<Arc<IndexedLogHgIdDataStore>> {
        self.indexedlog_cache.clone()
    }

    pub fn with_content_store(&self, cs: Arc<ContentStore>) -> Self {
        let mut clone = self.clone();
        clone.contentstore = Some(cs);
        clone
    }
}

fn use_memcache(creation_time: Instant) -> bool {
    // Only use memcache if the process has been around a while. It takes 2s to setup, which
    // hurts responiveness for short commands.
    creation_time.elapsed() > MEMCACHE_DELAY
}

impl LegacyStore for FileStore {
    /// Returns only the local cache / shared stores, in place of the local-only stores, such that writes will go directly to the local cache.
    /// For compatibility with ContentStore::get_shared_mutable
    fn get_shared_mutable(&self) -> Arc<dyn HgIdMutableDeltaStore> {
        // this is infallible in ContentStore so panic if there are no shared/cache stores.
        assert!(
            self.indexedlog_cache.is_some() || self.lfs_cache.is_some(),
            "cannot get shared_mutable, no shared / local cache stores available"
        );
        Arc::new(FileStore {
            extstored_policy: self.extstored_policy.clone(),
            lfs_threshold_bytes: self.lfs_threshold_bytes.clone(),
            edenapi_retries: self.edenapi_retries.clone(),
            allow_write_lfs_ptrs: self.allow_write_lfs_ptrs,
            prefer_computing_aux_data: self.prefer_computing_aux_data,

            indexedlog_local: self.indexedlog_cache.clone(),
            lfs_local: self.lfs_cache.clone(),

            indexedlog_cache: None,
            lfs_cache: None,

            memcache: None,
            cache_to_memcache: false,

            edenapi: None,
            lfs_remote: None,

            contentstore: None,
            fetch_logger: self.fetch_logger.clone(),
            metrics: self.metrics.clone(),
            activity_logger: self.activity_logger.clone(),

            aux_local: None,
            aux_cache: None,

            creation_time: Instant::now(),
            lfs_progress: self.lfs_progress.clone(),

            // Conservatively flushing on drop here, didn't see perf problems and might be needed by Python
            flush_on_drop: true,
        })
    }

    fn get_logged_fetches(&self) -> HashSet<RepoPathBuf> {
        let mut seen = self
            .fetch_logger
            .as_ref()
            .map(|fl| fl.take_seen())
            .unwrap_or_default();
        if let Some(contentstore) = self.contentstore.as_ref() {
            seen.extend(contentstore.get_logged_fetches());
        }
        seen
    }

    fn get_file_content(&self, key: &Key) -> Result<Option<Bytes>> {
        self.metrics.write().api.hg_getfilecontent.call(0);
        self.fetch(
            std::iter::once(key.clone()),
            FileAttributes::CONTENT,
            FetchMode::AllowRemote,
        )
        .single()?
        .map(|entry| entry.content.unwrap().file_content())
        .transpose()
    }

    // If ContentStore is available, these call into ContentStore. Otherwise, implement these
    // methods on top of scmstore (though they should still eventaully be removed).
    fn add_pending(
        &self,
        key: &Key,
        data: Bytes,
        meta: Metadata,
        location: RepackLocation,
    ) -> Result<()> {
        self.metrics.write().api.hg_addpending.call(0);
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.add_pending(key, data, meta, location)
        } else {
            let delta = Delta {
                data,
                base: None,
                key: key.clone(),
            };

            match location {
                RepackLocation::Local => self.add(&delta, &meta),
                RepackLocation::Shared => self.get_shared_mutable().add(&delta, &meta),
            }
        }
    }

    fn commit_pending(&self, location: RepackLocation) -> Result<Option<Vec<PathBuf>>> {
        self.metrics.write().api.hg_commitpending.call(0);
        if let Some(contentstore) = self.contentstore.as_ref() {
            contentstore.commit_pending(location)
        } else {
            self.flush()?;
            Ok(None)
        }
    }
}

impl HgIdDataStore for FileStore {
    // Fetch the raw content of a single TreeManifest blob
    fn get(&self, key: StoreKey) -> Result<StoreResult<Vec<u8>>> {
        self.metrics.write().api.hg_get.call(0);
        Ok(
            match self
                .fetch(
                    std::iter::once(key.clone()).filter_map(|sk| sk.maybe_into_key()),
                    FileAttributes::CONTENT,
                    FetchMode::AllowRemote,
                )
                .single()?
            {
                Some(entry) => StoreResult::Found(entry.content.unwrap().hg_content()?.into_vec()),
                None => StoreResult::NotFound(key),
            },
        )
    }

    fn get_meta(&self, key: StoreKey) -> Result<StoreResult<Metadata>> {
        self.metrics.write().api.hg_getmeta.call(0);
        Ok(
            match self
                .fetch(
                    std::iter::once(key.clone()).filter_map(|sk| sk.maybe_into_key()),
                    FileAttributes::CONTENT,
                    FetchMode::AllowRemote,
                )
                .single()?
            {
                Some(entry) => StoreResult::Found(entry.content.unwrap().metadata()?),
                None => StoreResult::NotFound(key),
            },
        )
    }

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

impl RemoteDataStore for FileStore {
    fn prefetch(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.metrics.write().api.hg_prefetch.call(keys.len());
        let missing = self
            .fetch(
                keys.iter().cloned().filter_map(|sk| sk.maybe_into_key()),
                FileAttributes::CONTENT,
                FetchMode::AllowRemote,
            )
            .missing()?
            .into_iter()
            .map(StoreKey::HgId)
            .collect();
        Ok(missing)
    }

    fn upload(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.metrics.write().api.hg_upload.call(keys.len());
        // TODO(meyer): Eliminate usage of legacy API, or at least minimize it (do we really need memcache + multiplex, etc)
        if let Some(ref lfs_remote) = self.lfs_remote {
            let mut multiplex = MultiplexDeltaStore::new();
            multiplex.add_store(self.get_shared_mutable());
            if use_memcache(self.creation_time) {
                if let Some(ref memcache) = self.memcache {
                    multiplex.add_store(memcache.clone());
                }
            }
            lfs_remote
                .clone()
                .datastore(Arc::new(multiplex))
                .upload(keys)
        } else {
            Ok(keys.to_vec())
        }
    }
}

impl LocalStore for FileStore {
    fn get_missing(&self, keys: &[StoreKey]) -> Result<Vec<StoreKey>> {
        self.metrics.write().api.hg_getmissing.call(keys.len());
        Ok(self
            .fetch(
                keys.iter().cloned().filter_map(|sk| sk.maybe_into_key()),
                FileAttributes::CONTENT,
                FetchMode::LocalOnly,
            )
            .missing()?
            .into_iter()
            .map(StoreKey::HgId)
            .collect())
    }
}

impl HgIdMutableDeltaStore for FileStore {
    fn add(&self, delta: &Delta, metadata: &Metadata) -> Result<()> {
        self.metrics.write().api.hg_add.call(0);
        if let Delta {
            data,
            base: None,
            key,
        } = delta.clone()
        {
            self.write_batch(std::iter::once((key, data, metadata.clone())))
        } else {
            bail!("Deltas with non-None base are not supported")
        }
    }

    fn flush(&self) -> Result<Option<Vec<PathBuf>>> {
        self.metrics.write().api.hg_flush.call(0);
        self.flush()?;
        Ok(None)
    }
}

// TODO(meyer): Content addressing not supported at all for trees. I could look for HgIds present here and fetch with
// that if available, but I feel like there's probably something wrong if this is called for trees.
impl ContentDataStore for FileStore {
    fn blob(&self, key: StoreKey) -> Result<StoreResult<Bytes>> {
        self.metrics.write().api.contentdatastore_blob.call(0);
        Ok(
            match self
                .fetch(
                    std::iter::once(key.clone()).filter_map(|sk| sk.maybe_into_key()),
                    FileAttributes::CONTENT,
                    FetchMode::LocalOnly,
                )
                .single()?
            {
                Some(entry) => StoreResult::Found(entry.content.unwrap().file_content()?),
                None => StoreResult::NotFound(key),
            },
        )
    }

    fn metadata(&self, key: StoreKey) -> Result<StoreResult<ContentMetadata>> {
        self.metrics.write().api.contentdatastore_metadata.call(0);
        Ok(
            match self
                .fetch(
                    std::iter::once(key.clone()).filter_map(|sk| sk.maybe_into_key()),
                    FileAttributes::CONTENT,
                    FetchMode::LocalOnly,
                )
                .single()?
            {
                Some(StoreFile {
                    content: Some(LazyFile::Lfs(_blob, pointer)),
                    ..
                }) => StoreResult::Found(pointer.into()),
                Some(_) => StoreResult::NotFound(key),
                None => StoreResult::NotFound(key),
            },
        )
    }
}


/*
 * 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::ops::BitAnd;
use std::ops::BitOr;
use std::ops::Not;
use std::ops::Sub;

use edenapi_types::FileAttributes as EdenApiFileAttributes;
use serde::Deserialize;
use serde::Serialize;

use crate::scmstore::attrs::StoreAttrs;

#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct FileAttributes {
    pub content: bool,
    pub aux_data: bool,
}

impl From<FileAttributes> for EdenApiFileAttributes {
    fn from(v: FileAttributes) -> Self {
        EdenApiFileAttributes {
            content: v.content,
            aux_data: v.aux_data,
        }
    }
}

impl StoreAttrs for FileAttributes {
    const NONE: Self = FileAttributes {
        content: false,
        aux_data: false,
    };

    /// Returns all the attributes which are present or can be computed from present attributes.
    fn with_computable(&self) -> FileAttributes {
        if self.content {
            *self | FileAttributes::AUX
        } else {
            *self
        }
    }
}

impl FileAttributes {
    pub const CONTENT: Self = FileAttributes {
        content: true,
        aux_data: false,
    };

    pub const AUX: Self = FileAttributes {
        content: false,
        aux_data: true,
    };
}

impl Not for FileAttributes {
    type Output = Self;

    fn not(self) -> Self::Output {
        FileAttributes {
            content: !self.content,
            aux_data: !self.aux_data,
        }
    }
}

impl BitAnd for FileAttributes {
    type Output = Self;

    fn bitand(self, rhs: Self) -> Self::Output {
        FileAttributes {
            content: self.content & rhs.content,
            aux_data: self.aux_data & rhs.aux_data,
        }
    }
}

impl BitOr for FileAttributes {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        FileAttributes {
            content: self.content | rhs.content,
            aux_data: self.aux_data | rhs.aux_data,
        }
    }
}

/// The subtraction operator is implemented here to mean "set difference" aka relative complement.
impl Sub for FileAttributes {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        self & !rhs
    }
}


/*
 * 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 edenapi_types::Blake3;
use edenapi_types::ContentId;
use edenapi_types::FileAuxData as EdenApiFileAuxData;
use edenapi_types::Sha1;
use serde::Deserialize;
use serde::Serialize;
use types::Sha256;

use crate::indexedlogauxstore::Entry as AuxDataEntry;

#[derive(Debug, Copy, Clone, Serialize, Deserialize, PartialEq, Eq, Default)]
pub struct FileAuxData {
    pub total_size: u64,
    pub content_id: ContentId,
    pub content_sha1: Sha1,
    pub content_sha256: Sha256,
    pub content_seeded_blake3: Option<Blake3>,
}

impl From<AuxDataEntry> for FileAuxData {
    fn from(v: AuxDataEntry) -> Self {
        FileAuxData {
            total_size: v.total_size() as u64,
            content_id: v.content_id(),
            content_sha1: v.content_sha1(),
            content_sha256: Sha256::from_byte_array(v.content_sha256().into()),
            content_seeded_blake3: v.content_seeded_blake3(),
        }
    }
}

impl From<FileAuxData> for AuxDataEntry {
    fn from(v: FileAuxData) -> Self {
        AuxDataEntry {
            total_size: v.total_size,
            content_id: v.content_id,
            content_sha1: v.content_sha1,
            content_sha256: v.content_sha256.into_inner().into(),
            content_seeded_blake3: v.content_seeded_blake3,
        }
    }
}

impl From<EdenApiFileAuxData> for FileAuxData {
    fn from(v: EdenApiFileAuxData) -> Self {
        FileAuxData {
            total_size: v.total_size,
            content_id: v.content_id,
            content_sha1: v.sha1,
            content_sha256: Sha256::from_byte_array(v.sha256.into()),
            content_seeded_blake3: v.seeded_blake3,
        }
    }
}

impl From<FileAuxData> for EdenApiFileAuxData {
    fn from(v: FileAuxData) -> Self {
        EdenApiFileAuxData {
            total_size: v.total_size,
            content_id: v.content_id,
            sha1: v.content_sha1,
            sha256: v.content_sha256.into_inner().into(),
            seeded_blake3: v.content_seeded_blake3,
        }
    }
}


/*
 * 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 anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use edenapi_types::FileEntry;
use hgstore::strip_metadata;
use minibytes::Bytes;
use types::HgId;
use types::Key;

use crate::indexedlogdatastore::Entry;
use crate::lfs::rebuild_metadata;
use crate::lfs::LfsPointersEntry;
use crate::memcache::McData;
use crate::scmstore::file::FileAuxData;
use crate::ContentHash;
use crate::Metadata;

/// A minimal file enum that simply wraps the possible underlying file types,
/// with no processing (so Entry might have the wrong Key.path, etc.)
#[derive(Debug)]
pub(crate) enum LazyFile {
    /// A response from calling into the legacy storage API
    ContentStore(Bytes, Metadata),

    /// An entry from a local IndexedLog. The contained Key's path might not match the requested Key's path.
    IndexedLog(Entry),

    /// A local LfsStore entry.
    Lfs(Bytes, LfsPointersEntry),

    /// An EdenApi FileEntry.
    EdenApi(FileEntry),

    /// A memcache entry, convertable to Entry. In this case the Key's path should match the requested Key's path.
    Memcache(McData),
}

impl LazyFile {
    #[allow(dead_code)]
    fn hgid(&self) -> Option<HgId> {
        use LazyFile::*;
        match self {
            ContentStore(_, _) => None,
            IndexedLog(ref entry) => Some(entry.key().hgid),
            Lfs(_, ref ptr) => Some(ptr.hgid()),
            EdenApi(ref entry) => Some(entry.key().hgid),
            Memcache(ref entry) => Some(entry.key.hgid),
        }
    }

    /// Compute's the aux data associated with this file from the content.
    pub(crate) fn aux_data(&mut self) -> Result<FileAuxData> {
        // TODO(meyer): Implement the rest of the aux data fields
        let aux_data = match self {
            LazyFile::Lfs(content, ref ptr) => FileAuxData {
                total_size: content.len() as u64,
                content_id: ContentHash::content_id(&content),
                content_sha1: ContentHash::sha1(&content),
                content_sha256: ptr.sha256(),
                content_seeded_blake3: Some(ContentHash::seeded_blake3(content)),
            },
            LazyFile::EdenApi(entry) if entry.aux_data.is_some() => entry
                .aux_data()
                .cloned()
                .ok_or_else(|| {
                    anyhow::anyhow!("Invalid EdenAPI entry in LazyFile. Aux data is empty")
                })?
                .into(),
            _ => {
                let content = self.file_content()?;
                FileAuxData {
                    total_size: content.len() as u64,
                    content_id: ContentHash::content_id(&content),
                    content_sha1: ContentHash::sha1(&content),
                    content_sha256: ContentHash::sha256(&content).unwrap_sha256(),
                    content_seeded_blake3: Some(ContentHash::seeded_blake3(&content)),
                }
            }
        };
        Ok(aux_data)
    }

    /// The file content, as would be found in the working copy (stripped of copy header)
    pub(crate) fn file_content(&mut self) -> Result<Bytes> {
        use LazyFile::*;
        Ok(match self {
            IndexedLog(ref mut entry) => strip_metadata(&entry.content()?)?.0,
            Lfs(ref blob, _) => blob.clone(),
            ContentStore(ref blob, _) => strip_metadata(blob)?.0,
            // TODO(meyer): Convert EdenApi to use minibytes
            EdenApi(ref entry) => strip_metadata(&entry.data()?.into())?.0,
            Memcache(ref entry) => strip_metadata(&entry.data)?.0,
        })
    }

    /// The file content, as would be found in the working copy, and also with copy info
    pub(crate) fn file_content_with_copy_info(&mut self) -> Result<(Bytes, Option<Key>)> {
        use LazyFile::*;
        Ok(match self {
            IndexedLog(ref mut entry) => strip_metadata(&entry.content()?)?,
            Lfs(ref blob, ref ptr) => (blob.clone(), ptr.copy_from().clone()),
            ContentStore(ref blob, _) => strip_metadata(blob)?,
            EdenApi(ref entry) => strip_metadata(&entry.data()?.into())?,
            Memcache(ref entry) => strip_metadata(&entry.data)?,
        })
    }

    /// The file content, as would be encoded in the Mercurial blob (with copy header)
    pub(crate) fn hg_content(&mut self) -> Result<Bytes> {
        use LazyFile::*;
        Ok(match self {
            IndexedLog(ref mut entry) => entry.content()?,
            Lfs(ref blob, ref ptr) => rebuild_metadata(blob.clone(), ptr),
            ContentStore(ref blob, _) => blob.clone(),
            EdenApi(ref entry) => entry.data()?.into(),
            Memcache(ref entry) => entry.data.clone(),
        })
    }

    pub(crate) fn metadata(&self) -> Result<Metadata> {
        use LazyFile::*;
        Ok(match self {
            IndexedLog(ref entry) => entry.metadata().clone(),
            Lfs(_, ref ptr) => Metadata {
                size: Some(ptr.size()),
                flags: None,
            },
            ContentStore(_, ref meta) => meta.clone(),
            EdenApi(ref entry) => entry.metadata()?.clone(),
            Memcache(ref entry) => entry.metadata.clone(),
        })
    }

    /// Convert the LazyFile to an indexedlog Entry, if it should ever be written to IndexedLog cache
    pub(crate) fn indexedlog_cache_entry(&self, key: Key) -> Result<Option<Entry>> {
        use LazyFile::*;
        Ok(match self {
            IndexedLog(ref entry) => Some(entry.clone().with_key(key)),
            EdenApi(ref entry) => Some(Entry::new(
                key,
                entry.data()?.into(),
                entry.metadata()?.clone(),
            )),
            // TODO(meyer): We shouldn't ever need to replace the key with Memcache, can probably just clone this.
            Memcache(ref entry) => Some({
                let entry: Entry = entry.clone().into();
                entry.with_key(key)
            }),
            // LFS Files should be written to LfsCache instead
            Lfs(_, _) => None,
            // ContentStore handles caching internally
            ContentStore(_, _) => None,
        })
    }
}

impl TryFrom<McData> for LfsPointersEntry {
    type Error = Error;

    fn try_from(e: McData) -> Result<Self, Self::Error> {
        if e.metadata.is_lfs() {
            Ok(LfsPointersEntry::from_bytes(e.data, e.key.hgid)?)
        } else {
            bail!("failed to convert McData entry to LFS pointer, is_lfs is false")
        }
    }
}

impl TryFrom<Entry> for LfsPointersEntry {
    type Error = Error;

    fn try_from(mut e: Entry) -> Result<Self, Self::Error> {
        if e.metadata().is_lfs() {
            Ok(LfsPointersEntry::from_bytes(e.content()?, e.key().hgid)?)
        } else {
            bail!("failed to convert entry to LFS pointer, is_lfs is false")
        }
    }
}

impl TryFrom<FileEntry> for LfsPointersEntry {
    type Error = Error;

    fn try_from(e: FileEntry) -> Result<Self, Self::Error> {
        if e.metadata()?.is_lfs() {
            Ok(LfsPointersEntry::from_bytes(e.data()?, e.key().hgid)?)
        } else {
            bail!("failed to convert EdenApi FileEntry to LFS pointer, is_lfs is false")
        }
    }
}


/*
 * 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.
 */

mod attrs;
mod auxdata;
mod lazy_file;
mod store_file;

pub use self::attrs::FileAttributes;
pub use self::auxdata::FileAuxData;
pub(crate) use self::lazy_file::LazyFile;
pub use self::store_file::StoreFile;


/*
 * 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::ops::BitOr;

use anyhow::anyhow;
use anyhow::Result;
use minibytes::Bytes;
use types::Key;

use crate::scmstore::file::LazyFile;
use crate::scmstore::value::StoreValue;
use crate::scmstore::FileAttributes;
use crate::scmstore::FileAuxData;

#[derive(Debug, Default)]
pub struct StoreFile {
    // TODO(meyer): We'll probably eventually need a better "canonical lazy file" abstraction, since EdenApi FileEntry won't always carry content
    pub(crate) content: Option<LazyFile>,
    pub(crate) aux_data: Option<FileAuxData>,
}

impl StoreValue for StoreFile {
    type Attrs = FileAttributes;

    /// Returns which attributes are present in this StoreFile
    fn attrs(&self) -> FileAttributes {
        FileAttributes {
            content: self.content.is_some(),
            aux_data: self.aux_data.is_some(),
        }
    }

    /// Return a StoreFile with only the specified subset of attributes
    fn mask(self, attrs: FileAttributes) -> Self {
        StoreFile {
            content: if attrs.content { self.content } else { None },
            aux_data: if attrs.aux_data { self.aux_data } else { None },
        }
    }
}

impl StoreFile {
    pub fn aux_data(&self) -> Result<FileAuxData> {
        self.aux_data
            .ok_or_else(|| anyhow!("no aux data available"))
    }

    pub(crate) fn compute_aux_data(&mut self) -> Result<()> {
        self.aux_data = Some(
            self.content
                .as_mut()
                .ok_or_else(|| anyhow!("failed to compute aux data, no content available"))?
                .aux_data()?,
        );
        Ok(())
    }

    pub fn file_content(&mut self) -> Result<Bytes> {
        self.content
            .as_mut()
            .ok_or_else(|| anyhow!("no content available"))?
            .file_content()
    }

    pub fn file_content_with_copy_info(&mut self) -> Result<(Bytes, Option<Key>)> {
        self.content
            .as_mut()
            .ok_or_else(|| anyhow!("no content available"))?
            .file_content_with_copy_info()
    }
}

impl BitOr for StoreFile {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        StoreFile {
            content: self.content.or(rhs.content),
            aux_data: self.aux_data.or(rhs.aux_data),
        }
    }
}

impl From<FileAuxData> for StoreFile {
    fn from(v: FileAuxData) -> Self {
        StoreFile {
            content: None,
            aux_data: Some(v),
        }
    }
}

impl From<LazyFile> for StoreFile {
    fn from(v: LazyFile) -> Self {
        StoreFile {
            content: Some(v),
            aux_data: None,
        }
    }
}