/*
 * 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 dice::UserComputationData;
use dupe::Dupe;

/// Knobs controlling how RunAction works.
#[derive(Copy, Clone, Dupe, Default)]
pub struct RunActionKnobs {
    /// Process dep files as they are generated.
    pub eager_dep_files: bool,

    /// Hash all commands using the same mechanism as dep files. This allows us to skip
    /// re-executing commands if their inputs and outputs haven't changed.
    pub hash_all_commands: bool,

    /// Whether to try reading from the action output cache (in buck-out/*/offline-cache)
    /// for network actions (download_file, cas_artifact). Used to support offline
    /// builds.
    pub use_network_action_output_cache: bool,

    /// Whether to enforce timeouts when running things on RE.
    pub enforce_re_timeouts: bool,
}

pub trait HasRunActionKnobs {
    fn set_run_action_knobs(&mut self, knobs: RunActionKnobs);

    fn get_run_action_knobs(&self) -> RunActionKnobs;
}

impl HasRunActionKnobs for UserComputationData {
    fn set_run_action_knobs(&mut self, knobs: RunActionKnobs) {
        self.data.set(knobs);
    }

    fn get_run_action_knobs(&self) -> RunActionKnobs {
        *self
            .data
            .get::<RunActionKnobs>()
            .expect("RunActionKnobs should be set")
    }
}


/*
 * 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 async_trait::async_trait;
use dice::DiceComputations;
use starlark::eval::Evaluator;

#[async_trait(?Send)]
pub trait AnonPromisesDyn<'v>: 'v {
    async fn run_promises(
        self: Box<Self>,
        dice: &DiceComputations,
        eval: &mut Evaluator<'v, '_>,
        description: String,
    ) -> anyhow::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::fmt::Debug;
use std::marker::PhantomData;

use allocative::Allocative;
use buck2_interpreter::starlark_promise::StarlarkPromise;
use buck2_node::configuration::execution::ExecutionPlatformResolution;
use buck2_util::late_binding::LateBinding;
use starlark::any::AnyLifetime;
use starlark::values::dict::DictOf;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueTyped;

use crate::analysis::anon_promises_dyn::AnonPromisesDyn;

pub static ANON_TARGET_REGISTRY_NEW: LateBinding<
    for<'v> fn(
        PhantomData<Value<'v>>,
        ExecutionPlatformResolution,
    ) -> Box<dyn AnonTargetsRegistryDyn<'v> + 'v>,
> = LateBinding::new("ANON_TARGET_REGISTRY_NEW");

pub trait AnonTargetsRegistryDyn<'v>:
    Debug + Allocative + Trace<'v> + AnyLifetime<'v> + 'v
{
    fn as_any_mut(&mut self) -> &mut dyn AnyLifetime<'v>;
    fn take_promises(&mut self) -> Option<Box<dyn AnonPromisesDyn<'v>>>;
    fn assert_no_promises(&self) -> anyhow::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.
 */

//! Rule analysis related Dice calculations
use std::collections::HashMap;
use std::future::Future;
use std::pin::Pin;

use async_trait::async_trait;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::configuration::pair::ConfigurationNoExec;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_ref::ConfiguredGraphNodeRef;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use buck2_util::late_binding::LateBinding;
use dice::DiceComputations;
use dupe::Dupe;

use crate::analysis::AnalysisResult;
use crate::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;

pub static EVAL_ANALYSIS_QUERY: LateBinding<
    for<'a> fn(
        &'a DiceComputations,
        &'a str,
        HashMap<String, ConfiguredTargetNode>,
    ) -> Pin<
        Box<dyn Future<Output = anyhow::Result<TargetSet<ConfiguredGraphNodeRef>>> + Send + 'a>,
    >,
> = LateBinding::new("EVAL_ANALYSIS_QUERY");

#[async_trait]
pub trait RuleAnalsysisCalculationImpl: Send + Sync + 'static {
    async fn get_analysis_result(
        &self,
        ctx: &DiceComputations,
        target: &ConfiguredTargetLabel,
    ) -> anyhow::Result<MaybeCompatible<AnalysisResult>>;
}

pub static RULE_ANALYSIS_CALCULATION: LateBinding<&'static dyn RuleAnalsysisCalculationImpl> =
    LateBinding::new("RULE_ANALYSIS_CALCULATION");

#[async_trait]
pub trait RuleAnalysisCalculation {
    /// Returns the analysis result for a ConfiguredTargetLabel. This is the full set of Providers
    /// returned by the target's rule implementation function.
    async fn get_analysis_result(
        &self,
        target: &ConfiguredTargetLabel,
    ) -> anyhow::Result<MaybeCompatible<AnalysisResult>>;

    /// Return the analysis result for a configuration rule `TargetLabel`
    /// (e. g. `constraint_value`).
    async fn get_configuration_analysis_result(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<AnalysisResult>;

    /// Returns the provider collection for a ConfiguredProvidersLabel. This is the full set of Providers
    /// returned by the target's rule implementation function.
    async fn get_providers(
        &self,
        target: &ConfiguredProvidersLabel,
    ) -> anyhow::Result<MaybeCompatible<FrozenProviderCollectionValue>>;
}

#[async_trait]
impl RuleAnalysisCalculation for DiceComputations {
    async fn get_analysis_result(
        &self,
        target: &ConfiguredTargetLabel,
    ) -> anyhow::Result<MaybeCompatible<AnalysisResult>> {
        RULE_ANALYSIS_CALCULATION
            .get()?
            .get_analysis_result(self, target)
            .await
    }

    async fn get_configuration_analysis_result(
        &self,
        target: &TargetLabel,
    ) -> anyhow::Result<AnalysisResult> {
        // Analysis for configuration nodes is always done with the unbound configuration.
        let target = target.configure_pair(ConfigurationNoExec::unbound().cfg_pair().dupe());
        Ok(self
            .get_analysis_result(&target)
            .await?
            .require_compatible()?)
    }

    async fn get_providers(
        &self,
        target: &ConfiguredProvidersLabel,
    ) -> anyhow::Result<MaybeCompatible<FrozenProviderCollectionValue>> {
        let analysis = self.get_analysis_result(target.target()).await?;

        analysis.try_map(|analysis| analysis.lookup_inner(target))
    }
}


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

use buck2_artifact::deferred::id::DeferredId;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_interpreter::starlark_profiler::StarlarkProfileDataAndStats;

use crate::deferred::types::DeferredLookup;
use crate::deferred::types::DeferredTable;

// TODO(@wendyy) move into `buck2_node`
pub mod anon_promises_dyn;
#[allow(unused)] // TODO(@wendyy) temporary
// TODO(@wendyy) move out of buck2_build_api and into `buck2_interpreter_for_build`
pub mod anon_targets_registry;
pub mod calculation;
pub(crate) mod promise_artifacts;
pub mod registry;

use allocative::Allocative;
use dupe::Dupe;

use crate::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;

#[derive(Debug, Clone, Dupe, Allocative)]
pub struct AnalysisResult {
    /// The actual provider collection, validated to be the correct type (`FrozenProviderCollection`)
    pub provider_collection: FrozenProviderCollectionValue,
    deferred: DeferredTable,
    pub profile_data: Option<Arc<StarlarkProfileDataAndStats>>,
}

impl AnalysisResult {
    /// Create a new AnalysisResult
    pub fn new(
        provider_collection: FrozenProviderCollectionValue,
        deferred: DeferredTable,
        profile_data: Option<Arc<StarlarkProfileDataAndStats>>,
    ) -> Self {
        Self {
            provider_collection,
            deferred,
            profile_data,
        }
    }

    pub fn providers(&self) -> &FrozenProviderCollectionValue {
        &self.provider_collection
    }

    /// Used to lookup an inner named provider result.
    pub fn lookup_inner(
        &self,
        label: &ConfiguredProvidersLabel,
    ) -> anyhow::Result<FrozenProviderCollectionValue> {
        self.provider_collection.lookup_inner(label)
    }

    pub fn lookup_deferred(&self, id: DeferredId) -> anyhow::Result<DeferredLookup<'_>> {
        self.deferred.lookup_deferred(id)
    }

    pub fn iter_deferreds(&self) -> impl Iterator<Item = DeferredLookup<'_>> {
        self.deferred.iter()
    }

    pub fn testing_deferred(&self) -> &DeferredTable {
        &self.deferred
    }
}


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

use allocative::Allocative;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_interpreter::starlark_promise::StarlarkPromise;
use dupe::Dupe;
use once_cell::sync::OnceCell;
use starlark::codemap::FileSpan;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::ValueTyped;

use crate::artifact_groups::promise::PromiseArtifact;
use crate::artifact_groups::promise::PromiseArtifactId;
use crate::artifact_groups::promise::PromiseArtifactResolveError;
use crate::interpreter::rule_defs::artifact::ValueAsArtifactLike;

#[derive(Debug, Trace, Allocative)]
struct PromiseArtifactEntry {
    location: Option<FileSpan>,
    artifact: PromiseArtifact,
    short_path: Option<ForwardRelativePathBuf>,
}

/// The PromiseArtifactRegistry stores promises registered with ctx.actions.promise_artifact() and their
/// corresponding internal PromiseArtifact. At the end of analysis (after promises have been resolved),
/// all PromiseArtifact will be updated to have the resolved artifact from the corresponding starlark promise.
#[derive(Debug, Trace, Allocative)]
pub struct PromiseArtifactRegistry<'v> {
    promises: Vec<ValueTyped<'v, StarlarkPromise<'v>>>,
    artifacts: Vec<PromiseArtifactEntry>,
    owner: BaseDeferredKey,
}
impl<'v> PromiseArtifactRegistry<'v> {
    pub(crate) fn new(owner: BaseDeferredKey) -> Self {
        Self {
            owner,
            promises: Vec::new(),
            artifacts: Vec::new(),
        }
    }

    pub(crate) fn resolve_all(&self) -> anyhow::Result<()> {
        for (promise, artifact_entry) in std::iter::zip(&self.promises, &self.artifacts) {
            match promise.get() {
                Some(v) => match ValueAsArtifactLike::unpack_value(v) {
                    Some(v) => {
                        artifact_entry
                            .artifact
                            .resolve(v.0, &artifact_entry.short_path)?;
                    }
                    None => {
                        return Err(PromiseArtifactResolveError::NotAnArtifact(
                            artifact_entry.location.clone(),
                            v.to_repr(),
                        )
                        .into());
                    }
                },
                None => {
                    return Err(PromiseArtifactResolveError::PromiseNotResolved(
                        artifact_entry.location.clone(),
                        promise.to_string(),
                    )
                    .into());
                }
            }
        }
        Ok(())
    }

    pub(crate) fn register(
        &mut self,
        promise: ValueTyped<'v, StarlarkPromise<'v>>,
        location: Option<FileSpan>,
        short_path: Option<ForwardRelativePathBuf>,
    ) -> anyhow::Result<PromiseArtifact> {
        let id = PromiseArtifactId::new(self.owner.dupe(), self.promises.len());
        let artifact = PromiseArtifact::new(Arc::new(OnceCell::new()), Arc::new(id));

        self.promises.push(promise);
        self.artifacts.push(PromiseArtifactEntry {
            location,
            artifact: artifact.dupe(),
            short_path,
        });
        Ok(artifact)
    }
}


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

use allocative::Allocative;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::artifact_type::DeclaredArtifact;
use buck2_artifact::artifact::artifact_type::OutputArtifact;
use buck2_artifact::deferred::id::DeferredId;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_execute::execute::request::OutputType;
use buck2_interpreter::starlark_promise::StarlarkPromise;
use buck2_node::configuration::execution::ExecutionPlatformResolution;
use derivative::Derivative;
use dupe::Dupe;
use indexmap::IndexSet;
use starlark::codemap::FileSpan;
use starlark::environment::FrozenModule;
use starlark::environment::Module;
use starlark::eval::Evaluator;
use starlark::values::Heap;
use starlark::values::OwnedFrozenValue;
use starlark::values::Trace;
use starlark::values::Tracer;
use starlark::values::Value;
use starlark::values::ValueTyped;

use crate::actions::registry::ActionsRegistry;
use crate::actions::UnregisteredAction;
use crate::analysis::anon_promises_dyn::AnonPromisesDyn;
use crate::analysis::anon_targets_registry::AnonTargetsRegistryDyn;
use crate::analysis::anon_targets_registry::ANON_TARGET_REGISTRY_NEW;
use crate::analysis::promise_artifacts::PromiseArtifactRegistry;
use crate::artifact_groups::promise::PromiseArtifact;
use crate::artifact_groups::registry::ArtifactGroupRegistry;
use crate::artifact_groups::ArtifactGroup;
use crate::deferred::types::BaseKey;
use crate::deferred::types::DeferredRegistry;
use crate::dynamic::registry::DynamicRegistry;
use crate::interpreter::rule_defs::artifact::associated::AssociatedArtifacts;
use crate::interpreter::rule_defs::artifact::output_artifact_like::OutputArtifactArg;
use crate::interpreter::rule_defs::artifact::StarlarkDeclaredArtifact;

#[derive(Derivative, Trace, Allocative)]
#[derivative(Debug)]
pub struct AnalysisRegistry<'v> {
    #[derivative(Debug = "ignore")]
    deferred: DeferredRegistry,
    #[derivative(Debug = "ignore")]
    actions: ActionsRegistry,
    #[derivative(Debug = "ignore")]
    artifact_groups: ArtifactGroupRegistry,
    #[derivative(Debug = "ignore")]
    dynamic: DynamicRegistry,
    pub anon_targets: Box<dyn AnonTargetsRegistryDyn<'v>>,
    artifact_promises: PromiseArtifactRegistry<'v>,
    analysis_value_storage: AnalysisValueStorage<'v>,
}

#[derive(thiserror::Error, Debug)]
enum DeclaredArtifactError {
    #[error("Can't declare an artifact with an empty filename component")]
    DeclaredEmptyFileName,
}

impl<'v> AnalysisRegistry<'v> {
    pub fn new_from_owner(
        owner: BaseDeferredKey,
        execution_platform: ExecutionPlatformResolution,
    ) -> anyhow::Result<AnalysisRegistry<'v>> {
        Self::new_from_owner_and_deferred(
            owner.dupe(),
            execution_platform,
            DeferredRegistry::new(BaseKey::Base(owner)),
        )
    }

    pub(crate) fn new_from_owner_and_deferred(
        owner: BaseDeferredKey,
        execution_platform: ExecutionPlatformResolution,
        deferred: DeferredRegistry,
    ) -> anyhow::Result<Self> {
        Ok(AnalysisRegistry {
            deferred,
            actions: ActionsRegistry::new(owner.dupe(), execution_platform.dupe()),
            artifact_groups: ArtifactGroupRegistry::new(),
            dynamic: DynamicRegistry::new(owner.dupe()),
            anon_targets: (ANON_TARGET_REGISTRY_NEW.get()?)(PhantomData, execution_platform),
            analysis_value_storage: AnalysisValueStorage::new(),
            artifact_promises: PromiseArtifactRegistry::new(owner),
        })
    }

    pub(crate) fn set_action_key(&mut self, action_key: Arc<str>) {
        self.actions.set_action_key(action_key);
    }

    /// Reserves a path in an output directory. Doesn't declare artifact,
    /// but checks that there is no previously declared artifact with a path
    /// which is in conflict with claimed `path`.
    pub fn claim_output_path(
        &mut self,
        eval: &Evaluator<'_, '_>,
        path: &ForwardRelativePath,
    ) -> anyhow::Result<()> {
        let declaration_location = eval.call_stack_top_location();
        self.actions.claim_output_path(path, declaration_location)
    }

    pub(crate) fn declare_dynamic_output(
        &mut self,
        path: BuckOutPath,
        output_type: OutputType,
    ) -> DeclaredArtifact {
        self.actions.declare_dynamic_output(path, output_type)
    }

    pub fn declare_output(
        &mut self,
        prefix: Option<&str>,
        filename: &str,
        output_type: OutputType,
        declaration_location: Option<FileSpan>,
    ) -> anyhow::Result<DeclaredArtifact> {
        // We want this artifact to be a file/directory inside the current context, which means
        // things like `..` and the empty path `.` can be bad ideas. The `::new` method checks for those
        // things and fails if they are present.

        if filename == "." || filename.is_empty() {
            return Err(DeclaredArtifactError::DeclaredEmptyFileName.into());
        }

        let path = ForwardRelativePath::new(filename)?.to_owned();
        let prefix = match prefix {
            None => None,
            Some(x) => Some(ForwardRelativePath::new(x)?.to_owned()),
        };
        self.actions
            .declare_artifact(prefix, path, output_type, declaration_location)
    }

    /// Takes a string or artifact/output artifact and converts it into an output artifact
    ///
    /// This is handy for functions like `ctx.actions.write` where it's nice to just let
    /// the user give us a string if they want as the output name.
    ///
    /// This function can declare new artifacts depending on the input.
    /// If there is no error, it returns a wrapper around the artifact (ArtifactDeclaration) and the corresponding OutputArtifact
    ///
    /// The valid types for `value` and subsequent actions are as follows:
    ///  - `str`: A new file is declared with this name.
    ///  - `StarlarkOutputArtifact`: The original artifact is returned
    ///  - `StarlarkArtifact`/`StarlarkDeclaredArtifact`: If the artifact is already bound, an error is raised. Otherwise we proceed with the original artifact.
    pub fn get_or_declare_output<'v2>(
        &mut self,
        eval: &Evaluator<'v2, '_>,
        value: OutputArtifactArg<'v2>,
        output_type: OutputType,
    ) -> anyhow::Result<(ArtifactDeclaration<'v2>, OutputArtifact)> {
        let declaration_location = eval.call_stack_top_location();
        let heap = eval.heap();
        let declared_artifact = match value {
            OutputArtifactArg::Str(path) => {
                let artifact =
                    self.declare_output(None, path, output_type, declaration_location.dupe())?;
                heap.alloc_typed(StarlarkDeclaredArtifact::new(
                    declaration_location,
                    artifact,
                    AssociatedArtifacts::new(),
                ))
            }
            OutputArtifactArg::OutputArtifact(output) => output.inner(),
            OutputArtifactArg::DeclaredArtifact(artifact) => artifact,
            OutputArtifactArg::WrongArtifact(artifact) => {
                return Err(artifact.0.as_output_error());
            }
        };

        let output = declared_artifact.output_artifact();
        output.ensure_output_type(output_type)?;
        Ok((
            ArtifactDeclaration {
                artifact: declared_artifact,
                heap,
            },
            output,
        ))
    }

    pub fn register_action<A: UnregisteredAction + 'static>(
        &mut self,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<OutputArtifact>,
        action: A,
        associated_value: Option<Value<'v>>,
    ) -> anyhow::Result<()> {
        let id = self
            .actions
            .register(&mut self.deferred, inputs, outputs, action)?;
        if let Some(value) = associated_value {
            self.analysis_value_storage.set_value(id, value);
        }
        Ok(())
    }

    pub fn create_transitive_set(
        &mut self,
        definition: Value<'v>,
        value: Option<Value<'v>>,
        children: Option<Value<'v>>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Value<'v>> {
        let set = self.artifact_groups.create_transitive_set(
            definition,
            value,
            children,
            &mut self.deferred,
            eval,
        )?;

        let key = set.key().deferred_key().id();
        let set = eval.heap().alloc_complex(set);

        self.analysis_value_storage.set_value(key, set);

        Ok(set)
    }

    pub fn register_dynamic_output(
        &mut self,
        dynamic: IndexSet<Artifact>,
        inputs: IndexSet<Artifact>,
        outputs: IndexSet<OutputArtifact>,
        attributes_lambda: Value<'v>,
    ) -> anyhow::Result<()> {
        let id = self
            .dynamic
            .register(dynamic, inputs, outputs, &mut self.deferred)?;
        self.analysis_value_storage.set_value(id, attributes_lambda);
        Ok(())
    }

    pub(crate) fn take_promises(&mut self) -> Option<Box<dyn AnonPromisesDyn<'v>>> {
        self.anon_targets.take_promises()
    }

    pub(crate) fn assert_no_promises(&self) -> anyhow::Result<()> {
        self.anon_targets.assert_no_promises()
    }

    pub fn register_artifact_promise(
        &mut self,
        promise: ValueTyped<'v, StarlarkPromise<'v>>,
        location: Option<FileSpan>,
        option: Option<ForwardRelativePathBuf>,
    ) -> anyhow::Result<PromiseArtifact> {
        self.artifact_promises.register(promise, location, option)
    }

    /// You MUST pass the same module to both the first function and the second one.
    /// It requires both to get the lifetimes to line up.
    pub fn finalize(
        self,
        env: &'v Module,
    ) -> anyhow::Result<
        impl FnOnce(Module) -> anyhow::Result<(FrozenModule, DeferredRegistry)> + 'static,
    > {
        let AnalysisRegistry {
            mut deferred,
            dynamic,
            actions,
            artifact_groups,
            anon_targets: _,
            analysis_value_storage,
            artifact_promises,
        } = self;
        artifact_promises.resolve_all()?;

        analysis_value_storage.write_to_module(env);
        Ok(move |env: Module| {
            let frozen_env = env.freeze()?;
            let analysis_value_fetcher = AnalysisValueFetcher {
                frozen_module: Some(frozen_env.dupe()),
            };
            actions.ensure_bound(&mut deferred, &analysis_value_fetcher)?;
            artifact_groups.ensure_bound(&mut deferred, &analysis_value_fetcher)?;
            dynamic.ensure_bound(&mut deferred, &analysis_value_fetcher)?;
            Ok((frozen_env, deferred))
        })
    }
}

pub struct ArtifactDeclaration<'v> {
    artifact: ValueTyped<'v, StarlarkDeclaredArtifact>,
    heap: &'v Heap,
}

impl<'v> ArtifactDeclaration<'v> {
    pub fn into_declared_artifact(
        self,
        extra_associated_artifacts: AssociatedArtifacts,
    ) -> ValueTyped<'v, StarlarkDeclaredArtifact> {
        self.heap.alloc_typed(
            self.artifact
                .with_extended_associated_artifacts(extra_associated_artifacts),
        )
    }
}

/// Store `Value<'v>` values for actions registered in an implementation function
///
/// Threading lifetimes through the various action registries is kind of a pain. So instead,
/// store the starlark values in this struct, using the `DeferredId` as the key.
///
/// These values eventually are written into the mutable `Module`, and a wrapper is
/// made available to get the `OwnedFrozenValue` back out after that `Module` is frozen.
///
/// Note that this object has internal mutation and is only expected to live for the duration
/// of impl function execution.
///
/// At the end of impl function execution, `write_to_module` should be called to ensure
/// that the values are written the top level of the `Module`.
#[derive(Debug, Allocative)]
struct AnalysisValueStorage<'v> {
    values: HashMap<DeferredId, Value<'v>>,
}

unsafe impl<'v> Trace<'v> for AnalysisValueStorage<'v> {
    fn trace(&mut self, tracer: &Tracer<'v>) {
        for v in self.values.values_mut() {
            tracer.trace(v)
        }
    }
}

/// Simple fetcher that fetches the values written in `AnalysisValueStorage::write_to_module`
///
/// These values are pulled from the `FrozenModule` that results from `env.freeze()`.
/// This is used by the action registry to make an `OwnedFrozenValue` available to
/// Actions' register function.
#[derive(Default)]
pub struct AnalysisValueFetcher {
    frozen_module: Option<FrozenModule>,
}

impl<'v> AnalysisValueStorage<'v> {
    fn new() -> Self {
        Self {
            values: HashMap::new(),
        }
    }

    /// Write all of the values to `module` using an internal name
    fn write_to_module(&self, module: &'v Module) {
        for (id, v) in self.values.iter() {
            let starlark_key = format!("$action_key_{}", id);
            module.set(&starlark_key, *v);
        }
    }

    /// Add a value to the internal hash map that maps ids -> values
    fn set_value(&mut self, id: DeferredId, value: Value<'v>) {
        self.values.insert(id, value);
    }
}

impl AnalysisValueFetcher {
    /// Get the `OwnedFrozenValue` that corresponds to a `DeferredId`, if present
    pub(crate) fn get(&self, id: DeferredId) -> anyhow::Result<Option<OwnedFrozenValue>> {
        match &self.frozen_module {
            None => Ok(None),
            Some(module) => {
                let starlark_key = format!("$action_key_{}", id);
                // This return `Err` is the symbol is private.
                // It is never private, but error is better than panic.
                module.get_option(&starlark_key)
            }
        }
    }
}


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

use allocative::Allocative;
use anyhow::Context as _;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_core::directory::Directory;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_execute::artifact::artifact_dyn::ArtifactDyn;
use buck2_execute::artifact::group::artifact_group_values_dyn::ArtifactGroupValuesDyn;
use buck2_execute::artifact_value::ArtifactValue;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::insert_artifact;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::directory::ActionSharedDirectory;
use buck2_execute::directory::INTERNER;
use dupe::Dupe;
use smallvec::smallvec;
use smallvec::SmallVec;

/// The [`ArtifactValue`]s for an [`crate::artifact_groups::ArtifactGroup`].
#[derive(Clone, Dupe, Allocative)]
pub struct ArtifactGroupValues(pub(super) Arc<ArtifactGroupValuesData>);

impl ArtifactGroupValues {
    /// Create a new instance of ArtifactGroupValues for a TransitiveSetProjection. This expects
    /// that all the children *will* have a Directory.
    pub fn new(
        values: SmallVec<[(Artifact, ArtifactValue); 1]>,
        children: Vec<Self>,
        artifact_fs: &ArtifactFs,
        digest_config: DigestConfig,
    ) -> anyhow::Result<Self> {
        let mut builder = ActionDirectoryBuilder::empty();

        for (artifact, value) in values.iter() {
            let path = artifact
                .resolve_path(artifact_fs)
                .context("Invalid artifact")?;
            insert_artifact(&mut builder, path.as_ref(), value)?;
        }

        for child in children.iter() {
            // NOTE: Technically, we could fall back to iterating the artifacts in the
            // ArtifactGroupValues here, but we *do* rely on the fact that TransitiveSetProjections
            // produce intermediate directories, so if they don't, it is preferable to report it.
            let child_dir = child
                .0
                .directory
                .as_ref()
                .context("TransitiveSetProjection was missing directory!")?;

            builder
                .merge(child_dir.to_builder())
                .context("Merge failed")?;
        }

        let directory = builder
            .fingerprint(digest_config.as_directory_serializer())
            .shared(&*INTERNER);

        Ok(Self(Arc::new(ArtifactGroupValuesData {
            values,
            children,
            directory: Some(directory),
        })))
    }

    pub fn from_artifact(artifact: Artifact, value: ArtifactValue) -> Self {
        Self(Arc::new(ArtifactGroupValuesData {
            values: smallvec![(artifact, value)],
            children: Vec::new(),
            directory: None,
        }))
    }

    pub fn add_to_directory(
        &self,
        builder: &mut ActionDirectoryBuilder,
        artifact_fs: &ArtifactFs,
    ) -> anyhow::Result<()> {
        if let Some(d) = self.0.directory.as_ref() {
            builder.merge(d.to_builder())?;
            return Ok(());
        }

        for (artifact, value) in self.iter() {
            let projrel_path = artifact.resolve_path(artifact_fs)?;
            insert_artifact(builder, projrel_path.as_ref(), value)?;
        }

        Ok(())
    }

    pub fn iter(&self) -> impl Iterator<Item = &(Artifact, ArtifactValue)> {
        TransitiveSetIterator::new(self)
    }

    pub fn shallow_equals(&self, other: &Self) -> bool {
        let this = &self.0;
        let other = &other.0;

        this.values == other.values
            && this
                .children
                .iter()
                .eq_by(&other.children, |x, y| Arc::ptr_eq(&x.0, &y.0))
    }
}

#[derive(Allocative)]
pub struct ArtifactGroupValuesData {
    pub(super) values: SmallVec<[(Artifact, ArtifactValue); 1]>,
    pub(super) children: Vec<ArtifactGroupValues>,
    /// If set, a precomputed directory represented the union of all values in this
    /// ArtifactGroupValuesData.
    pub(super) directory: Option<ActionSharedDirectory>,
}

/// An opaque identifier for the identity of a ArtifactGroupValue. There is no operation on this
/// that makes sense except for comparison.
#[derive(Hash, Eq, PartialEq)]
pub struct ArtifactValueIdentity(usize);

impl TransitiveSetContainer for ArtifactGroupValues {
    type Value = (Artifact, ArtifactValue);
    type Identity = ArtifactValueIdentity;

    fn values(&self) -> &[Self::Value] {
        &self.0.values
    }

    fn children(&self) -> &[Self] {
        &self.0.children
    }

    fn identity(&self) -> Self::Identity {
        ArtifactValueIdentity(Arc::as_ptr(&self.0) as usize)
    }
}

pub trait TransitiveSetContainer: Sized {
    type Value: Sized;
    type Identity: Hash + Eq + PartialEq;

    fn values(&self) -> &[Self::Value];

    fn children(&self) -> &[Self];

    fn identity(&self) -> Self::Identity;
}

pub struct TransitiveSetIterator<'a, C, V, I> {
    values: &'a [V],
    queue: Vec<&'a C>,
    seen: HashSet<I>,
}

impl<'a, C>
    TransitiveSetIterator<
        'a,
        C,
        <C as TransitiveSetContainer>::Value,
        <C as TransitiveSetContainer>::Identity,
    >
where
    C: TransitiveSetContainer,
{
    pub fn new(container: &'a C) -> Self {
        let mut ret = Self {
            values: container.values(),
            queue: Vec::new(),
            seen: HashSet::new(),
        };
        ret.enqueue_children(container.children());
        ret
    }

    pub fn enqueue_children(&mut self, transitive: &'a [C]) {
        for t in transitive.iter().rev() {
            if self.seen.insert(t.identity()) {
                self.queue.push(t);
            }
        }
    }
}

impl<'a, C> Iterator
    for TransitiveSetIterator<
        'a,
        C,
        <C as TransitiveSetContainer>::Value,
        <C as TransitiveSetContainer>::Identity,
    >
where
    C: TransitiveSetContainer,
{
    type Item = &'a <C as TransitiveSetContainer>::Value;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some((v, rest)) = self.values.split_first() {
                self.values = rest;
                return Some(v);
            }

            let next = self.queue.pop()?;
            self.values = next.values();
            self.enqueue_children(next.children());
        }
    }
}

impl ArtifactGroupValuesDyn for ArtifactGroupValues {
    fn iter(&self) -> Box<dyn Iterator<Item = (&dyn ArtifactDyn, &ArtifactValue)> + '_> {
        Box::new(
            self.iter()
                .map(|(artifact, value)| (artifact as &dyn ArtifactDyn, value)),
        )
    }

    fn add_to_directory(
        &self,
        builder: &mut ActionDirectoryBuilder,
        artifact_fs: &ArtifactFs,
    ) -> anyhow::Result<()> {
        self.add_to_directory(builder, artifact_fs)
    }
}

#[cfg(test)]
mod tests {
    use buck2_artifact::artifact::artifact_type::testing::BuildArtifactTestingExt;
    use buck2_artifact::artifact::build_artifact::BuildArtifact;
    use buck2_artifact::deferred::id::DeferredId;
    use buck2_core::configuration::data::ConfigurationData;
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
    use buck2_core::target::label::ConfiguredTargetLabel;
    use buck2_execute::artifact_value::ArtifactValue;

    use super::*;

    fn artifact(name: &str) -> (Artifact, ArtifactValue) {
        let target =
            ConfiguredTargetLabel::testing_parse("cell//pkg:foo", ConfigurationData::testing_new());

        let artifact = BuildArtifact::testing_new(
            target.dupe(),
            ForwardRelativePathBuf::unchecked_new(name.to_owned()),
            DeferredId::testing_new(0),
        );

        let value = ArtifactValue::file(DigestConfig::testing_default().empty_file());

        (Artifact::from(artifact), value)
    }

    impl ArtifactGroupValuesData {
        fn value(mut self, v: &(Artifact, ArtifactValue)) -> Self {
            self.values.push((v.0.dupe(), v.1.dupe()));
            self
        }

        fn chain(mut self, child: &ArtifactGroupValues) -> Self {
            self.children.push(child.dupe());
            self
        }

        fn build(self) -> ArtifactGroupValues {
            ArtifactGroupValues(Arc::new(self))
        }
    }

    fn builder() -> ArtifactGroupValuesData {
        ArtifactGroupValuesData {
            values: Default::default(),
            children: Default::default(),
            directory: None,
        }
    }

    #[test]
    fn test_iter() {
        let a1 = artifact("a1");
        let a2 = artifact("a1");
        let a3 = artifact("a1");

        let v2 = builder().value(&a2).build();
        let v3 = builder().value(&a3).build();
        let values = builder().value(&a1).chain(&v2).chain(&v3).build();

        let mut iter = values.iter();
        assert_eq!(iter.next(), Some(&a1));
        assert_eq!(iter.next(), Some(&a2));
        assert_eq!(iter.next(), Some(&a3));
        assert_eq!(iter.next(), None);
    }

    #[test]
    fn test_shallow_eq() {
        let a1 = artifact("a1");
        let a2 = artifact("a1");
        let a3 = artifact("a1");

        let v2 = builder().value(&a2).build();
        let v3 = builder().value(&a3).build();

        {
            let s1 = builder().value(&a1).chain(&v2).chain(&v3).build();
            let s2 = builder().value(&a1).chain(&v2).chain(&v3).build();
            assert!(s1.shallow_equals(&s2));
        }

        {
            // Different artifacts
            let s1 = builder().value(&a1).chain(&v2).chain(&v3).build();
            let s2 = builder().chain(&v2).chain(&v3).build();
            assert!(!s1.shallow_equals(&s2));
        }

        {
            // Different children
            let s1 = builder().value(&a1).chain(&v2).chain(&v3).build();
            let s2 = builder().value(&a1).chain(&v2).build();
            assert!(!s1.shallow_equals(&s2));
        }
    }
}


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

use allocative::Allocative;
use anyhow::Context as _;
use async_recursion::async_recursion;
use async_trait::async_trait;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::artifact_type::ArtifactKind;
use buck2_artifact::artifact::artifact_type::BaseArtifactKind;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_artifact::artifact::projected_artifact::ProjectedArtifact;
use buck2_artifact::artifact::source_artifact::SourceArtifact;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::file_ops::FileOps;
use buck2_common::file_ops::PathMetadata;
use buck2_common::file_ops::PathMetadataOrRedirection;
use buck2_common::result::SharedResult;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::directory::DirectoryData;
use buck2_execute::artifact_value::ArtifactValue;
use buck2_execute::digest_config::HasDigestConfig;
use buck2_execute::directory::extract_artifact_value;
use buck2_execute::directory::insert_artifact;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::directory::ActionDirectoryEntry;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::directory::ActionSharedDirectory;
use buck2_execute::directory::INTERNER;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use futures::future;
use futures::stream::FuturesOrdered;
use futures::Future;
use futures::FutureExt;
use more_futures::cancellation::CancellationContext;
use ref_cast::RefCast;
use smallvec::SmallVec;
use thiserror::Error;

use crate::actions::artifact::get_artifact_fs::GetArtifactFs;
use crate::actions::calculation::ActionCalculation;
use crate::actions::execute::action_executor::ActionOutputs;
use crate::artifact_groups::ArtifactGroup;
use crate::artifact_groups::ArtifactGroupValues;
use crate::artifact_groups::ResolvedArtifactGroup;
use crate::artifact_groups::TransitiveSetProjectionKey;
use crate::deferred::calculation::DeferredCalculation;
use crate::keep_going;

#[async_trait]
pub trait ArtifactGroupCalculation {
    /// Makes an 'Artifact' available to be accessed
    async fn ensure_artifact_group(
        &self,
        input: &ArtifactGroup,
    ) -> anyhow::Result<ArtifactGroupValues>;
}

#[async_trait]
impl ArtifactGroupCalculation for DiceComputations {
    /// makes the 'Artifact' available to be accessed
    async fn ensure_artifact_group(
        &self,
        input: &ArtifactGroup,
    ) -> anyhow::Result<ArtifactGroupValues> {
        // TODO consider if we need to cache this
        ensure_artifact_group_staged(self, input)
            .await?
            .to_group_values(input)
    }
}

/// A large build may have many artifact dependency edges and so may have many of the
/// `ensure_build_artifact_*()` futures live at any time. To support this efficiently
/// we provide these `*_staged()` functions that provide an optimized Future implementation
/// for waiting on the dependency edge and an optimized form to represent the result (as
/// we also may have many of those results alive across await points as things need to wait
/// on all dependencies).
///
/// Performance sensitive things should use these staged functions, wait for all their results
/// and then synchronously process them and drop any intermediate data structures before their
/// next yield point.
///
/// Some of the optimizations this provides:
///  - The staged future is kept to a minimum size (which we track in an assertion below).
///  - The result of the staged future is kept to a minimum size (also tracked below).
///  - For the single Artifact case from ensure_artifact_group_staged, we defer allocation
///    of the ArtifactGroupValues until `to_group_values()` is called. For callers waiting
///    on many inputs, this allows them to only allocate those large values only after all
///    inputs are ready.
pub(crate) fn ensure_artifact_group_staged<'a>(
    ctx: &'a DiceComputations,
    input: &'a ArtifactGroup,
) -> impl Future<Output = anyhow::Result<EnsureArtifactGroupReady>> + 'a {
    match input.assert_resolved() {
        ResolvedArtifactGroup::Artifact(artifact) => {
            ensure_artifact_staged(ctx, artifact).left_future()
        }
        ResolvedArtifactGroup::TransitiveSetProjection(key) => ctx
            .compute(EnsureTransitiveSetProjectionKey::ref_cast(key))
            .map(|v| Ok(EnsureArtifactGroupReady::TransitiveSet(v??)))
            .right_future(),
    }
}

/// See [ensure_artifact_group_staged].
pub(super) fn ensure_base_artifact_staged<'a>(
    dice: &'a DiceComputations,
    artifact: &'a BaseArtifactKind,
) -> impl Future<Output = anyhow::Result<EnsureArtifactGroupReady>> + 'a {
    match artifact {
        BaseArtifactKind::Build(built) => ensure_build_artifact_staged(dice, built).left_future(),
        BaseArtifactKind::Source(source) => {
            ensure_source_artifact_staged(dice, source).right_future()
        }
    }
}

/// See [ensure_artifact_group_staged].
pub(super) fn ensure_artifact_staged<'a>(
    dice: &'a DiceComputations,
    artifact: &'a Artifact,
) -> impl Future<Output = anyhow::Result<EnsureArtifactGroupReady>> + 'a {
    match artifact.data() {
        ArtifactKind::Base(base) => ensure_base_artifact_staged(dice, base).left_future(),
        ArtifactKind::Projected(projected) => dice
            .compute(EnsureProjectedArtifactKey::ref_cast(projected))
            .map(|v| Ok(EnsureArtifactGroupReady::Single(v??)))
            .right_future(),
    }
}

fn ensure_build_artifact_staged<'a>(
    dice: &'a DiceComputations,
    built: &'a BuildArtifact,
) -> impl Future<Output = anyhow::Result<EnsureArtifactGroupReady>> + 'a {
    ActionCalculation::build_action(dice, built.key()).map(move |action_outputs| {
        let action_outputs = action_outputs?;
        if let Some(value) = action_outputs.get(built.get_path()) {
            Ok(EnsureArtifactGroupReady::Single(value.dupe()))
        } else {
            Err(
                EnsureArtifactStagedError::BuildArtifactMissing(built.clone(), action_outputs)
                    .into(),
            )
        }
    })
}

fn ensure_source_artifact_staged<'a>(
    dice: &'a DiceComputations,
    source: &'a SourceArtifact,
) -> impl Future<Output = anyhow::Result<EnsureArtifactGroupReady>> + 'a {
    async move {
        Ok(EnsureArtifactGroupReady::Single(
            path_artifact_value(dice, Arc::new(source.get_path().to_cell_path()))
                .await?
                .into(),
        ))
    }
    .boxed()
}

// These errors should be unreachable, they indicate misuse of the staged ensure artifact (or other buck
// invariant violations), but it's still better to propagate them as Error than to panic!().
#[derive(Debug, Error)]
pub enum EnsureArtifactStagedError {
    #[error("Tried to unpack single artifact, but got transitive set")]
    UnpackSingleTransitiveSet,
    #[error("Expected a transitive set, got a single artifact")]
    ExpectedTransitiveSet,
    // This one could probably be a panic! if DICE didn't eagerly re-evaluate all deps.
    #[error("Building an artifact didn't produce it. Expected `{0}` but only have `{1:?}`")]
    BuildArtifactMissing(BuildArtifact, ActionOutputs),
}

/// Represents the "ready" stage of an ensure_artifact_*() call. At this point the
/// ArtifactValue/ArtifactGroupValues can be synchronously accessed/constructed.
pub(crate) enum EnsureArtifactGroupReady {
    Single(ArtifactValue),
    TransitiveSet(ArtifactGroupValues),
}

impl EnsureArtifactGroupReady {
    /// Converts the ensured artifact to an ArtifactGroupValues. The caller must ensure that the passed in artifact
    /// is the same one that was used to ensure this.
    pub(crate) fn to_group_values(
        self,
        artifact: &ArtifactGroup,
    ) -> anyhow::Result<ArtifactGroupValues> {
        match self {
            EnsureArtifactGroupReady::TransitiveSet(values) => Ok(values),
            EnsureArtifactGroupReady::Single(value) => match artifact.assert_resolved() {
                ResolvedArtifactGroup::Artifact(artifact) => {
                    Ok(ArtifactGroupValues::from_artifact(artifact.clone(), value))
                }
                ResolvedArtifactGroup::TransitiveSetProjection(_) => {
                    Err(EnsureArtifactStagedError::ExpectedTransitiveSet.into())
                }
            },
        }
    }

    fn unpack_single(self) -> anyhow::Result<ArtifactValue> {
        match self {
            EnsureArtifactGroupReady::Single(value) => Ok(value),
            EnsureArtifactGroupReady::TransitiveSet(..) => {
                Err(EnsureArtifactStagedError::UnpackSingleTransitiveSet.into())
            }
        }
    }
}

static_assertions::assert_eq_size!(EnsureArtifactGroupReady, [usize; 3]);

// This assertion assures we don't unknowingly regress the size of this critical future.
// TODO(cjhopman): We should be able to wrap this in a convenient assertion macro.
#[allow(unused, clippy::diverging_sub_expression)]
fn _assert_ensure_artifact_group_future_size() {
    let v = ensure_artifact_group_staged(panic!(), panic!());
    let e = [0u8; 704 / 8];
    static_assertions::assert_eq_size_ptr!(&v, &e);
}

async fn dir_artifact_value(
    ctx: &DiceComputations,
    cell_path: Arc<CellPath>,
) -> anyhow::Result<ActionDirectoryEntry<ActionSharedDirectory>> {
    // We keep running into this performance footgun where a large directory is declared
    // as a source on a toolchain, and then every BuildKey using that toolchain ends up taking
    // a DICE edge on PathMetadataKey of every file inside that directory, blowing up Buck2's
    // memory use. This diff introduces an intermediate DICE key `DirArtifactValueKey` for
    // getting the artifact value of a source directory. Every BuildKey
    // using that directory now only depends on one DirArtifactValueKey, and that DirArtifactValueKey
    // depends on the PathMetadataKey of every member of the directory.
    #[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
    #[display(fmt = "dir_artifact_value({})", .0)]
    struct DirArtifactValueKey(Arc<CellPath>);

    #[async_trait]
    impl Key for DirArtifactValueKey {
        type Value = SharedResult<ActionSharedDirectory>;

        async fn compute(
            &self,
            ctx: &DiceComputations,
            _cancellation: &CancellationContext,
        ) -> Self::Value {
            let file_ops = ctx.file_ops();
            let files = file_ops.read_dir(self.0.as_ref().as_ref()).await?.included;

            let entries = files.iter().map(|x| async {
                // TODO(scottcao): This current creates a `DirArtifactValueKey` for each subdir of a source directory.
                // Instead, this should be 1 key for the entire top-level directory since there's almost
                // no chance of getting cache hit with a sub-directory.
                let value =
                    path_artifact_value(ctx, Arc::new(self.0.as_ref().join(&x.file_name))).await?;
                anyhow::Ok((x.file_name.clone(), value))
            });

            let entries = future::try_join_all(entries).await?;
            let entries = entries.into_iter().collect();

            let digest_config = ctx.global_data().get_digest_config();
            let d: DirectoryData<_, _, _> =
                DirectoryData::new(entries, digest_config.as_directory_serializer());
            Ok(INTERNER.intern(d))
        }

        fn equality(x: &Self::Value, y: &Self::Value) -> bool {
            match (x, y) {
                (Ok(x), Ok(y)) => x.fingerprint() == y.fingerprint(),
                _ => false,
            }
        }
    }

    let res = ctx.compute(&DirArtifactValueKey(cell_path)).await??;
    Ok(ActionDirectoryEntry::Dir(res))
}

#[async_recursion]
async fn path_artifact_value(
    ctx: &DiceComputations,
    cell_path: Arc<CellPath>,
) -> anyhow::Result<ActionDirectoryEntry<ActionSharedDirectory>> {
    let file_ops = &ctx.file_ops() as &dyn FileOps;
    let raw = file_ops
        .read_path_metadata(cell_path.as_ref().as_ref())
        .await?;
    match PathMetadataOrRedirection::from(raw) {
        PathMetadataOrRedirection::PathMetadata(meta) => match meta {
            PathMetadata::ExternalSymlink(symlink) => Ok(ActionDirectoryEntry::Leaf(
                ActionDirectoryMember::ExternalSymlink(symlink),
            )),
            PathMetadata::File(metadata) => Ok(ActionDirectoryEntry::Leaf(
                ActionDirectoryMember::File(metadata),
            )),
            PathMetadata::Directory => dir_artifact_value(ctx, cell_path).await,
        },
        PathMetadataOrRedirection::Redirection(r) => {
            // TODO (T126181780): This should have a limit on recursion.
            path_artifact_value(ctx, r).await
        }
    }
}

#[derive(Clone, Dupe, Eq, PartialEq, Hash, Display, Debug, Allocative, RefCast)]
#[repr(transparent)]
pub struct EnsureProjectedArtifactKey(pub(crate) ProjectedArtifact);

#[async_trait]
impl Key for EnsureProjectedArtifactKey {
    type Value = SharedResult<ArtifactValue>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellation: &CancellationContext,
    ) -> Self::Value {
        let base_value = ensure_base_artifact_staged(ctx, self.0.base())
            .await?
            .unpack_single()?;

        let artifact_fs = ctx.get_artifact_fs().await?;
        let digest_config = ctx.global_data().get_digest_config();

        let base_path = match self.0.base() {
            BaseArtifactKind::Build(built) => artifact_fs.resolve_build(built.get_path()),
            BaseArtifactKind::Source(source) => artifact_fs.resolve_source(source.get_path())?,
        };

        let mut builder = ActionDirectoryBuilder::empty();
        insert_artifact(&mut builder, base_path.as_ref(), &base_value)?;

        let value = extract_artifact_value(&builder, &base_path.join(self.0.path()), digest_config)
            .with_context(|| {
                format!(
                    "The path `{}` cannot be projected in the artifact `{}`",
                    self.0.path(),
                    self.0.base()
                )
            })?
            .with_context(|| {
                format!(
                    "The path `{}` does not exist in the artifact `{}`",
                    self.0.path(),
                    self.0.base()
                )
            })?;

        Ok(value)
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Ok(x), Ok(y)) => x == y,
            _ => false,
        }
    }
}

#[derive(Clone, Dupe, Eq, PartialEq, Hash, Display, Debug, Allocative, RefCast)]
#[repr(transparent)]
pub struct EnsureTransitiveSetProjectionKey(pub TransitiveSetProjectionKey);

#[async_trait]
impl Key for EnsureTransitiveSetProjectionKey {
    type Value = SharedResult<ArtifactGroupValues>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        _cancellation: &CancellationContext,
    ) -> Self::Value {
        let set = ctx
            .compute_deferred_data(&self.0.key)
            .await
            .context("Failed to compute deferred")?;

        let artifact_fs = ctx.get_artifact_fs().await?;

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

        let (values, children) = {
            // Compute the new inputs. Note that ordering here (and below) is important to ensure
            // stability of the ArtifactGroupValues we produce across executions, so we use
            // FuturesOrdered.

            let ensure_futs: FuturesOrdered<_> = sub_inputs
                .iter()
                .map(|v| ensure_artifact_group_staged(ctx, v))
                .collect();

            let ready_inputs: Vec<_> =
                tokio::task::unconstrained(keep_going::try_join_all(ctx, ensure_futs)).await?;

            // Partition our inputs in artifacts and projections.
            let mut values_count = 0;
            for input in sub_inputs.iter() {
                if let ArtifactGroup::Artifact(..) = input {
                    values_count += 1;
                }
            }

            let mut values = SmallVec::<[_; 1]>::with_capacity(values_count);
            let mut children = Vec::with_capacity(sub_inputs.len() - values_count);

            for (group, ready) in zip(sub_inputs.iter(), ready_inputs.into_iter()) {
                match group.assert_resolved() {
                    ResolvedArtifactGroup::Artifact(artifact) => {
                        values.push((artifact.dupe(), ready.unpack_single()?))
                    }
                    ResolvedArtifactGroup::TransitiveSetProjection(..) => {
                        children.push(ready.to_group_values(group)?)
                    }
                }
            }
            (values, children)
        };

        // At this point we're holding a lot of data and want to ensure that we don't hold that across any
        // .await, so move into a little sync closure and call that
        (move || {
            let digest_config = ctx.global_data().get_digest_config();

            let values = ArtifactGroupValues::new(values, children, &artifact_fs, digest_config)
                .context("Failed to construct ArtifactGroupValues")?;

            Ok(values)
        })()
    }

    fn equality(x: &Self::Value, y: &Self::Value) -> bool {
        match (x, y) {
            (Ok(x), Ok(y)) => x.shallow_equals(y),
            _ => false,
        }
    }
}


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

use allocative::Allocative;
use buck2_artifact::deferred::data::DeferredData;
use dupe::Dupe;
use starlark::values::OwnedFrozenValue;
use starlark::values::OwnedFrozenValueTyped;
use starlark::values::Value;

use crate::deferred::types::AnyValue;
use crate::deferred::types::TrivialDeferred;
use crate::interpreter::rule_defs::transitive_set::FrozenTransitiveSet;

pub type TransitiveSetKey = DeferredData<DeferredTransitiveSetData>;

#[derive(Clone, Dupe, Debug, Allocative)]
pub struct DeferredTransitiveSetData(
    #[allocative(skip)] // TODO(nga): visit heap.
    pub(super)  OwnedFrozenValueTyped<FrozenTransitiveSet>,
);

impl TrivialDeferred for DeferredTransitiveSetData {
    fn as_any_value(&self) -> &dyn AnyValue {
        self
    }

    fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
}

impl DeferredTransitiveSetData {
    pub fn testing_new(value: OwnedFrozenValue) -> DeferredTransitiveSetData {
        DeferredTransitiveSetData(value.downcast().unwrap())
    }

    pub fn as_transitive_set(&self) -> OwnedFrozenValueTyped<FrozenTransitiveSet> {
        self.0.dupe()
    }

    pub fn as_value(&self) -> Value {
        self.0.to_value()
    }
}


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

mod artifact_group_values;
pub mod calculation;
pub mod deferred;
pub mod promise;
pub mod registry;

use std::hash::Hash;

use allocative::Allocative;
pub use artifact_group_values::ArtifactGroupValues;
use buck2_artifact::artifact::artifact_type::Artifact;
use derive_more::Display;
use dupe::Dupe;
use gazebo::variants::UnpackVariants;

use crate::artifact_groups::deferred::TransitiveSetKey;
use crate::artifact_groups::promise::PromiseArtifact;

/// An [ArtifactGroup] can expand to one or more [Artifact]. Those Artifacts will be made available
/// to Actions when they execute.
#[derive(
    Clone,
    Debug,
    Display,
    Dupe,
    PartialEq,
    Eq,
    Hash,
    UnpackVariants,
    Allocative
)]
pub enum ArtifactGroup {
    Artifact(Artifact),
    TransitiveSetProjection(TransitiveSetProjectionKey),
    Promise(PromiseArtifact),
}

impl ArtifactGroup {
    pub fn assert_resolved(&self) -> ResolvedArtifactGroup {
        self.resolved().unwrap()
    }

    pub fn resolved(&self) -> anyhow::Result<ResolvedArtifactGroup> {
        Ok(match self {
            ArtifactGroup::Artifact(a) => ResolvedArtifactGroup::Artifact(a),
            ArtifactGroup::TransitiveSetProjection(a) => {
                ResolvedArtifactGroup::TransitiveSetProjection(a)
            }
            ArtifactGroup::Promise(p) => ResolvedArtifactGroup::Artifact(p.get_err()?),
        })
    }
}
pub enum ResolvedArtifactGroup<'a> {
    Artifact(&'a Artifact),
    TransitiveSetProjection(&'a TransitiveSetProjectionKey),
}

#[derive(Clone, Debug, Display, Dupe, PartialEq, Eq, Hash, Allocative)]
#[display(fmt = "TransitiveSetProjection({}, {})", key, projection)]
pub struct TransitiveSetProjectionKey {
    pub key: TransitiveSetKey,
    pub projection: usize,
}


/*
 * 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::fmt::Display;
use std::hash::Hash;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use dupe::Dupe;
use once_cell::sync::OnceCell;
use starlark::codemap::FileSpan;
use thiserror::Error;

use crate::interpreter::rule_defs::artifact::StarlarkArtifactLike;

#[derive(Debug, Error)]
pub(crate) enum PromiseArtifactResolveError {
    #[error("artifact_promise(){} resolved promise was not an artifact (was `{1}`)", maybe_declared_at(&_0))]
    NotAnArtifact(Option<FileSpan>, String),
    #[error("artifact_promise(){1}{} promise wasn't resolved", maybe_declared_at(&_0))]
    PromiseNotResolved(Option<FileSpan>, String),
    #[error("artifact_promise() resolved multiple times")]
    AlreadyResolved,
    #[error(
        "artifact_promise(){0} is not yet resolved, this indicates it was used as an artifact during the analysis where it was declared"
    )]
    PromiseNotYetResolved(PromiseArtifact),
    #[error(
        "artifact_promise() resolved to artifact with associated artifacts, this isn't allowed"
    )]
    HasAssociatedArtifacts,
    #[error(
        "artifact_promise() resolved to a source artifact, this isn't allowed: an artifact_promise() must resolve to a build artifact"
    )]
    SourceArtifact,
    #[error(
        "artifact_promise() was called with `short_path = {0}`, but it did not match the artifact's actual short path: `{1}`"
    )]
    ShortPathMismatch(ForwardRelativePathBuf, String),
}

fn maybe_declared_at(location: &Option<FileSpan>) -> String {
    match location {
        Some(v) => format!(" (declared at {})", v),
        None => String::new(),
    }
}

/// A PromiseArtifact is used to assert that a starlark promise will resolve to an artifact. While the analysis
/// that declares the promise is running, the promise may not yet be resolved. In that case, operations that require
/// the underlying artifact will fail. Once that analysis is complete, all promises will be resolved and this will
/// have a reference to the fully resolved artifact.
#[derive(Clone, Debug, Dupe, Allocative)]
pub struct PromiseArtifact {
    artifact: Arc<OnceCell<Artifact>>,
    id: Arc<PromiseArtifactId>,
}

#[derive(Clone, Debug, Dupe, Allocative, Hash, Eq, PartialEq)]
pub struct PromiseArtifactId {
    owner: BaseDeferredKey,
    id: usize,
}
impl PromiseArtifactId {
    pub(crate) fn new(owner: BaseDeferredKey, id: usize) -> PromiseArtifactId {
        Self { owner, id }
    }
}

impl PromiseArtifact {
    pub fn new(artifact: Arc<OnceCell<Artifact>>, id: Arc<PromiseArtifactId>) -> Self {
        Self { artifact, id }
    }

    pub fn get_err(&self) -> anyhow::Result<&Artifact> {
        match self.artifact.get() {
            Some(v) => Ok(v),
            None => Err(PromiseArtifactResolveError::PromiseNotYetResolved(self.clone()).into()),
        }
    }

    pub fn unwrap(&self) -> &Artifact {
        self.get_err().unwrap()
    }

    pub fn get(&self) -> Option<&Artifact> {
        self.artifact.get()
    }

    pub(crate) fn resolve(
        &self,
        artifact: &dyn StarlarkArtifactLike,
        expected_short_path: &Option<ForwardRelativePathBuf>,
    ) -> anyhow::Result<()> {
        if let Some(v) = artifact.get_associated_artifacts() {
            if !v.is_empty() {
                return Err(PromiseArtifactResolveError::HasAssociatedArtifacts.into());
            }
        }
        let bound = artifact
            .get_bound_artifact()
            .context("expected bound artifact for promise_artifact resolve")?;
        if bound.is_source() {
            return Err(PromiseArtifactResolveError::SourceArtifact.into());
        }
        if let Some(expected_short_path) = expected_short_path {
            bound.get_path().with_short_path(|artifact_short_path| {
                if artifact_short_path != expected_short_path {
                    Err(anyhow::Error::new(
                        PromiseArtifactResolveError::ShortPathMismatch(
                            expected_short_path.clone(),
                            artifact_short_path.to_string(),
                        ),
                    ))
                } else {
                    Ok(())
                }
            })?;
        }
        match self.artifact.set(bound) {
            Ok(_) => Ok(()),
            Err(_) => Err(PromiseArtifactResolveError::AlreadyResolved.into()),
        }
    }

    pub(crate) fn id(&self) -> &PromiseArtifactId {
        self.id.as_ref()
    }
}

impl Display for PromiseArtifact {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if let Some(artifact) = self.artifact.get() {
            write!(f, "PromiseArtifact(Resolved({}))", artifact)
        } else {
            write!(
                f,
                "PromiseArtifact(Unresolved({}, {}))",
                &self.id.owner, self.id.id
            )
        }
    }
}

impl Hash for PromiseArtifact {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.id.hash(state);
    }
}

impl PartialEq for PromiseArtifact {
    fn eq(&self, other: &Self) -> bool {
        self.id == other.id
    }
}

impl Eq for PromiseArtifact {}


/*
 * 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 allocative::Allocative;
use anyhow::Context as _;
use dupe::Dupe;
use starlark::eval::Evaluator;
use starlark::values::Value;

use crate::analysis::registry::AnalysisValueFetcher;
use crate::artifact_groups::deferred::DeferredTransitiveSetData;
use crate::deferred::types::DeferredRegistry;
use crate::deferred::types::ReservedTrivialDeferredData;
use crate::interpreter::rule_defs::transitive_set::FrozenTransitiveSet;
use crate::interpreter::rule_defs::transitive_set::TransitiveSet;

#[derive(Allocative)]
pub struct ArtifactGroupRegistry {
    pending: Vec<ReservedTrivialDeferredData<DeferredTransitiveSetData>>,
}

impl ArtifactGroupRegistry {
    pub fn new() -> Self {
        Self {
            pending: Vec::new(),
        }
    }

    pub fn create_transitive_set<'v>(
        &mut self,
        definition: Value<'v>,
        value: Option<Value<'v>>,
        children: Option<Value<'v>>,
        deferred: &mut DeferredRegistry,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<TransitiveSet<'v>> {
        let reserved = deferred.reserve_trivial::<DeferredTransitiveSetData>();
        let set = TransitiveSet::new_from_values(
            reserved.data().dupe(),
            definition,
            value,
            children,
            eval,
        )?;
        self.pending.push(reserved);
        Ok(set)
    }

    pub(crate) fn ensure_bound(
        self,
        registry: &mut DeferredRegistry,
        analysis_value_fetcher: &AnalysisValueFetcher,
    ) -> anyhow::Result<()> {
        for key in self.pending {
            let id = key.data().deferred_key().id();

            let set = analysis_value_fetcher
                .get(id)?
                .with_context(|| format!("Key is missing in AnalysisValueFetcher: {:?}", id))?;

            match set.downcast::<FrozenTransitiveSet>() {
                Ok(set) => {
                    registry.bind_trivial(key, DeferredTransitiveSetData(set));
                }
                Err(set) => {
                    return Err(anyhow::anyhow!(
                        "ArtifactGroupRegistry::ensure_bound found a value that was not a FrozenTransitiveSet: {:?}",
                        set
                    ));
                }
            }
        }

        Ok(())
    }
}


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

//! This module provides support for implementing and consuming attr
//! types.
//!
//! Support for implementing new attribute types is in the `val` module.
//!
//! Attribute values are present in 4 different states:
//!   1. initial starlark value - the value passed to the rule function
//!   2. coerced value - the value captured after processing the build file. At this point, the type has
//!      been checked (i.e. for an attrs.list(attrs.string()), we've confirmed that (1) was an iterable of strings).
//!      This is done when invoking a rule function, so it has no access to configuration or information from
//!      other build files.
//!   3. configured value - this is roughly (2) with a specific configuration attached to all configurable
//!      values. For example a dep or a source (where the source is an output) are configurable (they both are basically
//!      targets).
//!   4. resolved value - this is the "resolved" attribute once again as a starlark value and as provided to a rule
//!      implementation function.
//!
//! Attribute coercion happens immediately on declaring a build target (via
//! invoking a rule function). It will validate the types of the
//! attribute values and do some simple conversions (ex. parse strings to
//! target labels).
//!
//! Attribute configuration happens when the configuration for a target
//! becomes available. There are two primary operations that happen
//! here: select resolution and target configuration.
//!
//! Attribute resolution happens just before invoking the rule
//! implementation. At this point, the context has access to all the
//! providers that are needed (based on inspection of the configured value)
//! and anything that requires a provider can be resolved to its final
//! value.
//!
//! Generally an attribute type is specified as its "resolved" type that the
//! rule implementation requests. For example, a rule that wants a list of files
//! for its sources would specify `attrs.list(attrs.file())` and would
//! receive a list of files in the implementation. The intermediate form of that
//! may be strings or targets or some other thing (e.g. a lazy glob, maybe).


/*
 * 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::BTreeMap;
use std::collections::HashMap;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_artifact::artifact::artifact_type::BaseArtifactKind;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_cli_proto::build_request::Materializations;
use buck2_common::executor_config::PathSeparatorKind;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_events::dispatch::console_message;
use buck2_execute::artifact::fs::ExecutorFs;
use dashmap::mapref::entry::Entry;
use dashmap::DashMap;
use dice::DiceComputations;
use dice::UserComputationData;
use dupe::Dupe;
use futures::future;
use futures::stream::BoxStream;
use futures::stream::FuturesUnordered;
use futures::stream::Stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures::FutureExt;
use itertools::Itertools;
use tokio::sync::Mutex;

use crate::actions::artifact::get_artifact_fs::GetArtifactFs;
use crate::actions::artifact::materializer::ArtifactMaterializer;
use crate::analysis::calculation::RuleAnalysisCalculation;
use crate::artifact_groups::calculation::ArtifactGroupCalculation;
use crate::artifact_groups::ArtifactGroup;
use crate::artifact_groups::ArtifactGroupValues;
use crate::build_signals::HasBuildSignals;
use crate::interpreter::rule_defs::cmd_args::AbsCommandLineContext;
use crate::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use crate::interpreter::rule_defs::cmd_args::SimpleCommandLineArtifactVisitor;
use crate::interpreter::rule_defs::provider::builtin::run_info::RunInfo;
use crate::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;
use crate::interpreter::rule_defs::provider::test_provider::TestProvider;

/// The types of provider to build on the configured providers label
#[derive(Debug, Clone, Dupe, Allocative)]
pub enum BuildProviderType {
    Default,
    DefaultOther,
    Run,
    Test,
}

#[derive(Clone, Debug, Allocative)]
pub struct BuildTargetResultGen<T> {
    pub outputs: Vec<T>,
    pub providers: FrozenProviderCollectionValue,
    pub run_args: Option<Vec<String>>,
}

pub type BuildTargetResult = BuildTargetResultGen<SharedResult<ProviderArtifacts>>;

impl BuildTargetResult {
    pub async fn collect_stream(
        mut stream: impl Stream<Item = anyhow::Result<BuildEvent>> + Unpin,
        fail_fast: bool,
    ) -> anyhow::Result<BTreeMap<ConfiguredProvidersLabel, Option<Self>>> {
        // Create a map of labels to outputs, but retain the expected index of each output.
        let mut res = HashMap::<
            ConfiguredProvidersLabel,
            Option<BuildTargetResultGen<(usize, SharedResult<ProviderArtifacts>)>>,
        >::new();

        while let Some(BuildEvent { label, variant }) = stream.try_next().await? {
            match variant {
                BuildEventVariant::SkippedIncompatible => {
                    res.entry((*label).clone()).or_insert(None);
                }
                BuildEventVariant::Prepared {
                    providers,
                    run_args,
                } => {
                    res.entry((*label).clone())
                        .or_insert(Some(BuildTargetResultGen {
                            outputs: Vec::new(),
                            providers,
                            run_args,
                        }));
                }
                BuildEventVariant::Output { index, output } => {
                    let is_err = output.is_err();

                    res.get_mut(label.as_ref())
                        .with_context(|| format!("BuildEventVariant::Output before BuildEventVariant::Prepared for {} (internal error)", label))?
                        .as_mut()
                        .with_context(|| format!("BuildEventVariant::Output for a skipped target: `{}` (internal error)", label))?
                        .outputs
                        .push((index, output));

                    if is_err && fail_fast {
                        break;
                    }
                }
            }
        }

        // Sort our outputs within each individual BuildTargetResult, then return those.
        // Also, turn our HashMap into a BTreeMap.
        let res = res
            .into_iter()
            .map(|(label, result)| {
                let result = result.map(|result| {
                    let BuildTargetResultGen {
                        mut outputs,
                        providers,
                        run_args,
                    } = result;

                    // No need for a stable sort: the indices are unique (see below).
                    outputs.sort_unstable_by_key(|(index, _outputs)| *index);

                    // TODO: This whole building thing needs quite a bit of refactoring. We might
                    // request the same targets multiple times here, but since we know that
                    // ConfiguredTargetLabel -> Output is going to be deterministic, we just dedupe
                    // them using the index.
                    BuildTargetResult {
                        outputs: outputs
                            .into_iter()
                            .unique_by(|(index, _outputs)| *index)
                            .map(|(_index, outputs)| outputs)
                            .collect(),
                        providers,
                        run_args,
                    }
                });

                (label, result)
            })
            .collect();

        Ok(res)
    }
}

enum BuildEventVariant {
    SkippedIncompatible,
    Prepared {
        providers: FrozenProviderCollectionValue,
        run_args: Option<Vec<String>>,
    },
    Output {
        output: SharedResult<ProviderArtifacts>,
        /// Ensure a stable ordering of outputs.
        index: usize,
    },
}

/// Events to be accumulated using BuildTargetResult::collect_stream.
pub struct BuildEvent {
    label: Arc<ConfiguredProvidersLabel>,
    variant: BuildEventVariant,
}

pub async fn build_configured_label<'a>(
    ctx: &'a DiceComputations,
    materialization_context: &MaterializationContext,
    providers_label: ConfiguredProvidersLabel,
    providers_to_build: &ProvidersToBuild,
    skippable: bool,
) -> anyhow::Result<BoxStream<'a, BuildEvent>> {
    let providers_label = Arc::new(providers_label);

    let artifact_fs = ctx.get_artifact_fs().await?;

    let (providers, outputs, run_args) = {
        // A couple of these objects aren't Send and so scope them here so async transform doesn't get concerned.
        let providers = match ctx.get_providers(providers_label.as_ref()).await? {
            MaybeCompatible::Incompatible(reason) => {
                if skippable {
                    console_message(reason.skipping_message(providers_label.target()));
                    return Ok(futures::stream::once(futures::future::ready(BuildEvent {
                        label: providers_label.dupe(),
                        variant: BuildEventVariant::SkippedIncompatible,
                    }))
                    .boxed());
                } else {
                    return Err(reason.to_err());
                }
            }
            MaybeCompatible::Compatible(v) => v,
        };

        // Important we use an an ordered collections, so the order matches the order the rule
        // author wrote.
        let mut outputs = Vec::new();
        // Providers that produced each output, in the order of outputs above. We use a separate collection
        // otherwise we'd build the same output twice when it's both in DefaultInfo and RunInfo
        let collection = providers.provider_collection();

        let mut run_args: Option<Vec<String>> = None;

        if providers_to_build.default {
            collection
                .default_info()
                .for_each_default_output_artifact_only(&mut |o| {
                    outputs.push((ArtifactGroup::Artifact(o), BuildProviderType::Default));
                    Ok(())
                })?;
        }
        if providers_to_build.default_other {
            collection
                .default_info()
                .for_each_default_output_other_artifacts_only(&mut |o| {
                    outputs.push((o, BuildProviderType::DefaultOther));
                    Ok(())
                })?;
            // TODO(marwhal): We can remove this once we migrate all other outputs to be handled with Artifacts directly
            collection.default_info().for_each_other_output(&mut |o| {
                outputs.push((o, BuildProviderType::DefaultOther));
                Ok(())
            })?;
        }
        if providers_to_build.run {
            if let Some(runinfo) = RunInfo::from_providers(providers.provider_collection()) {
                let mut artifact_visitor = SimpleCommandLineArtifactVisitor::new();
                runinfo.visit_artifacts(&mut artifact_visitor)?;
                for input in artifact_visitor.inputs {
                    outputs.push((input, BuildProviderType::Run));
                }
                // Produce arguments to run on a local machine.
                let path_separator = if cfg!(windows) {
                    PathSeparatorKind::Windows
                } else {
                    PathSeparatorKind::Unix
                };
                let executor_fs = ExecutorFs::new(&artifact_fs, path_separator);
                let mut cli = Vec::<String>::new();
                let mut ctx = AbsCommandLineContext::new(&executor_fs);
                runinfo.add_to_command_line(&mut cli, &mut ctx)?;
                run_args = Some(cli);
            }
        }
        if providers_to_build.tests {
            if let Some(test_provider) = <dyn TestProvider>::from_collection(collection) {
                let mut artifact_visitor = SimpleCommandLineArtifactVisitor::new();
                test_provider.visit_artifacts(&mut artifact_visitor)?;
                for input in artifact_visitor.inputs {
                    outputs.push((input, BuildProviderType::Test));
                }
            }
        }

        (providers, outputs, run_args)
    };

    if let Some(signals) = ctx.per_transaction_data().get_build_signals() {
        signals.top_level_target(
            providers_label.target().dupe(),
            outputs
                .iter()
                .map(|(output, _type)| output.dupe())
                .collect(),
        );
    }

    if !skippable && outputs.is_empty() {
        console_message(format!(
            "target {} does not have any outputs: building it does nothing",
            providers_label.target()
        ));
    }

    let outputs = outputs
        .into_iter()
        .enumerate()
        .map({
            |(index, (output, provider_type))| {
                let materialization_context = materialization_context.dupe();
                materialize_artifact_group_owned(ctx, output, materialization_context).map(
                    move |res| {
                        let res = res.shared_error().map(|values| ProviderArtifacts {
                            values,
                            provider_type,
                        });

                        (index, res)
                    },
                )
            }
        })
        .collect::<FuturesUnordered<_>>()
        .map({
            let providers_label = providers_label.dupe();
            move |(index, output)| BuildEvent {
                label: providers_label.dupe(),
                variant: BuildEventVariant::Output { index, output },
            }
        });

    let stream = futures::stream::once(futures::future::ready(BuildEvent {
        label: providers_label.dupe(),
        variant: BuildEventVariant::Prepared {
            providers,
            run_args,
        },
    }))
    .chain(outputs)
    .boxed();

    Ok(stream)
}
pub async fn materialize_artifact_group_owned(
    ctx: &DiceComputations,
    artifact_group: ArtifactGroup,
    materialization_context: MaterializationContext,
) -> anyhow::Result<ArtifactGroupValues> {
    materialize_artifact_group(ctx, &artifact_group, &materialization_context).await
}

#[derive(Clone, Allocative)]
pub struct ProviderArtifacts {
    pub values: ArtifactGroupValues,
    pub provider_type: BuildProviderType,
}

// what type of artifacts to build based on the provider it came from
#[derive(Default, Clone)]
pub struct ProvidersToBuild {
    pub default: bool,
    pub default_other: bool,
    pub run: bool,
    pub tests: bool,
}

impl Debug for ProviderArtifacts {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ProviderArtifacts")
            .field("values", &self.values.iter().collect::<Vec<_>>())
            .field("provider_type", &self.provider_type)
            .finish()
    }
}

pub async fn materialize_artifact_group(
    ctx: &DiceComputations,
    artifact_group: &ArtifactGroup,
    materialization_context: &MaterializationContext,
) -> anyhow::Result<ArtifactGroupValues> {
    let values = ctx.ensure_artifact_group(artifact_group).await?;

    if let MaterializationContext::Materialize { map, force } = materialization_context {
        future::try_join_all(values.iter().filter_map(|(artifact, _value)| {
            match artifact.as_parts().0 {
                BaseArtifactKind::Build(artifact) => {
                    match map.entry(artifact.dupe()) {
                        Entry::Vacant(v) => {
                            // Ensure we won't request this artifact elsewhere, and proceed to request
                            // it.
                            v.insert(());
                        }
                        Entry::Occupied(..) => {
                            // We've already requested this artifact, no use requesting it again.
                            return None;
                        }
                    }

                    Some(ctx.try_materialize_requested_artifact(artifact, *force))
                }
                BaseArtifactKind::Source(..) => None,
            }
        }))
        .await
        .context("Failed to materialize artifacts")?;
    }

    Ok(values)
}

#[derive(Clone, Dupe)]
pub enum MaterializationContext {
    Skip,
    Materialize {
        /// This map contains all the artifacts that we enqueued for materialization. This ensures
        /// we don't enqueue the same thing more than once.
        map: Arc<DashMap<BuildArtifact, ()>>,
        /// Whether we should force the materialization of requested artifacts, or defer to the
        /// config.
        force: bool,
    },
}

impl MaterializationContext {
    /// Create a new MaterializationContext that will force all materializations.
    pub fn force_materializations() -> Self {
        Self::Materialize {
            map: Arc::new(DashMap::new()),
            force: true,
        }
    }
}

pub trait ConvertMaterializationContext {
    fn from(self) -> MaterializationContext;

    fn with_existing_map(self, map: &Arc<DashMap<BuildArtifact, ()>>) -> MaterializationContext;
}

impl ConvertMaterializationContext for Materializations {
    fn from(self) -> MaterializationContext {
        match self {
            Materializations::Skip => MaterializationContext::Skip,
            Materializations::Default => MaterializationContext::Materialize {
                map: Arc::new(DashMap::new()),
                force: false,
            },
            Materializations::Materialize => MaterializationContext::Materialize {
                map: Arc::new(DashMap::new()),
                force: true,
            },
        }
    }

    fn with_existing_map(self, map: &Arc<DashMap<BuildArtifact, ()>>) -> MaterializationContext {
        match self {
            Materializations::Skip => MaterializationContext::Skip,
            Materializations::Default => MaterializationContext::Materialize {
                map: map.dupe(),
                force: false,
            },
            Materializations::Materialize => MaterializationContext::Materialize {
                map: map.dupe(),
                force: true,
            },
        }
    }
}

pub trait HasCreateUnhashedSymlinkLock {
    fn set_create_unhashed_symlink_lock(&mut self, lock: Arc<Mutex<()>>);

    fn get_create_unhashed_symlink_lock(&self) -> Arc<Mutex<()>>;
}

impl HasCreateUnhashedSymlinkLock for UserComputationData {
    fn set_create_unhashed_symlink_lock(&mut self, lock: Arc<Mutex<()>>) {
        self.data.set(lock);
    }

    fn get_create_unhashed_symlink_lock(&self) -> Arc<Mutex<()>> {
        self.data
            .get::<Arc<Mutex<()>>>()
            .expect("Lock for creating unhashed symlinks should be set")
            .dupe()
    }
}


/*
 * 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 allocative::Allocative;
use gazebo::variants::UnpackVariants;

use crate::build::BuildTargetResult;

#[derive(Clone, Debug, derive_more::Display, UnpackVariants, Allocative)]
pub enum BxlBuildResult {
    None,
    #[display(fmt = "build result")]
    Built(BuildTargetResult),
}

impl BxlBuildResult {
    pub fn new(result: Option<BuildTargetResult>) -> Self {
        match result {
            Some(result) => Self::Built(result),
            None => Self::None,
        }
    }
}


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

//! DICE calculations for bxl

use std::fmt::Debug;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_core::base_deferred_key::BaseDeferredKeyDyn;
use buck2_util::late_binding::LateBinding;
use dashmap::DashMap;
use dice::DiceComputations;
use dupe::Dupe;

use crate::bxl::result::BxlResult;

#[async_trait]
pub trait BxlCalculationDyn: Debug + Send + Sync + 'static {
    async fn eval_bxl(
        &self,
        ctx: &DiceComputations,
        bxl: Arc<dyn BaseDeferredKeyDyn>,
    ) -> anyhow::Result<BxlComputeResult>;
}

#[derive(Allocative, Clone, Dupe)]
pub struct BxlComputeResult {
    pub bxl_result: Arc<BxlResult>,
    pub materializations: Arc<DashMap<BuildArtifact, ()>>,
}

/// Dependency injection for BXL.
///
/// BXL implementation lives in downstream crate.
/// This field is initialized at program start, so this crate can call BXL calculation.
pub static BXL_CALCULATION_IMPL: LateBinding<&'static dyn BxlCalculationDyn> =
    LateBinding::new("BXL_CALCULATION_IMPL");


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

//! Common utilities for bxl
use std::sync::Arc;

use buck2_common::executor_config::CommandExecutorConfig;
use buck2_common::executor_config::CommandGenerationOptions;
use buck2_common::executor_config::Executor;
use buck2_common::executor_config::LocalExecutorOptions;
use buck2_common::executor_config::PathSeparatorKind;
use buck2_core::configuration::pair::ConfigurationNoExec;
use buck2_node::configuration::execution::ExecutionPlatform;
use buck2_node::configuration::execution::ExecutionPlatformResolution;
use once_cell::sync::Lazy;

// TODO(bobyf) this should be configured by the bxl function similar to
// execution groups, which we don't have yet. Currently we hard code this to
// local since IDE will mostly be invoking local tools anyways
pub static EXECUTION_PLATFORM: Lazy<ExecutionPlatformResolution> = Lazy::new(|| {
    ExecutionPlatformResolution::new(
        Some(ExecutionPlatform::legacy_execution_platform(
            Arc::new(CommandExecutorConfig {
                executor: Executor::Local(LocalExecutorOptions::default()),
                options: CommandGenerationOptions {
                    path_separator: PathSeparatorKind::system_default(),
                    output_paths_behavior: Default::default(),
                },
            }),
            ConfigurationNoExec::unspecified(),
        )),
        vec![],
    )
});


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

//!
//! bxl is the Buck Extension Language, allowing any integrator to write Starlark code that
//! introspects buck2 internal graphs in a safe, incremental way to perform more complex operations
//!

pub mod build_result;

pub mod calculation;
pub mod execution_platform;
pub mod result;
pub mod types;


/*
 * 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 allocative::Allocative;
use buck2_artifact::deferred::id::DeferredId;
use buck2_core::fs::buck_out_path::BuckOutPath;
use indexmap::IndexSet;

use crate::artifact_groups::ArtifactGroup;
use crate::bxl::build_result::BxlBuildResult;
use crate::deferred::types::DeferredLookup;
use crate::deferred::types::DeferredTable;

/// The result of evaluating a bxl function
#[derive(Allocative)]
pub enum BxlResult {
    /// represents that the bxl function has no built results
    None {
        output_loc: BuckOutPath,
        error_loc: BuckOutPath,
    },
    /// a bxl that deals with builds
    BuildsArtifacts {
        output_loc: BuckOutPath,
        error_loc: BuckOutPath,
        built: Vec<BxlBuildResult>,
        artifacts: Vec<ArtifactGroup>,
        deferred: DeferredTable,
    },
}

impl BxlResult {
    pub fn new(
        output_loc: BuckOutPath,
        error_loc: BuckOutPath,
        ensured_artifacts: IndexSet<ArtifactGroup>,
        deferred: DeferredTable,
    ) -> Self {
        if ensured_artifacts.is_empty() {
            Self::None {
                output_loc,
                error_loc,
            }
        } else {
            Self::BuildsArtifacts {
                output_loc,
                error_loc,
                built: vec![],
                artifacts: ensured_artifacts.into_iter().collect(),
                deferred,
            }
        }
    }

    /// looks up an 'Deferred' given the id
    pub fn lookup_deferred(&self, id: DeferredId) -> anyhow::Result<DeferredLookup<'_>> {
        match self {
            BxlResult::None { .. } => Err(anyhow::anyhow!("Bxl never attempted to build anything")),
            BxlResult::BuildsArtifacts { deferred, .. } => deferred.lookup_deferred(id),
        }
    }

    pub fn get_output_loc(&self) -> &BuckOutPath {
        match self {
            BxlResult::None { output_loc, .. } => output_loc,
            BxlResult::BuildsArtifacts { output_loc, .. } => output_loc,
        }
    }

    pub fn get_error_loc(&self) -> &BuckOutPath {
        match self {
            BxlResult::None { error_loc, .. } => error_loc,
            BxlResult::BuildsArtifacts { error_loc, .. } => error_loc,
        }
    }
}


/*
 * 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 allocative::Allocative;
use buck2_data::ToProtoMessage;
use buck2_interpreter::path::BxlFilePath;
use derive_more::Display;
use serde::Serialize;
use serde::Serializer;

/// The identifier used to find the implementation function for this bxl. Should point at the output of `bxl()`
#[derive(
    Debug, Clone, Display, Eq, PartialEq, Hash, Ord, PartialOrd, Allocative
)]
#[display(fmt = "{}:{}", bxl_path, name)]
pub struct BxlFunctionLabel {
    /// The cell, package, and file that contains the output of `bxl()`
    pub bxl_path: BxlFilePath,
    /// The name of the symbol that is bound to the output of `bxl()`, e.g. `custom_query`
    pub name: String,
}

impl Serialize for BxlFunctionLabel {
    fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        s.serialize_str(&format!("{}", self))
    }
}

impl ToProtoMessage for BxlFunctionLabel {
    type Message = buck2_data::BxlFunctionLabel;

    fn as_proto(&self) -> Self::Message {
        buck2_data::BxlFunctionLabel {
            bxl_path: self.bxl_path.to_string(),
            name: self.name.clone(),
        }
    }
}