/*
 * 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::Any;
use std::collections::HashMap;
use std::fmt;
use std::hash::Hash;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_analysis::analysis::calculation::AnalysisKey;
use buck2_analysis::analysis::calculation::AnalysisKeyActivationData;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::calculation::BuildKey;
use buck2_build_api::actions::calculation::BuildKeyActivationData;
use buck2_build_api::actions::RegisteredAction;
use buck2_build_api::artifact_groups::calculation::EnsureProjectedArtifactKey;
use buck2_build_api::artifact_groups::calculation::EnsureTransitiveSetProjectionKey;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::artifact_groups::ResolvedArtifactGroup;
use buck2_build_api::build_signals::BuildSignals;
use buck2_build_api::build_signals::BuildSignalsInstaller;
use buck2_build_api::build_signals::CREATE_BUILD_SIGNALS;
use buck2_build_api::deferred::calculation::DeferredCompute;
use buck2_build_api::deferred::calculation::DeferredResolve;
use buck2_build_signals::CriticalPathBackendName;
use buck2_build_signals::DeferredBuildSignals;
use buck2_build_signals::FinishBuildSignals;
use buck2_build_signals::NodeDuration;
use buck2_configured::nodes::calculation::ConfiguredTargetNodeKey;
use buck2_core::package::PackageLabel;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_data::ToProtoMessage;
use buck2_events::dispatch::instant_event;
use buck2_events::dispatch::with_dispatcher_async;
use buck2_events::dispatch::EventDispatcher;
use buck2_events::metadata;
use buck2_events::span::SpanId;
use buck2_interpreter_for_build::interpreter::calculation::IntepreterResultsKeyActivationData;
use buck2_interpreter_for_build::interpreter::calculation::InterpreterResultsKey;
use buck2_node::nodes::eval_result::EvaluationResult;
use derive_more::From;
use dice::ActivationData;
use dice::ActivationTracker;
use dupe::Dupe;
use dupe::OptionDupedExt;
use itertools::Itertools;
use smallvec::SmallVec;
use static_assertions::assert_eq_size;
use tokio::sync::mpsc::UnboundedReceiver;
use tokio::sync::mpsc::UnboundedSender;
use tokio::task::JoinHandle;
use tokio_stream::wrappers::UnboundedReceiverStream;
use tokio_stream::StreamExt;

use crate::backend::backend::BuildListenerBackend;
use crate::backend::default::DefaultBackend;
use crate::backend::longest_path_graph::LongestPathGraphBackend;

mod backend;

/// A node in our critical path graph.
#[derive(Hash, Eq, PartialEq, Clone, Dupe, Debug, From)]
enum NodeKey {
    // Those are DICE keys.
    BuildKey(BuildKey),
    AnalysisKey(AnalysisKey),
    EnsureProjectedArtifactKey(EnsureProjectedArtifactKey),
    EnsureTransitiveSetProjectionKey(EnsureTransitiveSetProjectionKey),
    DeferredCompute(DeferredCompute),
    DeferredResolve(DeferredResolve),
    ConfiguredTargetNodeKey(ConfiguredTargetNodeKey),
    InterpreterResultsKey(InterpreterResultsKey),

    // This one is not a DICE key.
    Materialization(BuildArtifact),
}

// Explain the sizeof this struct (and avoid regressing it since we store it in the longest path
// graph implementation).

assert_eq_size!(BuildKey, [usize; 4]);
assert_eq_size!(AnalysisKey, [usize; 2]);
assert_eq_size!(EnsureTransitiveSetProjectionKey, [usize; 5]);
assert_eq_size!(EnsureProjectedArtifactKey, [usize; 7]);
assert_eq_size!(DeferredCompute, [usize; 4]);
assert_eq_size!(DeferredResolve, [usize; 4]);
assert_eq_size!(ConfiguredTargetNodeKey, [usize; 2]);
assert_eq_size!(InterpreterResultsKey, [usize; 1]);
assert_eq_size!(BuildArtifact, [usize; 6]);
assert_eq_size!(NodeKey, [usize; 7]);

impl NodeKey {
    fn from_any(key: &dyn Any) -> Option<Self> {
        let key = if let Some(key) = key.downcast_ref::<BuildKey>() {
            Self::BuildKey(key.dupe())
        } else if let Some(key) = key.downcast_ref::<AnalysisKey>() {
            Self::AnalysisKey(key.dupe())
        } else if let Some(key) = key.downcast_ref::<EnsureProjectedArtifactKey>() {
            Self::EnsureProjectedArtifactKey(key.dupe())
        } else if let Some(key) = key.downcast_ref::<EnsureTransitiveSetProjectionKey>() {
            Self::EnsureTransitiveSetProjectionKey(key.dupe())
        } else if let Some(key) = key.downcast_ref::<DeferredCompute>() {
            Self::DeferredCompute(key.dupe())
        } else if let Some(key) = key.downcast_ref::<DeferredResolve>() {
            Self::DeferredResolve(key.dupe())
        } else if let Some(key) = key.downcast_ref::<ConfiguredTargetNodeKey>() {
            Self::ConfiguredTargetNodeKey(key.dupe())
        } else if let Some(key) = key.downcast_ref::<InterpreterResultsKey>() {
            Self::InterpreterResultsKey(key.dupe())
        } else {
            return None;
        };

        Some(key)
    }
}

impl fmt::Display for NodeKey {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::BuildKey(k) => write!(f, "BuildKey({})", k),
            Self::AnalysisKey(k) => write!(f, "AnalysisKey({})", k),
            Self::EnsureProjectedArtifactKey(k) => write!(f, "EnsureProjectedArtifactKey({})", k),
            Self::EnsureTransitiveSetProjectionKey(k) => {
                write!(f, "EnsureTransitiveSetProjectionKey({})", k)
            }
            Self::DeferredCompute(k) => write!(f, "DeferredCompute({})", k),
            Self::DeferredResolve(k) => write!(f, "DeferredResolve({})", k),
            Self::ConfiguredTargetNodeKey(k) => write!(f, "ConfiguredTargetNodeKey({})", k),
            Self::InterpreterResultsKey(k) => write!(f, "InterpreterResultsKey({})", k),
            Self::Materialization(k) => write!(f, "Materialization({})", k),
        }
    }
}

struct TopLevelTargetSignal {
    pub label: ConfiguredTargetLabel,
    pub artifacts: Vec<ArtifactGroup>,
}

struct FinalMaterializationSignal {
    pub artifact: BuildArtifact,
    pub duration: NodeDuration,
    pub span_id: Option<SpanId>,
}

/* These signals are distinct from the main Buck event bus because some
 * analysis needs access to the entire build graph, and serializing the
 * entire build graph isn't feasible - therefore, we have these signals
 * with an unserializable but lightweight handle on a RegisteredAction.
 */
#[derive(From)]
enum BuildSignal {
    Evaluation(Evaluation),
    TopLevelTarget(TopLevelTargetSignal),
    FinalMaterialization(FinalMaterializationSignal),
    BuildFinished,
}

/// Data for a BuildSignal that is the result of a DICE key evaluation.
pub struct Evaluation {
    /// The key we evaluated.
    key: NodeKey,
    /// The duration. By default this'll be zero, unless activation data says otherwise.
    duration: NodeDuration,
    /// The dependencies.
    dep_keys: Vec<NodeKey>,
    /// Spans that correspond to this key. We use this when producing a chrome trace.
    spans: SmallVec<[SpanId; 1]>,

    // NOTE: The fields below aren't usually going to be both set, but it doesn't really hurt (for
    // now) to have them not tied to the right variant.
    /// The RegisteredAction that corresponds to this Evaluation (this will only be present for
    /// NodeKey::BuildKey).
    action: Option<Arc<RegisteredAction>>,

    /// The Load result that corresponds to this Evaluation (this will only be pesent for
    /// InterpreterResultsKey).
    load_result: Option<Arc<EvaluationResult>>,
}

pub struct BuildSignalSender {
    sender: UnboundedSender<BuildSignal>,
}

impl BuildSignals for BuildSignalSender {
    fn top_level_target(&self, label: ConfiguredTargetLabel, artifacts: Vec<ArtifactGroup>) {
        let _ignored = self
            .sender
            .send(TopLevelTargetSignal { label, artifacts }.into());
    }

    fn final_materialization(
        &self,
        artifact: BuildArtifact,
        duration: NodeDuration,
        span_id: Option<SpanId>,
    ) {
        let _ignored = self.sender.send(
            FinalMaterializationSignal {
                artifact,
                duration,
                span_id,
            }
            .into(),
        );
    }
}

impl ActivationTracker for BuildSignalSender {
    /// We received a DICE key. Check if it's one of the keys we care about (i.e. can we downcast
    /// it to NodeKey?), and then if that's the case, extract its dependencies and activation data
    /// (if any).
    fn key_activated(
        &self,
        key: &dyn Any,
        deps: &mut dyn Iterator<Item = &dyn Any>,
        activation_data: ActivationData,
    ) {
        let key = match NodeKey::from_any(key) {
            Some(key) => key,
            None => return,
        };

        let mut signal = Evaluation {
            key,
            action: None,
            duration: NodeDuration::zero(),
            dep_keys: deps.into_iter().filter_map(NodeKey::from_any).collect(),
            spans: Default::default(),
            load_result: None,
        };

        /// Given an Option containing an Any, take it if and only if it contains a T.
        fn downcast_and_take<T: 'static>(
            data: &mut Option<Box<dyn Any + Send + Sync + 'static>>,
        ) -> Option<T> {
            if data.as_ref().map(|d| d.is::<T>()) != Some(true) {
                return None;
            }

            // Unwrap safety: we just checked that the option is occupied and the type matches
            Some(*data.take().unwrap().downcast().ok().unwrap())
        }

        if let ActivationData::Evaluated(mut activation_data) = activation_data {
            if let Some(BuildKeyActivationData {
                action,
                duration,
                spans,
            }) = downcast_and_take(&mut activation_data)
            {
                signal.action = Some(action);
                signal.duration = duration;
                signal.spans = spans;
            } else if let Some(AnalysisKeyActivationData { duration, spans }) =
                downcast_and_take(&mut activation_data)
            {
                signal.duration = NodeDuration {
                    user: duration,
                    total: duration,
                };
                signal.spans = spans;
            } else if let Some(IntepreterResultsKeyActivationData {
                duration,
                result,
                spans,
            }) = downcast_and_take(&mut activation_data)
            {
                signal.duration = NodeDuration {
                    user: duration,
                    total: duration,
                };

                signal.load_result = result.ok();
                signal.spans = spans;
            }
        }

        let _ignored = self.sender.send(signal.into());
    }
}

pub struct DeferredBuildSignalsImpl {
    sender: Arc<BuildSignalSender>,
    receiver: UnboundedReceiver<BuildSignal>,
}

impl DeferredBuildSignals for DeferredBuildSignalsImpl {
    fn start(
        self: Box<Self>,
        events: EventDispatcher,
        backend: CriticalPathBackendName,
    ) -> Box<dyn FinishBuildSignals> {
        let handle = match backend {
            CriticalPathBackendName::LongestPathGraph => {
                start_backend(events, self.receiver, LongestPathGraphBackend::new())
            }
            CriticalPathBackendName::Default => {
                start_backend(events, self.receiver, DefaultBackend::new())
            }
        };

        Box::new(FinishBuildSignalsImpl {
            sender: self.sender,
            handle,
        }) as _
    }
}

pub struct FinishBuildSignalsImpl {
    sender: Arc<BuildSignalSender>,
    handle: JoinHandle<anyhow::Result<()>>,
}

#[async_trait]
impl FinishBuildSignals for FinishBuildSignalsImpl {
    async fn finish(self: Box<Self>) -> anyhow::Result<()> {
        let _ignored = self.sender.sender.send(BuildSignal::BuildFinished);

        self.handle
            .await
            .context("Error joining critical path task")?
    }
}

fn start_backend(
    events: EventDispatcher,
    receiver: UnboundedReceiver<BuildSignal>,
    backend: impl BuildListenerBackend + Send + 'static,
) -> JoinHandle<anyhow::Result<()>> {
    let listener = BuildSignalReceiver::new(receiver, backend);
    tokio::spawn(with_dispatcher_async(events.dupe(), async move {
        listener.run_and_log().await
    }))
}

struct BuildSignalReceiver<T> {
    receiver: UnboundedReceiverStream<BuildSignal>,
    // Maps a PackageLabel to the first PackageLabel that had an edge to it. When that PackageLabel
    // shows up, we'll give it a dependency on said first PackageLabel that had an edge to it, which
    // is how we discovered its existence.
    first_edge_to_load: HashMap<PackageLabel, PackageLabel>,
    backend: T,
}

impl<T> BuildSignalReceiver<T>
where
    T: BuildListenerBackend,
{
    fn new(receiver: UnboundedReceiver<BuildSignal>, backend: T) -> Self {
        Self {
            receiver: UnboundedReceiverStream::new(receiver),
            backend,
            first_edge_to_load: HashMap::new(),
        }
    }

    pub async fn run_and_log(mut self) -> anyhow::Result<()> {
        while let Some(event) = self.receiver.next().await {
            match event {
                BuildSignal::Evaluation(eval) => self.process_evaluation(eval),
                BuildSignal::TopLevelTarget(top_level) => {
                    self.process_top_level_target(top_level)?
                }
                BuildSignal::FinalMaterialization(final_materialization) => {
                    self.process_final_materialization(final_materialization)?
                }
                BuildSignal::BuildFinished => break,
            }
        }

        let now = Instant::now();

        let BuildInfo {
            critical_path,
            num_nodes,
            num_edges,
        } = self.backend.finish()?;

        let compute_elapsed = now.elapsed();

        let meta_entry_data = NodeData {
            action: None,
            duration: NodeDuration {
                user: Duration::ZERO,
                total: compute_elapsed,
            },
            span_ids: Default::default(),
        };

        let meta_entry = (
            buck2_data::critical_path_entry2::ComputeCriticalPath {}.into(),
            &meta_entry_data,
            &Some(compute_elapsed),
        );

        let critical_path2 = critical_path
            .iter()
            .filter_map(|(key, data, potential_improvement)| {
                let entry: buck2_data::critical_path_entry2::Entry = match key {
                    NodeKey::BuildKey(key) => {
                        let owner = key.0.owner().to_proto().into();

                        // If we have a NodeKey that's an ActionKey we'd expect to have an `action`
                        // in our data (unless we didn't actually run it because of e.g. early
                        // cutoff, in which case omitting it is what we want).
                        let action = data.action.as_ref()?;

                        buck2_data::critical_path_entry2::ActionExecution {
                            owner: Some(owner),
                            name: Some(buck2_data::ActionName {
                                category: action.category().as_str().to_owned(),
                                identifier: action.identifier().unwrap_or("").to_owned(),
                            }),
                        }
                        .into()
                    }
                    NodeKey::AnalysisKey(key) => buck2_data::critical_path_entry2::Analysis {
                        target: Some(key.0.as_proto().into()),
                    }
                    .into(),
                    NodeKey::Materialization(key) => {
                        let owner = key.key().owner().to_proto().into();

                        buck2_data::critical_path_entry2::Materialization {
                            owner: Some(owner),
                            path: key.get_path().path().to_string(),
                        }
                        .into()
                    }
                    NodeKey::InterpreterResultsKey(key) => buck2_data::critical_path_entry2::Load {
                        package: key.0.to_string(),
                    }
                    .into(),
                    NodeKey::EnsureProjectedArtifactKey(..) => return None,
                    NodeKey::EnsureTransitiveSetProjectionKey(..) => return None,
                    NodeKey::DeferredCompute(..) => return None,
                    NodeKey::DeferredResolve(..) => return None,
                    NodeKey::ConfiguredTargetNodeKey(..) => return None,
                };

                Some((entry, data, potential_improvement))
            })
            .chain(std::iter::once(meta_entry))
            .map(|(entry, data, potential_improvement)| {
                anyhow::Ok(buck2_data::CriticalPathEntry2 {
                    span_ids: data
                        .span_ids
                        .iter()
                        .map(|span_id| (*span_id).into())
                        .collect(),
                    duration: Some(data.duration.critical_path_duration().try_into()?),
                    user_duration: Some(data.duration.user.try_into()?),
                    total_duration: Some(data.duration.total.try_into()?),
                    potential_improvement_duration: potential_improvement
                        .map(|p| p.try_into())
                        .transpose()?,
                    entry: Some(entry),
                })
            })
            .collect::<Result<Vec<_>, _>>()?;

        instant_event(buck2_data::BuildGraphExecutionInfo {
            critical_path: Vec::new(),
            critical_path2,
            metadata: metadata::collect(),
            num_nodes,
            num_edges,
            uses_total_duration: true,
            backend_name: Some(T::name().to_string()),
        });
        Ok(())
    }

    /// Receive an Evaluation. Do a little enrichment if it's a load, then pass through to the
    /// underying backend.
    fn process_evaluation(&mut self, mut evaluation: Evaluation) {
        self.enrich_load(&mut evaluation);

        self.backend.process_node(
            evaluation.key,
            evaluation.action,
            evaluation.duration,
            evaluation.dep_keys.into_iter(),
            evaluation.spans,
        );
    }

    /// If the evaluation is a load (InterpreterResultsKey) and carries a load_result, then inject
    /// some extra edges that indicate which packages have now become visibile as a result of this
    /// load.
    fn enrich_load(&mut self, evaluation: &mut Evaluation) {
        let pkg = match &evaluation.key {
            NodeKey::InterpreterResultsKey(InterpreterResultsKey(pkg)) => pkg,
            _ => return,
        };

        if let Some(load_result) = &evaluation.load_result {
            let deps_pkg = load_result
                .targets()
                .values()
                .flat_map(|target| target.deps().map(|t| t.pkg()))
                .unique()
                .map(|pkg| pkg.dupe());

            for dep_pkg in deps_pkg {
                if dep_pkg == *pkg {
                    continue;
                }

                self.first_edge_to_load
                    .entry(dep_pkg)
                    .or_insert_with(|| pkg.dupe());
            }
        }

        let first_edge = self.first_edge_to_load.get(pkg);

        if let Some(first_edge) = first_edge {
            evaluation
                .dep_keys
                .push(NodeKey::InterpreterResultsKey(InterpreterResultsKey(
                    first_edge.dupe(),
                )));
        }
    }

    // TODO: We would need something similar with anon targets.
    fn process_top_level_target(
        &mut self,
        top_level: TopLevelTargetSignal,
    ) -> Result<(), anyhow::Error> {
        let artifact_keys =
            top_level
                .artifacts
                .into_iter()
                .filter_map(|dep| match dep.assert_resolved() {
                    ResolvedArtifactGroup::Artifact(artifact) => artifact
                        .action_key()
                        .duped()
                        .map(BuildKey)
                        .map(NodeKey::BuildKey),
                    ResolvedArtifactGroup::TransitiveSetProjection(key) => {
                        Some(NodeKey::EnsureTransitiveSetProjectionKey(
                            EnsureTransitiveSetProjectionKey(key.dupe()),
                        ))
                    }
                });

        self.backend.process_top_level_target(
            NodeKey::AnalysisKey(AnalysisKey(top_level.label)),
            artifact_keys,
        );

        Ok(())
    }

    fn process_final_materialization(
        &mut self,
        materialization: FinalMaterializationSignal,
    ) -> Result<(), anyhow::Error> {
        let dep = NodeKey::BuildKey(BuildKey(materialization.artifact.key().dupe()));

        self.backend.process_node(
            NodeKey::Materialization(materialization.artifact),
            None,
            materialization.duration,
            std::iter::once(dep),
            materialization.span_id.into_iter().collect(),
        );

        Ok(())
    }
}

pub struct BuildInfo {
    // Node, its data, and its potential for improvement
    critical_path: Vec<(NodeKey, NodeData, Option<Duration>)>,
    num_nodes: u64,
    num_edges: u64,
}

#[derive(Clone)]
struct NodeData {
    action: Option<Arc<RegisteredAction>>,
    duration: NodeDuration,
    span_ids: SmallVec<[SpanId; 1]>,
}

assert_eq_size!(NodeData, [usize; 8]);

fn create_build_signals() -> (BuildSignalsInstaller, Box<dyn DeferredBuildSignals>) {
    let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();

    let sender = Arc::new(BuildSignalSender { sender });
    let installer = BuildSignalsInstaller {
        build_signals: sender.dupe() as _,
        activation_tracker: sender.dupe() as _,
    };

    let deferred = Box::new(DeferredBuildSignalsImpl { receiver, sender });

    (installer, deferred as _)
}

pub fn init_late_bindings() {
    CREATE_BUILD_SIGNALS.init(create_build_signals)
}


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

use std::sync::Arc;

use buck2_build_api::actions::RegisteredAction;
use buck2_build_signals::CriticalPathBackendName;
use buck2_build_signals::NodeDuration;
use buck2_events::span::SpanId;
use smallvec::SmallVec;

use crate::BuildInfo;
use crate::NodeKey;

pub(crate) trait BuildListenerBackend {
    fn process_node(
        &mut self,
        key: NodeKey,
        value: Option<Arc<RegisteredAction>>,
        duration: NodeDuration,
        dep_keys: impl Iterator<Item = NodeKey>,
        span_ids: SmallVec<[SpanId; 1]>,
    );

    fn process_top_level_target(
        &mut self,
        analysis: NodeKey,
        artifacts: impl Iterator<Item = NodeKey>,
    );

    fn finish(self) -> anyhow::Result<BuildInfo>;

    fn name() -> CriticalPathBackendName;
}


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

use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Display;
use std::hash::Hash;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context as _;
use buck2_build_api::actions::RegisteredAction;
use buck2_build_signals::CriticalPathBackendName;
use buck2_build_signals::NodeDuration;
use buck2_events::span::SpanId;
use dupe::Dupe;
use gazebo::prelude::VecExt;
use itertools::Itertools;
use smallvec::SmallVec;

use crate::backend::backend::BuildListenerBackend;
use crate::BuildInfo;
use crate::NodeData;
use crate::NodeKey;

#[derive(Clone, Dupe)]
struct CriticalPathNode<TKey: Eq, TValue> {
    /// The aggregated duration of this critical path.
    pub duration: Duration,
    /// The value of this node. If None, this node just won't be included when displaying.
    pub value: TValue,
    pub prev: Option<TKey>,
}

fn extract_critical_path<TKey: Hash + Eq, TValue>(
    predecessors: &HashMap<TKey, CriticalPathNode<TKey, TValue>>,
) -> anyhow::Result<Vec<(&TKey, &TValue, Duration)>>
where
    TKey: Display,
{
    let mut tail = predecessors
        .iter()
        .max_by_key(|(_key, data)| data.duration)
        .map(|q| q.0);

    let mut path = vec![];
    let mut visited = HashSet::new();

    while let Some(v) = tail.take() {
        if !visited.insert(v) {
            return Err(anyhow::anyhow!(
                "Cycle in critical path: visited {} twice",
                v
            ));
        }

        tail = predecessors.get(v).and_then(|node| {
            path.push((v, &node.value, node.duration));
            node.prev.as_ref()
        });
    }

    // Take differences of adjacent elements to recover action time from cumulative sum.
    path.reverse();
    for i in (1..path.len()).rev() {
        path[i].2 = path[i].2.saturating_sub(path[i - 1].2);
    }

    Ok(path)
}

pub(crate) struct DefaultBackend {
    predecessors: HashMap<NodeKey, CriticalPathNode<NodeKey, NodeData>>,
    num_nodes: u64,
    num_edges: u64,
}

impl DefaultBackend {
    pub(crate) fn new() -> Self {
        Self {
            predecessors: HashMap::new(),
            num_nodes: 0,
            num_edges: 0,
        }
    }
}

impl BuildListenerBackend for DefaultBackend {
    fn process_node(
        &mut self,
        key: NodeKey,
        value: Option<Arc<RegisteredAction>>,
        duration: NodeDuration,
        dep_keys: impl Iterator<Item = NodeKey>,
        span_ids: SmallVec<[SpanId; 1]>,
    ) {
        let longest_ancestor = dep_keys
            .unique()
            .filter_map(|node_key| {
                self.num_edges += 1;
                let node_data = self.predecessors.get(&node_key)?;
                Some((node_key, node_data.duration))
            })
            .max_by_key(|d| d.1);

        let value = NodeData {
            action: value,
            duration,
            span_ids,
        };

        let node = match longest_ancestor {
            Some((key, ancestor_duration)) => CriticalPathNode {
                prev: Some(key.dupe()),
                value,
                duration: ancestor_duration + duration.critical_path_duration(),
            },
            None => CriticalPathNode {
                prev: None,
                value,
                duration: duration.critical_path_duration(),
            },
        };

        self.num_nodes += 1;
        self.predecessors.insert(key, node);
    }

    fn process_top_level_target(
        &mut self,
        _analysis: NodeKey,
        _artifacts: impl Iterator<Item = NodeKey>,
    ) {
    }

    fn finish(self) -> anyhow::Result<BuildInfo> {
        let critical_path = extract_critical_path(&self.predecessors)
            .context("Error extracting critical path")?
            .into_map(|(key, data, _duration)| (key.dupe(), data.clone(), None));

        Ok(BuildInfo {
            critical_path,
            num_nodes: self.num_nodes,
            num_edges: self.num_edges,
        })
    }

    fn name() -> CriticalPathBackendName {
        CriticalPathBackendName::Default
    }
}

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

    type CriticalPathMap = HashMap<i32, CriticalPathNode<i32, Option<i32>>>;

    fn cp_insert(
        predecessors: &mut CriticalPathMap,
        key: i32,
        prev: Option<i32>,
        duration: Duration,
    ) {
        predecessors.insert(
            key,
            CriticalPathNode {
                duration,
                value: Some(key),
                prev,
            },
        );
    }
    #[test]
    fn empty_path() {
        let predecessors = CriticalPathMap::new();
        assert_eq!(extract_critical_path(&predecessors).unwrap(), vec![]);
    }

    #[test]
    fn unit_path() {
        let mut predecessors = CriticalPathMap::new();
        cp_insert(&mut predecessors, 1, None, Duration::from_secs(3));
        assert_eq!(
            extract_critical_path(&predecessors).unwrap(),
            vec![(&1, &Some(1), Duration::from_secs(3))],
        );
    }

    #[test]
    fn long_path() {
        let mut predecessors = HashMap::new();
        /*   -> 1 -> 2 -> 3
         *   5s   6s   7s
         *
         *      1 -> 4
         *        9s
         */
        cp_insert(&mut predecessors, 1, None, Duration::from_secs(5));
        cp_insert(&mut predecessors, 2, Some(1), Duration::from_secs(11));
        cp_insert(&mut predecessors, 3, Some(2), Duration::from_secs(18));
        cp_insert(&mut predecessors, 4, Some(1), Duration::from_secs(14));
        assert_eq!(
            extract_critical_path(&predecessors).unwrap(),
            vec![
                (&1, &Some(1), Duration::from_secs(5)),
                (&2, &Some(2), Duration::from_secs(6)),
                (&3, &Some(3), Duration::from_secs(7)),
            ],
        );
    }

    #[test]
    fn cycle_path() {
        let mut predecessors = HashMap::new();
        cp_insert(&mut predecessors, 1, Some(2), Duration::from_secs(5));
        cp_insert(&mut predecessors, 2, Some(1), Duration::from_secs(11));
        assert!(extract_critical_path(&predecessors).is_err());
    }
}


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

use std::sync::Arc;
use std::time::Duration;

use anyhow::Context as _;
use buck2_build_api::actions::RegisteredAction;
use buck2_build_signals::CriticalPathBackendName;
use buck2_build_signals::NodeDuration;
use buck2_core::soft_error;
use buck2_critical_path::compute_critical_path_potentials;
use buck2_critical_path::GraphBuilder;
use buck2_critical_path::OptionalVertexId;
use buck2_critical_path::PushError;
use buck2_events::span::SpanId;
use dupe::Dupe;
use smallvec::SmallVec;

use crate::backend::backend::BuildListenerBackend;
use crate::BuildInfo;
use crate::NodeData;
use crate::NodeKey;

/// An implementation of critical path that uses a longest-paths graph in order to produce
/// potential savings in addition to the critical path.
pub(crate) struct LongestPathGraphBackend {
    builder: anyhow::Result<GraphBuilder<NodeKey, NodeData>>,
    top_level_analysis: Vec<VisibilityEdge>,
}

/// Represents nodes that block us "seeing" other parts of the graph until they finish evaluating.
struct VisibilityEdge {
    node: NodeKey,
    makes_visible: Vec<NodeKey>,
}

impl LongestPathGraphBackend {
    pub(crate) fn new() -> Self {
        Self {
            builder: Ok(GraphBuilder::new()),
            top_level_analysis: Vec::new(),
        }
    }
}

impl BuildListenerBackend for LongestPathGraphBackend {
    fn process_node(
        &mut self,
        key: NodeKey,
        action: Option<Arc<RegisteredAction>>,
        duration: NodeDuration,
        dep_keys: impl Iterator<Item = NodeKey>,
        span_ids: SmallVec<[SpanId; 1]>,
    ) {
        let builder = match self.builder.as_mut() {
            Ok(b) => b,
            Err(..) => return,
        };

        let res = builder.push(
            key,
            dep_keys,
            NodeData {
                action,
                duration,
                span_ids,
            },
        );

        let res = res.or_else(|err| match err {
            e @ PushError::Overflow => Err(e.into()),
            e @ PushError::DuplicateKey { .. } => {
                soft_error!("critical_path_duplicate_key", e.into(), quiet: true)?;
                anyhow::Ok(())
            }
        });

        match res {
            Ok(()) => {}
            Err(e) => self.builder = Err(e),
        }
    }

    fn process_top_level_target(
        &mut self,
        analysis: NodeKey,
        artifacts: impl Iterator<Item = NodeKey>,
    ) {
        self.top_level_analysis.push(VisibilityEdge {
            node: analysis,
            makes_visible: artifacts.collect(),
        })
    }

    fn finish(self) -> anyhow::Result<BuildInfo> {
        let (graph, keys, mut data) = {
            let (graph, keys, data) = self.builder?.finish();

            let mut first_analysis = graph.allocate_vertex_data(OptionalVertexId::none());
            let mut n = 0;

            for visibility in &self.top_level_analysis {
                let analysis = &visibility.node;
                let artifacts = &visibility.makes_visible;

                let analysis = match keys.get(analysis) {
                    Some(k) => k,
                    None => continue, // Nothing depends on this,
                };

                let mut queue = Vec::new();

                // We have an analysis and a set of artifacts that we decided to build after having
                // evaluated this analysis. So, traverse all of those artifacts' dependencies, and
                // label them as depending on this top level analysis (but we only do that once).
                // Concretely, this expresses the idea that we only started knowing we cared about
                // those artifacts once we finished that analysis.

                for artifact in artifacts {
                    let artifact = match keys.get(artifact) {
                        Some(a) => a,
                        None => {
                            // Not built. Unexpected, but we don't report signals in all failure cases so that can happen.
                            continue;
                        }
                    };

                    queue.push(artifact);

                    while let Some(i) = queue.pop() {
                        if first_analysis[i].is_some() {
                            continue;
                        }

                        // We only traverse edges to things that produce artifacts here.
                        match keys[i] {
                            NodeKey::BuildKey(..)
                            | NodeKey::EnsureTransitiveSetProjectionKey(..)
                            | NodeKey::EnsureProjectedArtifactKey(..) => {}
                            _ => {
                                continue;
                            }
                        };

                        first_analysis[i] = analysis.into();
                        queue.extend(graph.iter_edges(i));
                        n += 1;
                    }
                }
            }

            let graph = graph
                .add_edges(&first_analysis, n)
                .context("Error adding first_analysis edges to graph")?;

            (graph, keys, data)
        };

        let durations = data.try_map_ref(|d| {
            d.duration
                .critical_path_duration()
                .as_micros()
                .try_into()
                .context("Duration `as_micros()` exceeds u64")
        })?;

        let (critical_path, critical_path_cost, replacement_durations) =
            compute_critical_path_potentials(&graph, &durations)
                .context("Error computing critical path potentials")?;

        drop(durations);

        let critical_path = critical_path
            .iter()
            .map(|(cp_idx, vertex_idx)| {
                let vertex_idx = *vertex_idx;
                let key = keys[vertex_idx].dupe();

                // OK to replace `data` with empty things here because we know that we will not access
                // the same index twice.
                let data = std::mem::replace(
                    &mut data[vertex_idx],
                    NodeData {
                        action: None,
                        duration: NodeDuration::zero(),
                        span_ids: Default::default(),
                    },
                );

                let potential = critical_path_cost.runtime - replacement_durations[cp_idx].runtime;

                (key, data, Some(Duration::from_micros(potential)))
            })
            .collect();

        Ok(BuildInfo {
            critical_path,
            num_nodes: graph.vertices_count() as _,
            num_edges: graph.edges_count() as _,
        })
    }

    fn name() -> CriticalPathBackendName {
        CriticalPathBackendName::LongestPathGraph
    }
}


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

#![allow(clippy::module_inception)]

pub mod backend;
pub mod default;
pub mod longest_path_graph;


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

use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::actions::artifact::get_artifact_fs::GetArtifactFs;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::build::materialize_artifact_group;
use buck2_build_api::build::BuildTargetResult;
use buck2_build_api::build::ConvertMaterializationContext;
use buck2_build_api::build::MaterializationContext;
use buck2_build_api::bxl::build_result::BxlBuildResult;
use buck2_build_api::bxl::calculation::BxlComputeResult;
use buck2_build_api::bxl::types::BxlFunctionLabel;
use buck2_cli_proto::build_request::Materializations;
use buck2_cli_proto::BxlRequest;
use buck2_cli_proto::BxlResponse;
use buck2_cli_proto::HasClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::result::SharedError;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::CellResolver;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::fs_util;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::package::PackageLabel;
use buck2_core::soft_error;
use buck2_core::tag_result;
use buck2_data::BxlEnsureArtifactsEnd;
use buck2_data::BxlEnsureArtifactsStart;
use buck2_data::StarlarkFailNoStacktrace;
use buck2_events::dispatch::get_dispatcher;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_interpreter::parse_import::parse_import_with_config;
use buck2_interpreter::parse_import::ParseImportOptions;
use buck2_interpreter::path::BxlFilePath;
use buck2_interpreter::path::StarlarkModulePath;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;
use futures::FutureExt;
use itertools::Itertools;
use starlark::errors::Diagnostic;

use crate::bxl::calculation::eval_bxl;
use crate::bxl::eval::get_bxl_callable;
use crate::bxl::eval::resolve_cli_args;
use crate::bxl::eval::BxlResolvedCliArgs;
use crate::bxl::eval::CliResolutionCtx;
use crate::bxl::key::BxlKey;
use crate::bxl::starlark_defs::functions::BxlErrorWithoutStacktrace;

pub async fn bxl_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: BxlRequest,
) -> anyhow::Result<BxlResponse> {
    run_server_command(BxlServerCommand { req }, ctx, partial_result_dispatcher).await
}

struct BxlServerCommand {
    req: BxlRequest,
}

#[async_trait]
impl ServerCommandTemplate for BxlServerCommand {
    type StartEvent = buck2_data::BxlCommandStart;
    type EndEvent = buck2_data::BxlCommandEnd;
    type Response = buck2_cli_proto::BxlResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

    fn start_event(&self) -> Self::StartEvent {
        let bxl_label = self.req.bxl_label.clone();
        buck2_data::BxlCommandStart { bxl_label }
    }

    fn end_event(&self, _response: &anyhow::Result<Self::Response>) -> Self::EndEvent {
        let bxl_label = self.req.bxl_label.clone();
        buck2_data::BxlCommandEnd { bxl_label }
    }

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        bxl(
            server_ctx,
            partial_result_dispatcher.as_writer(),
            ctx,
            &self.req,
        )
        .await
    }

    fn is_success(&self, response: &Self::Response) -> bool {
        response.error_messages.is_empty()
    }
}

async fn bxl(
    server_ctx: &dyn ServerCommandContextTrait,
    stdout: impl Write,
    ctx: DiceTransaction,
    request: &BxlRequest,
) -> anyhow::Result<buck2_cli_proto::BxlResponse> {
    let cwd = server_ctx.working_dir();
    let cell_resolver = ctx.get_cell_resolver().await?;
    let bxl_label = parse_bxl_label_from_cli(cwd, &request.bxl_label, &cell_resolver)?;
    let project_root = server_ctx.project_root().to_string();

    let client_ctx = request.client_context()?;
    let global_target_platform =
        target_platform_from_client_context(client_ctx, server_ctx, &ctx).await?;

    let bxl_args =
        match get_bxl_cli_args(cwd, &ctx, &bxl_label, &request.bxl_args, &cell_resolver).await? {
            BxlResolvedCliArgs::Resolved(bxl_args) => Arc::new(bxl_args),
            // Return early if user passed in `--help`
            BxlResolvedCliArgs::Help => {
                return Ok(BxlResponse {
                    project_root,
                    error_messages: Vec::new(),
                });
            }
        };

    let final_artifact_materializations =
        Materializations::from_i32(request.final_artifact_materializations)
            .with_context(|| "Invalid final_artifact_materializations")
            .unwrap();

    let bxl_key = BxlKey::new(bxl_label.clone(), bxl_args, global_target_platform);

    let ctx = &ctx;

    let BxlComputeResult {
        bxl_result,
        materializations,
    } = match eval_bxl(ctx, bxl_key.clone()).await {
        Ok(result) => result,
        Err(e) => {
            if !request.print_stacktrace {
                if let Some(shared) = e.downcast_ref::<SharedError>() {
                    if let Some(diag) = shared.inner().downcast_ref::<Diagnostic>() {
                        if let Some(fail_no_stacktrace) =
                            diag.message.downcast_ref::<BxlErrorWithoutStacktrace>()
                        {
                            let dispatcher = get_dispatcher();
                            dispatcher.instant_event(StarlarkFailNoStacktrace {
                                trace: format!("{}", diag),
                            });
                            dispatcher.console_message(
                                "Re-run the script with `-v5` to show the full stacktrace"
                                    .to_owned(),
                            );
                            return Err((fail_no_stacktrace.clone()).into());
                        }
                    }
                }
            }
            return Err(e);
        }
    };

    let materialization_context = ConvertMaterializationContext::with_existing_map(
        final_artifact_materializations,
        // Note: even though we have an Arc of the materialization map, we must actually clone the map
        // so that we don't mutate the materialization state stored when materializing the ensured
        // artifacts. We need to clone it so that we don't re-materialize what was already done, but
        // in a separate instance of the map.
        &Arc::new((*materializations).clone()),
    );

    let build_result = ensure_artifacts(ctx, &materialization_context, &bxl_result).await;
    copy_output(stdout, ctx, bxl_result.get_output_loc()).await?;
    copy_output(server_ctx.stderr()?, ctx, bxl_result.get_error_loc()).await?;

    let error_messages = match build_result {
        Ok(_) => vec![],
        Err(errors) => errors.iter().map(|e| format!("{:#}", e)).unique().collect(),
    };
    Ok(BxlResponse {
        project_root,
        error_messages,
    })
}

pub(crate) async fn get_bxl_cli_args(
    cwd: &ProjectRelativePath,
    ctx: &DiceTransaction,
    bxl_label: &BxlFunctionLabel,
    bxl_args: &Vec<String>,
    cell_resolver: &CellResolver,
) -> anyhow::Result<BxlResolvedCliArgs> {
    let cur_package = PackageLabel::from_cell_path(cell_resolver.get_cell_path(&cwd)?.as_ref());
    let cell_name = cell_resolver.find(&cwd)?;

    let target_alias_resolver = ctx.target_alias_resolver_for_cell(cell_name).await?;

    let bxl_module = ctx
        .get_loaded_module(StarlarkModulePath::BxlFile(&bxl_label.bxl_path))
        .await?;

    let frozen_callable = get_bxl_callable(bxl_label, &bxl_module)?;
    let cli_ctx = CliResolutionCtx {
        target_alias_resolver,
        cell_resolver: cell_resolver.dupe(),
        relative_dir: cur_package,
        dice: ctx,
    };

    resolve_cli_args(bxl_label, &cli_ctx, bxl_args, &frozen_callable).await
}

async fn copy_output<W: Write>(
    mut output: W,
    dice: &DiceComputations,
    output_loc: &BuckOutPath,
) -> anyhow::Result<()> {
    let loc = dice.global_data().get_io_provider().project_root().resolve(
        &dice
            .get_artifact_fs()
            .await?
            .buck_out_path_resolver()
            .resolve_gen(output_loc),
    );

    // we write the output to a file in buck-out as cache so we don't use memory caching it in
    // DICE. So now we open the file and read it all into the destination stream.
    let mut file = tag_result!(
        "bxl_output_missing",
        fs_util::open_file(loc),
        quiet: true,
        daemon_in_memory_state_is_corrupted: true,
        task: false
    )?;
    io::copy(&mut file, &mut output)?;
    Ok(())
}

async fn ensure_artifacts(
    ctx: &DiceComputations,
    materialization_ctx: &MaterializationContext,
    bxl_result: &buck2_build_api::bxl::result::BxlResult,
) -> Result<(), Vec<SharedError>> {
    match bxl_result {
        buck2_build_api::bxl::result::BxlResult::None { .. } => Ok(()),
        buck2_build_api::bxl::result::BxlResult::BuildsArtifacts {
            built, artifacts, ..
        } => {
            get_dispatcher()
                .span_async(BxlEnsureArtifactsStart {}, async move {
                    (
                        ensure_artifacts_inner(ctx, materialization_ctx, built, artifacts).await,
                        BxlEnsureArtifactsEnd {},
                    )
                })
                .await
        }
    }
}

async fn ensure_artifacts_inner(
    ctx: &DiceComputations,
    materialization_ctx: &MaterializationContext,
    built: &[BxlBuildResult],
    artifacts: &[ArtifactGroup],
) -> Result<(), Vec<SharedError>> {
    let mut futs = vec![];

    built.iter().for_each(|res| match res {
        BxlBuildResult::Built(BuildTargetResult { outputs, .. }) => {
            outputs.iter().for_each(|res| match res {
                Ok(artifacts) => {
                    for (artifact, _value) in artifacts.values.iter() {
                        futs.push(
                            async {
                                materialize_artifact_group(
                                    ctx,
                                    &ArtifactGroup::Artifact(artifact.dupe()),
                                    materialization_ctx,
                                )
                                .await?;
                                Ok(())
                            }
                            .boxed(),
                        )
                    }
                }
                Err(e) => futs.push(futures::future::ready(Err(e.dupe())).boxed()),
            });
        }

        BxlBuildResult::None => {}
    });

    artifacts.iter().for_each(|a| {
        futs.push(
            async move {
                materialize_artifact_group(ctx, a, materialization_ctx).await?;
                Ok(())
            }
            .boxed(),
        );
    });

    let res = futures::future::join_all(futs)
        .await
        .into_iter()
        .filter_map(|res| res.err())
        .collect::<Vec<_>>();
    if res.is_empty() { Ok(()) } else { Err(res) }
}

#[derive(Debug, thiserror::Error)]
enum BxlLabelError {
    #[error(
        "bxl label should be of format `<cell>//path/to/file.bxl:function_name`, but got `{0}`"
    )]
    Format(String),
    #[error(
        "The bxl function path `{got}` should use the canonical name `{wanted}`. If your bxl changes aren't being detected, this is probably why"
    )]
    WrongCell { got: String, wanted: CellPath },
}

/// Parse the bxl function label out of cli pattern
pub(crate) fn parse_bxl_label_from_cli(
    cwd: &ProjectRelativePath,
    bxl_label: &str,
    cell_resolver: &CellResolver,
) -> anyhow::Result<BxlFunctionLabel> {
    let current_cell = cell_resolver.get_cell_path(cwd)?;

    // Targets with cell aliases should be resolved against the cell mapping
    // as defined the cell derived from the cwd.
    let cell_alias_resolver = cell_resolver
        .get(current_cell.cell())
        .unwrap()
        .cell_alias_resolver();

    let (bxl_path, bxl_fn) = bxl_label
        .rsplit_once(':')
        .ok_or_else(|| BxlLabelError::Format(bxl_label.to_owned()))?;

    const OPTS: ParseImportOptions = ParseImportOptions {
        allow_missing_at_symbol: true,
        allow_relative_imports: true,
    };
    let import_path =
        parse_import_with_config(cell_alias_resolver, &current_cell, bxl_path, &OPTS)?;

    let project_path = cell_resolver.resolve_path(import_path.as_ref())?;
    let reformed_path = cell_resolver.get_cell_path(&project_path)?;
    if reformed_path.cell() != import_path.cell() {
        soft_error!(
            "bxl_label_wrong_cell",
            BxlLabelError::WrongCell {
                got: bxl_path.to_owned(),
                wanted: reformed_path,
            }
            .into()
        )?;
    }

    Ok(BxlFunctionLabel {
        bxl_path: BxlFilePath::new(import_path)?,
        name: bxl_fn.to_owned(),
    })
}


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

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

use std::sync::Once;

pub mod bxl;
pub mod command;
pub mod profile_command;

pub fn init_late_bindings() {
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        bxl::starlark_defs::init_bxl_functions();
        bxl::calculation::init_bxl_calculation_impl();
    });
}

#[test]
fn init_late_bindings_for_test() {
    #[ctor::ctor]
    fn init() {
        init_late_bindings();
    }
}


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

use async_trait::async_trait;
use buck2_cli_proto::profile_request::ProfileOpts;
use buck2_cli_proto::HasClientContext;
use buck2_cli_proto::ProfileRequest;
use buck2_cli_proto::ProfileResponse;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_interpreter::dice::starlark_profiler::StarlarkProfilerConfiguration;
use buck2_interpreter::starlark_profiler::StarlarkProfileModeOrInstrumentation;
use buck2_profile::get_profile_response;
use buck2_profile::starlark_profiler_configuration_from_request;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::NoPartialResult;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceTransaction;
use futures::FutureExt;

use crate::bxl::eval::eval;
use crate::bxl::eval::BxlResolvedCliArgs;
use crate::bxl::key::BxlKey;
use crate::command::get_bxl_cli_args;
use crate::command::parse_bxl_label_from_cli;

pub async fn bxl_profile_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
    req: ProfileRequest,
) -> anyhow::Result<ProfileResponse> {
    run_server_command(
        BxlProfileServerCommand { req },
        ctx,
        partial_result_dispatcher,
    )
    .await
}

struct BxlProfileServerCommand {
    req: ProfileRequest,
}

#[async_trait]
impl ServerCommandTemplate for BxlProfileServerCommand {
    type StartEvent = buck2_data::ProfileCommandStart;
    type EndEvent = buck2_data::ProfileCommandEnd;
    type Response = buck2_cli_proto::ProfileResponse;
    type PartialResult = NoPartialResult;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        _partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        match self
            .req
            .profile_opts
            .as_ref()
            .expect("BXL profile opts not populated")
        {
            ProfileOpts::BxlProfile(opts) => {
                let output = AbsPath::new(Path::new(&self.req.destination_path))?;

                let profile_mode = starlark_profiler_configuration_from_request(&self.req)?;

                let profile_data = match profile_mode {
                    StarlarkProfilerConfiguration::ProfileBxl(profile_mode) => {
                        let cwd = server_ctx.working_dir();

                        let cell_resolver = ctx.get_cell_resolver().await?;
                        let bxl_label =
                            parse_bxl_label_from_cli(cwd, &opts.bxl_label, &cell_resolver)?;

                        let bxl_args = match get_bxl_cli_args(
                            cwd,
                            &ctx,
                            &bxl_label,
                            &opts.bxl_args,
                            &cell_resolver,
                        )
                        .await?
                        {
                            BxlResolvedCliArgs::Resolved(bxl_args) => Arc::new(bxl_args),
                            _ => {
                                return Err(anyhow::anyhow!(
                                    "Help docs were displayed. No profiler data available"
                                ));
                            }
                        };

                        let client_ctx = self.req.client_context()?;
                        let global_target_platform =
                            target_platform_from_client_context(client_ctx, server_ctx, &ctx)
                                .await?;

                        let bxl_key =
                            BxlKey::new(bxl_label.clone(), bxl_args, global_target_platform);

                        server_ctx
                            .cancellation_context()
                            .with_structured_cancellation(|observer| {
                                async move {
                                    anyhow::Ok(
                                        eval(
                                            &ctx,
                                            bxl_key,
                                            StarlarkProfileModeOrInstrumentation::Profile(
                                                profile_mode,
                                            ),
                                            observer,
                                        )
                                        .await?
                                        .1
                                        .map(Arc::new)
                                        .expect("No bxl profile data found"),
                                    )
                                }
                                .boxed()
                            })
                            .await?
                    }
                    _ => {
                        return Err(anyhow::anyhow!("Incorrect profile mode (internal error)"));
                    }
                };

                get_profile_response(profile_data, &self.req, output)
            }
            _ => {
                return Err(anyhow::anyhow!(
                    "Expected BXL profile opts, not target profile opts"
                ));
            }
        }
    }

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


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

use std::sync::Arc;

use async_trait::async_trait;
use buck2_build_api::bxl::calculation::BxlCalculationDyn;
use buck2_build_api::bxl::calculation::BxlComputeResult;
use buck2_build_api::bxl::calculation::BXL_CALCULATION_IMPL;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_common::result::ToUnsharedResultExt;
use buck2_core::base_deferred_key::BaseDeferredKeyDyn;
use buck2_interpreter::dice::starlark_profiler::GetStarlarkProfilerInstrumentation;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use futures::future::FutureExt;
use more_futures::cancellation::CancellationContext;

use crate::bxl::eval::eval;
use crate::bxl::key::BxlKey;

#[derive(Debug)]
struct BxlCalculationImpl;

#[async_trait]
impl BxlCalculationDyn for BxlCalculationImpl {
    async fn eval_bxl(
        &self,
        ctx: &DiceComputations,
        bxl: Arc<dyn BaseDeferredKeyDyn>,
    ) -> anyhow::Result<BxlComputeResult> {
        eval_bxl(ctx, BxlKey::from_base_deferred_key_dyn_impl_err(bxl)?).await
    }
}

pub(crate) fn init_bxl_calculation_impl() {
    BXL_CALCULATION_IMPL.init(&BxlCalculationImpl);
}

pub(crate) async fn eval_bxl(
    ctx: &DiceComputations,
    bxl: BxlKey,
) -> anyhow::Result<BxlComputeResult> {
    ctx.compute(&internal::BxlComputeKey(bxl))
        .await?
        .unshared_error()
}

#[async_trait]
impl Key for internal::BxlComputeKey {
    type Value = SharedResult<BxlComputeResult>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        cancellation: &CancellationContext,
    ) -> Self::Value {
        let key = self.0.dupe();

        let profiler = ctx.get_profile_mode_for_intermediate_analysis().await?;

        cancellation
            .with_structured_cancellation(|observer| {
                async move {
                    eval(ctx, key, profiler, observer).await.shared_error().map(
                        |(result, _, materializations)| BxlComputeResult {
                            bxl_result: Arc::new(result),
                            materializations,
                        },
                    )
                }
                .boxed()
            })
            .await
    }

    fn equality(_: &Self::Value, _: &Self::Value) -> bool {
        false
    }
}

mod internal {
    use allocative::Allocative;
    use derive_more::Display;
    use dupe::Dupe;

    use crate::bxl::key::BxlKey;

    #[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
    pub(crate) struct BxlComputeKey(pub BxlKey);
}

#[cfg(test)]
pub(crate) mod testing {
    pub(crate) use crate::bxl::calculation::internal::BxlComputeKey;
}


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

#[cfg(test)]
mod tests {
    use std::any;
    use std::any::Demand;
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering;
    use std::sync::Arc;

    use allocative::Allocative;
    use buck2_build_api::bxl::calculation::BxlComputeResult;
    use buck2_build_api::bxl::result::BxlResult;
    use buck2_build_api::bxl::types::BxlFunctionLabel;
    use buck2_build_api::deferred::calculation::DeferredCalculation;
    use buck2_build_api::deferred::types::BaseKey;
    use buck2_build_api::deferred::types::Deferred;
    use buck2_build_api::deferred::types::DeferredCtx;
    use buck2_build_api::deferred::types::DeferredInput;
    use buck2_build_api::deferred::types::DeferredRegistry;
    use buck2_build_api::deferred::types::DeferredTable;
    use buck2_build_api::deferred::types::DeferredValue;
    use buck2_common::dice::data::testing::SetTestingIoProvider;
    use buck2_common::executor_config::CommandExecutorConfig;
    use buck2_common::result::ToSharedResultExt;
    use buck2_core::base_deferred_key::BaseDeferredKey;
    use buck2_core::fs::buck_out_path::BuckOutPath;
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
    use buck2_core::fs::project::ProjectRootTemp;
    use buck2_execute::digest_config::DigestConfig;
    use buck2_execute::digest_config::SetDigestConfig;
    use buck2_execute::execute::dice_data::set_fallback_executor_config;
    use buck2_interpreter::path::BxlFilePath;
    use buck2_util::collections::ordered_map::OrderedMap;
    use dice::testing::DiceBuilder;
    use dice::UserComputationData;
    use dupe::Dupe;
    use indexmap::IndexSet;

    use crate::bxl::calculation::testing::BxlComputeKey;
    use crate::bxl::key::BxlKey;

    #[derive(Allocative)]
    struct FakeDeferred(usize, IndexSet<DeferredInput>, Arc<AtomicBool>);

    impl any::Provider for FakeDeferred {
        fn provide<'a>(&'a self, _demand: &mut Demand<'a>) {}
    }

    impl Deferred for FakeDeferred {
        type Output = usize;

        fn inputs(&self) -> &IndexSet<DeferredInput> {
            &self.1
        }

        fn execute(
            &self,
            _ctx: &mut dyn DeferredCtx,
        ) -> anyhow::Result<DeferredValue<Self::Output>> {
            self.2.store(true, Ordering::SeqCst);
            Ok(DeferredValue::Ready(self.0))
        }
    }

    #[tokio::test]
    async fn lookup_deferred_from_bxl() -> anyhow::Result<()> {
        let bxl = BxlKey::new(
            BxlFunctionLabel {
                bxl_path: BxlFilePath::testing_new("cell", "dir"),
                name: "foo".to_owned(),
            },
            Arc::new(OrderedMap::new()),
            None,
        );

        let mut deferred = DeferredRegistry::new(BaseKey::Base(BaseDeferredKey::BxlLabel(
            bxl.dupe().into_base_deferred_key_dyn_impl(),
        )));

        let executed0 = Arc::new(AtomicBool::new(false));
        let executed1 = Arc::new(AtomicBool::new(false));
        let data0 = deferred.defer(FakeDeferred(1, IndexSet::new(), executed0.dupe()));
        let data1 = deferred.defer(FakeDeferred(5, IndexSet::new(), executed1.dupe()));
        let deferred_result = DeferredTable::new(deferred.take_result()?);

        let fs = ProjectRootTemp::new()?;
        let dice = DiceBuilder::new()
            .set_data(|data| {
                data.set_testing_io_provider(&fs);
                data.set_digest_config(DigestConfig::testing_default());
            })
            .mock_and_return(
                BxlComputeKey(bxl.dupe()),
                anyhow::Ok(BxlComputeResult {
                    bxl_result: Arc::new(BxlResult::BuildsArtifacts {
                        output_loc: BuckOutPath::new(
                            BaseDeferredKey::BxlLabel(bxl.dupe().into_base_deferred_key_dyn_impl()),
                            ForwardRelativePathBuf::unchecked_new("test".to_owned()),
                        ),
                        error_loc: BuckOutPath::new(
                            BaseDeferredKey::BxlLabel(bxl.dupe().into_base_deferred_key_dyn_impl()),
                            ForwardRelativePathBuf::unchecked_new("error_test".to_owned()),
                        ),
                        built: vec![],
                        artifacts: vec![],
                        deferred: deferred_result,
                    }),
                    materializations: Arc::new(Default::default()),
                })
                .shared_error(),
            );

        let mut dice_data = UserComputationData::new();
        set_fallback_executor_config(&mut dice_data.data, CommandExecutorConfig::testing_local());

        let dice = dice.build(dice_data)?.commit().await;
        let deferred_result = dice.compute_deferred_data(&data0).await?;
        assert_eq!(*deferred_result, 1);
        assert!(executed0.load(Ordering::SeqCst));
        // we should cache deferred execution
        executed0.store(false, Ordering::SeqCst);
        let deferred_result = dice.compute_deferred_data(&data0).await?;
        assert_eq!(*deferred_result, 1);
        assert!(!executed0.load(Ordering::SeqCst));

        let deferred_result = dice.compute_deferred_data(&data1).await?;
        assert_eq!(*deferred_result, 5);
        assert!(executed1.load(Ordering::SeqCst));
        // we should cache deferred execution
        executed1.store(false, Ordering::SeqCst);
        assert_eq!(*deferred_result, 5);
        assert!(!executed1.load(Ordering::SeqCst));

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

use std::cell::RefCell;
use std::sync::Arc;

use anyhow::Context;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::artifact::get_artifact_fs::GetArtifactFs;
use buck2_build_api::bxl::result::BxlResult;
use buck2_build_api::bxl::types::BxlFunctionLabel;
use buck2_build_api::deferred::types::DeferredTable;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::events::HasEvents;
use buck2_common::target_aliases::BuckConfigTargetAliasResolver;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::cells::CellResolver;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_core::package::PackageLabel;
use buck2_data::BxlExecutionEnd;
use buck2_data::BxlExecutionStart;
use buck2_events::dispatch::console_message;
use buck2_events::dispatch::with_dispatcher;
use buck2_events::dispatch::with_dispatcher_async;
use buck2_execute::digest_config::HasDigestConfig;
use buck2_interpreter::dice::starlark_provider::with_starlark_eval_provider;
use buck2_interpreter::factory::StarlarkEvaluatorProvider;
use buck2_interpreter::file_loader::LoadedModule;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_interpreter::path::StarlarkModulePath;
use buck2_interpreter::print_handler::EventDispatcherPrintHandler;
use buck2_interpreter::starlark_profiler::StarlarkProfileDataAndStats;
use buck2_interpreter::starlark_profiler::StarlarkProfileModeOrInstrumentation;
use buck2_interpreter::starlark_profiler::StarlarkProfiler;
use buck2_interpreter::starlark_profiler::StarlarkProfilerOrInstrumentation;
use buck2_util::collections::ordered_map::OrderedMap;
use clap::ErrorKind;
use dashmap::DashMap;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;
use itertools::Itertools;
use more_futures::cancellable_future::CancellationObserver;
use starlark::environment::Module;
use starlark::eval::Evaluator;
use starlark::values::structs::AllocStruct;
use starlark::values::OwnedFrozenValueTyped;
use starlark::values::Value;
use starlark::values::ValueTyped;
use thiserror::Error;

use crate::bxl::key::BxlKey;
use crate::bxl::starlark_defs::cli_args::CliArgValue;
use crate::bxl::starlark_defs::context::starlark_async::BxlSafeDiceComputations;
use crate::bxl::starlark_defs::context::BxlContext;
use crate::bxl::starlark_defs::FrozenBxlFunction;

pub(crate) async fn eval(
    ctx: &DiceComputations,
    key: BxlKey,
    profile_mode_or_instrumentation: StarlarkProfileModeOrInstrumentation,
    liveness: CancellationObserver,
) -> anyhow::Result<(
    BxlResult,
    Option<StarlarkProfileDataAndStats>,
    Arc<DashMap<BuildArtifact, ()>>,
)> {
    // Note: because we use `block_in_place`, that will prevent the inner future from being polled
    // and yielded. So, for cancellation observers to work properly within the dice cancellable
    // future context, we need the future that it's attached to the cancellation context can
    // yield and be polled. To ensure that, we have to spawn the future that then enters block_in_place

    let dispatcher = ctx.per_transaction_data().get_dispatcher().dupe();

    let (_, futs) = unsafe {
        // SAFETY: as long as we don't `forget` the return object from `scope_and_collect`, it is safe

        // Additional cancellation notes:
        // the `scope_and_collect` will block on drop, but it will move the blocking to a tokio
        // blocking thread, freeing up the main worker threads. Additionally, the `spawn_cancellable`
        // on the scope will be dropped at the earliest await point. If we are within the blocking
        // section of bxl, the cancellation observer will be notified and cause the blocking calls
        // to terminate.
        async_scoped::TokioScope::scope_and_collect(|s| {
            s.spawn_cancellable(
                with_dispatcher_async(dispatcher.dupe(), async move {
                    let bxl_module = ctx
                        .get_loaded_module(StarlarkModulePath::BxlFile(&key.label().bxl_path))
                        .await?;

                    let cell_resolver = ctx.get_cell_resolver().await?;

                    let bxl_cell = cell_resolver
                        .get(key.label().bxl_path.cell())
                        .with_context(|| {
                            format!("Cell does not exist: `{}`", key.label().bxl_path.cell())
                        })?
                        .dupe();

                    let target_alias_resolver = ctx
                        .target_alias_resolver_for_cell(key.label().bxl_path.cell())
                        .await?;

                    let project_fs = ctx.global_data().get_io_provider().project_root().dupe();
                    let artifact_fs = ctx.get_artifact_fs().await?;

                    let digest_config = ctx.global_data().get_digest_config();

                    // The bxl function may trigger async operations like builds, analysis, parsing etc, but those
                    // will be blocking calls so that starlark can remain synchronous.
                    // So indicate to tokio that this may block in place to avoid starvation. Ideally we use
                    // spawn_blocking but that requires a static lifetime. There is no `join`s of multiple
                    // futures that requires work to be done on the current thread, so using block_in_place
                    // should have no noticeable different compared to spawn_blocking

                    // we put a file as our output stream cache. The file is associated with the `BxlKey`, which
                    // is super important, as it HAS to be the SAME as the DiceKey so that DICE is keeping
                    // the output file cache up to date.
                    let output_stream = BuckOutPath::new(
                        BaseDeferredKey::BxlLabel(key.dupe().into_base_deferred_key_dyn_impl()),
                        ForwardRelativePathBuf::unchecked_new(
                            "__bxl_internal__/outputstream_cache".to_owned(),
                        ),
                    );
                    let file_path = artifact_fs
                        .buck_out_path_resolver()
                        .resolve_gen(&output_stream);

                    let file = RefCell::new(Box::new(project_fs.create_file(&file_path, false)?));

                    let error_stream = BuckOutPath::new(
                        BaseDeferredKey::BxlLabel(key.dupe().into_base_deferred_key_dyn_impl()),
                        ForwardRelativePathBuf::unchecked_new(
                            "__bxl_internal__/errorstream_cache".to_owned(),
                        ),
                    );
                    let error_file_path = artifact_fs
                        .buck_out_path_resolver()
                        .resolve_gen(&error_stream);

                    let error_file =
                        RefCell::new(Box::new(project_fs.create_file(&error_file_path, false)?));

                    let print = EventDispatcherPrintHandler(dispatcher.clone());

                    let mut profiler_opt = profile_mode_or_instrumentation
                        .profile_mode()
                        .map(|profile_mode| StarlarkProfiler::new(profile_mode.dupe(), true));

                    let mut profiler = match &mut profiler_opt {
                        None => StarlarkProfilerOrInstrumentation::disabled(),
                        Some(profiler) => StarlarkProfilerOrInstrumentation::for_profiler(profiler),
                    };

                    let global_target_platform = key.global_target_platform().clone();

                    let (bxl_result, materializations) = with_starlark_eval_provider(
                        ctx,
                        &mut profiler,
                        format!("bxl:{}", key),
                        move |provider| {
                            let env = Module::new();

                            let resolved_args = env.heap().alloc(AllocStruct(
                                key.cli_args()
                                    .iter()
                                    .map(|(k, v)| (k, v.as_starlark(env.heap()))),
                            ));

                            let mut eval = provider.make(&env)?;
                            let bxl_function_name = key.label().name.clone();
                            let frozen_callable = get_bxl_callable(key.label(), &bxl_module)?;
                            eval.set_print_handler(&print);

                            let bxl_ctx = BxlContext::new(
                                eval.heap(),
                                key,
                                resolved_args,
                                target_alias_resolver,
                                project_fs,
                                artifact_fs,
                                cell_resolver,
                                bxl_cell.name(),
                                BxlSafeDiceComputations::new(ctx, liveness),
                                file,
                                error_file,
                                digest_config,
                                global_target_platform,
                            )?;

                            let bxl_ctx =
                                ValueTyped::<BxlContext>::new(env.heap().alloc(bxl_ctx)).unwrap();

                            let result = tokio::task::block_in_place(|| {
                                with_dispatcher(dispatcher.clone(), || {
                                    dispatcher.clone().span(
                                        BxlExecutionStart {
                                            name: bxl_function_name,
                                        },
                                        || {
                                            (
                                                eval_bxl(
                                                    &mut eval,
                                                    frozen_callable,
                                                    bxl_ctx.to_value(),
                                                    provider,
                                                ),
                                                BxlExecutionEnd {},
                                            )
                                        },
                                    )
                                })
                            })?;
                            if !result.is_none() {
                                return Err(anyhow::anyhow!(NotAValidReturnType(
                                    result.get_type()
                                )));
                            }

                            let (actions, ensured_artifacts, materializations) =
                                BxlContext::take_state(bxl_ctx)?;
                            std::mem::drop(eval);

                            let (actions_finalizer, ensured_artifacts, materializations) = {
                                // help rust understand that actions is consumed here.
                                let actions = actions;
                                match actions {
                                    Some(registry) => (
                                        Some(registry.finalize(&env)?),
                                        ensured_artifacts,
                                        materializations,
                                    ),
                                    None => (None, ensured_artifacts, materializations),
                                }
                            };

                            let (frozen_module, bxl_result) = match actions_finalizer {
                                Some(actions_finalizer) => {
                                    // this bxl registered actions, so extract the deferreds from it
                                    let (frozen_module, deferred) = actions_finalizer(env)?;

                                    let deferred_table =
                                        DeferredTable::new(deferred.take_result()?);

                                    (
                                        frozen_module,
                                        BxlResult::new(
                                            output_stream,
                                            error_stream,
                                            ensured_artifacts,
                                            deferred_table,
                                        ),
                                    )
                                }
                                None => {
                                    let frozen_module = env.freeze()?;

                                    // this bxl did not try to build anything, so we don't have any deferreds
                                    (
                                        frozen_module,
                                        BxlResult::new(
                                            output_stream,
                                            error_stream,
                                            ensured_artifacts,
                                            DeferredTable::new(Vec::new()),
                                        ),
                                    )
                                }
                            };

                            provider
                                .visit_frozen_module(Some(&frozen_module))
                                .context("Profiler heap visitation failed")?;

                            Ok((bxl_result, materializations))
                        },
                    )
                    .await?;

                    let profile_data = profiler_opt.map(|p| p.finish()).transpose()?;
                    Ok((bxl_result, profile_data, materializations))
                }),
                || Err(anyhow::anyhow!("cancelled")),
            )
        })
    }
    .await;

    match futs.into_iter().exactly_one() {
        Ok(res) => res?,
        Err(_) => panic!("only spawned one task"),
    }
}

fn eval_bxl<'a>(
    eval: &mut Evaluator<'a, '_>,
    frozen_callable: OwnedFrozenValueTyped<FrozenBxlFunction>,
    ctx: Value<'a>,
    provider: &mut dyn StarlarkEvaluatorProvider,
) -> anyhow::Result<Value<'a>> {
    let bxl_impl = frozen_callable.implementation();
    let result = eval.eval_function(bxl_impl.to_value(), &[ctx], &[]);

    provider
        .evaluation_complete(eval)
        .context("Profiler finalization failed")?;
    result
}

pub fn get_bxl_callable<'a>(
    spec: &BxlFunctionLabel,
    bxl_module: &'a LoadedModule,
) -> anyhow::Result<OwnedFrozenValueTyped<FrozenBxlFunction>> {
    let callable = bxl_module.env().get_any_visibility(&spec.name)?.0;

    Ok(callable
        .downcast::<FrozenBxlFunction>()
        .unwrap_or_else(|e| {
            panic!(
                "A bxl function should be a BxlFunction. It was a {}",
                e.value().get_type(),
            )
        }))
}

pub struct CliResolutionCtx<'a> {
    pub target_alias_resolver: BuckConfigTargetAliasResolver,
    pub cell_resolver: CellResolver,
    pub relative_dir: PackageLabel,
    pub dice: &'a DiceTransaction,
}

pub(crate) enum BxlResolvedCliArgs {
    Resolved(OrderedMap<String, CliArgValue>),
    Help,
}

pub(crate) async fn resolve_cli_args<'a>(
    spec: &BxlFunctionLabel,
    cli_ctx: &CliResolutionCtx<'a>,
    bxl_args: &Vec<String>,
    frozen_callable: &'a FrozenBxlFunction,
) -> anyhow::Result<BxlResolvedCliArgs> {
    match frozen_callable
        .to_clap(clap::Command::new(&spec.name).no_binary_name(true))
        .try_get_matches_from(bxl_args)
    {
        Ok(args) => Ok(BxlResolvedCliArgs::Resolved(
            frozen_callable.parse_clap(args, cli_ctx).await?,
        )),
        Err(e) => match e.kind() {
            ErrorKind::DisplayHelp => {
                let mut help_out = Vec::new();

                frozen_callable
                    .to_clap(clap::Command::new(&spec.name).no_binary_name(true))
                    .write_long_help(&mut help_out)
                    .unwrap();
                let help_msg = String::from_utf8(help_out)?;

                console_message(help_msg);

                Ok(BxlResolvedCliArgs::Help)
            }
            _ => Err(e.into()),
        },
    }
}

#[derive(Debug, Error)]
#[error("Expected `NoneType` to be returned from bxl. Got return value `{0}`")]
struct NotAValidReturnType(&'static str);


/*
 * 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::Any;
use std::collections::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_build_api::bxl::types::BxlFunctionLabel;
use buck2_core::base_deferred_key::BaseDeferredKeyDyn;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_core::target::label::TargetLabel;
use buck2_data::action_key_owner::BaseDeferredKeyProto;
use buck2_data::ToProtoMessage;
use buck2_util::collections::ordered_map::OrderedMap;
use cmp_any::PartialEqAny;
use dupe::Dupe;

use crate::bxl::starlark_defs::cli_args::CliArgValue;

#[derive(
    Clone,
    Dupe,
    derive_more::Display,
    Debug,
    Eq,
    Hash,
    PartialEq,
    Ord,
    PartialOrd,
    Allocative
)]
pub(crate) struct BxlKey(Arc<BxlKeyData>);

impl BxlKey {
    pub(crate) fn new(
        spec: BxlFunctionLabel,
        bxl_args: Arc<OrderedMap<String, CliArgValue>>,
        global_target_platform: Option<TargetLabel>,
    ) -> Self {
        Self(Arc::new(BxlKeyData {
            spec,
            bxl_args,
            global_target_platform,
        }))
    }

    pub(crate) fn label(&self) -> &BxlFunctionLabel {
        &self.0.spec
    }

    pub(crate) fn cli_args(&self) -> &Arc<OrderedMap<String, CliArgValue>> {
        &self.0.bxl_args
    }

    pub(crate) fn into_base_deferred_key_dyn_impl(self) -> Arc<dyn BaseDeferredKeyDyn> {
        self.0
    }

    pub(crate) fn from_base_deferred_key_dyn_impl(
        key: Arc<dyn BaseDeferredKeyDyn>,
    ) -> Option<Self> {
        key.into_any().downcast().ok().map(BxlKey)
    }

    pub(crate) fn from_base_deferred_key_dyn_impl_err(
        key: Arc<dyn BaseDeferredKeyDyn>,
    ) -> anyhow::Result<Self> {
        Self::from_base_deferred_key_dyn_impl(key).context("Not BxlKey (internal error)")
    }

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

#[derive(
    Clone,
    derive_more::Display,
    Debug,
    Eq,
    Hash,
    PartialEq,
    Ord,
    PartialOrd,
    Allocative
)]
#[display(fmt = "{}", "spec")]
struct BxlKeyData {
    spec: BxlFunctionLabel,
    bxl_args: Arc<OrderedMap<String, CliArgValue>>,
    global_target_platform: Option<TargetLabel>,
}

impl BxlKeyData {
    fn as_proto(&self) -> buck2_data::BxlFunctionKey {
        buck2_data::BxlFunctionKey {
            label: Some(self.spec.as_proto()),
        }
    }
}

impl BaseDeferredKeyDyn for BxlKeyData {
    fn eq_token(&self) -> PartialEqAny {
        PartialEqAny::new(self)
    }

    fn hash(&self) -> u64 {
        let mut hasher = DefaultHasher::new();
        Hash::hash(self, &mut hasher);
        hasher.finish()
    }

    fn make_hashed_path(
        &self,
        base: &ProjectRelativePath,
        prefix: &ForwardRelativePath,
        action_key: Option<&str>,
        path: &ForwardRelativePath,
    ) -> ProjectRelativePathBuf {
        let label = &self.spec;
        let cell_relative_path = label.bxl_path.path().path().as_str();

        let output_hash = {
            let mut hasher = DefaultHasher::new();
            self.bxl_args.hash(&mut hasher);
            let output_hash = hasher.finish();
            format!("{:x}", output_hash)
        };

        // It is performance critical that we use slices and allocate via `join` instead of
        // repeated calls to `join` on the path object because `join` allocates on each call,
        // which has a significant impact.
        let parts = [
            base.as_str(),
            "/",
            prefix.as_str(),
            "-bxl/",
            label.bxl_path.cell().as_str(),
            "/",
            output_hash.as_str(),
            "/",
            cell_relative_path,
            if cell_relative_path.is_empty() {
                ""
            } else {
                "/"
            },
            "__",
            label.name.as_str(),
            "__",
            action_key.unwrap_or_default(),
            if action_key.is_none() { "" } else { "__" },
            "/",
            path.as_str(),
        ];

        ProjectRelativePathBuf::unchecked_new(parts.concat())
    }

    fn configured_label(&self) -> Option<ConfiguredTargetLabel> {
        None
    }

    fn to_proto(&self) -> BaseDeferredKeyProto {
        BaseDeferredKeyProto::BxlKey(self.as_proto())
    }

    fn into_any(self: Arc<Self>) -> Arc<dyn Any + Send + Sync> {
        self
    }
}


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

pub mod calculation;
mod deferred;
pub mod eval;
pub(crate) mod key;
pub mod starlark_defs;
pub(crate) mod value_as_starlark_target_label;


/*
 * 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 buck2_common::target_aliases::BuckConfigTargetAliasResolver;
use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::name::CellName;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::cells::CellResolver;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::TargetLabel;
use buck2_interpreter::types::target_label::StarlarkTargetLabel;
use dupe::Dupe;
use starlark::values::Value;
use starlark::values::ValueLike;

#[derive(Debug, thiserror::Error)]
enum ValueAsTargetLabelError {
    #[error("Expected a single target like item, but was `{0}`")]
    NotATarget(String),
}

pub trait ValueAsStarlarkTargetLabel {
    fn parse_target_platforms(
        self,
        target_alias_resolver: &BuckConfigTargetAliasResolver,
        cell_resolver: &CellResolver,
        cell_name: CellName,
        default_target_platform: &Option<TargetLabel>,
    ) -> anyhow::Result<Option<TargetLabel>>;
}

impl<'v> ValueAsStarlarkTargetLabel for Value<'v> {
    fn parse_target_platforms(
        self,
        target_alias_resolver: &BuckConfigTargetAliasResolver,
        cell_resolver: &CellResolver,
        cell_name: CellName,
        default_target_platform: &Option<TargetLabel>,
    ) -> anyhow::Result<Option<TargetLabel>> {
        let target_platform = if self.is_none() {
            default_target_platform.clone()
        } else if let Some(s) = self.unpack_str() {
            Some(
                ParsedPattern::<TargetPatternExtra>::parse_relaxed(
                    target_alias_resolver,
                    // TODO(nga): Parse relaxed relative to cell root is incorrect.
                    CellPathRef::new(cell_name, CellRelativePath::empty()),
                    s,
                    cell_resolver,
                )?
                .as_target_label(s)?,
            )
        } else if let Some(target) = self.downcast_ref::<StarlarkTargetLabel>() {
            Some(target.label().dupe())
        } else {
            return Err(anyhow::anyhow!(ValueAsTargetLabelError::NotATarget(
                self.to_repr()
            )));
        };

        Ok(target_platform)
    }
}


/*
 * 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 buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::unconfigured::TargetNode;
use starlark::values::Heap;
use starlark::values::Value;

use crate::bxl::starlark_defs::nodes::configured::StarlarkConfiguredTargetNode;
use crate::bxl::starlark_defs::nodes::unconfigured::StarlarkTargetNode;

/// BXL-specific node allocator.
pub trait AllocNode {
    fn alloc(self, heap: &Heap) -> Value;
}

impl AllocNode for TargetNode {
    fn alloc(self, heap: &Heap) -> Value {
        heap.alloc(StarlarkTargetNode(self))
    }
}

impl AllocNode for ConfiguredTargetNode {
    fn alloc(self, heap: &Heap) -> Value {
        heap.alloc(StarlarkConfiguredTargetNode(self))
    }
}


/*
 * 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_build_api::analysis::AnalysisResult;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use derive_more::Display;
use starlark::any::ProvidesStaticType;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::starlark_module;
use starlark::starlark_simple_value;
use starlark::values::starlark_value;
use starlark::values::FrozenValue;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::StarlarkDocs;

#[derive(
    ProvidesStaticType,
    Debug,
    Display,
    NoSerialize,
    StarlarkDocs,
    Allocative
)]
#[display(fmt = "{:?}", self)]
#[starlark_docs(directory = "bxl")]
pub struct StarlarkAnalysisResult {
    analysis: AnalysisResult,
    label: ConfiguredProvidersLabel,
}

impl StarlarkAnalysisResult {
    pub(crate) fn new(analysis: AnalysisResult, label: ConfiguredProvidersLabel) -> Self {
        Self { analysis, label }
    }
}

starlark_simple_value!(StarlarkAnalysisResult);

#[starlark_value(type = "analysis_result")]
impl<'v> StarlarkValue<'v> for StarlarkAnalysisResult {
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(starlark_analysis_result_methods)
    }
}

/// The result of running an analysis in bxl.
#[starlark_module]
fn starlark_analysis_result_methods(builder: &mut MethodsBuilder) {
    /// Access the providers of the rule. Returns a `[ProviderCollection]` the same as accessing
    /// providers of dependencies within a rule implementation.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_providers(ctx):
    ///     node = ctx.configured_targets("root//bin:the_binary")
    ///     providers = ctx.analysis(node).providers()
    ///     ctx.output.print(providers[FooInfo])
    ///     providers = ctx.analysis("//:bin").providers()
    ///     ctx.output.print(providers[FooInfo])
    /// ```
    fn providers<'v>(this: &'v StarlarkAnalysisResult) -> anyhow::Result<FrozenValue> {
        unsafe {
            // SAFETY:: this actually just returns a FrozenValue from in the StarlarkAnalysisResult
            // which is kept alive for 'v
            Ok(this
                .analysis
                .lookup_inner(&this.label)?
                .value()
                .to_frozen_value())
        }
    }
}