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

#![feature(async_closure)]
#![feature(try_blocks)]
#![feature(provide_any)]

use std::sync::Once;

pub mod analysis;
pub mod aquery;
pub mod cquery;
pub mod dice;
pub mod frontend;
pub mod uquery;

pub fn init_late_bindings() {
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        analysis::environment::init_classpath_for_targets();
        analysis::environment::init_query_functions();
        analysis::eval::init_eval_analysis_query();
        aquery::find_matching_action::init_find_matching_action();
        frontend::init_query_frontend();
        cquery::bxl::init_new_bxl_cquery_functions();
        uquery::bxl::init_new_bxl_uquery_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.
 */

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

use allocative::Allocative;
use async_trait::async_trait;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_common::result::SharedResult;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_frontend::ConfiguredTargetNodeCalculation;
use buck2_node::nodes::configured_ref::ConfiguredGraphNodeRef;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use dupe::IterDupedExt;
use dupe::OptionDupedExt;
use futures::FutureExt;
use indexmap::IndexMap;
use more_futures::cancellation::CancellationContext;

use crate::analysis::environment::get_from_template_placeholder_info;
use crate::analysis::environment::ConfiguredGraphQueryEnvironmentDelegate;

pub struct AnalysisDiceQueryDelegate<'c> {
    pub ctx: &'c DiceComputations,
}

impl<'c> AnalysisDiceQueryDelegate<'c> {
    pub(crate) fn ctx(&self) -> &DiceComputations {
        self.ctx
    }
}

pub struct AnalysisConfiguredGraphQueryDelegate<'a> {
    pub dice_query_delegate: Arc<AnalysisDiceQueryDelegate<'a>>,
    pub resolved_literals: HashMap<String, ConfiguredTargetNode>,
}

#[async_trait]
impl<'a> ConfiguredGraphQueryEnvironmentDelegate for AnalysisConfiguredGraphQueryDelegate<'a> {
    fn eval_literal(&self, literal: &str) -> anyhow::Result<ConfiguredTargetNode> {
        self.resolved_literals
            .get(literal)
            .duped()
            .ok_or_else(|| anyhow::anyhow!(""))
    }

    async fn get_targets_from_template_placeholder_info(
        &self,
        template_name: &'static str,
        targets: TargetSet<ConfiguredGraphNodeRef>,
    ) -> anyhow::Result<TargetSet<ConfiguredGraphNodeRef>> {
        #[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
        #[display(fmt = "template_placeholder_info_query({})", template_name)]
        struct TemplatePlaceholderInfoQueryKey {
            template_name: &'static str,
            // Use `ConfiguredTargetLabel` instead of `ConfiguredGraphNodeRef` here because `ConfiguredGraphNodeRef`
            // only computes Hash and PartialEq based on the label. If we use `ConfiguredGraphNodeRef` directly we
            // may cache stale ConfiguredTargetNodes and end up with a bug like T133069783.
            targets: Arc<Vec<ConfiguredTargetLabel>>,
        }

        #[async_trait]
        impl Key for TemplatePlaceholderInfoQueryKey {
            type Value = SharedResult<Arc<TargetSet<ConfiguredGraphNodeRef>>>;

            async fn compute(
                &self,
                ctx: &DiceComputations,
                _cancellation: &CancellationContext,
            ) -> Self::Value {
                let (targets, label_to_artifact) = futures::future::try_join(
                    futures::future::try_join_all(self.targets.iter().map(|target| {
                        ctx.get_configured_target_node(target)
                            .map(|res| res?.require_compatible())
                    })),
                    get_from_template_placeholder_info(
                        ctx,
                        self.template_name,
                        self.targets.iter().duped(),
                    ),
                )
                .await?;

                let targets: TargetSet<_> =
                    targets.into_iter().map(ConfiguredGraphNodeRef).collect();
                let targets = find_target_nodes(targets, label_to_artifact)?;
                Ok(Arc::new(targets))
            }

            fn equality(_: &Self::Value, _: &Self::Value) -> bool {
                // result is not comparable
                false
            }
        }

        let targets: Vec<_> = targets
            .into_iter()
            .map(|target| target.label().dupe())
            .collect();
        let targets = self
            .dice_query_delegate
            .ctx()
            .compute(&TemplatePlaceholderInfoQueryKey {
                template_name,
                targets: Arc::new(targets),
            })
            .await??;

        // TODO(scottcao): Make all query functions return an Arc as an output so we can avoid making an unnecessary
        // clone here
        Ok(targets.as_ref().clone())
    }
}

/// Finds the nodes for a list of target labels within the deps of the provided targets.
///
/// It may seem like if we have ConfiguredTargetLabel we should just be able to lookup the
/// nodes directly, but that would require going through dice and then dice would record
/// dependencies on all the nodes that we lookup. It's common for these queries to operate
/// over inputs that are aggregated as data flows up the graph and it's important that we
/// don't inadvertently cause flattening of those sets.
fn find_target_nodes(
    targets: TargetSet<ConfiguredGraphNodeRef>,
    label_to_artifact: IndexMap<ConfiguredTargetLabel, Artifact>,
) -> anyhow::Result<TargetSet<ConfiguredGraphNodeRef>> {
    let mut queue: VecDeque<_> = targets.iter().duped().collect();
    let mut seen = targets;
    let mut result = TargetSet::new();

    while let Some(target) = queue.pop_front() {
        if label_to_artifact.contains_key(target.label()) {
            result.insert(target.dupe());
        }
        if result.len() == label_to_artifact.len() {
            return Ok(result);
        }
        for dep in target.0.target_deps() {
            let dep = ConfiguredGraphNodeRef(dep.dupe());
            if seen.insert(dep.dupe()) {
                queue.push_back(dep);
            }
        }
    }

    Ok(result)
}


/*
 * 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 std::collections::VecDeque;
use std::fmt;
use std::fmt::Debug;
use std::marker::PhantomData;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::artifact_type::OutputArtifact;
use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
use buck2_build_api::artifact_groups::deferred::DeferredTransitiveSetData;
use buck2_build_api::artifact_groups::deferred::TransitiveSetKey;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::artifact_groups::ResolvedArtifactGroup;
use buck2_build_api::deferred::calculation::DeferredCalculation;
use buck2_build_api::deferred::types::DeferredValueReady;
use buck2_build_api::interpreter::rule_defs::artifact_tagging::ArtifactTag;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use buck2_build_api::interpreter::rule_defs::cmd_args::SimpleCommandLineArtifactVisitor;
use buck2_build_api::interpreter::rule_defs::provider::builtin::template_placeholder_info::TemplatePlaceholderInfo;
use buck2_build_api::interpreter::rule_defs::transitive_set::TransitiveSet;
use buck2_build_api::query::analysis::CLASSPATH_FOR_TARGETS;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::ProvidersName;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_ref::ConfiguredGraphNodeRef;
use buck2_node::nodes::configured_ref::ConfiguredGraphNodeRefLookup;
use buck2_node::query::query_functions::CONFIGURED_GRAPH_QUERY_FUNCTIONS;
use buck2_query::query::environment::LabeledNode;
use buck2_query::query::environment::QueryEnvironment;
use buck2_query::query::syntax::simple::eval::error::QueryError;
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::eval::values::QueryValue;
use buck2_query::query::syntax::simple::functions::helpers::QueryBinaryOp;
use buck2_query::query::syntax::simple::functions::helpers::QueryFunction;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use buck2_query::query::syntax::simple::functions::QueryFunctions;
use buck2_query::query::traversal::async_depth_limited_traversal;
use buck2_query::query::traversal::async_fast_depth_first_postorder_traversal;
use buck2_query::query::traversal::AsyncTraversalDelegate;
use buck2_query::query_module;
use buck2_query_parser::BinaryOp;
use dice::DiceComputations;
use dupe::Dupe;
use indexmap::IndexMap;
use thiserror::Error;

#[derive(Debug, Error)]
enum AnalysisQueryError {
    #[error("file literals aren't supported in query attributes (got `{0}`)")]
    FileLiteralsNotAllowed(String),
    #[error(
        "template_placeholder_info `{0}` of target `{1}` used in query attributes had artifact (`{2}`) not produced by a target, only target-produced artifacts supported here"
    )]
    NonTargetBoundArtifact(String, ConfiguredTargetLabel, Artifact),
}

#[async_trait]
pub trait ConfiguredGraphQueryEnvironmentDelegate: Send + Sync {
    fn eval_literal(&self, literal: &str) -> anyhow::Result<ConfiguredTargetNode>;

    async fn get_targets_from_template_placeholder_info(
        &self,
        template_name: &'static str,
        targets: TargetSet<ConfiguredGraphNodeRef>,
    ) -> anyhow::Result<TargetSet<ConfiguredGraphNodeRef>>;
}

pub struct ConfiguredGraphQueryEnvironment<'a> {
    delegate: &'a dyn ConfiguredGraphQueryEnvironmentDelegate,
}

#[derive(Debug)]
struct ConfiguredGraphFunctions<'a>(PhantomData<&'a ()>);
#[query_module(ConfiguredGraphQueryEnvironment<'a>)]
impl<'a> ConfiguredGraphFunctions<'a> {
    async fn classpath(
        &self,
        env: &ConfiguredGraphQueryEnvironment<'a>,
        targets: TargetSet<ConfiguredGraphNodeRef>,
        depth: Option<u64>,
    ) -> Result<QueryValue<ConfiguredGraphNodeRef>, QueryError> {
        // if depth param is provided and it is not equal to 1, then it's not supported
        let mut run_first_order_classpath = false;
        if let Some(depth_int) = depth.map(|v| v as i32) {
            run_first_order_classpath = depth_int == 1;
            if !run_first_order_classpath {
                return Err(QueryError::InvalidDepth(depth_int));
            }
        }

        let template_name = if run_first_order_classpath {
            "first_order_classpath"
        } else {
            "classpath_including_targets_with_no_output"
        };

        let targets = env
            .get_targets_from_template_placeholder_info(template_name, targets)
            .await?;
        Ok(targets.into())
    }
}

impl<'a> ConfiguredGraphQueryEnvironment<'a> {
    pub fn new(delegate: &'a dyn ConfiguredGraphQueryEnvironmentDelegate) -> Self {
        Self { delegate }
    }

    pub fn functions() -> impl QueryFunctions<Env = ConfiguredGraphQueryEnvironment<'a>> {
        struct Functions<'a> {
            defaults: DefaultQueryFunctionsModule<ConfiguredGraphQueryEnvironment<'a>>,
            extra_functions: ConfiguredGraphFunctions<'a>,
        }

        impl Debug for Functions<'_> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                f.debug_struct("Functions").finish_non_exhaustive()
            }
        }

        impl<'a> QueryFunctions for Functions<'a> {
            type Env = ConfiguredGraphQueryEnvironment<'a>;
            fn get(
                &self,
                name: &str,
            ) -> Option<&dyn QueryFunction<ConfiguredGraphQueryEnvironment<'a>>> {
                if let Some(v) = self.extra_functions.get(name) {
                    Some(v)
                } else {
                    self.defaults.get(name)
                }
            }

            fn get_op(
                &self,
                op: BinaryOp,
            ) -> Option<&dyn QueryBinaryOp<ConfiguredGraphQueryEnvironment<'a>>> {
                if let Some(v) = self.extra_functions.get_op(op) {
                    Some(v)
                } else {
                    self.defaults.get_op(op)
                }
            }
        }

        Functions {
            defaults: DefaultQueryFunctionsModule::new(),
            extra_functions: ConfiguredGraphFunctions(PhantomData),
        }
    }

    async fn get_targets_from_template_placeholder_info(
        &self,
        template_name: &'static str,
        targets: TargetSet<ConfiguredGraphNodeRef>,
    ) -> anyhow::Result<TargetSet<ConfiguredGraphNodeRef>> {
        self.delegate
            .get_targets_from_template_placeholder_info(template_name, targets)
            .await
    }
}

pub(crate) fn init_query_functions() {
    CONFIGURED_GRAPH_QUERY_FUNCTIONS.init(Arc::new(ConfiguredGraphQueryEnvironment::functions()));
}

#[async_trait]
impl<'a> QueryEnvironment for ConfiguredGraphQueryEnvironment<'a> {
    type Target = ConfiguredGraphNodeRef;

    async fn get_node(&self, node_ref: &ConfiguredGraphNodeRef) -> anyhow::Result<Self::Target> {
        Ok(node_ref.dupe())
    }

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

    async fn eval_literals(&self, literal: &[&str]) -> anyhow::Result<TargetSet<Self::Target>> {
        let mut result = TargetSet::new();
        for lit in literal {
            result.insert(ConfiguredGraphNodeRef(self.delegate.eval_literal(lit)?));
        }
        Ok(result)
    }

    async fn eval_file_literal(&self, literal: &str) -> anyhow::Result<FileSet> {
        Err(AnalysisQueryError::FileLiteralsNotAllowed(literal.to_owned()).into())
    }

    async fn dfs_postorder(
        &self,
        root: &TargetSet<Self::Target>,
        delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
    ) -> anyhow::Result<()> {
        async_fast_depth_first_postorder_traversal(
            &ConfiguredGraphNodeRefLookup,
            root.iter().map(LabeledNode::node_ref),
            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(&ConfiguredGraphNodeRefLookup, root.iter(), delegate, depth)
            .await
    }

    async fn owner(&self, _paths: &FileSet) -> anyhow::Result<TargetSet<Self::Target>> {
        Err(QueryError::FunctionUnimplemented("owner").into())
    }
}

async fn dice_lookup_transitive_set(
    ctx: &DiceComputations,
    key: TransitiveSetKey,
) -> anyhow::Result<DeferredValueReady<DeferredTransitiveSetData>> {
    ctx.compute_deferred_data(&key).await
}

async fn get_template_info_provider_artifacts(
    ctx: &DiceComputations,
    configured_label: &ConfiguredTargetLabel,
    template_name: &str,
) -> anyhow::Result<Vec<ArtifactGroup>> {
    let providers_label =
        ConfiguredProvidersLabel::new(configured_label.dupe(), ProvidersName::Default);

    let providers = ctx.get_providers(&providers_label);

    let mut artifacts = vec![];

    match providers.await? {
        MaybeCompatible::Incompatible(reason) => {
            eprintln!("{}", reason.skipping_message(configured_label));
        }
        MaybeCompatible::Compatible(providers) => {
            let providers_collection = providers.provider_collection();

            if let Some(template_placeholder_info) =
                TemplatePlaceholderInfo::from_providers(providers_collection)
            {
                if let Some(template_info) = template_placeholder_info
                    .keyed_variables()
                    .get(template_name)
                {
                    let mut cmd_visitor = SimpleCommandLineArtifactVisitor::new();
                    if let either::Either::Left(command_line_arg) = template_info {
                        CommandLineArgLike::visit_artifacts(
                            command_line_arg.as_command_line_arg(),
                            &mut cmd_visitor,
                        )?;
                    } else if let either::Either::Right(map) = template_info {
                        for (_, command_line_arg) in map.iter() {
                            CommandLineArgLike::visit_artifacts(
                                command_line_arg.as_command_line_arg(),
                                &mut cmd_visitor,
                            )?;
                        }
                    }

                    for input in cmd_visitor.inputs {
                        artifacts.push(input);
                    }
                }
            }
        }
    }

    Ok(artifacts)
}

pub(crate) async fn get_from_template_placeholder_info<'x>(
    ctx: &'x DiceComputations,
    template_name: &'static str,
    targets: impl IntoIterator<Item = ConfiguredTargetLabel>,
) -> anyhow::Result<IndexMap<ConfiguredTargetLabel, Artifact>> {
    let mut label_to_artifact: IndexMap<ConfiguredTargetLabel, Artifact> = IndexMap::new();

    // Traversing tsets adds complexity here. Ideally, we could just do a normal traversal of these starlark values
    // we get from the template_info provider, but the cmdlinearglike interface only gives us access via ArtifactGroup
    // which only have the key for the tset projection, not the tset or projection itself.
    //
    // Then, next we would want to do a traversal over those ArtifactGroup and use the dice ctx to map the tset projection
    // keys back to their values. We can't do that because that would flatten the tset into the dice deps and rdeps storage.
    //
    // So, instead we need to extract out the ArtifactGroups from the template_info values and lookup the corresponding
    // tsets, manually traverse them and repeat (because the tset node values themselves return ArtifactGroup).
    //
    // This means that we will have unnessary dice nodes pointing to each tset value appearing in the template info and any
    // tset value appearing in another tset's nodes. In the common case, though, that set will be small and we won't have
    // flattened any full tset into the dice deps storage. We'll call those "top-level" tset nodes.

    // This will contain the ArtifactGroups we encounter during our traversal (so only artifacts and top-level tset nodes).
    // Artifacts are put here to keep them in the correct order in the output, tsets are top-level tset nodes that we need
    // to traverse.
    let artifacts = futures::future::try_join_all(targets.into_iter().map(|target| async move {
        let artifacts = get_template_info_provider_artifacts(ctx, &target, template_name).await?;
        anyhow::Ok(
            artifacts
                .into_iter()
                .map(move |artifact| (target.dupe(), artifact)),
        )
    }))
    .await?;
    let mut artifacts: VecDeque<_> = artifacts.into_iter().flatten().collect();

    // This will contain the TransitiveSetProjectionKey we encounter as top-level nodes and we will also put in TransitiveSetProjectionKey
    // for all the tset nodes that we encounter during our traversal of those top-level nodes. We don't need to track artifacts because
    // we just extract the targetlabel and put that in the output set and that can dedupe them (and we don't need to further
    // traverse artifacts).
    let mut seen = HashSet::new();

    while let Some((target, artifact)) = artifacts.pop_front() {
        let handle_artifact =
            |label_to_artifact: &mut IndexMap<ConfiguredTargetLabel, Artifact>,
             artifact: &Artifact|
             -> anyhow::Result<()> {
                if let Some(owner) = artifact.owner() {
                    let target_label = owner.unpack_target_label().ok_or_else(|| {
                        AnalysisQueryError::NonTargetBoundArtifact(
                            template_name.to_owned(),
                            target.dupe(),
                            artifact.dupe(),
                        )
                    })?;
                    label_to_artifact.insert(target_label.dupe(), artifact.dupe());
                }
                Ok(())
            };

        match artifact.resolved()? {
            ResolvedArtifactGroup::Artifact(artifact) => {
                handle_artifact(&mut label_to_artifact, artifact)?;
            }
            ResolvedArtifactGroup::TransitiveSetProjection(tset_key) => {
                // We've encountered a "top-level" tset node that we haven't yet seen (as either a top-level or intermediate node, doesn't matter).
                if seen.insert(tset_key.dupe()) {
                    let tset_value = dice_lookup_transitive_set(ctx, tset_key.key.dupe()).await?;

                    // Now we can traverse this tset from that node. This is a different traversal than our top-level one as we will
                    // be accessing tset internals directly and so we can actually traverse the starlark objects without going back through
                    // dice. We'll be working all with values with lifetimes from `tset_value`.
                    //
                    // We can't use tset's normal traverse because we need to avoid retraversing parts of the tset graph that we've already
                    // traversed (through other top-level tset nodes).
                    let mut queue = VecDeque::new();
                    queue.push_back(tset_value.as_value());
                    while let Some(v) = queue.pop_front() {
                        let as_tset =
                            TransitiveSet::from_value(v).context("invalid tset structure")?;

                        // Visit the projection value itself. As this is an opaque cmdargs-like thing, it may contain more top-level tset node
                        // references that need to be pushed into the outer queue.
                        if let Some(v) = as_tset.get_projection_value(tset_key.projection)? {
                            struct Visitor<'a>(
                                &'a mut VecDeque<(ConfiguredTargetLabel, ArtifactGroup)>,
                                ConfiguredTargetLabel,
                            );
                            impl<'a> CommandLineArtifactVisitor for Visitor<'a> {
                                fn visit_input(
                                    &mut self,
                                    input: ArtifactGroup,
                                    _tag: Option<&ArtifactTag>,
                                ) {
                                    self.0.push_back((self.1.dupe(), input));
                                }

                                fn visit_output(
                                    &mut self,
                                    _artifact: OutputArtifact,
                                    _tag: Option<&ArtifactTag>,
                                ) {
                                    // ignored
                                }
                            }
                            v.as_command_line_err()?
                                .visit_artifacts(&mut Visitor(&mut artifacts, target.dupe()))?;
                        }

                        // Enqueue any children we haven't yet seen (and mark them seen).
                        for child in as_tset.children.iter() {
                            let child_as_tset =
                                TransitiveSet::from_value(*child).context("Invalid deferred")?;
                            let projection_key =
                                child_as_tset.get_projection_key(tset_key.projection);
                            if seen.insert(projection_key) {
                                queue.push_back(*child);
                            }
                        }
                    }
                }
            }
        }
    }
    Ok(label_to_artifact)
}

/// Used by `audit classpath`
pub(crate) fn init_classpath_for_targets() {
    CLASSPATH_FOR_TARGETS.init(|ctx, targets| {
        Box::pin(get_from_template_placeholder_info(
            ctx,
            "classpath",
            targets,
        ))
    })
}


/*
 * 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 buck2_build_api::analysis::calculation::EVAL_ANALYSIS_QUERY;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_ref::ConfiguredGraphNodeRef;
use buck2_query::query::syntax::simple::eval::evaluator::QueryEvaluator;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use dice::DiceComputations;

use crate::analysis::configured_graph::AnalysisConfiguredGraphQueryDelegate;
use crate::analysis::configured_graph::AnalysisDiceQueryDelegate;
use crate::analysis::environment::ConfiguredGraphQueryEnvironment;

pub fn init_eval_analysis_query() {
    EVAL_ANALYSIS_QUERY.init(|ctx, query, resolved_literals| {
        Box::pin(eval_analysis_query(ctx, query, resolved_literals))
    });
}

async fn eval_analysis_query(
    ctx: &DiceComputations,
    query: &str,
    resolved_literals: HashMap<String, ConfiguredTargetNode>,
) -> anyhow::Result<TargetSet<ConfiguredGraphNodeRef>> {
    let dice_query_delegate = Arc::new(AnalysisDiceQueryDelegate { ctx });
    let delegate = AnalysisConfiguredGraphQueryDelegate {
        dice_query_delegate,
        resolved_literals,
    };

    let functions = ConfiguredGraphQueryEnvironment::functions();
    let env = ConfiguredGraphQueryEnvironment::new(&delegate);
    let evaluator = QueryEvaluator::new(&env, &functions);

    let result = evaluator.eval_query(query).await?;
    result.try_into_targets()
}


/*
 * 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 common cquery/uquery pieces.

use buck2_query::query::environment::QueryEnvironment;
use buck2_query::query::syntax::simple::eval::evaluator::QueryEvaluator;
use buck2_query::query::syntax::simple::eval::literals::extract_target_literals;
use buck2_query::query::syntax::simple::eval::multi_query::process_multi_query;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use buck2_query_parser::placeholder::QUERY_PERCENT_S_PLACEHOLDER;
use futures::Future;
use gazebo::prelude::*;
use starlark::collections::SmallSet;
use thiserror::Error;

#[derive(Debug, Error)]
enum EvalQueryError {
    #[error("Query args supplied without any `%s` placeholder in the query, got args {}", .0.map(|x| format!("`{}`", x)).join(", "))]
    ArgsWithoutPlaceholder(Vec<String>),
    #[error("Placeholder `%s` in query argument `{0}`")]
    PlaceholderInPattern(String),
}

pub async fn eval_query<
    Env: QueryEnvironment,
    Fut: Future<Output = anyhow::Result<Env>>,
    A: AsRef<str>,
>(
    functions: &DefaultQueryFunctionsModule<Env>,
    query: &str,
    query_args: &[A],
    environment: impl FnOnce(Vec<String>) -> Fut,
) -> anyhow::Result<QueryEvaluationResult<Env::Target>> {
    let mut literals = SmallSet::new();
    if query.contains(QUERY_PERCENT_S_PLACEHOLDER) {
        // We'd really like the query args to only be literals (file or target).
        // If that didn't work, we'd really like query args to be well-formed expressions.
        // Unfortunately Buck1 just substitutes in arbitrarily strings, where the query
        // or query_args may not form anything remotely valid.
        // We have to be backwards compatible :(
        for q in query_args {
            let q = q.as_ref();
            if q.contains(QUERY_PERCENT_S_PLACEHOLDER) {
                return Err(EvalQueryError::PlaceholderInPattern(q.to_owned()).into());
            }
            extract_target_literals(
                functions,
                &query.replace(QUERY_PERCENT_S_PLACEHOLDER, q),
                &mut literals,
            )?;
        }
        let env = environment(literals.into_iter().collect()).await?;
        let results = process_multi_query(query, query_args, |input, query| {
            let evaluator = QueryEvaluator::new(&env, functions);
            async move { (input, evaluator.eval_query(&query).await) }
        })
        .await;
        Ok(QueryEvaluationResult::Multiple(results))
    } else if !query_args.is_empty() {
        Err(
            EvalQueryError::ArgsWithoutPlaceholder(query_args.map(|s| s.as_ref().to_owned()))
                .into(),
        )
    } else {
        extract_target_literals(functions, query, &mut literals)?;
        let env = environment(literals.into_iter().collect()).await?;
        Ok(QueryEvaluationResult::Single(
            QueryEvaluator::new(&env, functions)
                .eval_query(query)
                .await?,
        ))
    }
}


/*
 * 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 configured_graph;
pub mod environment;
pub(crate) mod eval;
pub(crate) 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.
 */

use std::sync::Arc;

use async_trait::async_trait;
use buck2_artifact::actions::key::ActionKey;
use buck2_build_api::actions::query::ActionQueryNode;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_query::query::environment::QueryEnvironment;
use buck2_query::query::syntax::simple::eval::error::QueryError;
use buck2_query::query::syntax::simple::eval::file_set::FileSet;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
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 crate::cquery::environment::CqueryDelegate;
use crate::uquery::environment::QueryLiterals;

/// CqueryDelegate resolves information needed by the QueryEnvironment.
#[async_trait]
pub trait AqueryDelegate: Send + Sync {
    fn cquery_delegate(&self) -> &dyn CqueryDelegate;

    async fn get_node(&self, key: &ActionKey) -> anyhow::Result<ActionQueryNode>;
}

pub struct AqueryEnvironment<'c> {
    delegate: Arc<dyn AqueryDelegate + 'c>,
    literals: Arc<dyn QueryLiterals<ActionQueryNode> + 'c>,
}

impl<'c> AqueryEnvironment<'c> {
    pub fn new(
        delegate: Arc<dyn AqueryDelegate + 'c>,
        literals: Arc<dyn QueryLiterals<ActionQueryNode> + 'c>,
    ) -> Self {
        Self { delegate, literals }
    }

    async fn get_node(&self, label: &ActionKey) -> anyhow::Result<ActionQueryNode> {
        self.delegate.get_node(label).await
    }
}

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

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

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

    async fn get_node_for_default_configured_target(
        &self,
        _node_ref: &ActionKey,
    ) -> 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<Self::Target>> {
        self.literals.eval_literals(literals).await
    }

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

    async fn dfs_postorder(
        &self,
        root: &TargetSet<Self::Target>,
        traversal_delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
    ) -> anyhow::Result<()> {
        // TODO(cjhopman): The query nodes deps are going to flatten the tset structure for its deps. In a typical
        // build graph, a traversal over just the graph of ActionQueryNode ends up being an `O(n)` operation at each
        // node and ends up with an `O(n^2)` cost. If instead we were to not flatten the structure and traverse the
        // mixed graph of action nodes and tset nodes, we'd get closer to `O(n + e)` which in practice is much better
        // (hence the whole point of tsets). While we can't change the ActionQueryNode deps() function to not flatten
        // the tset, we aren't required to do these traversal's using that function.
        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<()> {
        // TODO(cjhopman): See above.
        async_depth_limited_traversal(self, root.iter_names(), delegate, depth).await
    }

    async fn owner(&self, _paths: &FileSet) -> anyhow::Result<TargetSet<Self::Target>> {
        Err(QueryError::NotAvailableInContext("owner").into())
    }
}


/*
 * 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_build_api::actions::query::ActionQueryNode;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
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::aquery::environment::AqueryEnvironment;
use crate::dice::aquery::DiceAqueryDelegate;
use crate::dice::get_dice_query_delegate;
use crate::uquery::environment::PreresolvedQueryLiterals;

pub(crate) struct AqueryEvaluator<'c> {
    dice_query_delegate: Arc<DiceAqueryDelegate<'c>>,
    functions: DefaultQueryFunctionsModule<AqueryEnvironment<'c>>,
}

impl AqueryEvaluator<'_> {
    pub(crate) async fn eval_query(
        &self,
        query: &str,
        query_args: &[String],
    ) -> anyhow::Result<QueryEvaluationResult<ActionQueryNode>> {
        eval_query(&self.functions, query, query_args, async move |literals| {
            let resolved_literals =
                PreresolvedQueryLiterals::pre_resolve(&*self.dice_query_delegate, &literals).await;
            Ok(AqueryEnvironment::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(crate) async fn get_aquery_evaluator<'a, 'c: 'a>(
    ctx: &'c DiceComputations,
    working_dir: &'a ProjectRelativePath,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<AqueryEvaluator<'c>> {
    let dice_query_delegate =
        get_dice_aquery_delegate(ctx, working_dir, global_target_platform).await?;
    let functions = DefaultQueryFunctionsModule::new();
    Ok(AqueryEvaluator {
        dice_query_delegate,
        functions,
    })
}

// Provides the dice query delegate for aquery evaluator
pub(crate) async fn get_dice_aquery_delegate<'a, 'c: 'a>(
    ctx: &'c DiceComputations,
    working_dir: &'a ProjectRelativePath,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<Arc<DiceAqueryDelegate<'c>>> {
    let dice_query_delegate =
        get_dice_query_delegate(ctx, working_dir, global_target_platform).await?;
    let dice_query_delegate = Arc::new(DiceAqueryDelegate::new(dice_query_delegate).await?);
    Ok(dice_query_delegate)
}


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

use buck2_artifact::actions::key::ActionKey;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_artifact::artifact::provide_outputs::ProvideOutputs;
use buck2_build_api::actions::query::ActionQueryNode;
use buck2_build_api::actions::query::FIND_MATCHING_ACTION;
use buck2_build_api::analysis::AnalysisResult;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use dice::DiceComputations;
use tracing::debug;

use crate::aquery::evaluator::get_dice_aquery_delegate;

fn check_output_path<'v>(
    build_artifact: &'v BuildArtifact,
    path_to_check: &'v ForwardRelativePathBuf,
) -> anyhow::Result<Option<&'v ActionKey>> {
    let path = build_artifact.get_path().path();

    debug!(
        "Checking action's output path: {}. Path to check: {}",
        path, path_to_check
    );

    if path_to_check.starts_with(path_to_check) {
        Ok(Some(build_artifact.key()))
    } else {
        Ok(None)
    }
}

async fn find_matching_action(
    ctx: &DiceComputations,
    working_dir: &ProjectRelativePath,
    global_target_platform: Option<TargetLabel>,
    analysis: &AnalysisResult,
    path_after_target_name: ForwardRelativePathBuf,
) -> anyhow::Result<Option<ActionQueryNode>> {
    let dice_aquery_delegate =
        get_dice_aquery_delegate(ctx, working_dir, global_target_platform.clone()).await?;

    for entry in analysis.iter_deferreds() {
        match any::request_value::<ProvideOutputs>(entry.as_complex()) {
            Some(outputs) => {
                let outputs = outputs.0?;
                for build_artifact in &outputs {
                    match check_output_path(build_artifact, &path_after_target_name)? {
                        Some(action_key) => {
                            return Ok(Some(
                                dice_aquery_delegate.get_action_node(action_key).await?,
                            ));
                        }
                        None => (),
                    }
                }
            }
            None => debug!("Could not extract outputs from deferred table entry"),
        }
    }
    Ok(None)
}

pub(crate) fn init_find_matching_action() {
    FIND_MATCHING_ACTION.init(
        |ctx, working_dir, global_target_platform, analysis, path_after_target_name| {
            Box::pin(find_matching_action(
                ctx,
                working_dir,
                global_target_platform,
                analysis,
                path_after_target_name,
            ))
        },
    );
}


/*
 * 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 mod environment;
pub mod evaluator;
pub(crate) mod find_matching_action;


/*
 * 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::sync::Arc;

use async_trait::async_trait;
use buck2_build_api::query::bxl::BxlCqueryFunctions;
use buck2_build_api::query::bxl::NEW_BXL_CQUERY_FUNCTIONS;
use buck2_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::package_boundary::HasPackageBoundaryExceptions;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
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::helpers::CapturedExpr;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctions;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use dice::DiceComputations;
use dupe::Dupe;

use crate::cquery::environment::CqueryEnvironment;
use crate::dice::DiceQueryDelegate;

fn cquery_functions<'v>() -> DefaultQueryFunctions<CqueryEnvironment<'v>> {
    DefaultQueryFunctions::new()
}

struct BxlCqueryFunctionsImpl<'c> {
    ctx: &'c DiceComputations,
    target_platform: Option<TargetLabel>,
    project_root: ProjectRoot,
    working_dir: ProjectRelativePathBuf,
}

impl<'c> BxlCqueryFunctionsImpl<'c> {
    async fn cquery_env(&self) -> anyhow::Result<CqueryEnvironment<'c>> {
        let cell_resolver = self.ctx.get_cell_resolver().await?;

        let package_boundary_exceptions = self.ctx.get_package_boundary_exceptions().await?;
        let target_alias_resolver = self
            .ctx
            .target_alias_resolver_for_working_dir(&self.working_dir)
            .await?;

        let dice_query_delegate = Arc::new(DiceQueryDelegate::new(
            self.ctx,
            &self.working_dir,
            self.project_root.dupe(),
            cell_resolver,
            self.target_platform.dupe(),
            package_boundary_exceptions,
            target_alias_resolver,
        )?);
        Ok(CqueryEnvironment::new(
            dice_query_delegate.dupe(),
            dice_query_delegate,
            // TODO(nga): add universe.
            None,
            // TODO(nga): use proper owner behavior.
            CqueryOwnerBehavior::Deprecated,
        ))
    }
}

#[async_trait]
impl<'c> BxlCqueryFunctions<'c> for BxlCqueryFunctionsImpl<'c> {
    async fn allpaths(
        &self,
        from: &TargetSet<ConfiguredTargetNode>,
        to: &TargetSet<ConfiguredTargetNode>,
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .allpaths(&self.cquery_env().await?, from, to)
            .await?)
    }

    async fn somepath(
        &self,
        from: &TargetSet<ConfiguredTargetNode>,
        to: &TargetSet<ConfiguredTargetNode>,
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .somepath(&self.cquery_env().await?, from, to)
            .await?)
    }

    async fn owner(&self, file_set: &FileSet) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .owner(&self.cquery_env().await?, file_set)
            .await?)
    }

    async fn deps(
        &self,
        targets: &TargetSet<ConfiguredTargetNode>,
        deps: Option<i32>,
        captured_expr: Option<&CapturedExpr>,
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .deps(
                &self.cquery_env().await?,
                &DefaultQueryFunctionsModule::new(),
                targets,
                deps,
                captured_expr,
            )
            .await?)
    }

    async fn rdeps(
        &self,
        universe: &TargetSet<ConfiguredTargetNode>,
        targets: &TargetSet<ConfiguredTargetNode>,
        depth: Option<i32>,
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .rdeps(&self.cquery_env().await?, universe, targets, depth)
            .await?)
    }

    async fn testsof(
        &self,
        targets: &TargetSet<ConfiguredTargetNode>,
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        Ok(cquery_functions()
            .testsof(&self.cquery_env().await?, targets)
            .await?)
    }

    async fn testsof_with_default_target_platform(
        &self,
        targets: &TargetSet<ConfiguredTargetNode>,
    ) -> anyhow::Result<Vec<MaybeCompatible<ConfiguredTargetNode>>> {
        Ok(cquery_functions()
            .testsof_with_default_target_platform(&self.cquery_env().await?, targets)
            .await?)
    }
}

pub(crate) fn init_new_bxl_cquery_functions() {
    NEW_BXL_CQUERY_FUNCTIONS.init(|ctx, target_platform, project_root, cell_name| {
        Box::pin(async move {
            let cell_resolver = ctx.get_cell_resolver().await?;
            let cell = cell_resolver.get(cell_name)?;
            // TODO(nga): working as as cell root is not right.
            //   Should be either the project root or user's current working directory.
            let working_dir = cell.path().as_project_relative_path().to_buf();

            Result::<Box<dyn BxlCqueryFunctions>, _>::Ok(Box::new(BxlCqueryFunctionsImpl {
                ctx,
                target_platform,
                project_root,
                working_dir,
            }))
        })
    })
}


/*
 * 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::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_core::target::label::TargetLabel;
use buck2_events::dispatch::console_message;
use buck2_node::configured_universe::CqueryUniverse;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_query::query::environment::QueryEnvironment;
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 dupe::Dupe;
use tracing::warn;

use crate::uquery::environment::allbuildfiles;
use crate::uquery::environment::rbuildfiles;
use crate::uquery::environment::QueryLiterals;
use crate::uquery::environment::UqueryDelegate;

#[derive(Debug, thiserror::Error)]
enum CqueryError {
    #[error("Target universe not specified (internal error)")]
    NoUniverse,
}

/// CqueryDelegate resolves information needed by the QueryEnvironment.
#[async_trait]
pub trait CqueryDelegate: Send + Sync {
    fn uquery_delegate(&self) -> &dyn UqueryDelegate;

    async fn get_node_for_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<MaybeCompatible<ConfiguredTargetNode>>;

    async fn get_node_for_configured_target(
        &self,
        target: &ConfiguredTargetLabel,
    ) -> anyhow::Result<ConfiguredTargetNode>;

    async fn get_configured_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<ConfiguredTargetLabel>;

    async fn get_node_for_default_configured_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<MaybeCompatible<ConfiguredTargetNode>>;
}

pub struct CqueryEnvironment<'c> {
    delegate: Arc<dyn CqueryDelegate + 'c>,
    literals: Arc<dyn QueryLiterals<ConfiguredTargetNode> + 'c>,
    // TODO(nga): BXL `cquery` function does not provides us the universe.
    // TODO(nga): do not compute the universe when we don't need it, because it is not free.
    //   For example, when evaluating
    //   ```
    //   buck2 cquery 'deps(//foo:bar)'
    //   ```
    universe: Option<CqueryUniverse>,
    owner_behavior: CqueryOwnerBehavior,
}

impl<'c> CqueryEnvironment<'c> {
    pub fn new(
        delegate: Arc<dyn CqueryDelegate + 'c>,
        literals: Arc<dyn QueryLiterals<ConfiguredTargetNode> + 'c>,
        universe: Option<CqueryUniverse>,
        owner_behavior: CqueryOwnerBehavior,
    ) -> Self {
        Self {
            delegate,
            literals,
            universe,
            owner_behavior,
        }
    }

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

    async fn get_node(
        &self,
        label: &ConfiguredTargetLabel,
    ) -> anyhow::Result<ConfiguredTargetNode> {
        self.delegate.get_node_for_configured_target(label).await
    }

    async fn get_node_for_default_configured_target(
        &self,
        label: &ConfiguredTargetLabel,
    ) -> anyhow::Result<MaybeCompatible<ConfiguredTargetNode>> {
        self.delegate
            .get_node_for_default_configured_target(label.unconfigured())
            .await
    }

    /// Deprecated `owner` function implementation.
    /// See [this post](https://fburl.com/0xv7u4bz) for details.
    async fn owner_deprecated(&self, path: &CellPath) -> anyhow::Result<Vec<ConfiguredTargetNode>> {
        // need to explicitly track this rather than checking for changes to result set since the owner might
        // already be in the set.
        let mut owners = Vec::new();
        match self
            .delegate
            .uquery_delegate()
            .get_enclosing_packages(path)
            .await
        {
            Ok(packages) => {
                let package_futs = packages.iter().map(|package| async move {
                    let mut result: Vec<ConfiguredTargetNode> = Vec::new();

                    // TODO(cjhopman): We should make sure that the file exists.
                    let targets = self
                        .delegate
                        .uquery_delegate()
                        .eval_build_file(package.dupe())
                        .await?;

                    for node in targets.targets().values() {
                        match self.delegate.get_node_for_target(node.label()).await? {
                            MaybeCompatible::Compatible(node) => {
                                for input in node.inputs() {
                                    if &input == path {
                                        result.push(node.dupe());
                                        // this intentionally only breaks out of the inner 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.
                                        break;
                                    }
                                }
                            }
                            MaybeCompatible::Incompatible(reason) => {
                                // TODO(scottcao): Add event for incompatible target skipping
                                console_message(reason.skipping_message(
                                    &self.delegate.get_configured_target(node.label()).await?,
                                ));
                            }
                        }
                    }

                    anyhow::Ok(result)
                });

                for nodes in futures::future::join_all(package_futs).await.into_iter() {
                    for node in nodes?.into_iter() {
                        owners.push(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.
            }
        };
        Ok(owners)
    }

    fn owner_correct(&self, path: &CellPath) -> anyhow::Result<Vec<ConfiguredTargetNode>> {
        let universe = self.universe.as_ref().context(CqueryError::NoUniverse)?;
        Ok(universe.owners(path))
    }
}

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

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

    async fn get_node_for_default_configured_target(
        &self,
        node_ref: &ConfiguredTargetLabel,
    ) -> anyhow::Result<MaybeCompatible<Self::Target>> {
        CqueryEnvironment::get_node_for_default_configured_target(self, node_ref).await
    }

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

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

    async fn dfs_postorder(
        &self,
        root: &TargetSet<ConfiguredTargetNode>,
        traversal_delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
    ) -> 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.uquery_delegate()).await;
    }

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

    async fn owner(&self, paths: &FileSet) -> anyhow::Result<TargetSet<Self::Target>> {
        let mut result = TargetSet::new();

        for path in paths.iter() {
            let owners = match &self.owner_behavior {
                CqueryOwnerBehavior::Deprecated => self.owner_deprecated(path).await?,
                CqueryOwnerBehavior::Correct => self.owner_correct(path)?,
            };
            if owners.is_empty() {
                warn!("No owner was found for {}", path);
            }
            result.extend(owners);
        }
        Ok(result)
    }
}

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


/*
 * 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_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_common::result::ToSharedResultExt;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use buck2_events::dispatch::console_message;
use buck2_node::configured_universe::CqueryUniverse;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use dice::DiceComputations;
use dupe::Dupe;
use futures::stream::FuturesUnordered;
use futures::StreamExt;
use gazebo::prelude::*;

use crate::analysis::evaluator::eval_query;
use crate::cquery::environment::CqueryEnvironment;
use crate::dice::get_dice_query_delegate;
use crate::dice::DiceQueryDelegate;
use crate::uquery::environment::PreresolvedQueryLiterals;
use crate::uquery::environment::QueryLiterals;
use crate::uquery::environment::UqueryDelegate;

pub struct CqueryEvaluator<'c> {
    dice_query_delegate: Arc<DiceQueryDelegate<'c>>,
    functions: DefaultQueryFunctionsModule<CqueryEnvironment<'c>>,
    owner_behavior: CqueryOwnerBehavior,
}

impl CqueryEvaluator<'_> {
    pub async fn eval_query<A: AsRef<str>, U: AsRef<str>>(
        &self,
        query: &str,
        query_args: &[A],
        target_universe: Option<&[U]>,
    ) -> anyhow::Result<QueryEvaluationResult<ConfiguredTargetNode>> {
        eval_query(&self.functions, query, query_args, async move |literals| {
            let (universe, resolved_literals) = match target_universe {
                None => {
                    if literals.is_empty() {
                        console_message(
                            "Query has no target literals and `--target-universe` is not specified.\n\
                            Such query is correct, but the result is always empty.\n\
                            Consider specifying `--target-universe` for this query\n\
                            or using `uquery` instead of `cquery`".to_owned());
                    }
                    // In the absence of a user-provided target universe, we use the target
                    // literals in the cquery as the universe.
                    resolve_literals_in_universe(&self.dice_query_delegate, &literals, &literals)
                        .await?
                }
                Some(universe) => {
                    resolve_literals_in_universe(&self.dice_query_delegate, &literals, universe)
                        .await?
                }
            };
            Ok(CqueryEnvironment::new(
                self.dice_query_delegate.dupe(),
                Arc::new(resolved_literals),
                Some(universe),
                self.owner_behavior,
            ))
        })
        .await
    }
}

pub(crate) async fn preresolve_literals_and_build_universe(
    dice_query_delegate: &DiceQueryDelegate<'_>,
    literals: &[String],
) -> anyhow::Result<(
    CqueryUniverse,
    PreresolvedQueryLiterals<ConfiguredTargetNode>,
)> {
    let resolved_literals =
        PreresolvedQueryLiterals::pre_resolve(dice_query_delegate, literals).await;
    let universe = CqueryUniverse::build(&resolved_literals.literals()?).await?;
    Ok((universe, resolved_literals))
}

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

// This will first resolve the universe to configured nodes and then gather all
// the deps. From there, it resolves the literals to any matching nodes in the universe deps.
async fn resolve_literals_in_universe<L: AsRef<str>, U: AsRef<str>>(
    dice_query_delegate: &DiceQueryDelegate<'_>,
    literals: &[L],
    universe: &[U],
) -> anyhow::Result<(
    CqueryUniverse,
    PreresolvedQueryLiterals<ConfiguredTargetNode>,
)> {
    // TODO(cjhopman): We should probably also resolve the literals to TargetNode so that
    // we can get errors for packages or targets that don't exist or fail to load.
    let refs: Vec<_> = universe.map(|v| v.as_ref());
    let universe_resolved = dice_query_delegate.eval_literals(&refs).await?;

    let universe = CqueryUniverse::build(&universe_resolved).await?;

    // capture a reference so the ref can be moved into the future below.
    let universe_ref = &universe;

    // TODO(cjhopman): Using the default resolution for recursive literals is inefficient.
    // If we can have a package-trie or cellpath-trie we can do the resolution directly
    // against the universe.
    let resolution_futs: FuturesUnordered<_> = literals
        .iter()
        .map(|lit| async move {
            let lit = lit.as_ref();
            let result: anyhow::Result<_> = try {
                let resolved_pattern = dice_query_delegate.resolve_target_patterns(&[lit]).await?;
                universe_ref.get(&resolved_pattern)
            };

            (lit.to_owned(), result.shared_error())
        })
        .collect();

    let resolved = resolution_futs.collect().await;
    Ok((universe, PreresolvedQueryLiterals::new(resolved)))
}


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

use std::hash::Hash;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_artifact::actions::key::ActionKey;
use buck2_build_api::actions::artifact::get_artifact_fs::GetArtifactFs;
use buck2_build_api::actions::calculation::ActionCalculation;
use buck2_build_api::actions::query::ActionInput;
use buck2_build_api::actions::query::ActionQueryNode;
use buck2_build_api::actions::query::SetProjectionInputs;
use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::artifact_groups::ResolvedArtifactGroup;
use buck2_build_api::artifact_groups::TransitiveSetProjectionKey;
use buck2_build_api::deferred::calculation::DeferredCalculation;
use buck2_common::result::SharedResult;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::pattern::ParsedPattern;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use dashmap::DashMap;
use dice::DiceComputations;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::future::Shared;
use futures::stream::FuturesOrdered;
use futures::StreamExt;
use gazebo::prelude::*;
use itertools::Either;
use itertools::Itertools;
use thiserror::Error;

use crate::aquery::environment::AqueryDelegate;
use crate::cquery::environment::CqueryDelegate;
use crate::dice::DiceQueryDelegate;
use crate::uquery::environment::QueryLiterals;

#[derive(Debug, Error)]
enum ActionQueryError {
    #[error(
        "`aquery` currently only supports literal target patterns, not package or recursive (got `{0}`)"
    )]
    UnsupportedTargetPattern(String),
}

/// A simple concurrent map with a `get_or_compute()` function
struct NodeCache<'c, K: Hash + Eq + PartialEq + Dupe, V: Dupe> {
    map: DashMap<K, Shared<BoxFuture<'c, V>>>,
}

impl<'c, K: Hash + Eq + PartialEq + Dupe, V: Dupe> NodeCache<'c, K, V> {
    fn new() -> Self {
        Self {
            map: DashMap::new(),
        }
    }

    /// Gets the value or computes it with the provided function. The function is called while holding
    /// a lock on the map and so should not do much work. The future returned by that function isn't
    /// polled until later so it's fine for it to do more work.
    fn get_or_compute<F>(self: Arc<Self>, key: K, f: F) -> Shared<BoxFuture<'c, V>>
    where
        F: FnOnce(K) -> BoxFuture<'c, V>,
    {
        if let Some(v) = self.map.get(&key) {
            return v.clone();
        }

        let entry = self.map.entry(key.dupe());
        entry.or_insert_with(move || f(key).shared()).clone()
    }
}

/// Cache for ActionQueryNode and things needed to construct it. This is required because
/// QueryTarget::deps() requires that deps are synchronously available and so we need to
/// be able to iterate the tset structure synchronously.
#[derive(Clone, Dupe)]
struct DiceAqueryNodesCache<'c> {
    action_nodes: Arc<NodeCache<'c, ActionKey, SharedResult<ActionQueryNode>>>,
    tset_nodes: Arc<NodeCache<'c, TransitiveSetProjectionKey, SharedResult<SetProjectionInputs>>>,
}

impl<'c> DiceAqueryNodesCache<'c> {
    fn new() -> Self {
        Self {
            action_nodes: Arc::new(NodeCache::new()),
            tset_nodes: Arc::new(NodeCache::new()),
        }
    }
}

pub(crate) struct DiceAqueryDelegate<'c> {
    base_delegate: DiceQueryDelegate<'c>,
    nodes_cache: DiceAqueryNodesCache<'c>,
    artifact_fs: Arc<ArtifactFs>,
}

/// Converts artifact inputs into aquery's ActionInput. This is mostly a matter of resolving the indirect
/// `TransitiveSetProjectionKey` to our direct shadow tset graph node `SetProjectionInputs`.
// TODO(cjhopman): I think we should change ArtifactGroup to hold a `(TransitiveSet, ProjectionIndex)` rather
// than `(TransitiveSetKey, ProjectionIndex)`. We already have that information when constructing it and the
// artifact side of it holds a starlark ref. That would allow someone with an ArtifactGroup to synchronously
// traverse the tset graph rather than needing to asynchronously resolve a TransitiveSetKey.
async fn convert_inputs<'c, 'a, Iter: IntoIterator<Item = &'a ArtifactGroup>>(
    ctx: &'c DiceComputations,
    node_cache: DiceAqueryNodesCache<'c>,
    inputs: Iter,
) -> anyhow::Result<Vec<ActionInput>> {
    let (artifacts, projections): (Vec<_>, Vec<_>) = Itertools::partition_map(
        inputs
            .into_iter()
            .filter_map(|input| match input.assert_resolved() {
                ResolvedArtifactGroup::Artifact(a) => a.action_key().map(Either::Left),
                ResolvedArtifactGroup::TransitiveSetProjection(key) => Some(Either::Right(key)),
            }),
        |v| v,
    );
    let mut deps = artifacts.into_map(|a| ActionInput::ActionKey(a.dupe()));
    let mut projection_deps: FuturesOrdered<_> = projections
        .into_iter()
        .map(|key| {
            let key = key.dupe();
            let node_cache = node_cache.dupe();
            get_tset_node(node_cache, ctx, key)
        })
        .collect();

    while let Some(node) = tokio::task::unconstrained(projection_deps.next()).await {
        deps.push(ActionInput::IndirectInputs(node?));
    }
    Ok(deps)
}

fn compute_tset_node<'c>(
    node_cache: DiceAqueryNodesCache<'c>,
    ctx: &'c DiceComputations,
    key: TransitiveSetProjectionKey,
) -> BoxFuture<'c, SharedResult<SetProjectionInputs>> {
    async move {
        let set = ctx
            .compute_deferred_data(&key.key)
            .await
            .context("Failed to compute deferred")?;

        let sub_inputs = set
            .as_transitive_set()
            .get_projection_sub_inputs(key.projection)?;

        let inputs = convert_inputs(&ctx, node_cache, sub_inputs.iter()).await?;

        let (direct, children) = inputs.into_iter().partition_map(|v| match v {
            ActionInput::ActionKey(action_key) => Either::Left(action_key),
            ActionInput::IndirectInputs(projection) => Either::Right(projection),
        });

        Ok(SetProjectionInputs::new(key.dupe(), direct, children))
    }
    .boxed()
}

async fn get_tset_node<'c>(
    node_cache: DiceAqueryNodesCache<'c>,
    ctx: &'c DiceComputations,
    key: TransitiveSetProjectionKey,
) -> anyhow::Result<SetProjectionInputs> {
    let copied_node_cache = node_cache.dupe();
    Ok(node_cache
        .tset_nodes
        .get_or_compute(key, move |key| {
            compute_tset_node(copied_node_cache, ctx, key)
        })
        .await?)
}

fn compute_action_node<'c>(
    node_cache: DiceAqueryNodesCache<'c>,
    ctx: &'c DiceComputations,
    key: ActionKey,
    fs: Arc<ArtifactFs>,
) -> BoxFuture<'c, SharedResult<ActionQueryNode>> {
    async move {
        let action = ActionCalculation::get_action(ctx, &key).await?;
        let deps = convert_inputs(&ctx, node_cache, action.inputs()?.iter()).await?;
        Ok(ActionQueryNode::new(action, deps, fs))
    }
    .boxed()
}

async fn get_action_node<'c>(
    node_cache: DiceAqueryNodesCache<'c>,
    ctx: &'c DiceComputations,
    key: ActionKey,
    fs: Arc<ArtifactFs>,
) -> anyhow::Result<ActionQueryNode> {
    let copied_node_cache = node_cache.dupe();
    Ok(node_cache
        .action_nodes
        .get_or_compute(key, move |key| {
            compute_action_node(copied_node_cache, ctx, key, fs)
        })
        .await?)
}

impl<'c> DiceAqueryDelegate<'c> {
    pub(crate) async fn new(
        base_delegate: DiceQueryDelegate<'c>,
    ) -> anyhow::Result<DiceAqueryDelegate<'c>> {
        let artifact_fs = Arc::new(base_delegate.ctx().get_artifact_fs().await?);
        Ok(DiceAqueryDelegate {
            base_delegate,
            nodes_cache: DiceAqueryNodesCache::new(),
            artifact_fs,
        })
    }

    pub async fn get_action_node(&self, key: &ActionKey) -> anyhow::Result<ActionQueryNode> {
        get_action_node(
            self.nodes_cache.dupe(),
            self.base_delegate.ctx(),
            key.dupe(),
            self.artifact_fs.dupe(),
        )
        .await
    }
}

#[async_trait]
impl<'c> AqueryDelegate for DiceAqueryDelegate<'c> {
    fn cquery_delegate(&self) -> &dyn CqueryDelegate {
        &self.base_delegate
    }

    async fn get_node(&self, key: &ActionKey) -> anyhow::Result<ActionQueryNode> {
        self.get_action_node(key).await
    }
}

#[async_trait]
impl<'c> QueryLiterals<ActionQueryNode> for DiceAqueryDelegate<'c> {
    async fn eval_literals(&self, literals: &[&str]) -> anyhow::Result<TargetSet<ActionQueryNode>> {
        // For literal evaluation, we resolve the providers pattern to the analysis result, pull out
        // the default outputs and look up the corresponding actions.
        // TODO(cjhopman): This is a common pattern and we should probably pull it out to a common
        // place, maybe with streaming results like we have for `target pattern -> target node`
        let mut result = TargetSet::new();
        for literal in literals {
            let label = self
                .base_delegate
                .literal_parser()
                .parse_providers_pattern(literal)?;
            match label {
                ParsedPattern::Target(package, target_name, providers) => {
                    let label = providers.into_providers_label(package, target_name.as_ref());
                    let configured_label = self
                        .base_delegate
                        .ctx()
                        .get_configured_provider_label(
                            &label,
                            self.base_delegate.global_target_platform(),
                        )
                        .await?;

                    match self
                        .base_delegate
                        .ctx()
                        .get_providers(&configured_label)
                        .await?
                    {
                        MaybeCompatible::Incompatible(_) => {
                            // ignored
                        }
                        MaybeCompatible::Compatible(providers) => {
                            for output in providers
                                .provider_collection()
                                .default_info()
                                .default_outputs()
                            {
                                if let Some(action_key) = output.artifact().action_key() {
                                    result.insert(self.get_action_node(action_key).await?);
                                }
                            }
                        }
                    }
                }
                ParsedPattern::Package(_) | ParsedPattern::Recursive(_) => {
                    return Err(
                        ActionQueryError::UnsupportedTargetPattern((*literal).to_owned()).into(),
                    );
                }
            }
        }
        Ok(result)
    }
}


/*
 * 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 async_trait::async_trait;
use buck2_build_api::configure_targets::load_compatible_patterns;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::package_boundary::HasPackageBoundaryExceptions;
use buck2_common::package_boundary::PackageBoundaryExceptions;
use buck2_common::package_listing::dice::HasPackageListingResolver;
use buck2_common::package_listing::resolver::PackageListingResolver;
use buck2_common::pattern::resolve::resolve_target_patterns;
use buck2_common::pattern::resolve::ResolvedPattern;
use buck2_common::target_aliases::BuckConfigTargetAliasResolver;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellAliasResolver;
use buck2_core::cells::CellResolver;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::ProvidersPatternExtra;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::query_file_literal::parse_query_file_literal;
use buck2_core::pattern::ParsedPattern;
use buck2_core::provider::label::ProvidersName;
use buck2_core::soft_error;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_core::target::label::TargetLabel;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_node::load_patterns::load_patterns;
use buck2_node::load_patterns::MissingTargetBehavior;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_frontend::ConfiguredTargetNodeCalculation;
use buck2_node::nodes::eval_result::EvaluationResult;
use buck2_node::nodes::frontend::TargetGraphCalculation;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
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 dice::DiceComputations;
use dupe::Dupe;
use gazebo::prelude::*;
use indexmap::indexset;

use crate::cquery::environment::CqueryDelegate;
use crate::uquery::environment::QueryLiterals;
use crate::uquery::environment::UqueryDelegate;

pub mod aquery;

#[derive(Debug, thiserror::Error)]
enum LiteralParserError {
    #[error("Expected a target pattern without providers, got: `{0}`")]
    ExpectingTargetPatternWithoutProviders(String),
}

pub(crate) struct LiteralParser {
    // file and target literals are resolved relative to the working dir.
    working_dir: CellPath,
    working_dir_abs: AbsNormPathBuf,
    project_root: ProjectRoot,
    cell_resolver: CellResolver,
    cell_alias_resolver: CellAliasResolver,
    target_alias_resolver: BuckConfigTargetAliasResolver,
}

impl LiteralParser {
    // We allow provider names and flavors in the value and it gets stripped out for the result as queries operate on the target graphs.
    fn parse_target_pattern(
        &self,
        value: &str,
    ) -> anyhow::Result<ParsedPattern<TargetPatternExtra>> {
        let providers_pattern = self.parse_providers_pattern(value)?;
        let target_pattern = match providers_pattern {
            ParsedPattern::Target(package, target_name, ProvidersPatternExtra { providers }) => {
                if providers != ProvidersName::Default {
                    // After converting this to hard error, replace this function body
                    // with direct call to `ParsedPattern::parse_relative`,
                    // as `parse_providers_pattern` does.
                    soft_error!(
                        "expecting_target_pattern_without_providers",
                        LiteralParserError::ExpectingTargetPatternWithoutProviders(
                            value.to_owned()
                        )
                        .into()
                    )?;
                }
                ParsedPattern::Target(package, target_name, TargetPatternExtra)
            }
            ParsedPattern::Package(package) => ParsedPattern::Package(package),
            ParsedPattern::Recursive(path) => ParsedPattern::Recursive(path),
        };
        Ok(target_pattern)
    }

    pub(crate) fn parse_providers_pattern(
        &self,
        value: &str,
    ) -> anyhow::Result<ParsedPattern<ProvidersPatternExtra>> {
        ParsedPattern::parse_relative(
            &self.target_alias_resolver,
            self.working_dir.as_ref(),
            value,
            &self.cell_resolver,
        )
    }

    fn parse_file_literal(&self, literal: &str) -> anyhow::Result<CellPath> {
        parse_query_file_literal(
            literal,
            &self.cell_alias_resolver,
            &self.cell_resolver,
            &self.working_dir_abs,
            &self.project_root,
        )
    }
}

/// A Uquery delegate that resolves TargetNodes with the provided
/// InterpreterCalculation.
pub struct DiceQueryDelegate<'c> {
    ctx: &'c DiceComputations,
    cell_resolver: CellResolver,
    literal_parser: Arc<LiteralParser>,
    global_target_platform: Option<TargetLabel>,
    package_boundary_exceptions: Arc<PackageBoundaryExceptions>,
}

impl<'c> DiceQueryDelegate<'c> {
    pub fn new<'a>(
        ctx: &'c DiceComputations,
        working_dir: &'a ProjectRelativePath,
        project_root: ProjectRoot,
        cell_resolver: CellResolver,
        global_target_platform: Option<TargetLabel>,
        package_boundary_exceptions: Arc<PackageBoundaryExceptions>,
        target_alias_resolver: BuckConfigTargetAliasResolver,
    ) -> anyhow::Result<Self> {
        let cell_path = cell_resolver.get_cell_path(working_dir)?;
        let cell_alias_resolver = cell_resolver
            .get(cell_path.cell())?
            .cell_alias_resolver()
            .dupe();
        let working_dir_abs = project_root.resolve(working_dir);

        Ok(Self {
            ctx,
            global_target_platform,
            cell_resolver: cell_resolver.dupe(),
            literal_parser: Arc::new(LiteralParser {
                working_dir_abs,
                working_dir: cell_path,
                project_root,
                cell_resolver,
                cell_alias_resolver,
                target_alias_resolver,
            }),
            package_boundary_exceptions,
        })
    }

    pub(crate) fn ctx(&self) -> &'c DiceComputations {
        self.ctx
    }

    pub(crate) fn literal_parser(&self) -> &LiteralParser {
        &self.literal_parser
    }

    pub(crate) fn global_target_platform(&self) -> Option<&TargetLabel> {
        self.global_target_platform.as_ref()
    }
}

#[async_trait]
impl<'c> UqueryDelegate for DiceQueryDelegate<'c> {
    async fn eval_build_file(
        &self,
        package: PackageLabel,
    ) -> anyhow::Result<Arc<EvaluationResult>> {
        self.ctx.get_interpreter_results(package).await
    }

    async fn eval_module_imports(&self, path: &ImportPath) -> anyhow::Result<Vec<ImportPath>> {
        //TODO(benfoxman): Don't need to get the whole module, just parse the imports.
        let module = self.ctx.get_loaded_module_from_import_path(path).await?;
        Ok(module.imports().cloned().collect())
    }

    // get the list of potential buildfile names for each cell
    fn get_buildfile_names_by_cell(&self) -> anyhow::Result<HashMap<CellName, &[FileNameBuf]>> {
        let resolver = &self.cell_resolver;
        let mut buildfile_names_by_cell = HashMap::<CellName, &[FileNameBuf]>::new();
        for (cell, _) in resolver.cells() {
            let prev = buildfile_names_by_cell.insert(cell, resolver.get(cell)?.buildfiles());
            assert!(prev.is_none());
        }
        Ok(buildfile_names_by_cell)
    }

    async fn resolve_target_patterns(
        &self,
        patterns: &[&str],
    ) -> anyhow::Result<ResolvedPattern<TargetPatternExtra>> {
        let parsed_patterns = patterns.try_map(|p| self.literal_parser.parse_target_pattern(p))?;
        let file_ops = self.ctx.file_ops();
        resolve_target_patterns(&self.cell_resolver, &parsed_patterns, &file_ops).await
    }

    // This returns 1 package normally but can return multiple packages if the path is covered under `self.package_boundary_exceptions`.
    async fn get_enclosing_packages(&self, path: &CellPath) -> anyhow::Result<Vec<PackageLabel>> {
        let package_listing_resolver = self.ctx.get_package_listing_resolver();

        // Without package boundary violations, there is only 1 owning package for a path.
        // However, with package boundary violations, all parent packages of the enclosing package can also be owners.
        if let Some(enclosing_violation_path) = self
            .package_boundary_exceptions
            .get_package_boundary_exception_path(path)
        {
            return Ok(package_listing_resolver
                .get_enclosing_packages(path.as_ref(), enclosing_violation_path.as_ref())
                .await?
                .into_iter()
                .collect());
        }

        let package = package_listing_resolver
            .get_enclosing_package(path.as_ref())
            .await?;
        Ok(vec![package])
    }

    async fn eval_file_literal(&self, literal: &str) -> anyhow::Result<FileSet> {
        let cell_path = self.literal_parser.parse_file_literal(literal)?;
        Ok(FileSet::new(indexset![FileNode(cell_path)]))
    }
}

#[async_trait]
impl<'c> CqueryDelegate for DiceQueryDelegate<'c> {
    fn uquery_delegate(&self) -> &dyn UqueryDelegate {
        self
    }

    async fn get_node_for_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<MaybeCompatible<ConfiguredTargetNode>> {
        let target = self
            .ctx
            .get_configured_target(target, self.global_target_platform.as_ref())
            .await?;
        Ok(self.ctx.get_configured_target_node(&target).await?)
    }

    async fn get_node_for_configured_target(
        &self,
        target: &ConfiguredTargetLabel,
    ) -> anyhow::Result<ConfiguredTargetNode> {
        Ok(self
            .ctx
            .get_configured_target_node(target)
            .await?
            .require_compatible()?)
    }

    async fn get_node_for_default_configured_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<MaybeCompatible<ConfiguredTargetNode>> {
        let target = self.ctx.get_default_configured_target(target).await?;
        self.ctx.get_configured_target_node(&target).await
    }

    async fn get_configured_target(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<ConfiguredTargetLabel> {
        self.ctx
            .get_configured_target(target, self.global_target_platform.as_ref())
            .await
    }
}

#[async_trait]
impl<'c> QueryLiterals<ConfiguredTargetNode> for DiceQueryDelegate<'c> {
    async fn eval_literals(
        &self,
        literals: &[&str],
    ) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
        let parsed_patterns = literals.try_map(|p| self.literal_parser.parse_target_pattern(p))?;
        load_compatible_patterns(
            self.ctx,
            parsed_patterns,
            self.global_target_platform.dupe(),
            MissingTargetBehavior::Fail,
        )
        .await
    }
}

#[async_trait]
impl<'c> QueryLiterals<TargetNode> for DiceQueryDelegate<'c> {
    async fn eval_literals(&self, literals: &[&str]) -> anyhow::Result<TargetSet<TargetNode>> {
        let parsed_patterns = literals.try_map(|p| self.literal_parser.parse_target_pattern(p))?;
        let loaded_patterns =
            load_patterns(self.ctx, parsed_patterns, MissingTargetBehavior::Fail).await?;
        let mut target_set = TargetSet::new();
        for (_package, results) in loaded_patterns.into_iter() {
            target_set.extend(results?.into_values());
        }
        Ok(target_set)
    }
}

pub(crate) async fn get_dice_query_delegate<'a, 'c: 'a>(
    ctx: &'c DiceComputations,
    working_dir: &'a ProjectRelativePath,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<DiceQueryDelegate<'c>> {
    let cell_resolver = ctx.get_cell_resolver().await?;
    let package_boundary_exceptions = ctx.get_package_boundary_exceptions().await?;
    let target_alias_resolver = ctx
        .target_alias_resolver_for_working_dir(working_dir)
        .await?;
    let project_root = ctx
        .global_data()
        .get_io_provider()
        .project_root()
        .to_owned();
    DiceQueryDelegate::new(
        ctx,
        working_dir,
        project_root,
        cell_resolver,
        global_target_platform,
        package_boundary_exceptions,
        target_alias_resolver,
    )
}


/*
 * 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::sync::Arc;

use async_trait::async_trait;
use buck2_build_api::query::bxl::BxlUqueryFunctions;
use buck2_build_api::query::bxl::NEW_BXL_UQUERY_FUNCTIONS;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::package_boundary::HasPackageBoundaryExceptions;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_node::nodes::unconfigured::TargetNode;
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::helpers::CapturedExpr;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctions;
use buck2_query::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use dice::DiceComputations;
use dupe::Dupe;

use crate::dice::DiceQueryDelegate;
use crate::uquery::environment::UqueryEnvironment;

fn uquery_functions<'v>() -> DefaultQueryFunctions<UqueryEnvironment<'v>> {
    DefaultQueryFunctions::new()
}

struct BxlUqueryFunctionsImpl<'c> {
    ctx: &'c DiceComputations,
    project_root: ProjectRoot,
    working_dir: ProjectRelativePathBuf,
}

impl<'c> BxlUqueryFunctionsImpl<'c> {
    async fn uquery_env(&self) -> anyhow::Result<UqueryEnvironment<'c>> {
        let cell_resolver = self.ctx.get_cell_resolver().await?;

        let package_boundary_exceptions = self.ctx.get_package_boundary_exceptions().await?;
        let target_alias_resolver = self
            .ctx
            .target_alias_resolver_for_working_dir(&self.working_dir)
            .await?;

        let dice_query_delegate = Arc::new(DiceQueryDelegate::new(
            self.ctx,
            &self.working_dir,
            self.project_root.dupe(),
            cell_resolver,
            None,
            package_boundary_exceptions,
            target_alias_resolver,
        )?);
        Ok(UqueryEnvironment::new(
            dice_query_delegate.dupe(),
            dice_query_delegate,
        ))
    }
}

#[async_trait]
impl<'c> BxlUqueryFunctions<'c> for BxlUqueryFunctionsImpl<'c> {
    async fn allpaths(
        &self,
        from: &TargetSet<TargetNode>,
        to: &TargetSet<TargetNode>,
    ) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .allpaths(&self.uquery_env().await?, from, to)
            .await?)
    }
    async fn somepath(
        &self,
        from: &TargetSet<TargetNode>,
        to: &TargetSet<TargetNode>,
    ) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .somepath(&self.uquery_env().await?, from, to)
            .await?)
    }
    async fn deps(
        &self,
        targets: &TargetSet<TargetNode>,
        deps: Option<i32>,
        captured_expr: Option<&CapturedExpr>,
    ) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .deps(
                &self.uquery_env().await?,
                &DefaultQueryFunctionsModule::new(),
                targets,
                deps,
                captured_expr,
            )
            .await?)
    }
    async fn rdeps(
        &self,
        universe: &TargetSet<TargetNode>,
        targets: &TargetSet<TargetNode>,
        depth: Option<i32>,
    ) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .rdeps(&self.uquery_env().await?, universe, targets, depth)
            .await?)
    }
    async fn testsof(
        &self,
        targets: &TargetSet<TargetNode>,
    ) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .testsof(&self.uquery_env().await?, targets)
            .await?)
    }
    async fn owner(&self, file_set: &FileSet) -> anyhow::Result<TargetSet<TargetNode>> {
        Ok(uquery_functions()
            .owner(&self.uquery_env().await?, file_set)
            .await?)
    }
}

pub(crate) fn init_new_bxl_uquery_functions() {
    NEW_BXL_UQUERY_FUNCTIONS.init(|ctx, project_root, cell_name| {
        Box::pin(async move {
            let cell_resolver = ctx.get_cell_resolver().await?;
            let cell = cell_resolver.get(cell_name)?;
            // TODO(nga): working as as cell root is not right.
            //   Should be either the project root or user's current working directory.
            let working_dir = cell.path().as_project_relative_path().to_buf();

            Result::<Box<dyn BxlUqueryFunctions>, _>::Ok(Box::new(BxlUqueryFunctionsImpl {
                ctx,
                project_root,
                working_dir,
            }))
        })
    })
}