/*
 * 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::ops::Add;
use std::sync::Arc;
use std::sync::Mutex;
use std::time::Duration;
use std::time::Instant;

use anyhow::Error;
use context::CoreContext;
use derive_more::Add;
use derive_more::Div;
use derive_more::Mul;
use derive_more::Sub;
use fbinit::FacebookInit;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::info;
use slog::Logger;
use stats::prelude::*;

use crate::detail::graph::Node;
use crate::detail::graph::NodeType;
use crate::detail::log;
use crate::detail::state::StepStats;

define_stats! {
    prefix = "mononoke.walker";
    walk_progress_walked: dynamic_timeseries("{}.progress.{}.walked", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_progress_queued: dynamic_timeseries("{}.progress.{}.queued", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_progress_errors: dynamic_timeseries("{}.progress.{}.errors", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_progress_missing: dynamic_timeseries("{}.progress.{}.missing", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_progress_hash_validation_failure: dynamic_timeseries("{}.progress.{}.hash_validation_failure", (subcommand: &'static str, repo: String); Rate, Sum),
}

pub trait ProgressRecorderUnprotected<SS> {
    fn record_step(&mut self, n: &Node, ss: Option<&SS>);

    fn set_sample_builder(&mut self, s: MononokeScubaSampleBuilder);
}

pub trait ProgressReporterUnprotected {
    /// Report progress unconditional. e.g. called at end of run to report final progress
    fn report_progress(&mut self);

    /// Apply your own throttling criteria, once called might chose to do nothing,
    /// in which case return None, otherwise return how long since last report.
    fn report_throttled(&mut self);
}

#[derive(Clone, Copy)]
pub struct ProgressOptions {
    pub sample_rate: u64,
    pub interval: Duration,
}

pub struct ProgressStateByTypeParams {
    pub fb: FacebookInit,
    pub logger: Logger,
    pub subcommand_stats_key: &'static str,
    pub repo_stats_key: String,
    pub types_sorted_by_name: Vec<NodeType>,
    options: ProgressOptions,
}

pub struct ProgressStateWorkByType<SS>
where
    SS: Add<SS, Output = SS> + Default,
{
    pub stats_by_type: HashMap<NodeType, (u64, SS)>,
    total_progress: u64,
}

impl<SS> ProgressStateWorkByType<SS>
where
    SS: Add<SS, Output = SS> + Copy + Default,
{
    fn record_step(&mut self, n: &Node, opt: Option<&SS>) {
        // Global stats
        self.total_progress += 1;
        // By type
        let k = n.get_type();
        let entry = self.stats_by_type.entry(k).or_insert((0, SS::default()));
        entry.0 += 1;
        match opt {
            Some(ss) => {
                entry.1 = entry.1 + *ss;
            }
            None => {}
        }
    }
}

#[derive(Add, Sub, Mul, Div, Clone, Copy, Default, Debug)]
pub struct ProgressSummary {
    walked: u64,
    checked: u64,
    queued: u64,
    errors: u64,
    missing: u64,
    hash_validation_failure: u64,
}

// Takes a summary type as a parameter. e.g. ProgressSummary
pub struct ProgressStateReporting<T> {
    pub start_time: Instant,
    pub last_summary_by_type: HashMap<NodeType, T>,
    pub last_summary: T,
    pub last_update: Instant,
}

// Can retain between runs to have cumulative progress reported
pub struct ProgressStateCountByType<SS, T>
where
    SS: Add<SS, Output = SS> + Default,
{
    pub params: ProgressStateByTypeParams,
    pub work_stats: ProgressStateWorkByType<SS>,
    pub reporting_stats: ProgressStateReporting<T>,
}

pub fn sort_by_string<C, T>(c: C) -> Vec<T>
where
    C: IntoIterator<Item = T>,
    T: ToString,
{
    let mut v: Vec<_> = c.into_iter().collect();
    v.sort_by_key(|k| k.to_string());
    v
}

impl<SS, T> ProgressStateCountByType<SS, T>
where
    SS: Add<SS, Output = SS> + Default,
    T: Default,
{
    pub fn new(
        fb: FacebookInit,
        logger: Logger,
        subcommand_stats_key: &'static str,
        repo_stats_key: String,
        included_types: HashSet<NodeType>,
        options: ProgressOptions,
    ) -> Self {
        let types_by_name = sort_by_string(included_types);

        let now = Instant::now();
        Self {
            params: ProgressStateByTypeParams {
                fb,
                logger,
                subcommand_stats_key,
                repo_stats_key,
                types_sorted_by_name: types_by_name,
                options,
            },
            // Updated by record_step
            work_stats: ProgressStateWorkByType::<SS> {
                stats_by_type: HashMap::new(),
                total_progress: 0,
            },
            // Updated by report_*
            reporting_stats: ProgressStateReporting::<T> {
                start_time: now,
                last_summary_by_type: HashMap::new(),
                last_summary: T::default(),
                last_update: now,
            },
        }
    }

    // Throttle by sample, then time
    pub fn should_log_throttled(&mut self) -> Option<Duration> {
        if self.work_stats.total_progress % self.params.options.sample_rate == 0 {
            let new_update = Instant::now();
            let delta_time = new_update.duration_since(self.reporting_stats.last_update);
            if delta_time >= self.params.options.interval {
                self.reporting_stats.last_update = new_update;
                return Some(delta_time);
            }
        }
        None
    }
}

impl ProgressStateCountByType<StepStats, ProgressSummary> {
    pub fn report_progress_log(&mut self, mut delta_time: Option<Duration>) {
        let summary_by_type: HashMap<NodeType, ProgressSummary> = self
            .work_stats
            .stats_by_type
            .iter()
            .map(|(k, (ps, ss))| {
                let s = ProgressSummary {
                    walked: *ps,
                    checked: ss.visited_of_type as u64,
                    // num_expanded_new is per type children which when summed == a top level queued stat
                    queued: ss.num_expanded_new as u64,
                    errors: ss.error_count as u64,
                    missing: ss.missing_count as u64,
                    hash_validation_failure: ss.hash_validation_failure_count as u64,
                };
                (*k, s)
            })
            .collect();

        let new_summary = summary_by_type
            .values()
            .fold(ProgressSummary::default(), |acc, v| acc + *v);
        let delta_summary = new_summary - self.reporting_stats.last_summary;

        let detail = &self
            .params
            .types_sorted_by_name
            .iter()
            .map(|t| {
                let s = summary_by_type.get(t).cloned().unwrap_or_default();
                format!("{}:{},{},{}", t, s.walked, s.checked, s.queued)
            })
            .collect::<Vec<_>>()
            .join(" ");

        if delta_time.is_none() {
            // Is the last log of a run or chunk, need to know the time
            let now = Instant::now();
            let t = now.duration_since(self.reporting_stats.last_update);
            delta_time = if t.as_millis() > 0 { Some(t) } else { None };
            self.reporting_stats.last_update = now;
        }

        let (delta_s, delta_summary_per_s) =
            delta_time.map_or((0, ProgressSummary::default()), |delta_time| {
                (
                    delta_time.as_secs(),
                    delta_summary * 1000 / (delta_time.as_millis() as u64),
                )
            });

        let total_time = self
            .reporting_stats
            .last_update
            .duration_since(self.reporting_stats.start_time);

        let total_summary_per_s = if total_time.as_millis() > 0 {
            new_summary * 1000 / (total_time.as_millis() as u64)
        } else {
            ProgressSummary::default()
        };

        info!(
            self.params.logger,
            #log::GRAPH,
            "Walked/s,Children/s,Walked,Errors,Missing,Children,Time; Delta {:06}/s,{:06}/s,{},{},{},{},{}s; Run {:06}/s,{:06}/s,{},{},{},{},{}s; Type:Walked,Checks,Children {}",
            delta_summary_per_s.walked,
            delta_summary_per_s.queued,
            delta_summary.walked,
            delta_summary.errors,
            delta_summary.missing,
            delta_summary.queued,
            delta_s,
            total_summary_per_s.walked,
            total_summary_per_s.queued,
            self.work_stats.total_progress,
            new_summary.errors,
            new_summary.missing,
            new_summary.queued,
            total_time.as_secs(),
            detail,
        );

        STATS::walk_progress_walked.add_value(
            delta_summary.walked as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );
        STATS::walk_progress_queued.add_value(
            delta_summary.queued as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );
        STATS::walk_progress_errors.add_value(
            delta_summary.errors as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );

        STATS::walk_progress_missing.add_value(
            delta_summary.missing as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );

        STATS::walk_progress_hash_validation_failure.add_value(
            delta_summary.hash_validation_failure as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );

        self.reporting_stats.last_summary_by_type = summary_by_type;
        self.reporting_stats.last_summary = new_summary;
    }
}

impl<SS, T> ProgressRecorderUnprotected<SS> for ProgressStateCountByType<SS, T>
where
    SS: Add<SS, Output = SS> + Copy + Default,
{
    fn record_step(&mut self, n: &Node, opt: Option<&SS>) {
        self.work_stats.record_step(n, opt);
    }

    fn set_sample_builder(&mut self, _s: MononokeScubaSampleBuilder) {
        // NOOP
    }
}

impl ProgressReporterUnprotected for ProgressStateCountByType<StepStats, ProgressSummary> {
    fn report_progress(&mut self) {
        self.report_progress_log(None);
    }

    fn report_throttled(&mut self) {
        if let Some(delta_time) = self.should_log_throttled() {
            self.report_progress_log(Some(delta_time));
        }
    }
}

pub trait ProgressRecorder<SS> {
    fn record_step(&self, n: &Node, ss: Option<&SS>);
    fn set_sample_builder(&self, s: MononokeScubaSampleBuilder);
}

pub trait ProgressReporter {
    fn report_progress(&self);
    fn report_throttled(&self);
}

#[derive(Debug)]
pub struct ProgressStateMutex<Inner> {
    inner: Arc<Mutex<Inner>>,
}

impl<Inner> ProgressStateMutex<Inner> {
    pub fn new(inner: Inner) -> Self {
        Self {
            inner: Arc::new(Mutex::new(inner)),
        }
    }
}

impl<Inner, SS> ProgressRecorder<SS> for ProgressStateMutex<Inner>
where
    Inner: ProgressRecorderUnprotected<SS>,
{
    fn record_step(&self, n: &Node, ss: Option<&SS>) {
        self.inner.lock().unwrap().record_step(n, ss)
    }

    fn set_sample_builder(&self, s: MononokeScubaSampleBuilder) {
        self.inner.lock().unwrap().set_sample_builder(s)
    }
}

impl<Inner> ProgressReporter for ProgressStateMutex<Inner>
where
    Inner: ProgressReporterUnprotected,
{
    fn report_progress(&self) {
        self.inner.lock().unwrap().report_progress()
    }

    fn report_throttled(&self) {
        self.inner.lock().unwrap().report_throttled()
    }
}

impl<Inner> Clone for ProgressStateMutex<Inner> {
    fn clone(&self) -> Self {
        Self {
            inner: self.inner.clone(),
        }
    }
}

// Log some status update, passing on all data unchanged
pub fn progress_stream<InStream, PS, Payload, SS, K>(
    quiet: bool,
    progress_state: &PS,
    s: InStream,
) -> impl Stream<Item = Result<(K, Payload, Option<SS>), Error>>
where
    InStream: Stream<Item = Result<(K, Payload, Option<SS>), Error>> + 'static + Send,
    PS: 'static + Send + Clone + ProgressRecorder<SS> + ProgressReporter,
    K: 'static,
    // Make sure we can convert from K reference to Node reference
    for<'b> &'b Node: From<&'b K>,
{
    s.map({
        let progress_state = progress_state.clone();
        move |r| {
            r.map(|(key, payload, stats_opt)| {
                {
                    let k: &K = &key;
                    let n: &Node = k.into();
                    progress_state.record_step(n, stats_opt.as_ref());
                    if !quiet {
                        progress_state.report_throttled();
                    }
                }
                (key, payload, stats_opt)
            })
        }
    })
}

// Final status summary, plus count of seen nodes
pub async fn report_state<InStream, ND, SS>(ctx: CoreContext, s: InStream) -> Result<(), Error>
where
    InStream: Stream<Item = Result<(Node, Option<ND>, Option<SS>), Error>> + 'static + Send,
{
    let (seen, loaded) = s
        .try_fold((0_usize, 0_usize), {
            async move |(mut seen, mut loaded), (_n, nd, _ss)| {
                seen += 1;
                if nd.is_some() {
                    loaded += 1;
                }
                Ok((seen, loaded))
            }
        })
        .await?;

    info!(ctx.logger(), #log::LOADED, "Seen,Loaded: {},{}", seen, loaded);
    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::collections::HashSet;
use std::fmt;
use std::hash;
use std::sync::Arc;

use anyhow::Error;
use async_trait::async_trait;
use bonsai_hg_mapping::BonsaiHgMapping;
use bonsai_hg_mapping::BonsaiHgMappingEntry;
use bulkops::Direction;
use context::CoreContext;
use context::SamplingKey;
use dashmap::DashMap;
use mercurial_types::HgChangesetId;
use mononoke_types::datetime::DateTime;
use mononoke_types::ChangesetId;
use phases::Phases;
use regex::Regex;
use slog::Logger;

use crate::detail::graph::EdgeType;
use crate::detail::graph::Node;
use crate::detail::graph::NodeData;
use crate::detail::graph::NodeType;
use crate::detail::graph::WrappedPathHash;
use crate::detail::graph::WrappedPathLike;
use crate::detail::state::InternedType;
use crate::detail::state::StepStats;
use crate::detail::state::WalkState;
use crate::detail::walk::EmptyRoute;
use crate::detail::walk::OutgoingEdge;
use crate::detail::walk::StepRoute;
use crate::detail::walk::TailingWalkVisitor;
use crate::detail::walk::VisitOne;
use crate::detail::walk::WalkVisitor;

pub trait SampleTrigger<K> {
    fn map_keys(&self, key: SamplingKey, walk_key: K);
}

#[derive(Clone, Debug, Default)]
pub struct SamplingOptions {
    pub sample_rate: u64,
    pub sample_offset: u64,
    pub node_types: HashSet<NodeType>,
    pub exclude_types: HashSet<NodeType>,
}

impl SamplingOptions {
    pub fn retain_or_default(&mut self, walk_include: &HashSet<NodeType>) {
        if self.node_types.is_empty() {
            self.node_types = walk_include
                .iter()
                .filter(|e| !self.exclude_types.contains(e))
                .cloned()
                .collect();
        } else {
            self.node_types.retain(|i| walk_include.contains(i));
        }
    }
}

pub struct SamplingWalkVisitor<T> {
    inner: WalkState,
    options: SamplingOptions,
    sample_path_regex: Option<Regex>,
    sampler: Arc<T>,
}

impl<T> SamplingWalkVisitor<T> {
    pub fn new(
        include_node_types: HashSet<NodeType>,
        include_edge_types: HashSet<EdgeType>,
        options: SamplingOptions,
        sample_path_regex: Option<Regex>,
        sampler: Arc<T>,
        enable_derive: bool,
        chunk_direction: Option<Direction>,
    ) -> Self {
        Self {
            inner: WalkState::new(
                include_node_types,
                include_edge_types,
                HashSet::new(),
                enable_derive,
                chunk_direction,
            ),
            options,
            sample_path_regex,
            sampler,
        }
    }
}

#[async_trait]
impl<T: Send + Sync> VisitOne for SamplingWalkVisitor<T> {
    fn in_chunk(&self, bcs_id: &ChangesetId) -> bool {
        self.inner.in_chunk(bcs_id)
    }
    fn needs_visit(&self, outgoing: &OutgoingEdge) -> bool {
        self.inner.needs_visit(outgoing)
    }
    async fn is_public(
        &self,
        ctx: &CoreContext,
        phases_store: &dyn Phases,
        bcs_id: &ChangesetId,
    ) -> Result<bool, Error> {
        self.inner.is_public(ctx, phases_store, bcs_id).await
    }
    fn get_hg_from_bonsai(&self, bcs_id: &ChangesetId) -> Option<HgChangesetId> {
        self.inner.get_hg_from_bonsai(bcs_id)
    }
    fn record_hg_from_bonsai(&self, bcs_id: &ChangesetId, hg_cs_id: HgChangesetId) {
        self.inner.record_hg_from_bonsai(bcs_id, hg_cs_id)
    }
    async fn get_bonsai_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<ChangesetId, Error> {
        self.inner
            .get_bonsai_from_hg(ctx, bonsai_hg_mapping, hg_cs_id)
            .await
    }
    async fn defer_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<Option<ChangesetId>, Error> {
        self.inner
            .defer_from_hg(ctx, bonsai_hg_mapping, hg_cs_id)
            .await
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Hash)]
pub struct PathTrackingRoute
 {
    // The path we reached this by
    pub path: Option
,
    /// When did this route see this path was updated.
    /// Taken from the last bonsai or hg changset stepped through.
    pub mtime: Option<DateTime>,
}

// We don't hold these tracking so as to keep memory usage down in scrub
impl
 StepRoute for PathTrackingRoute

where
    P: WrappedPathLike + fmt::Debug,
{
    fn source_node(&self) -> Option<&Node> {
        None
    }
    fn via_node(&self) -> Option<&Node> {
        None
    }
}

impl
 PathTrackingRoute

where
    P: WrappedPathLike + Eq + Clone,
{
    fn evolve(route: Option<Self>, walk_item: &OutgoingEdge, mtime: Option<&DateTime>) -> Self {
        let existing_path = route.as_ref().and_then(|r| r.path.as_ref());
        let existing_mtime = route.as_ref().and_then(|r| r.mtime.as_ref());
        let new_path = P::evolve_path(existing_path, walk_item);

        // reuse same route if possible
        if new_path == existing_path && (mtime.is_none() || mtime == existing_mtime) {
            if let Some(route) = route {
                return route;
            }
        }

        Self {
            path: new_path.cloned(),
            mtime: mtime.cloned().or_else(|| route.and_then(|r| r.mtime)),
        }
    }
}

// Name the stream output key type
#[derive(Debug, Eq, Hash, PartialEq)]
pub struct WalkKeyOptPath
 {
    pub node: Node,
    pub path: Option
,
}

// Map the key type so progress reporting works
impl<'a, P: WrappedPathLike> From<&'a WalkKeyOptPath
> for &'a Node {
    fn from(from: &'a WalkKeyOptPath
) -> &'a Node {
        &from.node
    }
}

// Name the stream output payload type
#[derive(Default)]
pub struct WalkPayloadMtime {
    pub data: Option<NodeData>,
    pub mtime: Option<DateTime>,
}

impl<T> TailingWalkVisitor for SamplingWalkVisitor<T> {
    fn start_chunk(
        &mut self,
        chunk_members: &HashSet<ChangesetId>,
        mapping_prepop: Vec<BonsaiHgMappingEntry>,
    ) -> Result<HashSet<OutgoingEdge>, Error> {
        self.inner.start_chunk(chunk_members, mapping_prepop)
    }

    fn clear_state(
        &mut self,
        node_types: &HashSet<NodeType>,
        interned_types: &HashSet<InternedType>,
    ) {
        self.inner.clear_state(node_types, interned_types)
    }

    fn end_chunks(&mut self, logger: &Logger, contiguous_bounds: bool) -> Result<(), Error> {
        self.inner.end_chunks(logger, contiguous_bounds)
    }

    fn num_deferred(&self) -> usize {
        self.inner.num_deferred()
    }
}

impl<T, P>
    WalkVisitor<(WalkKeyOptPath
, WalkPayloadMtime, Option<StepStats>), PathTrackingRoute
>
    for SamplingWalkVisitor<T>
where
    T: SampleTrigger<WalkKeyOptPath
> + Send + Sync,
    P: WrappedPathLike + Clone + Eq + fmt::Display,
{
    fn start_step(
        &self,
        mut ctx: CoreContext,
        route: Option<&PathTrackingRoute
>,
        step: &OutgoingEdge,
    ) -> Option<CoreContext> {
        if self.options.node_types.contains(&step.target.get_type()) {
            let repo_path = route.and_then(|r| P::evolve_path(r.path.as_ref(), step));
            if self.sample_path_regex.as_ref().map_or_else(
                || true,
                |re| match repo_path {
                    None => false,
                    Some(repo_path) => re.is_match(&repo_path.to_string()),
                },
            ) {
                let should_sample = match self.options.sample_rate {
                    0 => false,
                    1 => true,
                    sample_rate => {
                        let sampling_fingerprint = repo_path.map_or_else(
                            || step.target.sampling_fingerprint(),
                            |r| Some(r.sampling_fingerprint()),
                        );
                        sampling_fingerprint
                            .map_or(self.options.sample_offset % sample_rate == 0, |fp| {
                                (fp + self.options.sample_offset) % sample_rate == 0
                            })
                    }
                };

                if should_sample {
                    let sampling_key = SamplingKey::new();
                    ctx = ctx.clone_and_sample(sampling_key);
                    self.sampler.map_keys(
                        sampling_key,
                        WalkKeyOptPath {
                            node: step.target.clone(),
                            path: repo_path.cloned(),
                        },
                    );
                }
            }
        }
        self.inner
            .start_step(ctx, route.map(|_| &EmptyRoute {}), step)
    }

    fn visit(
        &self,
        ctx: &CoreContext,
        resolved: OutgoingEdge,
        node_data: Option<NodeData>,
        route: Option<PathTrackingRoute
>,
        outgoing: Vec<OutgoingEdge>,
    ) -> (
        (WalkKeyOptPath
, WalkPayloadMtime, Option<StepStats>),
        PathTrackingRoute
,
        Vec<OutgoingEdge>,
    ) {
        let inner_route = route.as_ref().map(|_| EmptyRoute {});

        let mtime = match &node_data {
            Some(NodeData::Changeset(bcs)) => {
                bcs.committer_date().or_else(|| Some(bcs.author_date()))
            }
            Some(NodeData::HgChangeset(hg_cs)) => Some(hg_cs.time()),
            _ => None,
        };

        let route = PathTrackingRoute::evolve(route, &resolved, mtime);
        let ((n, nd, stats), _inner_route, outgoing) =
            self.inner
                .visit(ctx, resolved, node_data, inner_route, outgoing);

        (
            (
                WalkKeyOptPath {
                    node: n,
                    path: route.path.clone(),
                },
                WalkPayloadMtime {
                    data: nd,
                    mtime: route.mtime.clone(),
                },
                stats,
            ),
            route,
            outgoing,
        )
    }

    fn defer_visit(
        &self,
        bcs_id: &ChangesetId,
        walk_item: &OutgoingEdge,
        route: Option<PathTrackingRoute
>,
    ) -> Result<
        (
            (WalkKeyOptPath
, WalkPayloadMtime, Option<StepStats>),
            PathTrackingRoute
,
        ),
        Error,
    > {
        let inner_route = route.as_ref().map(|_| EmptyRoute {});
        let route = PathTrackingRoute::evolve(route, walk_item, None);
        let ((n, _nd, stats), _inner_route) =
            self.inner.defer_visit(bcs_id, walk_item, inner_route)?;
        Ok((
            (
                WalkKeyOptPath {
                    node: n,
                    path: route.path.clone(),
                },
                WalkPayloadMtime::default(),
                stats,
            ),
            route,
        ))
    }
}

// Super simple sampling visitor impl for scrubbing
impl<T>
    WalkVisitor<
        (
            WalkKeyOptPath<WrappedPathHash>,
            WalkPayloadMtime,
            Option<StepStats>,
        ),
        EmptyRoute,
    > for SamplingWalkVisitor<T>
where
    T: SampleTrigger<WalkKeyOptPath<WrappedPathHash>> + Send + Sync,
{
    fn start_step(
        &self,
        mut ctx: CoreContext,
        route: Option<&EmptyRoute>,
        step: &OutgoingEdge,
    ) -> Option<CoreContext> {
        if self.options.node_types.contains(&step.target.get_type()) {
            let should_sample = match self.options.sample_rate {
                0 => false,
                1 => true,
                sample_rate => step
                    .target
                    .sampling_fingerprint()
                    .map_or(self.options.sample_offset % sample_rate == 0, |fp| {
                        (fp + self.options.sample_offset) % sample_rate == 0
                    }),
            };

            if should_sample {
                let sampling_key = SamplingKey::new();
                ctx = ctx.clone_and_sample(sampling_key);
                self.sampler.map_keys(
                    sampling_key,
                    WalkKeyOptPath {
                        node: step.target.clone(),
                        path: None,
                    },
                );
            }
        }
        self.inner.start_step(ctx, route, step)
    }

    fn visit(
        &self,
        ctx: &CoreContext,
        resolved: OutgoingEdge,
        node_data: Option<NodeData>,
        route: Option<EmptyRoute>,
        outgoing: Vec<OutgoingEdge>,
    ) -> (
        (
            WalkKeyOptPath<WrappedPathHash>,
            WalkPayloadMtime,
            Option<StepStats>,
        ),
        EmptyRoute,
        Vec<OutgoingEdge>,
    ) {
        let ((n, nd, stats), route, outgoing) =
            self.inner.visit(ctx, resolved, node_data, route, outgoing);
        let output = (
            WalkKeyOptPath {
                node: n,
                path: None,
            },
            WalkPayloadMtime {
                data: nd,
                mtime: None,
            },
            stats,
        );
        (output, route, outgoing)
    }

    fn defer_visit(
        &self,
        bcs_id: &ChangesetId,
        walk_item: &OutgoingEdge,
        route: Option<EmptyRoute>,
    ) -> Result<
        (
            (
                WalkKeyOptPath<WrappedPathHash>,
                WalkPayloadMtime,
                Option<StepStats>,
            ),
            EmptyRoute,
        ),
        Error,
    > {
        let ((n, nd, stats), route) = self.inner.defer_visit(bcs_id, walk_item, route)?;
        let output = (
            WalkKeyOptPath {
                node: n,
                path: None,
            },
            WalkPayloadMtime {
                data: nd,
                mtime: None,
            },
            stats,
        );
        Ok((output, route))
    }
}

// Map from a Sampling Key the sample type T
// And from a graph level step S to the sampling key
#[derive(Debug)]
pub struct WalkSampleMapping<S, T>
where
    S: Eq + fmt::Debug + hash::Hash,
{
    // T can keep a one to many mapping, e.g. some nodes like
    // chunked files have multiple blobstore keys
    inflight: DashMap<SamplingKey, T>,
    // 1:1 relationship, each step has one SamplingKey
    inflight_reverse: DashMap<S, SamplingKey>,
}

impl<T> SampleTrigger<Node> for WalkSampleMapping<Node, T>
where
    T: Default,
{
    fn map_keys(&self, sample_key: SamplingKey, walk_key: Node) {
        self.inflight.insert(sample_key, T::default());
        self.inflight_reverse.insert(walk_key, sample_key);
    }
}

impl<T, P> SampleTrigger<WalkKeyOptPath
> for WalkSampleMapping<Node, T>
where
    T: Default,
    P: WrappedPathLike + Eq + hash::Hash + fmt::Debug,
{
    fn map_keys(&self, sample_key: SamplingKey, walk_key: WalkKeyOptPath
) {
        self.inflight.insert(sample_key, T::default());
        self.inflight_reverse.insert(walk_key.node, sample_key);
    }
}

impl<T, P> SampleTrigger<WalkKeyOptPath
> for WalkSampleMapping<WalkKeyOptPath
, T>
where
    T: Default,
    P: WrappedPathLike + Eq + hash::Hash + fmt::Debug,
{
    fn map_keys(&self, sample_key: SamplingKey, walk_key: WalkKeyOptPath
) {
        self.inflight.insert(sample_key, T::default());
        self.inflight_reverse.insert(walk_key, sample_key);
    }
}

impl<S, T> WalkSampleMapping<S, T>
where
    S: Eq + fmt::Debug + hash::Hash,
{
    pub fn new() -> Self {
        Self {
            inflight: DashMap::new(),
            inflight_reverse: DashMap::new(),
        }
    }

    // Called from the blobstore sampling callback
    pub fn inflight(&self) -> &DashMap<SamplingKey, T> {
        &self.inflight
    }

    // Needs to be called to stop tracking the node and thus free memory.
    // Can be called from the vistor visit, or in the stream processing
    // walk output.
    pub fn complete_step(&self, s: &S) -> Option<T> {
        let reverse_mapping = self.inflight_reverse.remove(s);
        reverse_mapping
            .as_ref()
            .and_then(|(_k, sample_key)| self.inflight.remove(sample_key))
            .map(|(_k, v)| v)
    }

    pub fn is_sampling(&self, s: &S) -> bool {
        self.inflight_reverse.contains_key(s)
    }
}


/*
 * 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::fmt;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::format_err;
use anyhow::Error;
use blobstore::BlobstoreGetData;
use blobstore::SizeMetadata;
use cloned::cloned;
use context::CoreContext;
use derive_more::Add;
use derive_more::Div;
use derive_more::Mul;
use derive_more::Sub;
use fbinit::FacebookInit;
use futures::future;
use futures::future::try_join_all;
use futures::future::FutureExt;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use futures::TryFutureExt;
use metaconfig_types::BlobstoreId;
use mononoke_types::datetime::DateTime;
use repo_identity::RepoIdentityRef;
use samplingblob::ComponentSamplingHandler;
use slog::info;
use stats::prelude::*;

use crate::args::OutputFormat;
use crate::commands::JobParams;
use crate::commands::JobWalkParams;
use crate::commands::RepoSubcommandParams;
use crate::commands::SCRUB;
use crate::detail::graph::FileContentData;
use crate::detail::graph::Node;
use crate::detail::graph::NodeData;
use crate::detail::graph::NodeType;
use crate::detail::graph::WrappedPathHash;
use crate::detail::graph::WrappedPathLike;
use crate::detail::log;
use crate::detail::pack::PackInfo;
use crate::detail::pack::PackInfoLogOptions;
use crate::detail::pack::PackInfoLogger;
use crate::detail::progress::progress_stream;
use crate::detail::progress::report_state;
use crate::detail::progress::ProgressOptions;
use crate::detail::progress::ProgressReporter;
use crate::detail::progress::ProgressReporterUnprotected;
use crate::detail::progress::ProgressStateCountByType;
use crate::detail::progress::ProgressStateMutex;
use crate::detail::sampling::PathTrackingRoute;
use crate::detail::sampling::SamplingOptions;
use crate::detail::sampling::SamplingWalkVisitor;
use crate::detail::sampling::WalkKeyOptPath;
use crate::detail::sampling::WalkPayloadMtime;
use crate::detail::sampling::WalkSampleMapping;
use crate::detail::sizing::SizingSample;
use crate::detail::tail::walk_exact_tail;
use crate::detail::validate::TOTAL;
use crate::detail::walk::EmptyRoute;
use crate::detail::walk::RepoWalkParams;
use crate::detail::walk::RepoWalkTypeParams;

define_stats! {
    prefix = "mononoke.walker";
    walk_progress_keys: dynamic_timeseries("{}.progress.{}.blobstore_keys", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_progress_bytes: dynamic_timeseries("{}.progress.{}.blobstore_bytes", (subcommand: &'static str, repo: String); Rate, Sum),
    walk_last_completed_by_type: dynamic_singleton_counter("{}.last_completed.{}.{}.{}", (subcommand: &'static str, repo: String, node_type: &'static str, desc: &'static str)),
}

#[derive(Add, Div, Mul, Sub, Clone, Copy, Default, Debug)]
pub struct ScrubStats {
    pub blobstore_bytes: u64,
    pub blobstore_keys: u64,
}

impl From<Option<&ScrubSample>> for ScrubStats {
    fn from(sample: Option<&ScrubSample>) -> Self {
        sample
            .map(|sample| ScrubStats {
                blobstore_keys: sample.data.values().len() as u64,
                blobstore_bytes: sample
                    .data
                    .values()
                    // Uncompressed size is always the same between stores, so can take first value
                    .map(|v| v.values().next().map_or(0, |v| v.unique_uncompressed_size))
                    .sum(),
            })
            .unwrap_or_default()
    }
}

impl From<Option<&SizingSample>> for ScrubStats {
    fn from(sample: Option<&SizingSample>) -> Self {
        sample
            .map(|sample| ScrubStats {
                blobstore_keys: sample.data.values().len() as u64,
                blobstore_bytes: sample
                    .data
                    .values()
                    .by_ref()
                    .map(|bytes| bytes.len() as u64)
                    .sum(),
            })
            .unwrap_or_default()
    }
}

impl fmt::Display for ScrubStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "{},{}", self.blobstore_bytes, self.blobstore_keys,)
    }
}

// Force load of leaf data like file contents that graph traversal did not need
fn loading_stream<InStream, SS, L>(
    limit_data_fetch: bool,
    scheduled_max: usize,
    s: InStream,
    sampler: Arc<WalkSampleMapping<Node, ScrubSample>>,
    output_node_types: HashSet<NodeType>,
    output_format: OutputFormat,
    pack_info_logger: Option<L>,
) -> impl Stream<Item = Result<(Node, Option<NodeData>, Option<ScrubStats>), Error>>
where
    InStream: Stream<
            Item = Result<
                (
                    WalkKeyOptPath<WrappedPathHash>,
                    WalkPayloadMtime,
                    Option<SS>,
                ),
                Error,
            >,
        >
        + 'static
        + Send,
    L: PackInfoLogger + 'static + Send,
{
    s.map_ok(move |(walk_key, payload, _progress_stats)| {
        let mtime = payload.mtime;
        match payload.data {
            Some(NodeData::FileContent(FileContentData::ContentStream(file_bytes_stream)))
                if !limit_data_fetch =>
            {
                cloned!(sampler);
                file_bytes_stream
                    .try_fold(0, |acc, file_bytes| future::ok(acc + file_bytes.size()))
                    .map_ok(move |num_bytes| {
                        let sample = sampler.complete_step(&walk_key.node);
                        (
                            walk_key,
                            mtime,
                            Some(NodeData::FileContent(FileContentData::Consumed(num_bytes))),
                            Some(sample),
                        )
                    })
                    .map_err(|e| e.context(format_err!("While scrubbing file content stream")))
                    .left_future()
            }
            data_opt => {
                if output_node_types.contains(&walk_key.node.get_type()) {
                    match output_format {
                        OutputFormat::Debug => {
                            println!("Node {:?}: NodeData: {:?}", walk_key.node, data_opt)
                        }
                        // Keep Node as non-Pretty so its on same line
                        OutputFormat::PrettyDebug => {
                            println!("Node {:?}: NodeData: {:#?}", walk_key.node, data_opt)
                        }
                    }
                }
                let sample = data_opt
                    .as_ref()
                    .map(|_d| sampler.complete_step(&walk_key.node));
                future::ok((walk_key, mtime, data_opt, sample)).right_future()
            }
        }
    })
    .try_buffer_unordered(scheduled_max)
    .map_ok(move |(walk_key, mtime, data_opt, sample)| {
        let size = if let Some(sample) = sample {
            let size = ScrubStats::from(sample.as_ref());
            if let Some(logger) = pack_info_logger.as_ref() {
                record_for_packer(logger, &walk_key, mtime, sample);
            }
            Some(size)
        } else {
            None
        };
        (walk_key.node, data_opt, size)
    })
}

fn record_for_packer<L>(
    logger: &L,
    walk_key: &WalkKeyOptPath<WrappedPathHash>,
    mtime: Option<DateTime>,
    sample: Option<ScrubSample>,
) where
    L: PackInfoLogger,
{
    if let Some(sample) = sample {
        for (blobstore_key, store_to_key_sizes) in sample.data {
            for (blobstore_id, key_sample) in store_to_key_sizes {
                logger.log(PackInfo {
                    blobstore_id,
                    blobstore_key: blobstore_key.as_str(),
                    node_type: walk_key.node.get_type(),
                    node_fingerprint: walk_key.node.sampling_fingerprint(),
                    similarity_key: walk_key.path.map(|p| p.sampling_fingerprint()),
                    mtime: mtime.map(|mtime| mtime.timestamp_secs() as u64),
                    uncompressed_size: key_sample.unique_uncompressed_size,
                    sizes: key_sample.sizes,
                    ctime: key_sample.ctime,
                })
            }
        }
    }
}

// Sample for one blobstore key
#[derive(Debug)]
struct ScrubKeySample {
    // Every key/store can provide uncompressed size
    unique_uncompressed_size: u64,
    // Only keys accessed via a packblob store have SizeMetadata
    sizes: Option<SizeMetadata>,
    ctime: Option<i64>,
}

// Holds a map from blobstore keys to their samples per store
#[derive(Debug)]
pub struct ScrubSample {
    data: HashMap<String, HashMap<Option<BlobstoreId>, ScrubKeySample>>,
}

impl Default for ScrubSample {
    fn default() -> Self {
        Self {
            data: HashMap::with_capacity(1),
        }
    }
}

impl ComponentSamplingHandler for WalkSampleMapping<Node, ScrubSample> {
    fn sample_get(
        &self,
        ctx: &CoreContext,
        key: &str,
        value: Option<&BlobstoreGetData>,
        inner_id: Option<BlobstoreId>,
    ) -> Result<(), Error> {
        ctx.sampling_key().map(|sampling_key| {
            self.inflight().get_mut(sampling_key).map(|mut guard| {
                value.map(|value| {
                    let sample = ScrubKeySample {
                        unique_uncompressed_size: value.as_bytes().len() as u64,
                        sizes: value.as_meta().sizes().cloned(),
                        ctime: value.as_meta().ctime(),
                    };
                    guard
                        .data
                        .entry(key.to_owned())
                        .or_default()
                        .insert(inner_id, sample)
                })
            })
        });
        Ok(())
    }
}

impl ProgressStateCountByType<ScrubStats, ScrubStats> {
    fn report_completion_stat(&self, stat: &ScrubStats, stat_key: &'static str) {
        for (desc, value) in &[
            ("blobstore_bytes", stat.blobstore_bytes),
            ("blobstore_keys", stat.blobstore_keys),
        ] {
            STATS::walk_last_completed_by_type.set_value(
                self.params.fb,
                *value as i64,
                (
                    self.params.subcommand_stats_key,
                    self.params.repo_stats_key.clone(),
                    stat_key,
                    desc,
                ),
            );
        }
    }

    fn report_completion_stats(&self) {
        for (k, v) in self.reporting_stats.last_summary_by_type.iter() {
            self.report_completion_stat(v, k.into())
        }
        self.report_completion_stat(&self.reporting_stats.last_summary, TOTAL)
    }

    pub fn report_progress_log(&mut self, delta_time: Option<Duration>) {
        let summary_by_type: HashMap<NodeType, ScrubStats> = self
            .work_stats
            .stats_by_type
            .iter()
            .map(|(k, (_i, v))| (*k, *v))
            .collect();
        let new_summary = summary_by_type
            .values()
            .fold(ScrubStats::default(), |acc, v| acc + *v);
        let delta_summary = new_summary - self.reporting_stats.last_summary;

        let def = ScrubStats::default();
        let detail = &self
            .params
            .types_sorted_by_name
            .iter()
            .map(|t| {
                let s = summary_by_type.get(t).unwrap_or(&def);
                format!("{}:{}", t, s)
            })
            .collect::<Vec<_>>()
            .join(" ");

        let (delta_s, delta_summary_per_s) =
            delta_time.map_or((0, ScrubStats::default()), |delta_time| {
                (
                    delta_time.as_secs(),
                    delta_summary * 1000 / (delta_time.as_millis() as u64),
                )
            });

        let total_time = self
            .reporting_stats
            .last_update
            .duration_since(self.reporting_stats.start_time);

        let total_summary_per_s = if total_time.as_millis() > 0 {
            new_summary * 1000 / (total_time.as_millis() as u64)
        } else {
            ScrubStats::default()
        };

        info!(
            self.params.logger,
            #log::SIZING,
            "Bytes/s,Keys/s,Bytes,Keys; Delta {:06}/s,{:06}/s,{},{}s; Run {:06}/s,{:06}/s,{},{}s; Type:Raw,Compressed {}",
            delta_summary_per_s.blobstore_bytes,
            delta_summary_per_s.blobstore_keys,
            delta_summary,
            delta_s,
            total_summary_per_s.blobstore_bytes,
            total_summary_per_s.blobstore_keys,
            new_summary,
            total_time.as_secs(),
            detail,
        );

        STATS::walk_progress_bytes.add_value(
            delta_summary.blobstore_bytes as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );
        STATS::walk_progress_keys.add_value(
            delta_summary.blobstore_keys as i64,
            (
                self.params.subcommand_stats_key,
                self.params.repo_stats_key.clone(),
            ),
        );

        self.reporting_stats.last_summary_by_type = summary_by_type;
        self.reporting_stats.last_summary = new_summary;

        if delta_time.is_none() {
            self.report_completion_stats()
        }
    }
}

impl ProgressReporterUnprotected for ProgressStateCountByType<ScrubStats, ScrubStats> {
    fn report_progress(&mut self) {
        self.report_progress_log(None);
    }

    fn report_throttled(&mut self) {
        if let Some(delta_time) = self.should_log_throttled() {
            self.report_progress_log(Some(delta_time));
        }
    }
}

#[derive(Clone)]
pub struct ScrubCommand {
    pub limit_data_fetch: bool,
    pub output_format: OutputFormat,
    pub output_node_types: HashSet<NodeType>,
    pub progress_options: ProgressOptions,
    pub sampling_options: SamplingOptions,
    pub pack_info_log_options: Option<PackInfoLogOptions>,
    pub sampler: Arc<WalkSampleMapping<Node, ScrubSample>>,
}

impl ScrubCommand {
    fn apply_repo(&mut self, repo_params: &RepoWalkParams) {
        self.sampling_options
            .retain_or_default(&repo_params.include_node_types);
    }
}

// Starts from the graph, (as opposed to walking from blobstore enumeration)
pub async fn scrub_objects(
    fb: FacebookInit,
    job_params: JobParams,
    command: ScrubCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let JobParams {
        walk_params,
        per_repo,
    } = job_params;

    let mut all_walks = Vec::new();
    for (sub_params, repo_params) in per_repo {
        cloned!(mut command, walk_params);

        command.apply_repo(&repo_params);

        let walk = run_one(
            fb,
            walk_params,
            sub_params,
            repo_params,
            command,
            Arc::clone(&cancellation_requested),
        );
        all_walks.push(walk);
    }
    try_join_all(all_walks).await.map(|_| ())
}

async fn run_one(
    fb: FacebookInit,
    job_params: JobWalkParams,
    sub_params: RepoSubcommandParams,
    repo_params: RepoWalkParams,
    command: ScrubCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let sizing_progress_state =
        ProgressStateMutex::new(ProgressStateCountByType::<ScrubStats, ScrubStats>::new(
            fb,
            repo_params.logger.clone(),
            SCRUB,
            repo_params.repo.repo_identity().name().to_string(),
            command.sampling_options.node_types.clone(),
            command.progress_options,
        ));

    let make_sink = {
        cloned!(command, job_params.quiet, sub_params.progress_state,);
        move |ctx: &CoreContext, repo_params: &RepoWalkParams| {
            let repo_name = repo_params.repo.repo_identity().name().to_string();
            cloned!(ctx, repo_params.scheduled_max);
            async move |walk_output, run_start, chunk_num, checkpoint_name| {
                let walk_progress = progress_stream(quiet, &progress_state, walk_output);
                let loading = loading_stream(
                    command.limit_data_fetch,
                    scheduled_max,
                    walk_progress,
                    command.sampler,
                    command.output_node_types,
                    command.output_format,
                    command
                        .pack_info_log_options
                        .map(|o| o.make_logger(repo_name, run_start, chunk_num, checkpoint_name)),
                );
                let report_sizing = progress_stream(quiet, &sizing_progress_state, loading);

                report_state(ctx, report_sizing).await?;
                sizing_progress_state.report_progress();
                progress_state.report_progress();
                Ok(())
            }
        }
    };

    let mut stream_node_types = command.output_node_types.clone();
    if !command.limit_data_fetch {
        stream_node_types.insert(NodeType::FileContent);
    }
    if command.pack_info_log_options.is_some() {
        // Need these to be able to see and log the commit time stamps
        stream_node_types.insert(NodeType::Changeset);
        stream_node_types.insert(NodeType::HgChangeset);
    }
    let required_node_data_types: HashSet<NodeType> = stream_node_types.into_iter().collect();

    let walk_state = SamplingWalkVisitor::new(
        repo_params.include_node_types.clone(),
        repo_params.include_edge_types.clone(),
        command.sampling_options,
        None,
        command.sampler,
        job_params.enable_derive,
        sub_params
            .tail_params
            .chunking
            .as_ref()
            .map(|v| v.direction),
    );

    let type_params = RepoWalkTypeParams {
        required_node_data_types,
        always_emit_edge_types: HashSet::new(),
        keep_edge_paths: command.pack_info_log_options.is_some(),
    };

    if command.pack_info_log_options.is_some() {
        walk_exact_tail::<_, _, _, _, _, PathTrackingRoute<WrappedPathHash>>(
            fb,
            job_params,
            repo_params,
            type_params,
            sub_params.tail_params,
            walk_state,
            make_sink,
            cancellation_requested,
        )
        .await
    } else {
        walk_exact_tail::<_, _, _, _, _, EmptyRoute>(
            fb,
            job_params,
            repo_params,
            type_params,
            sub_params.tail_params,
            walk_state,
            make_sink,
            cancellation_requested,
        )
        .await
    }
}


/*
 * 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::cmp::min;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::io::Cursor;
use std::io::Write;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Error;
use async_compression::metered::MeteredWrite;
use async_compression::Compressor;
use async_compression::CompressorType;
use blobstore::BlobstoreGetData;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use derive_more::Add;
use derive_more::Div;
use derive_more::Mul;
use derive_more::Sub;
use fbinit::FacebookInit;
use futures::future;
use futures::future::try_join_all;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::stream::Stream;
use futures::stream::TryStreamExt;
use maplit::hashset;
use mononoke_types::BlobstoreBytes;
use repo_identity::RepoIdentityRef;
use samplingblob::SamplingHandler;
use slog::info;

use crate::commands::JobParams;
use crate::commands::JobWalkParams;
use crate::commands::RepoSubcommandParams;
use crate::commands::COMPRESSION_BENEFIT;
use crate::detail::graph::FileContentData;
use crate::detail::graph::Node;
use crate::detail::graph::NodeData;
use crate::detail::graph::NodeType;
use crate::detail::graph::WrappedPath;
use crate::detail::progress::progress_stream;
use crate::detail::progress::report_state;
use crate::detail::progress::ProgressOptions;
use crate::detail::progress::ProgressReporter;
use crate::detail::progress::ProgressReporterUnprotected;
use crate::detail::progress::ProgressStateCountByType;
use crate::detail::progress::ProgressStateMutex;
use crate::detail::sampling::PathTrackingRoute;
use crate::detail::sampling::SamplingOptions;
use crate::detail::sampling::SamplingWalkVisitor;
use crate::detail::sampling::WalkKeyOptPath;
use crate::detail::sampling::WalkPayloadMtime;
use crate::detail::sampling::WalkSampleMapping;
use crate::detail::tail::walk_exact_tail;
use crate::detail::walk::RepoWalkParams;
use crate::detail::walk::RepoWalkTypeParams;

#[derive(Add, Div, Mul, Sub, Clone, Copy, Default, Debug)]
struct SizingStats {
    raw: u64,
    compressed: u64,
}

impl SizingStats {
    fn compression_benefit_pct(&self) -> u64 {
        if self.raw == 0 {
            0
        } else {
            100 * (self.raw - self.compressed) / self.raw
        }
    }
}

impl fmt::Display for SizingStats {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(
            fmt,
            "{},{},{}%",
            self.raw,
            self.compressed,
            self.compression_benefit_pct()
        )
    }
}

fn try_compress(raw_data: &Bytes, compressor_type: CompressorType) -> Result<SizingStats, Error> {
    let raw = raw_data.len() as u64;
    let compressed_buf = MeteredWrite::new(Cursor::new(Vec::with_capacity(4 * 1024)));
    let mut compressor = Compressor::new(compressed_buf, compressor_type);
    compressor.write_all(raw_data)?;
    let compressed_buf = compressor.try_finish().map_err(|(_encoder, e)| e)?;
    // Assume we wouldn't compress if its bigger
    let compressed = min(raw, compressed_buf.total_thru());
    Ok(SizingStats { raw, compressed })
}

// Force load of leaf data and check compression ratio
fn size_sampling_stream<InStream, InStats>(
    scheduled_max: usize,
    s: InStream,
    compressor_type: CompressorType,
    sampler: Arc<WalkSampleMapping<Node, SizingSample>>,
) -> impl Stream<Item = Result<(Node, Option<NodeData>, Option<SizingStats>), Error>>
where
    InStream: Stream<
            Item = Result<
                (
                    WalkKeyOptPath<WrappedPath>,
                    Option<NodeData>,
                    Option<InStats>,
                ),
                Error,
            >,
        >
        + 'static
        + Send,
    InStats: 'static + Send,
{
    s.map_ok(move |(walk_key, data_opt, _stats_opt)| {
        match (&walk_key.node, data_opt) {
            (Node::FileContent(_content_id), Some(NodeData::FileContent(fc)))
                if sampler.is_sampling(&walk_key.node) =>
            {
                match fc {
                    FileContentData::Consumed(_num_loaded_bytes) => {
                        future::ok(_num_loaded_bytes).left_future()
                    }
                    // Consume the stream to make sure we loaded all blobs
                    FileContentData::ContentStream(file_bytes_stream) => file_bytes_stream
                        .try_fold(0, |acc, file_bytes| future::ok(acc + file_bytes.size()))
                        .right_future(),
                }
                .and_then({
                    cloned!(sampler);
                    move |fs_stream_size| {
                        // Report the blobstore sizes in sizing stats, more accurate than stream sizes, as headers included
                        let sizes = sampler
                            .complete_step(&walk_key.node)
                            .map(|sizing_sample| {
                                sizing_sample.data.values().try_fold(
                                    SizingStats::default(),
                                    |acc, v| {
                                        try_compress(v.as_bytes(), compressor_type)
                                            .map(|sizes| acc + sizes)
                                    },
                                )
                            })
                            .transpose();

                        future::ready(sizes.map(|sizes| {
                            // Report the filestore stream's bytes size in the Consumed node
                            (
                                walk_key.node,
                                Some(NodeData::FileContent(FileContentData::Consumed(
                                    fs_stream_size,
                                ))),
                                sizes,
                            )
                        }))
                    }
                })
                .left_future()
            }
            (_, data_opt) => {
                // Report the blobstore sizes in sizing stats, more accurate than stream sizes, as headers included
                let sizes = sampler
                    .complete_step(&walk_key.node)
                    .map(|sizing_sample| {
                        sizing_sample
                            .data
                            .values()
                            .try_fold(SizingStats::default(), |acc, v| {
                                try_compress(v.as_bytes(), compressor_type).map(|sizes| acc + sizes)
                            })
                    })
                    .transpose();

                future::ready(sizes.map(|sizes| (walk_key.node, data_opt, sizes))).right_future()
            }
        }
    })
    .try_buffer_unordered(scheduled_max)
}

impl ProgressStateCountByType<SizingStats, SizingStats> {
    pub fn report_progress_log(&mut self, delta_time: Option<Duration>) {
        let summary_by_type: HashMap<NodeType, SizingStats> = self
            .work_stats
            .stats_by_type
            .iter()
            .map(|(k, (_i, v))| (*k, *v))
            .collect();
        let new_summary = summary_by_type
            .values()
            .fold(SizingStats::default(), |acc, v| acc + *v);
        let delta_summary = new_summary - self.reporting_stats.last_summary;

        let def = SizingStats::default();
        let detail = &self
            .params
            .types_sorted_by_name
            .iter()
            .map(|t| {
                let s = summary_by_type.get(t).unwrap_or(&def);
                format!("{}:{}", t, s)
            })
            .collect::<Vec<_>>()
            .join(" ");

        let (delta_s, delta_summary_per_s) =
            delta_time.map_or((0, SizingStats::default()), |delta_time| {
                (
                    delta_time.as_secs(),
                    delta_summary * 1000 / (delta_time.as_millis() as u64),
                )
            });

        let total_time = self
            .reporting_stats
            .last_update
            .duration_since(self.reporting_stats.start_time);

        let total_summary_per_s = if total_time.as_millis() > 0 {
            new_summary * 1000 / (total_time.as_millis() as u64)
        } else {
            SizingStats::default()
        };

        info!(
            self.params.logger,
            "Raw/s,Compressed/s,Raw,Compressed,%Saving; Delta {:06}/s,{:06}/s,{},{}s; Run {:06}/s,{:06}/s,{},{}s; Type:Raw,Compressed,%Saving {}",
            delta_summary_per_s.raw,
            delta_summary_per_s.compressed,
            delta_summary,
            delta_s,
            total_summary_per_s.raw,
            total_summary_per_s.compressed,
            new_summary,
            total_time.as_secs(),
            detail,
        );

        self.reporting_stats.last_summary_by_type = summary_by_type;
        self.reporting_stats.last_summary = new_summary;
    }
}

impl ProgressReporterUnprotected for ProgressStateCountByType<SizingStats, SizingStats> {
    fn report_progress(&mut self) {
        self.report_progress_log(None);
    }

    fn report_throttled(&mut self) {
        if let Some(delta_time) = self.should_log_throttled() {
            self.report_progress_log(Some(delta_time));
        }
    }
}

#[derive(Debug)]
pub struct SizingSample {
    pub data: HashMap<String, BlobstoreBytes>,
}

impl Default for SizingSample {
    fn default() -> Self {
        Self {
            data: HashMap::with_capacity(1),
        }
    }
}

impl SamplingHandler for WalkSampleMapping<Node, SizingSample> {
    fn sample_get(
        &self,
        ctx: &CoreContext,
        key: &str,
        value: Option<&BlobstoreGetData>,
    ) -> Result<(), Error> {
        ctx.sampling_key().map(|sampling_key| {
            self.inflight().get_mut(sampling_key).map(|mut guard| {
                value.map(|value| guard.data.insert(key.to_owned(), value.as_bytes().clone()))
            })
        });
        Ok(())
    }
}

#[derive(Clone)]
pub struct SizingCommand {
    pub compression_level: i32,
    pub progress_options: ProgressOptions,
    pub sampling_options: SamplingOptions,
    pub sampler: Arc<WalkSampleMapping<Node, SizingSample>>,
}

impl SizingCommand {
    fn apply_repo(&mut self, repo_params: &RepoWalkParams) {
        self.sampling_options
            .retain_or_default(&repo_params.include_node_types);
    }
}

// Subcommand entry point for estimate of file compression benefit
pub async fn compression_benefit(
    fb: FacebookInit,
    job_params: JobParams,
    command: SizingCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let JobParams {
        walk_params,
        per_repo,
    } = job_params;

    let mut all_walks = Vec::new();
    for (sub_params, repo_params) in per_repo {
        cloned!(mut command, walk_params);

        command.apply_repo(&repo_params);

        let walk = run_one(
            fb,
            walk_params,
            sub_params,
            repo_params,
            command,
            Arc::clone(&cancellation_requested),
        );
        all_walks.push(walk);
    }
    try_join_all(all_walks).await.map(|_| ())
}

async fn run_one(
    fb: FacebookInit,
    job_params: JobWalkParams,
    sub_params: RepoSubcommandParams,
    repo_params: RepoWalkParams,
    command: SizingCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let sizing_progress_state =
        ProgressStateMutex::new(ProgressStateCountByType::<SizingStats, SizingStats>::new(
            fb,
            repo_params.logger.clone(),
            COMPRESSION_BENEFIT,
            repo_params.repo.repo_identity().name().to_string(),
            command.sampling_options.node_types.clone(),
            command.progress_options,
        ));

    let make_sink = {
        cloned!(command, job_params.quiet, sub_params.progress_state,);
        move |ctx: &CoreContext, repo_params: &RepoWalkParams| {
            cloned!(ctx, repo_params.scheduled_max);
            async move |walk_output, _run_start, _chunk_num, _checkpoint_name| {
                cloned!(ctx, sizing_progress_state);
                // Sizing doesn't use mtime, so remove it from payload
                let walk_progress = progress_stream(quiet, &progress_state, walk_output).map_ok(
                    |(key, payload, stats): (_, WalkPayloadMtime, _)| (key, payload.data, stats),
                );

                let compressor = size_sampling_stream(
                    scheduled_max,
                    walk_progress,
                    CompressorType::Zstd {
                        level: command.compression_level,
                    },
                    command.sampler,
                );
                let report_sizing = progress_stream(quiet, &sizing_progress_state, compressor);

                report_state(ctx, report_sizing).await?;
                sizing_progress_state.report_progress();
                progress_state.report_progress();
                Ok(())
            }
        }
    };

    let walk_state = SamplingWalkVisitor::new(
        repo_params.include_node_types.clone(),
        repo_params.include_edge_types.clone(),
        command.sampling_options,
        None,
        command.sampler,
        job_params.enable_derive,
        sub_params
            .tail_params
            .chunking
            .as_ref()
            .map(|v| v.direction),
    );

    let type_params = RepoWalkTypeParams {
        required_node_data_types: hashset![NodeType::FileContent],
        always_emit_edge_types: HashSet::new(),
        keep_edge_paths: true,
    };

    walk_exact_tail::<_, _, _, _, _, PathTrackingRoute<WrappedPath>>(
        fb,
        job_params,
        repo_params,
        type_params,
        sub_params.tail_params,
        walk_state,
        make_sink,
        cancellation_requested,
    )
    .await
}