/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_common::pattern::resolve::ResolvedPattern;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::name::CellName;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::eval_result::EvaluationResult;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_query::query::environment::LabeledNode;
use buck2_query::query::environment::QueryEnvironment;
use buck2_query::query::environment::QueryTarget;
use buck2_query::query::syntax::simple::eval::error::QueryError;
use buck2_query::query::syntax::simple::eval::file_set::FileNode;
use buck2_query::query::syntax::simple::eval::file_set::FileSet;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use buck2_query::query::syntax::simple::functions::docs::QueryEnvironmentDescription;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use buck2_query::query::syntax::simple::functions::HasModuleDescription;
use buck2_query::query::traversal::async_depth_first_postorder_traversal;
use buck2_query::query::traversal::async_depth_limited_traversal;
use buck2_query::query::traversal::AsyncNodeLookup;
use buck2_query::query::traversal::AsyncTraversalDelegate;
use buck2_query::query::traversal::ChildVisitor;
use derive_more::Display;
use dupe::Dupe;
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use gazebo::prelude::*;
use indexmap::IndexSet;
use itertools::Itertools;
use ref_cast::RefCast;
use thiserror::Error;
use tracing::warn;

type ArcCellPath = Arc<CellPath>;

#[derive(Debug, Error)]
enum QueryLiteralResolutionError {
    #[error("literal `{0}` missing in pre-resolved literals")]
    LiteralMissing(String),
}

#[derive(Debug, Error)]
enum RBuildFilesError {
    #[error("no parent found for the file `{0}`")]
    ParentDoesNotExist(ArcCellPath),
    #[error("internal error: no buildfile names found for the cell `{0}`")]
    CellMissingBuildFileNames(CellName),
}

pub enum SpecialAttr {
    String(String),
}

/// UqueryDelegate resolves information needed by the QueryEnvironment.
#[async_trait]
pub trait UqueryDelegate: Send + Sync {
    /// Returns the EvaluationResult for evaluation of the buildfile.
    async fn eval_build_file(
        &self,
        packages: PackageLabel,
    ) -> anyhow::Result<Arc<EvaluationResult>>;

    /// Get the imports from a LoadedModule corresponding to some path.
    async fn eval_module_imports(&self, path: &ImportPath) -> anyhow::Result<Vec<ImportPath>>;

    fn get_buildfile_names_by_cell(&self) -> anyhow::Result<HashMap<CellName, &[FileNameBuf]>>;

    /// Resolves a target pattern.
    async fn resolve_target_patterns(
        &self,
        pattern: &[&str],
    ) -> anyhow::Result<ResolvedPattern<TargetPatternExtra>>;

    async fn eval_file_literal(&self, literal: &str) -> anyhow::Result<FileSet>;

    // Get all enclosing packages needed to compute owner function.
    // This always includes the immediate enclosing package of the path but can also include
    // all parent packages if the package matches `project.package_boundary_exceptions` buckconfig.
    async fn get_enclosing_packages(&self, path: &CellPath) -> anyhow::Result<Vec<PackageLabel>>;
}

#[async_trait]
pub trait QueryLiterals<T: QueryTarget>: Send + Sync {
    async fn eval_literals(&self, literals: &[&str]) -> anyhow::Result<TargetSet<T>>;
}

pub struct UqueryEnvironment<'c> {
    delegate: Arc<dyn UqueryDelegate + 'c>,
    literals: Arc<dyn QueryLiterals<TargetNode> + 'c>,
}

pub struct PreresolvedQueryLiterals<T: QueryTarget> {
    resolved_literals: HashMap<String, SharedResult<TargetSet<T>>>,
}

impl<T: QueryTarget> PreresolvedQueryLiterals<T> {
    pub fn new(resolved_literals: HashMap<String, SharedResult<TargetSet<T>>>) -> Self {
        Self { resolved_literals }
    }

    pub async fn pre_resolve(base: &dyn QueryLiterals<T>, literals: &[String]) -> Self {
        let futs = literals
            .iter()
            .map(|lit| async move { (lit.to_owned(), base.eval_literals(&[lit]).await) });
        let mut resolved_literals = HashMap::new();
        for (literal, result) in futures::future::join_all(futs).await {
            resolved_literals.insert(literal, result.shared_error());
        }
        Self { resolved_literals }
    }

    /// All the literals, or error if resolution of any failed.
    pub fn literals(&self) -> anyhow::Result<TargetSet<T>> {
        let mut literals = TargetSet::new();
        for result in self.resolved_literals.values() {
            literals.extend(result.as_ref().map_err(|e| e.dupe())?);
        }
        Ok(literals)
    }
}

#[async_trait]
impl<T: QueryTarget> QueryLiterals<T> for PreresolvedQueryLiterals<T> {
    async fn eval_literals(&self, literals: &[&str]) -> anyhow::Result<TargetSet<T>> {
        let mut targets = TargetSet::new();
        for lit in literals {
            let resolved = match self
                .resolved_literals
                .get(*lit)
                .ok_or_else(|| QueryLiteralResolutionError::LiteralMissing((*lit).to_owned()))?
            {
                Ok(v) => v,
                Err(e) => return Err(e.dupe().into()),
            };
            targets.extend(resolved);
        }
        Ok(targets)
    }
}

impl<'c> UqueryEnvironment<'c> {
    pub fn new(
        delegate: Arc<dyn UqueryDelegate + 'c>,
        literals: Arc<dyn QueryLiterals<TargetNode> + 'c>,
    ) -> Self {
        Self { delegate, literals }
    }

    pub fn describe() -> QueryEnvironmentDescription {
        QueryEnvironmentDescription {
            name: "Uquery Environment".to_owned(),
            mods: vec![DefaultQueryFunctionsModule::<Self>::describe()],
        }
    }

    async fn get_node(&self, target: &TargetLabel) -> anyhow::Result<TargetNode> {
        let package = self
            .delegate
            .eval_build_file(target.pkg())
            .await
            .with_context(|| format!("Error looking up `{}``", target))?;
        let node = package.resolve_target(target.name())?;
        Ok(node.dupe())
    }
}

#[async_trait]
impl<'c> QueryEnvironment for UqueryEnvironment<'c> {
    type Target = TargetNode;

    async fn get_node(&self, node_ref: &TargetLabel) -> anyhow::Result<Self::Target> {
        UqueryEnvironment::get_node(self, node_ref).await
    }

    async fn get_node_for_default_configured_target(
        &self,
        _node_ref: &TargetLabel,
    ) -> anyhow::Result<MaybeCompatible<Self::Target>> {
        Err(QueryError::FunctionUnimplemented(
            "get_node_for_default_configured_target() only for CqueryEnvironment",
        )
        .into())
    }

    async fn eval_literals(&self, literals: &[&str]) -> anyhow::Result<TargetSet<TargetNode>> {
        self.literals.eval_literals(literals).await
    }

    async fn eval_file_literal(&self, literal: &str) -> anyhow::Result<FileSet> {
        self.delegate.eval_file_literal(literal).await
    }

    async fn dfs_postorder(
        &self,
        root: &TargetSet<TargetNode>,
        traversal_delegate: &mut dyn AsyncTraversalDelegate<TargetNode>,
    ) -> anyhow::Result<()> {
        async_depth_first_postorder_traversal(self, root.iter_names(), traversal_delegate).await
    }

    async fn depth_limited_traversal(
        &self,
        root: &TargetSet<Self::Target>,
        delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
        depth: u32,
    ) -> anyhow::Result<()> {
        async_depth_limited_traversal(self, root.iter_names(), delegate, depth).await
    }

    async fn allbuildfiles(&self, universe: &TargetSet<Self::Target>) -> anyhow::Result<FileSet> {
        return allbuildfiles(universe, &*self.delegate).await;
    }

    async fn rbuildfiles(&self, universe: &FileSet, argset: &FileSet) -> anyhow::Result<FileSet> {
        return rbuildfiles(universe, argset, &*self.delegate).await;
    }

    async fn owner(&self, paths: &FileSet) -> anyhow::Result<TargetSet<Self::Target>> {
        let mut result: TargetSet<Self::Target> = TargetSet::new();
        for path in paths.iter() {
            // need to explicitly track this rather than checking for changes to result set since the owner might
            // already be in the set.
            let mut found_owner = false;
            match self.delegate.get_enclosing_packages(path).await {
                Ok(packages) => {
                    let package_futs = packages.map(|package| async move {
                        // TODO(cjhopman): We should make sure that the file exists.
                        let targets = self.delegate.eval_build_file(package.dupe()).await?;

                        let owner_targets: Vec<Self::Target> = targets
                            .targets()
                            .values()
                            .filter_map(|node| {
                                for input in node.inputs() {
                                    if &input == path {
                                        return Some(node.dupe());
                                        // this intentionally breaks out of the loop. We don't need to look at the
                                        // other inputs of this target, but it's possible for a single file to be owned by
                                        // multiple targets.
                                    }
                                }
                                None
                            })
                            .collect();
                        anyhow::Ok(owner_targets)
                    });

                    for nodes in futures::future::join_all(package_futs).await.into_iter() {
                        for node in nodes?.into_iter() {
                            found_owner = true;
                            result.insert(node);
                        }
                    }
                }
                Err(_) => {
                    // we don't consider this an error, it's usually the case that the user
                    // just wants to know the target owning the file if it exists.
                }
            };
            if !found_owner {
                warn!("No owner was found for {}", path);
            }
        }
        Ok(result)
    }
}

#[async_trait]
impl<'a> AsyncNodeLookup<TargetNode> for UqueryEnvironment<'a> {
    async fn get(&self, label: &TargetLabel) -> anyhow::Result<TargetNode> {
        self.get_node(label).await
    }
}

pub(crate) async fn allbuildfiles<'c, T: QueryTarget>(
    universe: &TargetSet<T>,
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<FileSet> {
    let mut paths = IndexSet::<FileNode>::new();

    let mut top_level_imports = Vec::<ImportPath>::new();

    for target in universe.iter() {
        paths.insert(FileNode(target.dupe().buildfile_path().path()));

        let eval_result = delegate
            .eval_build_file(target.buildfile_path().package())
            .await?; // TODO: no longer use eval_build_file, just parse imports directly (will solve async issue too)

        top_level_imports.extend(eval_result.imports().iter().cloned());
    }

    let loads = get_transitive_loads(top_level_imports, delegate).await?;

    let mut new_paths = IndexSet::<FileNode>::new();
    for load in &loads {
        new_paths.insert(FileNode(load.path().clone()));
    }

    Ok(FileSet::new(paths).union(&FileSet::new(new_paths)))
}

pub(crate) async fn rbuildfiles<'c>(
    universe: &FileSet,
    argset: &FileSet,
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<FileSet> {
    let universe_paths: Vec<ArcCellPath> =
        universe.iter().map(|file| Arc::new(file.clone())).collect();
    // step 1: split the build files and bzl files
    let (buildfiles, bzlfiles) = split_universe_files(&universe_paths, delegate)?;

    // step 2: get all top level imports accordingly
    let top_level_import_by_build_file =
        top_level_imports_by_build_file(&buildfiles, &bzlfiles, delegate).await?;

    // step 3: get the first order imports for every loaded file, to lookup during traversal
    // TODO(benfoxman) this is actually unnecessary, since we can get the imports while traversing.
    // Will fix this in a separate diff.
    let all_top_level_imports: Vec<ImportPath> = top_level_import_by_build_file
        .iter()
        .flat_map(|(_, imports)| imports.iter().cloned())
        .collect();

    let first_order_import_map = first_order_imports(&all_top_level_imports, delegate).await?;

    // step 4: do rdeps()-like traversal from the top level imports down to bzls

    #[derive(Clone, Dupe)]
    struct Node(Arc<ImportPath>);

    impl Node {
        fn import_path(&self) -> &ImportPath {
            &self.0
        }
    }

    #[derive(Display, Debug, Hash, Eq, PartialEq, Clone, RefCast)]
    #[repr(transparent)]
    struct NodeRef(ImportPath);

    impl LabeledNode for Node {
        type NodeRef = NodeRef;

        fn node_ref(&self) -> &NodeRef {
            NodeRef::ref_cast(self.import_path())
        }
    }

    struct Lookup {}

    #[async_trait]
    impl AsyncNodeLookup<Node> for Lookup {
        async fn get(&self, label: &NodeRef) -> anyhow::Result<Node> {
            Ok(Node(Arc::new(label.0.clone())))
        }
    }

    struct Delegate<'c> {
        output_paths: Vec<ImportPath>,
        argset: &'c FileSet,
        first_order_import_map: HashMap<ImportPath, Vec<ImportPath>>,
    }

    #[async_trait]
    impl AsyncTraversalDelegate<Node> for Delegate<'_> {
        fn visit(&mut self, node: Node) -> anyhow::Result<()> {
            let node_import = node.import_path();
            if self.argset.iter().contains(node_import.path()) {
                self.output_paths.push(node_import.clone());
            } else {
                let loads = self
                    .first_order_import_map
                    .get(node_import)
                    .expect("import path should exist");
                for load in loads.iter() {
                    for arg in self.argset.iter() {
                        if load.path() == arg {
                            self.output_paths.push(node_import.clone());
                            return Ok(());
                        }
                    }
                }
            }
            Ok(())
        }
        async fn for_each_child(
            &mut self,
            node: &Node,
            func: &mut dyn ChildVisitor<Node>,
        ) -> anyhow::Result<()> {
            for import in self
                .first_order_import_map
                .get(node.import_path())
                .expect("import path should exist")
            {
                func.visit(NodeRef(import.clone()))?;
            }
            Ok(())
        }
    }

    let lookup = Lookup {};
    let mut delegate = Delegate {
        output_paths: vec![],
        argset,
        first_order_import_map,
    };

    // step 5: do traversal, get all modified imports
    async_depth_first_postorder_traversal(
        &lookup,
        all_top_level_imports.iter().map(NodeRef::ref_cast),
        &mut delegate,
    )
    .await?;

    let mut output_files = IndexSet::<FileNode>::new();
    for file in &delegate.output_paths {
        output_files.insert(FileNode(file.path().clone()));
    }

    // finally, since we didn't check the top level imports of buildfiles,
    // we do this right now
    for file in buildfiles {
        if argset.iter().contains(&*file) {
            output_files.insert(FileNode((*file).clone()));
        } else {
            let imports = top_level_import_by_build_file
                .get(&file)
                .expect("should have stored a universe file's input");

            for import in imports.iter() {
                if output_files
                    .iter()
                    .contains(&FileNode(import.path().clone()))
                {
                    output_files.insert(FileNode((*file).clone()));
                    break;
                }
            }
        }
    }

    Ok(FileSet::new(output_files))
}

fn split_universe_files<'c>(
    universe: &[ArcCellPath],
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<(Vec<ArcCellPath>, Vec<ArcCellPath>)> {
    let mut buildfiles = Vec::<ArcCellPath>::new();
    let mut bzlfiles = Vec::<ArcCellPath>::new();
    let buildfile_names_by_cell = delegate.get_buildfile_names_by_cell()?;
    for file in universe {
        let buildfile_names_for_file =
            buildfile_names_by_cell.get(&file.cell()).ok_or_else(|| {
                anyhow::anyhow!(RBuildFilesError::CellMissingBuildFileNames(file.cell()))
            })?;

        if let Some(name) = file.path().file_name() {
            if buildfile_names_for_file
                .iter()
                .map(<FileNameBuf as AsRef<FileName>>::as_ref)
                .contains(&name)
            {
                buildfiles.push(file.dupe());
            }
        } else {
            // TODO: right now we assume non-buildfiles are bzl's - we might want to handle error cases later.
            bzlfiles.push(file.dupe());
        }
    }
    Ok((buildfiles, bzlfiles))
}

async fn top_level_imports_by_build_file<'c>(
    buildfiles: &[ArcCellPath],
    bzlfiles: &[ArcCellPath],
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<HashMap<ArcCellPath, Vec<ImportPath>>> {
    let mut top_level_import_by_build_file = HashMap::<ArcCellPath, Vec<ImportPath>>::new();

    for file in bzlfiles {
        let imports = vec![ImportPath::new_same_cell(file.as_ref().clone())?];
        top_level_import_by_build_file.insert(file.dupe(), imports);
    }

    let mut buildfile_futs: FuturesUnordered<_> = buildfiles
        .iter()
        .map(|file| async move {
            if let Some(parent) = file.parent() {
                (
                    file.dupe(),
                    delegate
                        .eval_build_file(PackageLabel::from_cell_path(parent))
                        .await,
                )
            } else {
                (
                    file.dupe(),
                    Err(RBuildFilesError::ParentDoesNotExist(file.dupe()).into()),
                )
            }
        })
        .collect();

    while let Some((file, eval_result)) = tokio::task::unconstrained(buildfile_futs.next()).await {
        top_level_import_by_build_file.insert(file.dupe(), eval_result?.imports().to_owned());
    }

    Ok(top_level_import_by_build_file)
}

// TODO: no need to get all the first order imports prior to the traversal - we can do this
// at the same time as the actual traversal.
async fn first_order_imports<'c>(
    all_top_level_imports: &[ImportPath],
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<HashMap<ImportPath, Vec<ImportPath>>> {
    let all_imports = get_transitive_loads(all_top_level_imports.to_vec(), delegate).await?;

    let mut all_first_order_futs: FuturesUnordered<_> = all_imports
        .iter()
        .map(|node| async move { (node, delegate.eval_module_imports(node).await) })
        .collect();

    let mut first_order_import_map = HashMap::<ImportPath, Vec<ImportPath>>::new();

    while let Some((import, first_order_imports)) =
        tokio::task::unconstrained(all_first_order_futs.next()).await
    {
        let prev = first_order_import_map.insert(import.clone(), first_order_imports?);
        assert!(prev.is_none());
    }
    Ok(first_order_import_map)
}

// Uquery and Cquery share ImportPath traversal logic, so we move the logic to this function.
pub(crate) async fn get_transitive_loads<'c>(
    top_level_imports: Vec<ImportPath>,
    delegate: &'c dyn UqueryDelegate,
) -> anyhow::Result<Vec<ImportPath>> {
    #[derive(Clone, Dupe)]
    struct Node(Arc<ImportPath>);

    impl Node {
        fn import_path(&self) -> &ImportPath {
            &self.0
        }
    }

    #[derive(Display, Debug, Hash, Eq, PartialEq, Clone, RefCast)]
    #[repr(transparent)]
    struct NodeRef(ImportPath);

    impl LabeledNode for Node {
        type NodeRef = NodeRef;

        fn node_ref(&self) -> &NodeRef {
            NodeRef::ref_cast(self.import_path())
        }
    }

    struct Lookup {}

    #[async_trait]
    impl AsyncNodeLookup<Node> for Lookup {
        async fn get(&self, label: &NodeRef) -> anyhow::Result<Node> {
            Ok(Node(Arc::new(label.0.clone())))
        }
    }

    struct Delegate<'c> {
        imports: Vec<ImportPath>,
        delegate: &'c dyn UqueryDelegate,
    }

    #[async_trait]
    impl AsyncTraversalDelegate<Node> for Delegate<'_> {
        fn visit(&mut self, target: Node) -> anyhow::Result<()> {
            self.imports.push(target.import_path().clone());
            Ok(())
        }

        async fn for_each_child(
            &mut self,
            target: &Node,
            func: &mut dyn ChildVisitor<Node>,
        ) -> anyhow::Result<()> {
            for import in self
                .delegate
                .eval_module_imports(target.import_path())
                .await?
            {
                func.visit(NodeRef(import.clone()))?;
            }
            Ok(())
        }
    }

    let mut traversal_delegate = Delegate {
        imports: vec![],
        delegate: delegate.dupe(),
    };
    let lookup = Lookup {};

    let import_nodes = top_level_imports.iter().map(NodeRef::ref_cast);

    async_depth_first_postorder_traversal(&lookup, import_nodes, &mut traversal_delegate).await?;

    Ok(traversal_delegate.imports)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Implementation of the cli and query_* attr query language.
use std::sync::Arc;

use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use dice::DiceComputations;
use dupe::Dupe;

use crate::analysis::evaluator::eval_query;
use crate::dice::get_dice_query_delegate;
use crate::dice::DiceQueryDelegate;
use crate::uquery::environment::PreresolvedQueryLiterals;
use crate::uquery::environment::UqueryEnvironment;

pub struct UqueryEvaluator<'c> {
    dice_query_delegate: Arc<DiceQueryDelegate<'c>>,
    functions: DefaultQueryFunctionsModule<UqueryEnvironment<'c>>,
}

impl UqueryEvaluator<'_> {
    pub async fn eval_query(
        &self,
        query: &str,
        query_args: &[String],
    ) -> anyhow::Result<QueryEvaluationResult<TargetNode>> {
        eval_query(&self.functions, query, query_args, async move |literals| {
            let resolved_literals =
                PreresolvedQueryLiterals::pre_resolve(&*self.dice_query_delegate, &literals).await;
            Ok(UqueryEnvironment::new(
                self.dice_query_delegate.dupe(),
                Arc::new(resolved_literals),
            ))
        })
        .await
    }
}

/// Evaluates some query expression. TargetNodes are resolved via the interpreter from
/// the provided DiceCtx.
pub async fn get_uquery_evaluator<'a, 'c: 'a>(
    ctx: &'c DiceComputations,
    working_dir: &'a ProjectRelativePath,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<UqueryEvaluator<'c>> {
    let dice_query_delegate =
        Arc::new(get_dice_query_delegate(ctx, working_dir, global_target_platform).await?);
    let functions = DefaultQueryFunctionsModule::new();

    Ok(UqueryEvaluator {
        dice_query_delegate,
        functions,
    })
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub(crate) mod bxl;
pub mod environment;
pub mod evaluator;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Parsing query expressions.
//!
//! The parser doesn't do any validation on the function names or function arguments, that is handled
//! during evaluation. The returned parsed tree has locations attached so that the evaluation can
//! produce good error messages.
//!
//! The grammar is something like:
//!
//! ```text
//!
//! # note that set's args are space-separated, not comma-separated and so cannot be treated as a function
//! EXPR ::=
//!          WORD
//!        | INTEGER
//!        | '(' EXPR ')'
//!        | 'set(' WORD * ')'
//!        | FUNCTION_NAME '(' EXPR ( ',' EXPR ) * ')'
//!        | EXPR 'intersect' EXPR
//!        | EXPR ' ^ ' EXPR
//!        | EXPR ' union ' EXPR
//!        | EXPR ' + ' EXPR
//!        | EXPR ' except ' EXPR
//!        | EXPR ' - ' EXPR
//!
//! # word is much broader than a normal identifier-like thing would allow since we don't want to require
//! # quoting targets "@fbcode//some:target" or common regexes ".*" or filenames "Foo.java".
//! WORD ::=
//!          '\'' not_single_quote * '\''
//!        | '"' not_double_quote * '"'
//!        | "a-zA-Z0-9*/@._:$#%-" *
//!
//! INTEGER ::= "0" | ("1-9" "0-9"*)
//!
//! FUNCTION_NAME ::= "a-zA-Z_" "a-zA-Z0-9_" *
//!
//! ```

pub mod placeholder;
pub mod span;
pub mod spanned;

use std::fmt::Display;

use dupe::Dupe;
use enum_map::Enum;
use gazebo::prelude::*;
use gazebo::variants::VariantName;
use nom::branch::alt;
use nom::bytes::complete::is_a;
use nom::bytes::complete::tag;
use nom::bytes::complete::take_till;
use nom::character::complete::alpha1;
use nom::character::complete::alphanumeric1;
use nom::character::complete::char;
use nom::character::complete::digit1;
use nom::character::complete::multispace0;
use nom::character::complete::multispace1;
use nom::combinator::all_consuming;
use nom::combinator::cut;
use nom::combinator::recognize;
use nom::error::context;
use nom::error::convert_error;
use nom::error::ErrorKind;
use nom::error::VerboseError;
use nom::multi::many0;
use nom::multi::many1;
use nom::multi::separated_list0;
use nom::sequence::delimited;
use nom::sequence::pair;
use nom::sequence::preceded;
use nom::sequence::terminated;
use nom::IResult;
use thiserror::Error;

use crate::span::Span;
use crate::spanned::Spanned;

// TODO(cjhopman): Add `LET WORD = expr IN expr`

// TODO(cjhopman): We should switch to our own error type here. VerboseError doesn't even allow us to construct
// our own error messages (so, for example, we can't have a good error message for too large integers) and doesn't
// support propagating anyhow or std errors (and since we can't do a custom message, we can't even capture them as a string).
type NomResult<'a, O, E> = IResult<Span<'a>, O, E>;

// Type alias.
trait NomParseError<'a>:
    nom::error::ParseError<Span<'a>> + nom::error::ContextError<Span<'a>> + 'a
{
}

impl<'a, E: nom::error::ParseError<Span<'a>> + nom::error::ContextError<Span<'a>> + 'a>
    NomParseError<'a> for E
{
}

#[derive(Debug, Error)]
enum ParseError {
    #[error("{0}")]
    NomError(String),
}

/// This is the main output type of the query parser.
pub type SpannedExpr<'a> = Spanned<Expr<'a>>;

#[derive(Debug, VariantName)]
pub enum Expr<'a> {
    // doesn't need to be Span since the Expr itself is always a SpannedExpr.
    String(&'a str),
    Integer(u64),
    Function {
        function_name: Span<'a>,
        args: Vec<SpannedExpr<'a>>,
    },
    BinaryOpSequence(Box<SpannedExpr<'a>>, Vec<(BinaryOp, SpannedExpr<'a>)>),
    Set(Vec<Span<'a>>),
    FileSet(Vec<Span<'a>>),
}

impl Display for Expr<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Expr::String(word) => {
                let quote = if word.contains('\'') { '"' } else { '\'' };
                // The grammar shouldn't allow a word that contains both quote-types, but
                // this assert should help ensure this gets updated if the grammar changes
                // to allow that.
                assert!(!word.contains(quote));
                write!(f, "{0}{1}{0}", quote, word)?;
            }
            Expr::Integer(val) => write!(f, "{}", val)?,
            Expr::Function {
                function_name,
                args,
            } => {
                f.write_str(function_name.fragment())?;
                f.write_str("(")?;
                for (i, v) in args.iter().enumerate() {
                    if i != 0 {
                        f.write_str(", ")?;
                    }
                    v.fmt(f)?;
                }
                f.write_str(")")?;
            }
            Expr::BinaryOpSequence(left, exprs) => {
                // prints exprs.len() leading '('
                // see https://stackoverflow.com/questions/35280798/printing-a-character-a-variable-number-of-times-with-println
                write!(f, "{1:(<0$} {2}", exprs.len(), "", left)?;
                for e in exprs {
                    write!(f, " {} {})", e.0, e.1)?;
                }
            }
            Expr::Set(vals) => {
                f.write_str("set(")?;
                for (i, v) in vals.iter().enumerate() {
                    if i != 0 {
                        f.write_str(" ")?;
                    }
                    f.write_str(v.fragment())?;
                }
                f.write_str(")")?;
            }
            Expr::FileSet(vals) => {
                f.write_str("fileset(")?;
                for (i, v) in vals.iter().enumerate() {
                    if i != 0 {
                        f.write_str(" ")?;
                    }
                    f.write_str(v.fragment())?;
                }
                f.write_str(")")?;
            }
        }
        Ok(())
    }
}

const INTERSECT: &str = "^";
const EXCEPT: &str = "-";
const UNION: &str = "+";

#[derive(Debug, Enum, Copy, Dupe, Clone)]
pub enum BinaryOp {
    Intersect,
    Except,
    Union,
}

impl Display for BinaryOp {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            BinaryOp::Intersect => INTERSECT,
            BinaryOp::Except => EXCEPT,
            BinaryOp::Union => UNION,
        })
    }
}

/// Wraps a parser producing `O` with one producing `Spanned<O>` by marking the output with a span covering all of the consumed input.
fn spanned<'a, O, E: NomParseError<'a>, F: FnMut(Span<'a>) -> IResult<Span<'a>, O, E>>(
    mut func: F,
) -> impl FnMut(Span<'a>) -> IResult<Span<'a>, Spanned<O>, E>
where
    O: 'a,
    F: 'a,
{
    move |original_input| {
        let start = original_input.location_offset();
        let (input, value) = func(original_input)?;
        let end = input.location_offset();
        Ok((
            input,
            Spanned {
                position: start..end,
                value,
            },
        ))
    }
}

/// nom's convert_error requires that the error's type be very str-like and Span isn't. So, we convert the Span error into a &str error.
fn convert_to_str_error(err: VerboseError<Span>) -> VerboseError<&str> {
    VerboseError {
        errors: err.errors.into_map(|(span, kind)| (span.fragment(), kind)),
    }
}

/// Parses a query string into a SpannedExpr. Requires that the entire input is consumed.
pub fn parse_expr(input: &str) -> anyhow::Result<SpannedExpr> {
    let span = Span::new(input);
    // Parse with fast error (`()`) first,
    // and on error reparse again with `VerboseError` to get detailed errors.
    match all_consuming(expr)(span) {
        Ok((_, value)) => Ok(value),
        Err(nom::Err::Failure(())) | Err(nom::Err::Error(())) => {
            match all_consuming(expr)(span) {
                Ok(..) => unreachable!(
                    "if fast parse didn't succeed, slow parse should not succeed as well"
                ),
                Err(nom::Err::Failure(err)) | Err(nom::Err::Error(err)) => {
                    // TODO(cjhopman): Do better error message conversion. The default nom one is rather verbose and references line
                    // numbers (which generally are meaningless in this context where we have a single line). It'd also be nice if we
                    // could keep the error information rather than just stringifying it here.
                    Err(
                        ParseError::NomError(convert_error(input, convert_to_str_error(err)))
                            .into(),
                    )
                }
                Err(nom::Err::Incomplete(..)) => unreachable!(),
            }
        }
        Err(nom::Err::Incomplete(..)) => unreachable!(),
    }
}

// Parses a non-infix op expression. This is split out so that we can parse a sequence of infix operators without recursion.
fn single_expr<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    // The ordering here is a little important, the first three of these all have a pattern of identifying
    // some prefix (ex, expr_set will identify "set(") and then failing if it doesn't match.
    //
    // The important bit there is probably around set/function/word since we don't want to, for example, parse out a
    //  word from "set(" (we'd only then get an error by the caller requiring that we consumed everything).
    //
    // The infix binary operators require left-recursion, so we can't just handle that like the others. Instead we
    // parse an expression from the beginning of the input and check after if there's a "trailing" infix operator.
    let (input, left_expr) = alt((
        preceded(char('('), cut(terminated(expr, char(')')))),
        expr_set,
        expr_fileset,
        expr_function,
        expr_int,
        expr_word,
    ))(input)?;

    Ok((input, left_expr))
}

/// Tries to parse an Expr.
fn expr<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    context(
        "expression",
        delimited(
            multispace0,
            spanned(|input| {
                let (input, left_expr) = single_expr(input)?;

                if let Ok((input, binary_ops)) = trailing_infix(input) {
                    Ok((
                        input,
                        Expr::BinaryOpSequence(Box::new(left_expr), binary_ops),
                    ))
                } else {
                    Ok((input, left_expr.value))
                }
            }),
            multispace0,
        ),
    )(input)
}

// trailing infix parses a sequence of trailing infix operators. It's important that we don't recurse for each item in such
// a sequence as in practice we get really long such chains.
fn trailing_infix(input: Span) -> NomResult<Vec<(BinaryOp, SpannedExpr)>, ()> {
    fn single_infix(input: Span) -> NomResult<(BinaryOp, SpannedExpr), ()> {
        let (input, operator) = binary_op(input)?;
        let (input, right) = cut(single_expr)(input)?;
        Ok((input, (operator, right)))
    }
    many1(single_infix)(input)
}

/// Tries to parse an Expr::Word
fn expr_word<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    spanned(|input| {
        let (remaining, word) = word(input)?;
        Ok((remaining, Expr::String(word.fragment())))
    })(input)
}

/// Tries to parse an Expr::Integer
fn expr_int<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    spanned(|input| {
        let (remaining, word) = digit1(input)?;

        let value = match word.fragment() {
            "0" => 0,
            val if val.starts_with('0') => {
                return Err(nom::Err::Failure(nom::error::ParseError::from_error_kind(
                    input,
                    ErrorKind::Digit,
                )));
            }
            val => val.parse::<u64>().map_err(|_| {
                nom::Err::Failure(nom::error::ParseError::from_error_kind(
                    input,
                    ErrorKind::Digit,
                ))
            })?,
        };

        Ok((remaining, Expr::Integer(value)))
    })(input)
}

fn word<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, Span<'a>, E> {
    fn non_quoted_word<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, Span<'a>, E> {
        recognize(many1(alt((alphanumeric1, is_a("*/@.-_:$#%")))))(input)
    }

    alt((
        preceded(
            char('\''),
            cut(terminated(take_till(|c| c == '\''), char('\''))),
        ),
        preceded(
            char('"'),
            cut(terminated(take_till(|c| c == '"'), char('"'))),
        ),
        non_quoted_word,
    ))(input)
}

fn binary_op(input: Span) -> NomResult<BinaryOp, ()> {
    fn keyword(long: &'static str) -> impl Fn(Span) -> NomResult<Span, ()> {
        // keywords require spaces separating from the exprs
        move |span| delimited(multispace1, tag(long), multispace1)(span)
    }

    fn symbol(sym: &'static str) -> impl FnMut(Span) -> NomResult<Span, ()> {
        // symbols can have spaces separating from the exprs, but don't require them
        move |span: Span| delimited(multispace0, tag(sym), multispace0)(span)
    }

    fn op(
        sym: &'static str,
        word: &'static str,
        op: BinaryOp,
    ) -> impl Fn(Span) -> NomResult<BinaryOp, ()> {
        move |span| Ok((alt((symbol(sym), keyword(word)))(span)?.0, op))
    }

    alt((
        op(EXCEPT, "except", BinaryOp::Except),
        op(INTERSECT, "intersect", BinaryOp::Intersect),
        op(UNION, "union", BinaryOp::Union),
    ))(input)
}

/// Tries to parse an Expr::Set. Will fail if it detects an unfinished "set("
fn expr_set<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    fn set_args<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, Vec<Span<'a>>, E> {
        separated_list0(multispace1, word)(input)
    }

    spanned(|input| {
        let (input, _) = tag("set(")(input)?;
        cut(move |input| {
            let (input, args) = delimited(multispace0, set_args, multispace0)(input)?;
            let (input, _) = char(')')(input)?;
            Ok((input, Expr::Set(args)))
        })(input)
    })(input)
}

/// Tries to parse an Expr::FileSet. Will fail if it detects an unfinished "fileset("
fn expr_fileset<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    fn set_args<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, Vec<Span<'a>>, E> {
        separated_list0(multispace1, word)(input)
    }

    spanned(|input| {
        let (input, _) = tag("fileset(")(input)?;
        cut(move |input| {
            let (input, args) = delimited(multispace0, set_args, multispace0)(input)?;
            let (input, _) = char(')')(input)?;
            Ok((input, Expr::FileSet(args)))
        })(input)
    })(input)
}

/// Tries to parse an Expr::Function. Will fail if it detects an unfinished "func("
// We don't need to worry about "set(" as the outermost expr() ensures that never gets to here.
fn expr_function<'a, E: NomParseError<'a>>(input: Span<'a>) -> NomResult<'a, SpannedExpr<'a>, E> {
    fn function_args<'a, E: NomParseError<'a>>(
        input: Span<'a>,
    ) -> NomResult<'a, Vec<SpannedExpr<'a>>, E> {
        separated_list0(terminated(tag(","), multispace0), expr)(input)
    }

    spanned(|input| {
        let (input, function_name) = recognize(pair(
            alt((alpha1, tag("_"))),
            many0(alt((alphanumeric1, tag("_")))),
        ))(input)?;
        let (input, _) = char('(')(input)?;
        cut(move |input| {
            let (input, args) = terminated(function_args, char(')'))(input)?;
            Ok((
                input,
                Expr::Function {
                    function_name,
                    args,
                },
            ))
        })(input)
    })(input)
}

#[cfg(test)]
mod tests {
    use std::fmt::Debug;

    use super::*;

    fn assert_err<'a, E, O: Debug>(res: IResult<Span, O, E>, ctx: &str) -> E
    where
        E: NomParseError<'a> + Debug,
    {
        match res {
            Err(nom::Err::Incomplete(..)) => {
                unreachable!("Only streaming parsing should return 'Err::Incomplete'");
            }
            Err(nom::Err::Error(e)) => e,
            _ => {
                panic!("{} Got `{:?}`", ctx, res);
            }
        }
    }

    fn assert_fail<'a, E, O: Debug>(res: IResult<Span, O, E>, ctx: &str) -> E
    where
        E: NomParseError<'a> + Debug,
    {
        match res {
            Err(nom::Err::Incomplete(..)) => {
                unreachable!("Only streaming parsing should return 'Err::Incomplete'");
            }
            Err(nom::Err::Failure(e)) => e,
            _ => {
                panic!("{} Got `{:?}`", ctx, res);
            }
        }
    }

    #[test]
    fn test_expr() -> anyhow::Result<()> {
        run_tests(
            expr,
            &[
                "f",
                "func(f, a)",
                "deps(%Ss) + set(a b c) ^ func(o(a(b, c, d)))",
                "a + b",
                "(a - (b))",
                "123",
            ],
            &[],
            &["func(", "set(", "(a", "01234"],
        );

        match parse_expr("set(a b c)") {
            Ok(Spanned {
                value: Expr::Set(..),
                ..
            }) => {}
            v => panic!("expected set expr, got `{:?}`", v),
        }

        match parse_expr("set(a b c) ^ func()") {
            Ok(Spanned {
                value: Expr::BinaryOpSequence(..),
                ..
            }) => {}
            v => panic!("expected intersect expr, got `{:?}`", v),
        }

        match parse_expr("func(set(a b c))") {
            Ok(Spanned {
                value: Expr::Function { .. },
                ..
            }) => {}
            v => panic!("expected function expr, got `{:?}`", v),
        }

        match parse_expr("//:tgt") {
            Ok(Spanned {
                value: Expr::String("//:tgt"),
                ..
            }) => {}
            v => panic!("expected '//:tgt', got `{:?}`", v),
        }

        Ok(())
    }

    fn run_tests<'a, O: Debug, F: FnMut(Span<'a>) -> IResult<Span<'a>, O, ()> + Copy>(
        mut parser: F,
        good_cases: &'a [&'a str],
        recoverable_cases: &'a [&'a str],
        nonrecoverable_cases: &'a [&'a str],
    ) {
        for case in good_cases {
            // good_cases the parser should consume the whole thing
            // maybe should consider splitting it into two sets of cases for the caller to
            // explicitly indicate the ones that should be all consuming
            all_consuming(parser)(Span::new(case)).unwrap_or_else(|err| {
                panic!(
                    "Expected successful, complete parse for input `{}`. Got error `{:?}`",
                    case, err
                )
            });
        }

        for case in recoverable_cases {
            assert_err(
                parser(Span::new(case)),
                &format!("Expected recoverable error for input `{}`.", case),
            );
        }

        for case in nonrecoverable_cases {
            assert_fail(
                parser(Span::new(case)),
                &format!("Expected hard failure for input `{}`.", case),
            );
        }
    }

    #[test]
    fn test_function() -> anyhow::Result<()> {
        run_tests(
            expr_function,
            &["a()", "a(a)", "a(a, b, c)", "a(a,   b,   c)"],
            // As long as we don't match "name(", it should be recoverable
            &["xyz", "", "(", ",a(", " a()", "a,a()"],
            // An error after the name + opening paren is non-recoverable
            &["a(", "a(a()", "a(abc, ", "a(abc, ()"],
        );
        Ok(())
    }

    #[test]
    fn test_set() -> anyhow::Result<()> {
        run_tests(
            expr_set,
            &["set(a)", "set(a b c)", "set(a   b   c)", "set(%s)", "set()"],
            // As long as we don't match "set(", it should be recoverable
            &["set", "", "(", "abc(", " set()", ",set()"],
            // An error after "set" + opening paren is non-recoverable
            &["set(", "set(a,)", "set(abc ", "set(abc ()"],
        );
        Ok(())
    }

    #[test]
    fn test_word() -> anyhow::Result<()> {
        run_tests(
            expr_word,
            &[
                "@cell//package:target",
                "'quoted'",
                "regex.*$#",
                r#""double^  -quoted""#,
                "%s",
                "%Ss",
            ],
            // there's not a lot of errors
            &["^word", "", ",word"],
            // unfinished quotes are unrecoverable.
            &["'unfinished quote", "'wrong quote end\""],
        );
        Ok(())
    }

    #[test]
    fn test_integer() -> anyhow::Result<()> {
        run_tests(expr_int, &["0", "1234"], &["w123", ".1", ""], &["0123"]);
        Ok(())
    }

    #[test]
    fn test_trailing_infix() -> anyhow::Result<()> {
        run_tests(
            trailing_infix,
            &[
                " + b",
                " + (b)",
                " ^ (b)",
                " - (b)- (c) -(d)-(e)",
                " intersect b",
                " union b",
                " except b",
            ],
            // there's not a lot of errors
            &["", "| b", " |b"],
            &[],
        );
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

/// Replace `%s` with remaining command line arguments which contain query literals.
pub const QUERY_PERCENT_S_PLACEHOLDER: &str = "%s";
/// Replace `%Ss` with query literals read from files from remaining command line arguments.
pub const QUERY_PERCENT_SS_PLACEHOLDER: &str = "%Ss";


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::ops::Deref;
use std::ops::RangeFrom;
use std::ops::RangeTo;
use std::str::CharIndices;
use std::str::Chars;

use dupe::Dupe;
use nom::Compare;
use nom::CompareResult;
use nom::InputIter;
use nom::InputLength;
use nom::InputTake;
use nom::Needed;
use nom::Offset;
use nom::Slice;
use nom::UnspecializedInput;

/// Similar to `LocatedSpan` from `nom_locate`, but without line number.
#[derive(Clone, Debug, PartialEq, Copy, Dupe)]
pub struct Span<'a> {
    offset: usize,
    fragment: &'a str,
}

impl<'a> Deref for Span<'a> {
    type Target = str;

    fn deref(&self) -> &str {
        self.fragment
    }
}

impl<'a> Span<'a> {
    pub(crate) fn new(s: &'a str) -> Span<'a> {
        Span {
            offset: 0,
            fragment: s,
        }
    }

    pub fn fragment(&self) -> &'a str {
        self.fragment
    }

    pub(crate) fn location_offset(&self) -> usize {
        self.offset
    }
}

impl<'a> UnspecializedInput for Span<'a> {}

impl<'a> InputLength for Span<'a> {
    fn input_len(&self) -> usize {
        self.fragment.input_len()
    }
}

impl<'a> Slice<RangeFrom<usize>> for Span<'a> {
    fn slice(&self, range: RangeFrom<usize>) -> Self {
        Span {
            offset: self.offset + range.start,
            fragment: &self.fragment[range.start..],
        }
    }
}

impl<'a> Slice<RangeTo<usize>> for Span<'a> {
    fn slice(&self, range: RangeTo<usize>) -> Self {
        Span {
            offset: self.offset,
            fragment: &self.fragment[..range.end],
        }
    }
}

impl<'a> InputTake for Span<'a> {
    fn take(&self, count: usize) -> Self {
        Span {
            offset: self.offset,
            fragment: &self.fragment[..count],
        }
    }

    fn take_split(&self, count: usize) -> (Self, Self) {
        let (a, b) = self.fragment.split_at(count);
        (
            Span {
                offset: self.offset + count,
                fragment: b,
            },
            Span {
                offset: self.offset,
                fragment: a,
            },
        )
    }
}

impl<'a> InputIter for Span<'a> {
    type Item = char;
    type Iter = CharIndices<'a>;
    type IterElem = Chars<'a>;

    fn iter_indices(&self) -> Self::Iter {
        self.fragment.char_indices()
    }

    fn iter_elements(&self) -> Self::IterElem {
        self.fragment.chars()
    }

    fn position<P>(&self, predicate: P) -> Option<usize>
    where
        P: Fn(Self::Item) -> bool,
    {
        self.fragment.position(predicate)
    }

    fn slice_index(&self, count: usize) -> Result<usize, Needed> {
        self.fragment.slice_index(count)
    }
}

impl<'a> Offset for Span<'a> {
    fn offset(&self, second: &Self) -> usize {
        let fst = self.offset;
        let snd = second.offset;

        snd - fst
    }
}

impl<'a> Compare<&'a str> for Span<'a> {
    fn compare(&self, t: &'a str) -> CompareResult {
        self.fragment.compare(t)
    }

    fn compare_no_case(&self, t: &'a str) -> CompareResult {
        self.fragment.compare_no_case(t)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::ops::Range;

use buck2_util::truncate::truncate;
use derive_more::Display;

/// Pairs some output value with a span covering where the value came from. Unlike nom_locate, we drop
/// the references into the input so that this has a 'static lifetime. To be useful, it'll need to be
/// recombined with the input string.
#[derive(Debug, Display)]
#[display(fmt = "{}", value)]
pub struct Spanned<T> {
    pub position: Range<usize>,
    pub value: T,
}

impl<T> Spanned<T> {
    pub fn map_res<L, R, F: FnOnce(&T) -> Result<L, R>>(
        &self,
        func: F,
    ) -> Result<Spanned<L>, Spanned<R>> {
        self.span(func(&self.value))
    }

    pub fn span<O, E>(&self, res: Result<O, E>) -> Result<Spanned<O>, Spanned<E>> {
        match res {
            Ok(value) => Ok(Spanned {
                position: self.position.clone(),
                value,
            }),
            Err(value) => Err(Spanned {
                position: self.position.clone(),
                value,
            }),
        }
    }

    pub async fn async_map_res<
        'a,
        L,
        R,
        Fut: std::future::Future<Output = Result<L, R>>,
        F: FnOnce(&'a T) -> Fut + 'a,
    >(
        &'a self,
        func: F,
    ) -> Result<Spanned<L>, Spanned<R>> {
        let position = self.position.clone();
        match func(&self.value).await {
            Ok(value) => Ok(Spanned { position, value }),
            Err(value) => Err(Spanned { position, value }),
        }
    }

    pub fn into_map_res<L, R, F: FnOnce(T) -> Result<L, R>>(
        self,
        func: F,
    ) -> Result<Spanned<L>, Spanned<R>> {
        match func(self.value) {
            Ok(value) => Ok(Spanned {
                position: self.position,
                value,
            }),
            Err(value) => Err(Spanned {
                position: self.position,
                value,
            }),
        }
    }

    pub async fn async_into_map_res<
        L,
        R,
        Fut: std::future::Future<Output = Result<L, R>>,
        F: FnOnce(T) -> Fut,
    >(
        self,
        func: F,
    ) -> Result<Spanned<L>, Spanned<R>> {
        match func(self.value).await {
            Ok(value) => Ok(Spanned {
                position: self.position,
                value,
            }),
            Err(value) => Err(Spanned {
                position: self.position,
                value,
            }),
        }
    }

    // Constructs a 2-line string that prints the line where the span occurs and then a line below that identifying the span.
    pub fn get_err_context(&self, input: &str) -> String {
        // TODO(cjhopman): This assumes that the input is a single line, it should (maybe) handle multi-line inputs.
        // TODO(cjhopman): This should cut off the beginning and/or end of long lines and focus around the span.
        // TODO(cjhopman): Consider using annotate-snippets like we do in starlark.
        let (rest, end) = input.split_at(self.position.end);
        let (start, inner) = rest.split_at(self.position.start);

        let inner = truncate(inner, 80);

        let pointer = match inner.len() {
            l if l < 3 => format!("{}^", " ".repeat(start.len())),
            l => format!("{}^{}^", " ".repeat(start.len()), "-".repeat(l - 2)),
        };
        format!("\n    {}{}{}\n    {}\n", start, inner, end, pointer)
    }
}

impl<T> std::ops::Deref for Spanned<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_truncate_span_in_middle() {
        let error = "rdeps(set(fbcode//buck2/... fbsource//xplat/buck2/..., fbsource//fbobjc/buck2/...), fbcode//buck2/cli:buck2)";
        let span = Spanned {
            position: 0..error.len(),
            value: false,
        };
        let context = span.get_err_context(error);
        let context_lines: Vec<&str> = context.split('\n').collect();
        assert_eq!(
            context_lines,
            [
                "",
                "    rdeps(set(fbcode//buck2/... fbsour<<omitted>>ck2/...), fbcode//buck2/cli:buck2)",
                "    ^-----------------------------------------------------------------------------^",
                "",
            ]
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::str::FromStr;

use allocative::Allocative;
use buck2_common::legacy_configs::LegacyBuckConfig;

static BUCK2_RE_CLIENT_CFG_SECTION: &str = "buck2_re_client";

/// We put functions here that both things need to implement for code that isn't gated behind a
/// fbcode_build or not(fbcode_build)
pub trait RemoteExecutionStaticMetadataImpl: Sized {
    fn from_legacy_config(legacy_config: &LegacyBuckConfig) -> anyhow::Result<Self>;
    fn cas_semaphore_size(&self) -> usize;
}

#[allow(unused)]
mod fbcode {
    use super::*;

    /// Metadata that doesn't change between executions
    #[derive(Clone, Debug, Default, Allocative)]
    pub struct RemoteExecutionStaticMetadata {
        pub cas_address: Option<String>,
        pub cas_connection_count: i32,
        pub action_cache_address: Option<String>,
        pub action_cache_connection_count: i32,
        pub engine_address: Option<String>,
        pub engine_connection_count: i32,

        pub verbose_logging: bool,

        pub use_manifold_rich_client: bool,
        pub use_zippy_rich_client: bool,
        pub use_p2p: bool,

        pub cas_thread_count: i32,

        pub rich_client_channels_per_blob: Option<i32>,
        pub rich_client_attempt_timeout_ms: Option<i32>,
        pub rich_client_retries_count: Option<i32>,
        pub force_enable_deduplicate_find_missing: Option<bool>,

        pub features_config_path: Option<String>,

        // curl reactor
        pub curl_reactor_max_number_of_retries: Option<i32>,
        pub curl_reactor_connection_timeout_ms: Option<i32>,
        pub curl_reactor_request_timeout_ms: Option<i32>,
    }

    impl RemoteExecutionStaticMetadataImpl for RemoteExecutionStaticMetadata {
        fn from_legacy_config(legacy_config: &LegacyBuckConfig) -> anyhow::Result<Self> {
            Ok(Self {
                cas_address: legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "cas_address")?,
                cas_connection_count: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "cas_connection_count")?
                    .unwrap_or(16),
                action_cache_address: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "action_cache_address")?,
                action_cache_connection_count: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "action_cache_connection_count")?
                    .unwrap_or(4),
                engine_address: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "engine_address")?,
                engine_connection_count: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "engine_connection_count")?
                    .unwrap_or(4),
                verbose_logging: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "verbose_logging")?
                    .unwrap_or(false),
                cas_thread_count: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "cas_thread_count")?
                    .unwrap_or(4),
                use_manifold_rich_client: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "use_manifold_rich_client_new")?
                    .unwrap_or(true),
                use_zippy_rich_client: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "use_zippy_rich_client")?
                    .unwrap_or(false),
                use_p2p: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "use_p2p")?
                    .unwrap_or(false),
                rich_client_channels_per_blob: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "rich_client_channels_per_blob")?,
                rich_client_attempt_timeout_ms: legacy_config.parse(
                    BUCK2_RE_CLIENT_CFG_SECTION,
                    "rich_client_attempt_timeout_ms",
                )?,
                rich_client_retries_count: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "rich_client_retries_count")?,
                force_enable_deduplicate_find_missing: legacy_config.parse(
                    BUCK2_RE_CLIENT_CFG_SECTION,
                    "force_enable_deduplicate_find_missing",
                )?,
                features_config_path: legacy_config
                    .parse(BUCK2_RE_CLIENT_CFG_SECTION, "features_config_path")?,
                curl_reactor_max_number_of_retries: legacy_config.parse(
                    BUCK2_RE_CLIENT_CFG_SECTION,
                    "curl_reactor_max_number_of_retries",
                )?,
                curl_reactor_connection_timeout_ms: legacy_config.parse(
                    BUCK2_RE_CLIENT_CFG_SECTION,
                    "curl_reactor_connection_timeout_ms",
                )?,
                curl_reactor_request_timeout_ms: legacy_config.parse(
                    BUCK2_RE_CLIENT_CFG_SECTION,
                    "curl_reactor_request_timeout_ms",
                )?,
            })
        }

        fn cas_semaphore_size(&self) -> usize {
            self.cas_connection_count as usize * 30
        }
    }
}

#[allow(unused)]
mod not_fbcode {
    use super::*;

    /// Metadata that doesn't change between executions
    #[derive(Clone, Debug, Default, Allocative)]
    pub struct RemoteExecutionStaticMetadata(pub Buck2OssReConfiguration);

    impl RemoteExecutionStaticMetadataImpl for RemoteExecutionStaticMetadata {
        fn from_legacy_config(legacy_config: &LegacyBuckConfig) -> anyhow::Result<Self> {
            Ok(Self(Buck2OssReConfiguration::from_legacy_config(
                legacy_config,
            )?))
        }

        fn cas_semaphore_size(&self) -> usize {
            // FIXME: make this configurable?
            1024
        }
    }
}

/// A configuration used only in our OSS builds. We still compile this always, which lets us
/// gate less code behind fbcode_build.
#[derive(Clone, Debug, Default, Allocative)]
pub struct Buck2OssReConfiguration {
    /// Address for RBE Content Addresable Storage service (including bytestream uploads service).
    pub cas_address: Option<String>,
    /// Address for RBE Engine service (including capabilities service).
    pub engine_address: Option<String>,
    /// Address for RBE Action Cache service.
    pub action_cache_address: Option<String>,
    /// Whether to use TLS to interact with remote execution.
    pub tls: bool,
    /// Path to a CA certificates bundle. This must be PEM-encoded. If none is set, a default
    /// bundle will be used.
    ///
    /// This can contain environment variables using shell interpolation syntax (i.e. $VAR). They
    /// will be substituted before using the value.
    pub tls_ca_certs: Option<String>,
    /// Path to a client certificate (and intermediate chain), as well as its associated private
    /// key. This must be PEM-encoded.
    ///
    /// This can contain environment variables using shell interpolation syntax (i.e. $VAR). They
    /// will be substituted before using the value.
    pub tls_client_cert: Option<String>,
    /// HTTP headers to inject in all requests to RE. This is a comma-separated list of `Header:
    /// Value` pairs. Minimal validation of those headers is done here.
    ///
    /// This can contain environment variables using shell interpolation syntax (i.e. $VAR). They
    /// will be substituted before using the value.
    pub http_headers: Vec<HttpHeader>,
    /// Whether to query capabilities from the RBE backend.
    pub capabilities: Option<bool>,
    /// The instance name to use in requests.
    pub instance_name: Option<String>,
}

#[derive(Clone, Debug, Default, Allocative)]
pub struct HttpHeader {
    pub key: String,
    pub value: String,
}

impl FromStr for HttpHeader {
    type Err = anyhow::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let mut iter = s.split(':');
        match (iter.next(), iter.next(), iter.next()) {
            (Some(key), Some(value), None) => Ok(Self {
                key: key.trim().to_owned(),
                value: value.trim().to_owned(),
            }),
            _ => Err(anyhow::anyhow!(
                "Invalid header (expect exactly one `:`): `{}`",
                s
            )),
        }
    }
}

impl Buck2OssReConfiguration {
    pub fn from_legacy_config(legacy_config: &LegacyBuckConfig) -> anyhow::Result<Self> {
        // this is used for all three services by default, if given; if one of
        // them has an explicit address given as well though, use that instead
        let default_address: Option<String> =
            legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "address")?;

        Ok(Self {
            cas_address: legacy_config
                .parse(BUCK2_RE_CLIENT_CFG_SECTION, "cas_address")?
                .or(default_address.clone()),
            engine_address: legacy_config
                .parse(BUCK2_RE_CLIENT_CFG_SECTION, "engine_address")?
                .or(default_address.clone()),
            action_cache_address: legacy_config
                .parse(BUCK2_RE_CLIENT_CFG_SECTION, "action_cache_address")?
                .or(default_address),
            tls: legacy_config
                .parse(BUCK2_RE_CLIENT_CFG_SECTION, "tls")?
                .unwrap_or(true),
            tls_ca_certs: legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "tls_ca_certs")?,
            tls_client_cert: legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "tls_client_cert")?,
            http_headers: legacy_config
                .parse_list(BUCK2_RE_CLIENT_CFG_SECTION, "http_headers")?
                .unwrap_or_default(), // Empty list is as good None.
            capabilities: legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "capabilities")?,
            instance_name: legacy_config.parse(BUCK2_RE_CLIENT_CFG_SECTION, "instance_name")?,
        })
    }
}

#[cfg(fbcode_build)]
pub use fbcode::RemoteExecutionStaticMetadata;
#[cfg(not(fbcode_build))]
pub use not_fbcode::RemoteExecutionStaticMetadata;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::sync::Arc;

use buck2_cli_proto::ClientContext;
use buck2_event_observer::dice_state::DiceState;
use buck2_event_observer::pending_estimate::pending_estimate;
use buck2_event_observer::span_tracker;
use buck2_event_observer::span_tracker::RootData;
use buck2_event_observer::span_tracker::Roots;
use buck2_events::dispatch::EventDispatcher;
use buck2_events::span::SpanId;
use buck2_events::BuckEvent;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use parking_lot::MutexGuard;
use tokio::sync::oneshot;

static ACTIVE_COMMANDS: Lazy<Mutex<HashMap<TraceId, ActiveCommandHandle>>> =
    Lazy::new(|| Mutex::new(HashMap::new()));

/// Return the active commands, if you can access them.
pub fn try_active_commands() -> Option<HashMap<TraceId, ActiveCommandHandle>> {
    // Note that this function is accessed during panic, so have to be super careful
    Some(ACTIVE_COMMANDS.try_lock()?.clone())
}

pub fn active_commands() -> MutexGuard<'static, HashMap<TraceId, ActiveCommandHandle>> {
    ACTIVE_COMMANDS.lock()
}

/// Broadcasts an instant event, returns whether any subscribers were connected.
pub fn broadcast_instant_event<E: Into<buck2_data::instant_event::Data> + Clone>(
    event: &E,
) -> bool {
    let mut has_subscribers = false;

    for cmd in ACTIVE_COMMANDS.lock().values() {
        cmd.dispatcher.instant_event(event.clone());
        has_subscribers = true;
    }

    has_subscribers
}

pub fn broadcast_shutdown(shutdown: &buck2_data::DaemonShutdown) {
    for cmd in ACTIVE_COMMANDS.lock().values() {
        cmd.notify_shutdown(shutdown.clone());
    }
}

/// Allows interactions with commands found via active_commands().
#[derive(Clone, Dupe)]
pub struct ActiveCommandHandle {
    /// A channel to send notifications to this command.
    dispatcher: EventDispatcher,

    /// A separate channel to broadcast shutdown events. This is separate from the EventDispatcher
    /// because we want to allow shutdown events to jump the queue.
    daemon_shutdown_channel: Arc<Mutex<Option<oneshot::Sender<buck2_data::DaemonShutdown>>>>,

    /// State for this command. This is used to expose what this command is doing to other clients.
    state: Arc<ActiveCommandState>,
}

impl ActiveCommandHandle {
    fn notify_shutdown(&self, shutdown: buck2_data::DaemonShutdown) {
        let channel = self.daemon_shutdown_channel.lock().take();

        if let Some(channel) = channel {
            let _ignored = channel.send(shutdown); // Nothing to do if receiver hung up.
        }
    }

    pub fn state(&self) -> &ActiveCommandState {
        self.state.as_ref()
    }
}

pub struct ActiveCommandDropGuard {
    trace_id: TraceId,
}

impl Drop for ActiveCommandDropGuard {
    fn drop(&mut self) {
        ACTIVE_COMMANDS.lock().remove(&self.trace_id);
    }
}

/// A handle to the stats for this command. We use this to broadcast state about this command.
pub struct ActiveCommandState {
    #[allow(unused)]
    pub argv: Vec<String>,

    spans: Mutex<SpansSnapshot>,
}

impl ActiveCommandState {
    pub fn spans(&self) -> SpansSnapshot {
        *self.spans.lock()
    }

    fn new(argv: Vec<String>) -> Self {
        Self {
            argv,
            spans: Mutex::new(SpansSnapshot::default()),
        }
    }
}

#[derive(PartialEq, Debug, Default, Copy, Clone, Dupe)]
pub struct SpansSnapshot {
    pub open: u64,
    pub closed: u64,
    pub pending: u64,
}

/// A wrapper around ActiveCommandState that allows 1 client to write to it.
pub struct ActiveCommandStateWriter {
    /// Maps a SpanId to whether it is a root (i.e. no parent)
    roots: Roots<Arc<BuckEvent>>,
    non_roots: HashSet<SpanId>,
    dice_state: DiceState,
    closed: u64,
    shared: Arc<ActiveCommandState>,
}

impl ActiveCommandStateWriter {
    fn new(shared: Arc<ActiveCommandState>) -> Self {
        Self {
            roots: Roots::default(),
            non_roots: HashSet::new(),
            dice_state: DiceState::new(),
            closed: 0,
            shared,
        }
    }

    pub fn peek_event(&mut self, buck_event: &BuckEvent) {
        use buck2_data::buck_event::Data::*;

        let mut changed = false;

        match buck_event.data() {
            SpanStart(..) => {
                let span_id = match buck_event.span_id() {
                    Some(id) => id,
                    None => return,
                };

                if !span_tracker::is_span_shown(buck_event) {
                    return;
                }

                let is_root = buck_event.parent_id().map_or(true, |id| {
                    !self.roots.contains(id) && !self.non_roots.contains(&id)
                });

                if is_root {
                    self.roots.insert(span_id, false, RootData::new(buck_event));
                    changed = true;
                } else {
                    self.non_roots.insert(span_id);
                }
            }
            SpanEnd(..) => {
                let span_id = match buck_event.span_id() {
                    Some(id) => id,
                    None => return,
                };

                // If it's a root, then we increment closed.
                if self.roots.remove(span_id).is_some() {
                    self.closed += 1;
                    changed = true;
                } else {
                    self.non_roots.remove(&span_id);
                }
            }
            Instant(instant) => {
                use buck2_data::instant_event::Data::*;

                match instant.data.as_ref() {
                    Some(DiceStateSnapshot(snapshot)) => {
                        self.dice_state.update(snapshot);
                        changed = true;
                    }
                    _ => {}
                }
            }
            _ => {}
        }

        if changed {
            let open = self.roots.len() as u64;
            let pending = pending_estimate(&self.roots, &self.dice_state);

            *self.shared.spans.lock() = SpansSnapshot {
                open,
                closed: self.closed,
                pending,
            };
        }
    }
}

pub struct ActiveCommand {
    pub guard: ActiveCommandDropGuard,
    pub state: ActiveCommandStateWriter,
    pub daemon_shutdown_channel: oneshot::Receiver<buck2_data::DaemonShutdown>,
}

impl ActiveCommand {
    pub fn new(event_dispatcher: &EventDispatcher, client_ctx: &ClientContext) -> Self {
        let (sender, receiver) = oneshot::channel();

        let state = Arc::new(ActiveCommandState::new(client_ctx.sanitized_argv.clone()));

        let trace_id = event_dispatcher.trace_id().dupe();
        let result = {
            // Scope the guard so it's locked as little as possible
            let mut active_commands = ACTIVE_COMMANDS.lock();

            let existing_active_commands = if active_commands.len() > 1 {
                Some(active_commands.clone())
            } else {
                None
            };

            active_commands.insert(
                trace_id.dupe(),
                ActiveCommandHandle {
                    dispatcher: event_dispatcher.dupe(),
                    daemon_shutdown_channel: Arc::new(Mutex::new(Some(sender))),
                    state: state.dupe(),
                },
            );

            existing_active_commands
        };

        if let Some(commands) = result {
            // Notify our command it is running concurrently with others.
            event_dispatcher.instant_event(buck2_data::TagEvent {
                tags: commands
                    .keys()
                    .map(|cmd| format!("concurrent_commands:{}", cmd))
                    .collect(),
            });

            // Notify other commands that they are concurrent with ours.
            for cmd in commands.values() {
                cmd.dispatcher.instant_event(buck2_data::TagEvent {
                    tags: vec![format!("concurrent_commands:{}", trace_id)],
                });
            }
        }

        Self {
            guard: ActiveCommandDropGuard { trace_id },
            daemon_shutdown_channel: receiver,
            state: ActiveCommandStateWriter::new(state),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::time::SystemTime;

    use super::*;

    #[test]
    fn test_active_command_state() {
        let mut writer =
            ActiveCommandStateWriter::new(Arc::new(ActiveCommandState::new(Vec::new())));

        let root = SpanId::new();
        let child = SpanId::new();
        let trace = TraceId::new();

        writer.peek_event(&BuckEvent::new(
            SystemTime::now(),
            trace.clone(),
            Some(root),
            None,
            buck2_data::SpanStartEvent {
                data: Some(buck2_data::AnalysisStart::default().into()),
            }
            .into(),
        ));

        assert_eq!(
            writer.shared.spans(),
            SpansSnapshot {
                open: 1,
                closed: 0,
                pending: 0
            }
        );

        writer.peek_event(&BuckEvent::new(
            SystemTime::now(),
            trace.clone(),
            Some(child),
            Some(root),
            buck2_data::SpanStartEvent {
                data: Some(buck2_data::AnalysisStageStart::default().into()),
            }
            .into(),
        ));

        assert_eq!(
            writer.shared.spans(),
            SpansSnapshot {
                open: 1,
                closed: 0,
                pending: 0
            }
        );

        writer.peek_event(&BuckEvent::new(
            SystemTime::now(),
            trace.clone(),
            Some(child),
            Some(root),
            buck2_data::SpanEndEvent {
                data: Some(buck2_data::AnalysisStageEnd::default().into()),
                ..Default::default()
            }
            .into(),
        ));

        assert_eq!(
            writer.shared.spans(),
            SpansSnapshot {
                open: 1,
                closed: 0,
                pending: 0
            }
        );

        writer.peek_event(&BuckEvent::new(
            SystemTime::now(),
            trace.clone(),
            Some(root),
            None,
            buck2_data::SpanEndEvent {
                data: Some(buck2_data::AnalysisEnd::default().into()),
                ..Default::default()
            }
            .into(),
        ));

        assert_eq!(
            writer.shared.spans(),
            SpansSnapshot {
                open: 0,
                closed: 1,
                pending: 0
            }
        );

        writer.peek_event(&BuckEvent::new(
            SystemTime::now(),
            trace,
            None,
            None,
            buck2_data::InstantEvent {
                data: Some(
                    buck2_data::DiceStateSnapshot {
                        key_states: {
                            let mut map = HashMap::new();
                            map.insert(
                                "BuildKey".to_owned(),
                                buck2_data::DiceKeyState {
                                    started: 4,
                                    finished: 2,
                                    check_deps_started: 0,
                                    check_deps_finished: 0,
                                },
                            );
                            map
                        },
                    }
                    .into(),
                ),
            }
            .into(),
        ));

        assert_eq!(
            writer.shared.spans(),
            SpansSnapshot {
                open: 0,
                closed: 1,
                pending: 2
            }
        );
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use anyhow::Context;
use async_trait::async_trait;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::NoPartialResult;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use chrono::TimeZone;
use chrono::Utc;
use dice::DiceTransaction;

use crate::ctx::ServerCommandContext;

pub(crate) async fn clean_stale_command(
    ctx: &ServerCommandContext<'_>,
    partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
    req: buck2_cli_proto::CleanStaleRequest,
) -> anyhow::Result<buck2_cli_proto::CleanStaleResponse> {
    run_server_command(
        CleanStaleServerCommand { req },
        ctx,
        partial_result_dispatcher,
    )
    .await
}

struct CleanStaleServerCommand {
    req: buck2_cli_proto::CleanStaleRequest,
}

#[async_trait]
impl ServerCommandTemplate for CleanStaleServerCommand {
    type StartEvent = buck2_data::CleanCommandStart;
    type EndEvent = buck2_data::CleanCommandEnd;
    type Response = buck2_cli_proto::CleanStaleResponse;
    type PartialResult = NoPartialResult;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        _partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        _ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        server_ctx
            .cancellation_context()
            .critical_section(async move || {
                let deferred_materializer = server_ctx.materializer();
                deferred_materializer.log_materializer_state(server_ctx.events());

                let extension = deferred_materializer
                    .as_deferred_materializer_extension()
                    .context("Deferred materializer is not in use")?;

                let keep_since_time = Utc
                    .timestamp_opt(self.req.keep_since_time, 0)
                    .single()
                    .context("Invalid timestamp")?;

                extension
                    .clean_stale_artifacts(keep_since_time, self.req.dry_run, self.req.tracked_only)
                    .await
                    .context("Failed to clean stale artifacts.")
            })
            .await
    }

    fn is_success(&self, _response: &Self::Response) -> bool {
        // No response if we failed.
        true
    }

    fn end_event(&self, response: &anyhow::Result<Self::Response>) -> Self::EndEvent {
        let clean_stale_stats = if let Ok(res) = response {
            res.stats.clone()
        } else {
            None
        };
        buck2_data::CleanCommandEnd { clean_stale_stats }
    }

    fn exclusive_command_name(&self) -> Option<String> {
        Some("clean --stale".to_owned())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashSet;

use anyhow::Context;
use buck2_cli_proto::config_override::ConfigType;
use buck2_cli_proto::ConfigOverride;
use buck2_common::legacy_configs::cells::BuckConfigBasedCells;
use buck2_common::legacy_configs::LegacyBuckConfigs;
use buck2_common::legacy_configs::LegacyConfigCmdArg;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;

fn config_type_from_i32(value: i32) -> anyhow::Result<ConfigType> {
    ConfigType::from_i32(value).with_context(|| {
        format!(
            "Unknown ConfigType enum value `{}` when trying to deserialize",
            value
        )
    })
}

fn get_legacy_config_args<'a, Iter: IntoIterator<Item = &'a ConfigOverride>>(
    config_overrides: Iter,
) -> anyhow::Result<Vec<LegacyConfigCmdArg>> {
    config_overrides
        .into_iter()
        .map(
            |config_arg| match config_type_from_i32(config_arg.config_type)? {
                ConfigType::Value => Ok(LegacyConfigCmdArg::Flag(
                    config_arg.config_override.to_owned(),
                )),
                ConfigType::File => Ok(LegacyConfigCmdArg::UnresolvedFile(
                    config_arg.config_override.to_owned(),
                )),
            },
        )
        .collect::<anyhow::Result<Vec<LegacyConfigCmdArg>>>()
}

/// Read the configs, returning the cell resolver and the legacy configs
pub fn parse_legacy_cells<'a, Iter: IntoIterator<Item = &'a ConfigOverride>>(
    config_overrides: Iter,
    cwd: &ProjectRelativePath,
    fs: &ProjectRoot,
) -> anyhow::Result<(CellResolver, LegacyBuckConfigs, HashSet<AbsNormPathBuf>)> {
    let config_values = get_legacy_config_args(config_overrides)?;
    // TODO: We do not need to reparse _all_ configs, instead we just need to
    // overlay any custom configs for the current build command on top of
    // the base configs derived from the config files. This requires us to
    // store the base configs + overlaid ones separately, so we can cheaply
    // recompose.
    let res = BuckConfigBasedCells::parse_with_config_args(fs, &config_values, cwd)?;
    Ok((res.cell_resolver, res.configs_by_name, res.config_paths))
}