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

use std::iter::Peekable;

use anyhow::Context;
use buck2_build_api::bxl::types::BxlFunctionLabel;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::name::CellName;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_core::package::PackageLabel;
use buck2_core::target::label::TargetLabel;
use buck2_core::target::name::TargetNameRef;
use buck2_core::target::name::EQ_SIGN_SUBST;
use buck2_interpreter::path::BxlFilePath;
use dupe::Dupe;
use itertools::Itertools;

#[derive(Debug, thiserror::Error)]
enum BuckOutPathParserError {
    #[error(
        "Malformed buck-out path. Expected format: `buck-out/<isolation_prefix>/<gen|tmp|test|gen-anon|gen-bxl>/<cell_name>/<cfg_hash>/<target_path?>/__<target_name>__/<__action__id__?>/<outputs>`. Actual path was: `{0}`"
    )]
    MalformedOutputPath(String),
}

/// The common attributes of each `buck-out` path type,
pub(crate) struct BuckOutPathTypeCommon {
    /// Configuration hash within the `buck-out` path.
    pub(crate) config_hash: String,
    /// The path starting from cell to the artifact, without the configuration hash. For example, in
    /// `buck-out/v2/gen/cell/<CONFIG_HASH>/path/to/__target_name__/target`, it would be `cell/path/to/__target_name__/target`.
    pub(crate) raw_path_to_output: ForwardRelativePathBuf,
}

/// The types of the `buck-out` path.
pub(crate) enum BuckOutPathType {
    BxlOutput {
        // `BxlFunctionLabel` contains the `CellPath` to the bxl function.
        bxl_function_label: BxlFunctionLabel,
        common_attrs: BuckOutPathTypeCommon,
    },
    AnonOutput {
        path: CellPath,
        target_label: TargetLabel,
        // Rule attr hash is part of anonymous target buck-outs.
        attr_hash: String,
        common_attrs: BuckOutPathTypeCommon,
    },
    RuleOutput {
        path: CellPath,
        target_label: TargetLabel,
        // This is the part of the buck-out after target name. For example, it would `artifact` in  `gen/path/to/__target_name__/artifact`
        path_after_target_name: ForwardRelativePathBuf,
        common_attrs: BuckOutPathTypeCommon,
    },
    TestOutput {
        path: CellPath,
        common_attrs: BuckOutPathTypeCommon,
    },
    TmpOutput {
        path: CellPath,
        target_label: TargetLabel,
        common_attrs: BuckOutPathTypeCommon,
    },
}

pub(crate) struct BuckOutPathParser<'v> {
    cell_resolver: &'v CellResolver,
}

fn validate_buck_out_and_isolation_prefix<'v>(
    iter: &mut Peekable<impl Iterator<Item = &'v FileName>>,
) -> anyhow::Result<()> {
    // Validate path starts with buck-out.
    match iter.next() {
        Some(buck_out) => {
            if buck_out != "buck-out" {
                return Err(anyhow::anyhow!("Path does not start with buck-out"));
            }
        }
        None => return Err(anyhow::anyhow!("Path does not start with buck-out")),
    }

    // Advance the iterator to isolation dir.
    match iter.next() {
        Some(_) => Ok(()),
        None => Err(anyhow::anyhow!("Path does not have an isolation dir")),
    }
}

struct BuckOutPathData {
    // Cell path of the target label that created the artifact.
    cell_path: CellPath,
    config_hash: String,
    anon_hash: Option<String>,
    /// The path starting from cell to the artifact, without the configuration hash. For example, in
    /// `buck-out/v2/gen/cell/<CONFIG_HASH>/path/to/__target_name__/target`, it would be `cell/path/to/__target_name__/target`.
    raw_path_to_output: ForwardRelativePathBuf,
}

fn get_cell_path<'v>(
    iter: &mut Peekable<impl Iterator<Item = &'v FileName> + Clone>,
    cell_resolver: &'v CellResolver,
    generated_prefix: &'v str,
) -> anyhow::Result<BuckOutPathData> {
    let is_anon = generated_prefix == "gen-anon";
    let is_test = generated_prefix == "test";
    // Get cell name and validate it exists
    match iter.next() {
        Some(cell_name) => {
            let cell_name = CellName::unchecked_new(cell_name.as_str())?;
            let mut raw_path_to_output = ForwardRelativePath::new(cell_name.as_str())?.to_buf();

            cell_resolver.get(cell_name)?;

            // Advance iterator to the config hash
            let config_hash = match iter.next() {
                Some(config_hash) => config_hash,
                None => {
                    return Err(anyhow::anyhow!(
                        "Path does not have a platform configuration"
                    ));
                }
            };

            iter.clone().for_each(|f| {
                raw_path_to_output.push(f);
            });

            // Get cell relative path and construct the cell path
            let mut cell_relative_path = CellRelativePath::unchecked_new("").to_owned();

            while let Some(part) = iter.next() {
                cell_relative_path = cell_relative_path.join(part).to_owned();

                // We make sure not to consume the target name part via the iterator.
                match iter.peek() {
                    Some(maybe_target_name) => {
                        let maybe_target_name = maybe_target_name.as_str();
                        // TODO(@wendyy) We assume that the first string that we find that starts with "__"
                        // is the target name. There is a small risk of naming collisions (ex: if there's a directory
                        // name that follows this convention that contains a build file), but I will fix this at a
                        // later date.
                        if (*maybe_target_name).starts_with("__") {
                            // If it's an anonymous target, then the last part before the target name is actually the
                            // hash, and not part of the cell relative path.
                            let cell_path = if is_anon {
                                CellPath::new(
                                    cell_name,
                                    cell_relative_path
                                        .parent()
                                        .with_context(|| "Invalid path for anonymous target")?
                                        .to_buf(),
                                )
                            } else {
                                CellPath::new(cell_name, cell_relative_path.to_buf())
                            };

                            let anon_hash = if is_anon {
                                // Iterator is pointing to the part right before the target name, aka the attr
                                // hash for the anonymous target.
                                Some(part.to_string())
                            } else {
                                None
                            };

                            let buck_out_path_data = BuckOutPathData {
                                cell_path,
                                config_hash: config_hash.to_string(),
                                anon_hash,
                                raw_path_to_output: raw_path_to_output.to_buf(),
                            };

                            return Ok(buck_out_path_data);
                        }
                    }
                    None => (),
                }
            }

            if is_test {
                let buck_out_path_data = BuckOutPathData {
                    cell_path: CellPath::new(cell_name, cell_relative_path.to_buf()),
                    config_hash: config_hash.to_string(),
                    anon_hash: None,
                    raw_path_to_output: raw_path_to_output.to_buf(),
                };
                Ok(buck_out_path_data)
            } else {
                Err(anyhow::anyhow!("Invalid target name"))
            }
        }
        None => Err(anyhow::anyhow!("Invalid cell name")),
    }
}

fn get_target_name<'v>(
    iter: &mut Peekable<impl Iterator<Item = &'v FileName>>,
) -> anyhow::Result<String> {
    // Get target name, which is prefixed and suffixed with "__"
    match iter.next() {
        Some(raw_target_name) => {
            let mut target_name_with_underscores =
                <&ForwardRelativePath>::from(raw_target_name).to_owned();

            while !target_name_with_underscores.as_str().ends_with("__") {
                match iter.next() {
                    Some(next) => {
                        target_name_with_underscores = target_name_with_underscores.join(next);
                    }
                    None => return Err(anyhow::anyhow!("Invalid target name")),
                }
            }

            let target_name_with_underscores = target_name_with_underscores.as_str();
            let target_name =
                &target_name_with_underscores[2..(target_name_with_underscores.len() - 2)];
            Ok(target_name.replace(EQ_SIGN_SUBST, "="))
        }
        None => Err(anyhow::anyhow!("Invalid target name")),
    }
}

fn get_target_label<'v>(
    iter: &mut Peekable<impl Iterator<Item = &'v FileName>>,
    path: CellPath,
) -> anyhow::Result<TargetLabel> {
    let target_name = get_target_name(iter)?;
    let package = PackageLabel::from_cell_path(path.as_ref());
    let target = TargetNameRef::new(target_name.as_str())?;
    let target_label = TargetLabel::new(package.dupe(), target);
    Ok(target_label)
}

fn get_bxl_function_label<'v>(
    iter: &mut Peekable<impl Iterator<Item = &'v FileName>>,
    path: CellPath,
) -> anyhow::Result<BxlFunctionLabel> {
    let target_name = get_target_name(iter)?;
    let bxl_path = BxlFilePath::new(path)?;
    let bxl_function_label = BxlFunctionLabel {
        bxl_path,
        name: target_name,
    };

    Ok(bxl_function_label)
}

impl<'v> BuckOutPathParser<'v> {
    pub(crate) fn new(cell_resolver: &'v CellResolver) -> BuckOutPathParser {
        BuckOutPathParser { cell_resolver }
    }

    // Validates and parses the buck-out path, returning the `BuckOutPathType`. Assumes
    // that the inputted path is not a symlink.
    pub(crate) fn parse(&self, output_path: &str) -> anyhow::Result<BuckOutPathType> {
        self.parse_inner(output_path)
            .with_context(|| BuckOutPathParserError::MalformedOutputPath(output_path.to_owned()))
    }

    fn parse_inner(&self, output_path: &str) -> anyhow::Result<BuckOutPathType> {
        let path_as_forward_rel_path = ForwardRelativePathBuf::new(output_path.to_owned())?;
        let mut iter = path_as_forward_rel_path.iter().peekable();

        validate_buck_out_and_isolation_prefix(&mut iter)?;

        // Advance the iterator to the prefix (tmp, test, gen, gen-anon, or gen-bxl)
        match iter.next() {
            Some(part) => {
                let result = match part.as_str() {
                    "tmp" => {
                        let buck_out_path_data =
                            get_cell_path(&mut iter, self.cell_resolver, "tmp")?;
                        let target_label =
                            get_target_label(&mut iter, buck_out_path_data.cell_path.clone())?;

                        let common_attrs = BuckOutPathTypeCommon {
                            config_hash: buck_out_path_data.config_hash,
                            raw_path_to_output: buck_out_path_data.raw_path_to_output,
                        };

                        Ok(BuckOutPathType::TmpOutput {
                            path: buck_out_path_data.cell_path,
                            target_label,
                            common_attrs,
                        })
                    }
                    "test" => {
                        let buck_out_path_data =
                            get_cell_path(&mut iter, self.cell_resolver, "test")?;

                        let common_attrs = BuckOutPathTypeCommon {
                            config_hash: buck_out_path_data.config_hash,
                            raw_path_to_output: buck_out_path_data.raw_path_to_output,
                        };

                        Ok(BuckOutPathType::TestOutput {
                            path: buck_out_path_data.cell_path,
                            common_attrs,
                        })
                    }
                    "gen" => {
                        let buck_out_path_data =
                            get_cell_path(&mut iter, self.cell_resolver, "gen")?;
                        let target_label =
                            get_target_label(&mut iter, buck_out_path_data.cell_path.clone())?;
                        let path_after_target_name =
                            ForwardRelativePathBuf::new(iter.clone().join("/"))?;
                        let common_attrs = BuckOutPathTypeCommon {
                            config_hash: buck_out_path_data.config_hash,
                            raw_path_to_output: buck_out_path_data.raw_path_to_output,
                        };

                        Ok(BuckOutPathType::RuleOutput {
                            path: buck_out_path_data.cell_path,
                            target_label,
                            path_after_target_name,
                            common_attrs,
                        })
                    }
                    "gen-anon" => {
                        let buck_out_path_data =
                            get_cell_path(&mut iter, self.cell_resolver, "gen-anon")?;
                        let target_label =
                            get_target_label(&mut iter, buck_out_path_data.cell_path.clone())?;
                        let common_attrs = BuckOutPathTypeCommon {
                            config_hash: buck_out_path_data.config_hash,
                            raw_path_to_output: buck_out_path_data.raw_path_to_output,
                        };

                        Ok(BuckOutPathType::AnonOutput {
                            path: buck_out_path_data.cell_path,
                            target_label,
                            attr_hash: buck_out_path_data
                                .anon_hash
                                .expect("No hash found in anonymous artifact buck-out"),
                            common_attrs,
                        })
                    }
                    "gen-bxl" => {
                        let buck_out_path_data =
                            get_cell_path(&mut iter, self.cell_resolver, "gen-bxl")?;
                        let bxl_function_label =
                            get_bxl_function_label(&mut iter, buck_out_path_data.cell_path)?;
                        let common_attrs = BuckOutPathTypeCommon {
                            config_hash: buck_out_path_data.config_hash,
                            raw_path_to_output: buck_out_path_data.raw_path_to_output,
                        };

                        Ok(BuckOutPathType::BxlOutput {
                            bxl_function_label,
                            common_attrs,
                        })
                    }
                    _ => Err(anyhow::anyhow!(
                        "Directory after isolation dir is invalid (should be gen, gen-bxl, gen-anon, tmp, or test)"
                    )),
                };

                // Validate for non-test outputs that the target name is not the last element in the path
                if part != "test" && iter.peek().is_none() {
                    Err(anyhow::anyhow!("No output artifacts found"))
                } else {
                    result
                }
            }
            None => Err(anyhow::anyhow!("Path is empty")),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use buck2_build_api::bxl::types::BxlFunctionLabel;
    use buck2_core::cells::cell_path::CellPath;
    use buck2_core::cells::cell_root_path::CellRootPath;
    use buck2_core::cells::name::CellName;
    use buck2_core::cells::paths::CellRelativePath;
    use buck2_core::cells::CellResolver;
    use buck2_core::configuration::data::ConfigurationData;
    use buck2_core::configuration::data::ConfigurationDataData;
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
    use buck2_core::fs::project_rel_path::ProjectRelativePath;
    use buck2_core::package::PackageLabel;
    use buck2_core::target::label::TargetLabel;
    use buck2_core::target::name::TargetNameRef;
    use buck2_interpreter::path::BxlFilePath;

    use crate::output::buck_out_path_parser::BuckOutPathParser;
    use crate::output::buck_out_path_parser::BuckOutPathType;

    fn get_parse_test_cell_resolver() -> anyhow::Result<CellResolver> {
        let cell_path = CellRootPath::new(ProjectRelativePath::new("foo/bar")?);

        let cell_resolver = CellResolver::testing_with_name_and_path(
            CellName::testing_new("bar"),
            cell_path.to_buf(),
        );

        Ok(cell_resolver)
    }

    #[test]
    fn test_buck_path_parser_validation() -> anyhow::Result<()> {
        let configuration = ConfigurationData::from_platform(
            "cfg_for//:testing_exec".to_owned(),
            ConfigurationDataData {
                constraints: BTreeMap::new(),
            },
        )
        .unwrap();
        let cell_resolver = get_parse_test_cell_resolver()?;
        let buck_out_parser = BuckOutPathParser::new(&cell_resolver);

        let malformed_path1 = "does/not/start/with/buck-out/blah/blah";
        let malformed_path2 = "buck-out/v2/invalid_buck_prefix/blah/blah/blah/blah";
        let malformed_path3 = "buck-out/v2/gen/bar/no/target/name/found";
        let malformed_path4 = "buck-out/v2/gen/bar/path/to/target/__but_no_artifacts__";

        let res = buck_out_parser.parse(malformed_path1);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        let res = buck_out_parser.parse(malformed_path2);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        let res = buck_out_parser.parse(malformed_path3);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        let res = buck_out_parser.parse(malformed_path4);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        let cell_does_not_exist =
            "buck-out/v2/gen/nonexistent_cell/cfg_hash/path/to/target/__target_name__/output";

        let res = buck_out_parser.parse(cell_does_not_exist);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        let config_hash = configuration.output_hash();
        let no_artifacts_after_target_name = &format!(
            "buck-out/v2/gen/bar/{}/path/to/target/__target_name__",
            config_hash
        );
        let res = buck_out_parser.parse(no_artifacts_after_target_name);
        assert!(res.err().unwrap().to_string().contains("Malformed"));

        Ok(())
    }

    #[test]
    fn test_buck_path_parser() -> anyhow::Result<()> {
        let configuration = ConfigurationData::from_platform(
            "cfg_for//:testing_exec".to_owned(),
            ConfigurationDataData {
                constraints: BTreeMap::new(),
            },
        )
        .unwrap();
        let cell_resolver = get_parse_test_cell_resolver()?;
        let buck_out_parser = BuckOutPathParser::new(&cell_resolver);

        let pkg = PackageLabel::new(
            CellName::testing_new("bar"),
            CellRelativePath::unchecked_new("path/to/target"),
        );

        let expected_target_label =
            TargetLabel::new(pkg.clone(), TargetNameRef::new("target_name")?);

        let expected_cell_path = CellPath::new(
            CellName::testing_new("bar"),
            CellRelativePath::unchecked_new("path/to/target").to_owned(),
        );

        let expected_config_hash = configuration.output_hash();

        let rule_path = format!(
            "buck-out/v2/gen/bar/{}/path/to/target/__target_name__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&rule_path)?;

        match res {
            BuckOutPathType::RuleOutput {
                path,
                target_label,
                path_after_target_name,
                common_attrs,
            } => {
                assert_eq!(
                    path_after_target_name,
                    ForwardRelativePathBuf::new("output".to_owned())?,
                );
                assert_eq!(target_label, expected_target_label,);
                assert_eq!(path, expected_cell_path,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/__target_name__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let rule_path_target_label_with_slashes = format!(
            "buck-out/v2/gen/bar/{}/path/to/target/__target_name_start/target_name_end__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&rule_path_target_label_with_slashes)?;

        let expected_target_label_with_slashes = TargetLabel::new(
            pkg.clone(),
            TargetNameRef::new("target_name_start/target_name_end")?,
        );

        match res {
            BuckOutPathType::RuleOutput {
                path,
                target_label,
                path_after_target_name,
                common_attrs,
            } => {
                assert_eq!(
                    path_after_target_name,
                    ForwardRelativePathBuf::new("output".to_owned())?,
                );
                assert_eq!(target_label, expected_target_label_with_slashes,);
                assert_eq!(path, expected_cell_path,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/__target_name_start/target_name_end__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let rule_path_with_equal_sign = format!(
            "buck-out/v2/gen/bar/{}/path/to/target/__target_name_eqsb_out__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&rule_path_with_equal_sign)?;

        let expected_target_label_with_equal_sign =
            TargetLabel::new(pkg, TargetNameRef::new("target_name=out")?);

        match res {
            BuckOutPathType::RuleOutput {
                path,
                target_label,
                path_after_target_name,
                common_attrs,
            } => {
                assert_eq!(
                    path_after_target_name,
                    ForwardRelativePathBuf::new("output".to_owned())?,
                );
                assert_eq!(target_label, expected_target_label_with_equal_sign,);
                assert_eq!(path, expected_cell_path,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/__target_name_eqsb_out__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let tmp_path = format!(
            "buck-out/v2/tmp/bar/{}/path/to/target/__target_name__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&tmp_path)?;

        match res {
            BuckOutPathType::TmpOutput {
                path,
                target_label,
                common_attrs,
            } => {
                assert_eq!(path, expected_cell_path,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(target_label, expected_target_label,);
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/__target_name__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let test_path = format!(
            "buck-out/v2/test/bar/{}/path/to/target/test/output",
            expected_config_hash
        );

        let expected_test_cell_path = CellPath::new(
            CellName::testing_new("bar"),
            CellRelativePath::unchecked_new("path/to/target/test/output").to_owned(),
        );

        let res = buck_out_parser.parse(&test_path)?;

        match res {
            BuckOutPathType::TestOutput { path, common_attrs } => {
                assert_eq!(path, expected_test_cell_path,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/test/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let anon_path = format!(
            "buck-out/v2/gen-anon/bar/{}/path/to/target/anon_hash/__target_name__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&anon_path)?;

        match res {
            BuckOutPathType::AnonOutput {
                path,
                target_label,
                attr_hash,
                common_attrs,
            } => {
                assert_eq!(target_label, expected_target_label,);
                assert_eq!(path, expected_cell_path,);
                assert_eq!(attr_hash, "anon_hash",);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/target/anon_hash/__target_name__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        let path = format!(
            "buck-out/v2/gen-bxl/bar/{}/path/to/function.bxl/__function_name__/output",
            expected_config_hash
        );

        let res = buck_out_parser.parse(&path)?;

        match res {
            BuckOutPathType::BxlOutput {
                bxl_function_label,
                common_attrs,
            } => {
                let path = CellPath::new(
                    CellName::testing_new("bar"),
                    CellRelativePath::unchecked_new("path/to/function.bxl").to_owned(),
                );

                let bxl_path = BxlFilePath::new(path)?;
                let expected_bxl_function_label = BxlFunctionLabel {
                    bxl_path,
                    name: "function_name".to_owned(),
                };

                assert_eq!(bxl_function_label, expected_bxl_function_label,);
                assert_eq!(common_attrs.config_hash, expected_config_hash.as_str());
                assert_eq!(
                    common_attrs.raw_path_to_output.as_str(),
                    "bar/path/to/function.bxl/__function_name__/output"
                )
            }
            _ => panic!("Should have parsed buck-out path successfully"),
        }

        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::io::Write;

use indexmap::IndexMap;
use regex::RegexSet;

use super::buck_out_path_parser::BuckOutPathType;

pub(crate) struct BuckOutPathTypePrinter {
    json: bool,
    attributes: Option<RegexSet>,
}

impl BuckOutPathTypePrinter {
    pub(crate) fn new(json: bool, attributes: &Vec<String>) -> anyhow::Result<Self> {
        let attributes = if attributes.is_empty() {
            None
        } else {
            Some(RegexSet::new(attributes)?)
        };

        Ok(BuckOutPathTypePrinter { json, attributes })
    }

    pub(crate) fn print(
        &self,
        path_type: &BuckOutPathType,
        mut stdout: impl Write,
    ) -> anyhow::Result<()> {
        if self.json {
            writeln!(
                &mut stdout,
                "{}",
                serde_json::to_string_pretty(&self.printable_attributes(path_type))?
            )?;
        } else {
            self.printable_attributes(path_type)
                .values()
                .try_for_each(|a| writeln!(&mut stdout, "{}", a))?;
        }
        Ok(())
    }

    fn printable_attributes(&self, path_type: &BuckOutPathType) -> IndexMap<String, String> {
        let all_attributes = self.all_attributes(path_type);

        if let Some(attributes) = &self.attributes {
            all_attributes
                .into_iter()
                .filter(|(k, _)| attributes.is_match(k))
                .collect()
        } else {
            all_attributes
        }
    }

    fn all_attributes(&self, path_type: &BuckOutPathType) -> IndexMap<String, String> {
        // Deterministic order
        let mut attributes = IndexMap::new();

        match path_type {
            BuckOutPathType::BxlOutput {
                bxl_function_label,
                common_attrs,
            } => {
                attributes.insert(
                    "bxl_function_label".to_owned(),
                    bxl_function_label.to_string(),
                );
                attributes.insert("config_hash".to_owned(), common_attrs.config_hash.clone());
                attributes.insert(
                    "full_artifact_path_no_hash".to_owned(),
                    common_attrs.raw_path_to_output.to_string(),
                );
            }
            BuckOutPathType::AnonOutput {
                path,
                target_label,
                attr_hash,
                common_attrs,
            } => {
                attributes.insert("cell_path".to_owned(), path.to_string());
                attributes.insert("target_label".to_owned(), target_label.to_string());
                attributes.insert("attr_hash".to_owned(), attr_hash.clone());
                attributes.insert("config_hash".to_owned(), common_attrs.config_hash.clone());
                attributes.insert(
                    "full_artifact_path_no_hash".to_owned(),
                    common_attrs.raw_path_to_output.to_string(),
                );
            }
            BuckOutPathType::RuleOutput {
                path,
                target_label,
                path_after_target_name,
                common_attrs,
            } => {
                attributes.insert("cell_path".to_owned(), path.to_string());
                attributes.insert("target_label".to_owned(), target_label.to_string());
                attributes.insert(
                    "short_artifact_path".to_owned(),
                    path_after_target_name.to_string(),
                );
                attributes.insert("config_hash".to_owned(), common_attrs.config_hash.clone());
                attributes.insert(
                    "full_artifact_path_no_hash".to_owned(),
                    common_attrs.raw_path_to_output.to_string(),
                );
            }
            BuckOutPathType::TestOutput { path, common_attrs } => {
                attributes.insert("cell_path".to_owned(), path.to_string());
                attributes.insert("config_hash".to_owned(), common_attrs.config_hash.clone());
                attributes.insert(
                    "full_artifact_path_no_hash".to_owned(),
                    common_attrs.raw_path_to_output.to_string(),
                );
            }
            BuckOutPathType::TmpOutput {
                path,
                target_label,
                common_attrs,
            } => {
                attributes.insert("cell_path".to_owned(), path.to_string());
                attributes.insert("target_label".to_owned(), target_label.to_string());
                attributes.insert("config_hash".to_owned(), common_attrs.config_hash.clone());
                attributes.insert(
                    "full_artifact_path_no_hash".to_owned(),
                    common_attrs.raw_path_to_output.to_string(),
                );
            }
        }

        attributes
    }
}


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

use async_trait::async_trait;
use buck2_audit::output::command::AuditOutputCommand;
use buck2_build_api::actions::query::FIND_MATCHING_ACTION;
use buck2_build_api::actions::query::PRINT_ACTION_NODE;
use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
use buck2_build_api::audit_output::AuditOutputResult;
use buck2_build_api::audit_output::AUDIT_OUTPUT;
use buck2_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::cells::CellResolver;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use dice::DiceComputations;
use thiserror::Error;

use crate::output::buck_out_path_parser::BuckOutPathParser;
use crate::output::buck_out_path_parser::BuckOutPathType;
use crate::AuditSubcommand;

#[derive(Debug, Error)]
pub(crate) enum AuditOutputError {
    #[error(
        "BXL, anonymous target, test, and tmp artifacts are not supported for audit output. Only rule output artifacts are supported. Path: `{0}`"
    )]
    UnsupportedPathType(String),
}

async fn audit_output<'v>(
    output_path: &'v str,
    working_dir: &'v ProjectRelativePath,
    cell_resolver: &'v CellResolver,
    dice_ctx: &'v DiceComputations,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<Option<AuditOutputResult>> {
    let buck_out_parser = BuckOutPathParser::new(cell_resolver);
    let parsed = buck_out_parser.parse(output_path)?;

    let (target_label, config_hash, path_after_target_name) = match parsed {
        BuckOutPathType::RuleOutput {
            target_label,
            common_attrs,
            path_after_target_name,
            ..
        } => (
            target_label,
            common_attrs.config_hash,
            path_after_target_name,
        ),
        _ => {
            return Err(anyhow::anyhow!(AuditOutputError::UnsupportedPathType(
                output_path.to_owned()
            )));
        }
    };

    let configured_target_label = dice_ctx
        .get_configured_target(&target_label, global_target_platform.as_ref())
        .await?;

    let command_config = configured_target_label.cfg();
    let command_config_hash = command_config.output_hash();
    if command_config_hash.as_str() != config_hash {
        return Ok(Some(AuditOutputResult::MaybeRelevant(target_label)));
    }

    let analysis = dice_ctx
        .get_analysis_result(&configured_target_label)
        .await?
        .require_compatible()?;

    Ok(FIND_MATCHING_ACTION.get()?(
        dice_ctx,
        working_dir,
        global_target_platform,
        &analysis,
        path_after_target_name,
    )
    .await?
    .map(AuditOutputResult::Match))
}

pub(crate) fn init_audit_output() {
    AUDIT_OUTPUT.init(
        |output_path, working_dir, cell_resolver, dice_ctx, global_target_platform| {
            Box::pin(audit_output(
                output_path,
                working_dir,
                cell_resolver,
                dice_ctx,
                global_target_platform,
            ))
        },
    );
}

#[async_trait]
impl AuditSubcommand for AuditOutputCommand {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        client_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        server_ctx
            .with_dice_ctx(async move |server_ctx, dice_ctx| {
                // First, we parse the buck-out path to get a target label. Next, we configure the target
                // label and run analysis on it to get the `DeferredTable`. Then, we iterate through the
                // deferred table's entries and look at their build outputs (if they have any) to try to
                // match the inputted buck-out path with the build output's buck-out path. Once we find
                // a matching path, we create the `ActionQueryNode` from the action key associated with the
                // matching build output, and print out the result.

                let working_dir = server_ctx.working_dir();
                let cell_resolver = dice_ctx.get_cell_resolver().await?;

                let global_target_platform = target_platform_from_client_context(
                    &client_ctx,
                    server_ctx,
                    &dice_ctx,
                )
                .await?;

                let result = audit_output(&self.output_path, working_dir, &cell_resolver, &dice_ctx, global_target_platform).await?;

                let mut stdout = stdout.as_writer();

                match result {
                    Some(result) => {
                        match result {
                            AuditOutputResult::Match(action) => {
                                (PRINT_ACTION_NODE.get()?)(&mut stdout, action, self.json, &self.query_attributes.get()?, &cell_resolver).await?
                            },
                            AuditOutputResult::MaybeRelevant(label) => {
                                writeln!(
                                    stdout,
                                    "Platform configuration of the buck-out path did not match the one used to invoke this command. Returning the most relevant unconfigured target label for the buck-out path: {}",
                                    label
                                )?;
                            }
                        }
                    },
                    None => {
                        // If we get here, that means we failed to find any matching actions
                        writeln!(
                            stdout,
                            "Failed to find an action that produced the output path. Make sure that you did not input a symlinked buck-out path."
                        )?;
                    }
                }

                Ok(())
            })
            .await
    }
}


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

pub(crate) mod buck_out_path_parser;
pub(crate) mod buck_out_path_type_printer;
pub mod command;
pub mod parse;


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

use async_trait::async_trait;
use buck2_audit::output::parse::AuditParseCommand;
use buck2_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

use super::buck_out_path_parser::BuckOutPathParser;
use super::buck_out_path_type_printer::BuckOutPathTypePrinter;
use crate::AuditSubcommand;

#[async_trait]
impl AuditSubcommand for AuditParseCommand {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        _client_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        server_ctx
            .with_dice_ctx(async move |_server_ctx, dice_ctx| {
                let cell_resolver = dice_ctx.get_cell_resolver().await?;
                let buck_out_parser = BuckOutPathParser::new(&cell_resolver);
                let parsed_path = buck_out_parser.parse(&self.output_path)?;

                let printer = BuckOutPathTypePrinter::new(self.json, &self.output_attribute)?;

                let stdout = stdout.as_writer();

                printer.print(&parsed_path, stdout)
            })
            .await
    }
}


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

//! Starlark debugging.

mod module;
mod package_deps;

use async_trait::async_trait;
use buck2_audit::starlark::StarlarkCommand;
use buck2_cli_proto::ClientContext;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

use crate::AuditSubcommand;

#[async_trait]
impl AuditSubcommand for StarlarkCommand {
    async fn server_execute(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        client_ctx: ClientContext,
    ) -> anyhow::Result<()> {
        match self {
            StarlarkCommand::Module(cmd) => {
                module::server_execute(cmd, server_ctx, stdout, client_ctx).await
            }
            StarlarkCommand::PackageDeps(cmd) => {
                package_deps::server_execute(cmd, server_ctx, stdout, client_ctx).await
            }
        }
    }
}


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

use std::io::Write;

use buck2_audit::starlark::module::StarlarkModuleCommand;
use buck2_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
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::StarlarkModulePath;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

pub(crate) async fn server_execute(
    command: &StarlarkModuleCommand,
    server_ctx: &dyn ServerCommandContextTrait,
    mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    _client_ctx: ClientContext,
) -> anyhow::Result<()> {
    server_ctx
        .with_dice_ctx(async move |server_ctx, dice_ctx| {
            let cell_resolver = dice_ctx.get_cell_resolver().await?;
            let current_cell_path = cell_resolver.get_cell_path(server_ctx.working_dir())?;
            let current_cell = BuildFileCell::new(current_cell_path.cell());

            let cell_alias_resolver = cell_resolver
                .get(current_cell_path.cell())?
                .cell_alias_resolver();

            let path = parse_import_with_config(
                cell_alias_resolver,
                &current_cell_path,
                &command.import_path,
                &ParseImportOptions {
                    allow_relative_imports: true,
                    // Otherwise `@arg` is expanded as mode file.
                    allow_missing_at_symbol: true,
                },
            )?;

            let import_path = ImportPath::new_with_build_file_cells(path, current_cell)?;

            let loaded_module = dice_ctx
                .get_loaded_module(StarlarkModulePath::LoadFile(&import_path))
                .await?;

            let mut stdout = stdout.as_writer();
            writeln!(stdout, "{}", loaded_module.path())?;
            writeln!(stdout)?;
            writeln!(stdout, "Imports:")?;
            for import in loaded_module.imports() {
                writeln!(stdout, "  {}", import)?;
            }
            writeln!(stdout)?;
            write!(stdout, "{}", loaded_module.env().dump_debug())?;
            Ok(())
        })
        .await
}


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

use std::collections::HashSet;
use std::io::Write;

use buck2_audit::starlark::package_deps::StarlarkPackageDepsCommand;
use buck2_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::pattern::parse_package::parse_package;
use buck2_interpreter::file_loader::LoadedModule;
use buck2_interpreter::load_module::INTERPRETER_CALCULATION_IMPL;
use buck2_interpreter::path::StarlarkModulePath;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;

pub(crate) async fn server_execute(
    command: &StarlarkPackageDepsCommand,
    server_ctx: &dyn ServerCommandContextTrait,
    mut stdout: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    _client_ctx: ClientContext,
) -> anyhow::Result<()> {
    server_ctx
        .with_dice_ctx(async move |server_ctx, dice_ctx| {
            let cell_resolver = dice_ctx.get_cell_resolver().await?;
            let current_cell_path = cell_resolver.get_cell_path(server_ctx.working_dir())?;
            let current_cell = BuildFileCell::new(current_cell_path.cell());

            let cell_alias_resolver = cell_resolver
                .get(current_cell_path.cell())?
                .cell_alias_resolver();

            let package = parse_package(&command.package, cell_alias_resolver)?;

            let module_deps = INTERPRETER_CALCULATION_IMPL
                .get()?
                .get_module_deps(&dice_ctx, package, current_cell)
                .await?;

            let mut stdout = stdout.as_writer();

            struct Printer {
                first: bool,
                visited: HashSet<ImportPath>,
            }

            impl Printer {
                fn print_module_and_deps(
                    &mut self,
                    module: &LoadedModule,
                    stdout: &mut dyn Write,
                ) -> anyhow::Result<()> {
                    let path = match module.path() {
                        StarlarkModulePath::LoadFile(path) => path,
                        StarlarkModulePath::BxlFile(_) => {
                            return Err(anyhow::anyhow!("bxl be here"));
                        }
                    };

                    if !self.visited.insert(path.clone()) {
                        return Ok(());
                    }

                    for import in module.loaded_modules().map.values() {
                        self.print_module_and_deps(import, stdout)?;
                    }

                    if !self.first {
                        writeln!(stdout)?;
                        writeln!(stdout)?;
                    }
                    self.first = false;

                    writeln!(stdout, "# {}", path)?;
                    writeln!(stdout)?;
                    write!(stdout, "{}", module.env().dump_debug())?;

                    Ok(())
                }
            }

            let mut printer = Printer {
                first: true,
                visited: HashSet::new(),
            };

            for module in module_deps.0.into_iter() {
                printer.print_module_and_deps(&module, &mut stdout)?;
            }

            Ok(())
        })
        .await
}


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

use std::future::Future;
use std::io::Write;
use std::pin::Pin;

use buck2_core::category::Category;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_util::late_binding::LateBinding;
use dice::DiceTransaction;

/// Implementation of `audit dep-files`.
pub static AUDIT_DEP_FILES: LateBinding<
    for<'a> fn(
        ctx: &'a DiceTransaction,
        ConfiguredTargetLabel,
        Category,
        Option<String>,
        &'a mut (dyn Write + Send),
    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'a>>,
> = LateBinding::new("AUDIT_DEP_FILES");


/*
 * 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::future::Future;
use std::pin::Pin;

use buck2_core::cells::CellResolver;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use buck2_util::late_binding::LateBinding;
use dice::DiceComputations;

use crate::actions::query::ActionQueryNode;

/// The result of audit output.
pub enum AuditOutputResult {
    /// The exact action that matched the buck-out path.
    Match(ActionQueryNode),
    /// If the platform configuration of the buck-out path doesn't match the platform used when calling
    /// audit output, then we return the unconfigured target label.
    MaybeRelevant(TargetLabel),
}

pub static AUDIT_OUTPUT: LateBinding<
    for<'v> fn(
        &'v str,
        &'v ProjectRelativePath,
        &'v CellResolver,
        &'v DiceComputations,
        global_target_platform: Option<TargetLabel>,
    )
        -> Pin<Box<dyn Future<Output = anyhow::Result<Option<AuditOutputResult>>> + 'v>>,
> = LateBinding::new("AUDIT_OUTPUT");

pub async fn audit_output<'v>(
    output_path: &'v str,
    working_dir: &'v ProjectRelativePath,
    cell_resolver: &'v CellResolver,
    dice_ctx: &'v DiceComputations,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<Option<AuditOutputResult>> {
    (AUDIT_OUTPUT.get()?)(
        output_path,
        working_dir,
        cell_resolver,
        dice_ctx,
        global_target_platform,
    )
    .await
}


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

use std::sync::Arc;

use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_signals::DeferredBuildSignals;
use buck2_build_signals::NodeDuration;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_events::span::SpanId;
use buck2_util::late_binding::LateBinding;
use dice::ActivationTracker;
use dice::UserComputationData;
use dupe::Dupe;

use crate::artifact_groups::ArtifactGroup;

pub static CREATE_BUILD_SIGNALS: LateBinding<
    fn() -> (BuildSignalsInstaller, Box<dyn DeferredBuildSignals>),
> = LateBinding::new("CREATE_BUILD_SIGNALS");

pub fn create_build_signals() -> (BuildSignalsInstaller, Box<dyn DeferredBuildSignals>) {
    (CREATE_BUILD_SIGNALS.get().unwrap())()
}

/// Everything we need to setup build signals when starting a command.
#[derive(Clone, Dupe)]
pub struct BuildSignalsInstaller {
    pub build_signals: Arc<dyn BuildSignals>,
    pub activation_tracker: Arc<dyn ActivationTracker>,
}

/// Send notifications to the build signals backend.
pub trait BuildSignals: Send + Sync + 'static {
    fn top_level_target(&self, label: ConfiguredTargetLabel, artifacts: Vec<ArtifactGroup>);

    fn final_materialization(
        &self,
        artifact: BuildArtifact,
        duration: NodeDuration,
        span_id: Option<SpanId>,
    );
}

pub trait SetBuildSignals {
    fn set_build_signals(&mut self, sender: Arc<dyn BuildSignals>);
}

impl SetBuildSignals for UserComputationData {
    fn set_build_signals(&mut self, sender: Arc<dyn BuildSignals>) {
        self.data.set(sender);
    }
}

pub trait HasBuildSignals {
    fn get_build_signals(&self) -> Option<&dyn BuildSignals>;
}

impl HasBuildSignals for UserComputationData {
    fn get_build_signals(&self) -> Option<&dyn BuildSignals> {
        self.data
            .get::<Arc<dyn BuildSignals>>()
            .ok()
            .map(|build_signals| build_signals.as_ref())
    }
}


/*
 * 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_common::dice::cells::SetCellResolver;
use buck2_common::dice::data::SetIoProvider;
use buck2_common::io::IoProvider;
use buck2_common::legacy_configs::dice::SetLegacyConfigs;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::digest_config::SetDigestConfig;
use dice::DetectCycles;
use dice::Dice;
use dice::WhichDice;
use dice::WhichSpawner;

/// Utility to configure the dice globals.
/// One place to not forget to initialize something in all places.
pub async fn configure_dice_for_buck(
    io: Arc<dyn IoProvider>,
    digest_config: DigestConfig,
    root_config: Option<&LegacyBuckConfig>,
    detect_cycles: Option<DetectCycles>,
    which_dice: Option<WhichDice>,
) -> anyhow::Result<Arc<Dice>> {
    let detect_cycles = detect_cycles.map_or_else(
        || {
            root_config
                .and_then(|c| {
                    c.parse::<DetectCycles>("buck2", "detect_cycles")
                        .transpose()
                })
                .unwrap_or(Ok(DetectCycles::Enabled))
        },
        Ok,
    )?;

    let which_dice = which_dice.map_or_else(
        || {
            root_config
                .and_then(|c| c.parse::<WhichDice>("buck2", "dice").transpose())
                .unwrap_or(Ok(WhichDice::Legacy))
        },
        Ok,
    )?;

    let which_spawner = root_config
        .and_then(|c| c.parse::<WhichSpawner>("buck2", "dice_spawner").transpose())
        .unwrap_or(Ok(WhichSpawner::DropCancel))?;

    let mut dice = match which_dice {
        WhichDice::Legacy => Dice::builder(),
        WhichDice::Modern => Dice::modern(),
    };
    dice.set_io_provider(io);
    dice.set_digest_config(digest_config);

    let dice = dice.build_with_which_spawner(detect_cycles, which_spawner);
    let mut dice_ctx = dice.updater();
    dice_ctx.set_none_cell_resolver()?;
    dice_ctx.set_none_legacy_configs()?;
    dice_ctx.commit().await;

    Ok(dice)
}


/*
 * 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_core::configuration::compatibility::IncompatiblePlatformReason;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_core::target::label::TargetLabel;
use buck2_events::dispatch::console_message;
use buck2_node::load_patterns::load_patterns;
use buck2_node::load_patterns::MissingTargetBehavior;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_frontend::ConfiguredTargetNodeCalculation;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use dice::DiceComputations;
use dupe::Dupe;
use gazebo::prelude::VecExt;
use starlark_map::small_set::SmallSet;

// Returns a tuple of compatible and incompatible targets.
fn split_compatible_incompatible(
    targets: impl Iterator<Item = anyhow::Result<MaybeCompatible<ConfiguredTargetNode>>>,
) -> anyhow::Result<(
    TargetSet<ConfiguredTargetNode>,
    SmallSet<ConfiguredTargetLabel>,
)> {
    let mut target_set = TargetSet::new();
    let mut incompatible_targets = SmallSet::new();

    for res in targets {
        match res? {
            MaybeCompatible::Incompatible(reason) => {
                incompatible_targets.insert(reason.target.dupe());
            }
            MaybeCompatible::Compatible(target) => {
                target_set.insert(target);
            }
        }
    }
    Ok((target_set, incompatible_targets))
}

pub async fn get_maybe_compatible_targets(
    ctx: &DiceComputations,
    loaded_targets: impl IntoIterator<Item = (PackageLabel, anyhow::Result<Vec<TargetNode>>)>,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<impl Iterator<Item = anyhow::Result<MaybeCompatible<ConfiguredTargetNode>>>> {
    let mut by_package_futs: Vec<_> = Vec::new();
    for (_package, result) in loaded_targets {
        let targets = result?;

        by_package_futs.extend({
            let global_target_platform = global_target_platform.as_ref();
            let target_futs: Vec<_> = targets.into_map(|target| async move {
                let target = ctx
                    .get_configured_target(target.label(), global_target_platform)
                    .await?;
                anyhow::Ok(ctx.get_configured_target_node(&target).await?)
            });

            target_futs
        });
    }

    Ok(futures::future::join_all(by_package_futs).await.into_iter())
}

/// Converts target nodes to a set of compatible configured target nodes.
pub async fn get_compatible_targets(
    ctx: &DiceComputations,
    loaded_targets: impl IntoIterator<Item = (PackageLabel, anyhow::Result<Vec<TargetNode>>)>,
    global_target_platform: Option<TargetLabel>,
) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
    let maybe_compatible_targets =
        get_maybe_compatible_targets(ctx, loaded_targets, global_target_platform).await?;

    let (compatible_targets, incompatible_targets) =
        split_compatible_incompatible(maybe_compatible_targets)?;

    if !incompatible_targets.is_empty() {
        console_message(IncompatiblePlatformReason::skipping_message_for_multiple(
            incompatible_targets.iter(),
        ));
    }

    Ok(compatible_targets)
}

pub async fn load_compatible_patterns(
    ctx: &DiceComputations,
    parsed_patterns: Vec<ParsedPattern<TargetPatternExtra>>,
    global_target_platform: Option<TargetLabel>,
    skip_missing_targets: MissingTargetBehavior,
) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
    let loaded_patterns = load_patterns(ctx, parsed_patterns, skip_missing_targets).await?;
    get_compatible_targets(
        ctx,
        loaded_patterns.iter_loaded_targets_by_package(),
        global_target_platform,
    )
    .await
}


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

//! Injects data used for build onto dice

use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use buck2_core::fs::buck_out_path::BuckOutPathResolver;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use derive_more::Display;
use dice::DiceComputations;
use dice::DiceTransactionUpdater;
use dice::InjectedKey;
use dupe::Dupe;

#[async_trait]
pub trait HasBuildContextData {
    async fn get_buck_out_path(&self) -> anyhow::Result<BuckOutPathResolver>;
}

pub trait SetBuildContextData {
    fn set_buck_out_path(&mut self, path: Option<ProjectRelativePathBuf>) -> anyhow::Result<()>;
}

#[derive(PartialEq, Eq, Allocative)]
pub struct BuildData {
    buck_out_path: ProjectRelativePathBuf,
}

#[derive(Clone, Dupe, Display, Debug, Eq, Hash, PartialEq, Allocative)]
#[display(fmt = "{:?}", self)]
struct BuildDataKey;

impl InjectedKey for BuildDataKey {
    type Value = Arc<BuildData>;

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

#[async_trait]
impl HasBuildContextData for DiceComputations {
    async fn get_buck_out_path(&self) -> anyhow::Result<BuckOutPathResolver> {
        let data = self.compute(&BuildDataKey).await?;
        Ok(BuckOutPathResolver::new(data.buck_out_path.to_buf()))
    }
}

impl SetBuildContextData for DiceTransactionUpdater {
    fn set_buck_out_path(&mut self, path: Option<ProjectRelativePathBuf>) -> anyhow::Result<()> {
        Ok(self.changed_to(vec![(
            BuildDataKey,
            Arc::new(BuildData {
                buck_out_path: path.unwrap_or_else(|| {
                    ProjectRelativePathBuf::unchecked_new("buck-out/v2".to_owned())
                }),
            }),
        )])?)
    }
}


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

use std::collections::HashMap;
use std::hash::Hash;

use dice::DiceComputations;
use dice::UserComputationData;
use futures::Stream;
use futures::StreamExt;
use indexmap::IndexMap;
use smallvec::SmallVec;

/// Evaluate a series of futures, returning a series of results.
/// If any future fails, it will fail.
/// If KEEP_GOING is true, it will first make all others continue.
pub async fn try_join_all<C, R, E>(
    ctx: &DiceComputations,
    mut inputs: impl Stream<Item = Result<R, E>> + Unpin,
) -> Result<C, E>
where
    C: KeepGoingCollectable<R>,
{
    let keep_going = ctx.per_transaction_data().get_keep_going();

    let size = inputs.size_hint().0;
    let mut res = C::with_capacity(size);
    let mut err = None;
    while let Some(x) = inputs.next().await {
        match x {
            Ok(x) => res.push(x),
            Err(e) => {
                if keep_going {
                    err = Some(e);
                } else {
                    return Err(e);
                }
            }
        }
    }

    if let Some(err) = err {
        return Err(err);
    }

    Ok(res)
}

pub trait KeepGoingCollectable<I> {
    fn with_capacity(cap: usize) -> Self;

    fn push(&mut self, item: I);
}

impl<K, V> KeepGoingCollectable<(K, V)> for IndexMap<K, V>
where
    K: PartialEq + Eq + Hash,
{
    fn with_capacity(cap: usize) -> Self {
        IndexMap::with_capacity(cap)
    }

    fn push(&mut self, item: (K, V)) {
        let (k, v) = item;
        IndexMap::insert(self, k, v);
    }
}

impl<K, V> KeepGoingCollectable<(K, V)> for HashMap<K, V>
where
    K: PartialEq + Eq + Hash,
{
    fn with_capacity(cap: usize) -> Self {
        HashMap::with_capacity(cap)
    }

    fn push(&mut self, item: (K, V)) {
        let (k, v) = item;
        HashMap::insert(self, k, v);
    }
}

impl<I> KeepGoingCollectable<I> for Vec<I> {
    fn with_capacity(cap: usize) -> Self {
        Vec::with_capacity(cap)
    }

    fn push(&mut self, item: I) {
        Vec::push(self, item);
    }
}

impl<I> KeepGoingCollectable<I> for SmallVec<[I; 1]> {
    fn with_capacity(cap: usize) -> Self {
        SmallVec::with_capacity(cap)
    }

    fn push(&mut self, item: I) {
        SmallVec::push(self, item);
    }
}

pub struct KeepGoingHolder(bool);

pub trait HasKeepGoing {
    fn set_keep_going(&mut self, keep_going: bool);

    fn get_keep_going(&self) -> bool;
}

impl HasKeepGoing for UserComputationData {
    fn set_keep_going(&mut self, keep_going: bool) {
        self.data.set(KeepGoingHolder(keep_going));
    }

    fn get_keep_going(&self) -> bool {
        self.data
            .get::<KeepGoingHolder>()
            .expect("KeepGoing should be set")
            .0
    }
}


/*
 * 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(box_patterns)]
#![feature(iter_order_by)]
#![feature(try_blocks)]
#![feature(try_trait_v2)]
#![feature(never_type)]
#![feature(provide_any)]
// Plugins
#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

#[macro_use]
extern crate starlark;

use std::sync::Once;

pub mod actions;
pub mod analysis;
pub mod artifact_groups;
pub mod attrs;
pub mod audit_dep_files;
pub mod audit_output;
pub mod build;
pub mod build_signals;
pub mod bxl;
pub mod configure_dice;
pub mod configure_targets;
pub mod context;
pub mod deferred;
pub mod dynamic;
pub mod interpreter;
pub mod keep_going;
pub mod query;
pub mod spawner;
pub mod transition;

pub fn init_late_bindings() {
    static ONCE: Once = Once::new();
    ONCE.call_once(|| {
        interpreter::more::init_more_functions();
    });
}


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

use std::any::Any;

use buck2_common::events::HasEvents;
use buck2_events::dispatch::with_dispatcher_async;
use dupe::Dupe;
use futures::future::BoxFuture;
use more_futures::spawner::Spawner;
use tokio::task::JoinHandle;

#[derive(Default)]
pub struct BuckSpawner;

impl<T: HasEvents> Spawner<T> for BuckSpawner {
    fn spawn(
        &self,
        ctx: &T,
        fut: BoxFuture<'static, Box<dyn Any + Send + 'static>>,
    ) -> JoinHandle<Box<dyn Any + Send + 'static>> {
        let dispatcher = ctx.get_dispatcher().dupe();
        let task = async move { with_dispatcher_async(dispatcher, fut).await };
        tokio::spawn(task)
    }
}

#[cfg(test)]
mod tests {
    use std::sync::Arc;

    use buck2_data::CommandEnd;
    use buck2_data::CommandStart;
    use buck2_events::create_source_sink_pair;
    use buck2_events::dispatch::span;
    use buck2_events::dispatch::EventDispatcher;
    use buck2_events::BuckEvent;
    use buck2_events::EventSource;
    use buck2_wrapper_common::invocation_id::TraceId;
    use dice::DiceData;
    use dice::UserComputationData;
    use futures::future::FutureExt;
    use more_futures::spawn::spawn_dropcancel_with_preamble;

    use super::*;

    async fn next_event<T>(source: &mut T) -> BuckEvent
    where
        T: EventSource,
    {
        source.receive().unwrap().unpack_buck().unwrap().clone()
    }

    fn create_ctx(dispatcher: EventDispatcher) -> UserComputationData {
        let dice_data = {
            let mut data = DiceData::new();
            data.set(dispatcher);
            data
        };
        UserComputationData {
            data: dice_data,
            ..Default::default()
        }
    }

    fn create_start_end_events() -> (CommandStart, CommandEnd) {
        let start = CommandStart {
            data: Default::default(),
            metadata: Default::default(),
        };
        let end = CommandEnd {
            data: Default::default(),
            metadata: Default::default(),
            is_success: true,
            error_messages: vec![],
        };

        (start, end)
    }

    #[tokio::test]
    async fn test_spawn() {
        let sp = BuckSpawner;

        // Create dispatcher
        let (mut events, sink) = create_source_sink_pair();
        let trace_id = TraceId::new();
        let dispatcher = EventDispatcher::new(trace_id.dupe(), sink);

        let ctx = create_ctx(dispatcher);
        let (start, end) = create_start_end_events();

        let task = async { span(start, || (Box::new(()) as _, end)) }.boxed();

        sp.spawn(&ctx, task).await.expect("Task panicked");

        // Check that the events are received successfully & correctly.
        let event = next_event(&mut events).await;
        let span_id = event.span_id().unwrap();
        assert_eq!(event.trace_id().unwrap(), trace_id);
        assert!(event.parent_id().is_none());

        let end = next_event(&mut events).await;
        assert_eq!(end.span_id().unwrap(), span_id);
    }

    #[tokio::test]
    async fn test_spawn_task() {
        let sp = Arc::new(BuckSpawner);

        // Create dispatchers
        let (mut events1, sink1) = create_source_sink_pair();
        let trace_id1 = TraceId::new();
        let dispatcher1 = EventDispatcher::new(trace_id1.dupe(), sink1);

        let (mut events2, sink2) = create_source_sink_pair();
        let trace_id2 = TraceId::new();
        let dispatcher2 = EventDispatcher::new(trace_id2.dupe(), sink2);

        let ctx1 = create_ctx(dispatcher1);
        let ctx2 = create_ctx(dispatcher2);
        let (start1, end1) = create_start_end_events();
        let (start2, end2) = create_start_end_events();

        // Create tasks with separate dispatchers
        let task1 = async {
            span(start1, || ((), end1));
            "Hello!"
        }
        .boxed();

        let task2 = async {
            span(start2, || ((), end2));
            "World!"
        }
        .boxed();

        let (_, poll1) = spawn_dropcancel_with_preamble(
            task1,
            futures::future::ready(()),
            sp.as_ref(),
            &ctx1,
            tracing::debug_span!("test"),
        );
        let (_, poll2) = spawn_dropcancel_with_preamble(
            task2,
            futures::future::ready(()),
            sp.as_ref(),
            &ctx2,
            tracing::debug_span!("test"),
        );
        let joins = vec![poll1, poll2];

        assert_eq!(
            futures::future::join_all(joins.into_iter()).await,
            ["Hello!", "World!"]
        );

        // Check that the events are received successfully from the correct event sources.
        let event = next_event(&mut events1).await;
        let span_id = event.span_id().unwrap();
        assert_eq!(event.trace_id().unwrap(), trace_id1);
        assert!(event.parent_id().is_none());

        let end = next_event(&mut events1).await;
        assert_eq!(end.span_id().unwrap(), span_id);

        let event = next_event(&mut events2).await;
        let span_id = event.span_id().unwrap();
        assert_eq!(event.trace_id().unwrap(), trace_id2);
        assert!(event.parent_id().is_none());

        let end = next_event(&mut events2).await;
        assert_eq!(end.span_id().unwrap(), span_id);
    }
}


/*
 * 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_core::configuration::data::ConfigurationData;
use buck2_core::configuration::transition::applied::TransitionApplied;
use buck2_core::configuration::transition::id::TransitionId;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_util::late_binding::LateBinding;
use dice::DiceComputations;

#[async_trait]
pub trait TransitionCalculation: Send + Sync + 'static {
    /// Apply transition function to configuration and cache the result.
    async fn apply_transition(
        &self,
        ctx: &DiceComputations,
        target_node: &TargetNode,
        conf: &ConfigurationData,
        transition_id: &TransitionId,
    ) -> anyhow::Result<Arc<TransitionApplied>>;
}

pub static TRANSITION_CALCULATION: LateBinding<&'static dyn TransitionCalculation> =
    LateBinding::new("TRANSITION_CALCULATION");


/*
 * 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::slice;
use std::vec;

use allocative::Allocative;
use indexmap::IndexSet;

/// It is a boxed slice, where all elements are unique.
#[derive(Debug, Allocative)]
pub struct BoxSliceSet<T>(Box<[T]>);

impl<T> BoxSliceSet<T> {
    #[inline]
    pub fn iter(&self) -> slice::Iter<'_, T> {
        self.0.iter()
    }

    #[inline]
    pub fn as_slice(&self) -> &[T] {
        &self.0
    }
}

impl<T> From<IndexSet<T>> for BoxSliceSet<T> {
    #[inline]
    fn from(set: IndexSet<T>) -> BoxSliceSet<T> {
        BoxSliceSet(set.into_iter().collect())
    }
}

impl<T> IntoIterator for BoxSliceSet<T> {
    type Item = T;
    type IntoIter = vec::IntoIter<T>;

    #[inline]
    fn into_iter(self) -> vec::IntoIter<T> {
        self.0.into_vec().into_iter()
    }
}

impl<'a, T> IntoIterator for &'a BoxSliceSet<T> {
    type Item = &'a T;
    type IntoIter = slice::Iter<'a, T>;

    #[inline]
    fn into_iter(self) -> slice::Iter<'a, T> {
        self.iter()
    }
}


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

use std::iter::zip;
use std::sync::Arc;
use std::time::Instant;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_artifact::actions::key::ActionKey;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_signals::NodeDuration;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_common::result::ToUnsharedResultExt;
use buck2_data::ToProtoMessage;
use buck2_events::dispatch::async_record_root_spans;
use buck2_events::dispatch::span_async;
use buck2_events::span::SpanId;
use buck2_execute::execute::kind::CommandExecutionKind;
use buck2_execute::execute::result::CommandExecutionReport;
use buck2_execute::execute::result::CommandExecutionStatus;
use buck2_execute::output_size::OutputSize;
use derive_more::Display;
use dice::DiceComputations;
use dice::Key;
use dupe::Dupe;
use futures::future;
use futures::stream::FuturesOrdered;
use futures::FutureExt;
use indexmap::IndexMap;
use more_futures::cancellation::CancellationContext;
use ref_cast::RefCast;
use smallvec::SmallVec;
use tracing::debug;

use crate::actions::execute::action_executor::ActionOutputs;
use crate::actions::execute::action_executor::HasActionExecutor;
use crate::actions::key::ActionKeyExt;
use crate::actions::RegisteredAction;
use crate::artifact_groups::calculation::ensure_artifact_group_staged;
use crate::deferred::calculation::DeferredCalculation;
use crate::keep_going;

#[async_trait]
pub trait ActionCalculation {
    async fn get_action(&self, action_key: &ActionKey) -> anyhow::Result<Arc<RegisteredAction>>;
    async fn build_action(&self, action_key: &ActionKey) -> anyhow::Result<ActionOutputs>;
    async fn build_artifact(&self, artifact: &BuildArtifact) -> anyhow::Result<ActionOutputs>;
}

async fn build_action_impl(
    ctx: &DiceComputations,
    cancellation: &CancellationContext<'_>,
    key: &ActionKey,
) -> anyhow::Result<ActionOutputs> {
    // Compute is only called if we have cache miss
    debug!("compute {}", key);

    let action = ActionCalculation::get_action(ctx, key).await?;

    if action.key() != key {
        // The action key we start with is on the DICE graph, and thus cached
        // and properly deduplicated. But if the underlying has a different key,
        // e.g. due to dynamic_output, then we might have two different action keys
        // pointing at the same underlying action. We need to make sure that
        // underlying action only gets called once, so call build_action once
        // again with the new key to get DICE deduplication.
        let res = ActionCalculation::build_action(ctx, action.key()).await;
        return res;
    }

    build_action_no_redirect(ctx, cancellation, action).await
}

async fn build_action_no_redirect(
    ctx: &DiceComputations,
    cancellation: &CancellationContext<'_>,
    action: Arc<RegisteredAction>,
) -> anyhow::Result<ActionOutputs> {
    let materialized_inputs = {
        let inputs = action.inputs()?;
        let ensure_futs: FuturesOrdered<_> = inputs
            .iter()
            .map(|v| ensure_artifact_group_staged(ctx, v))
            .collect();

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

        let mut results = IndexMap::with_capacity(inputs.len());
        for (artifact, ready) in zip(inputs.iter(), ready_inputs.into_iter()) {
            results.insert(artifact.clone(), ready.to_group_values(artifact)?);
        }
        results
    };

    let start_event = buck2_data::ActionExecutionStart {
        key: Some(action.key().as_proto()),
        kind: action.kind().into(),
        name: Some(buck2_data::ActionName {
            category: action.category().as_str().to_owned(),
            identifier: action.identifier().unwrap_or("").to_owned(),
        }),
    };

    let executor = ctx
        .get_action_executor(action.execution_config())
        .await
        .context(format!("for action `{}`", action))?;

    let now = Instant::now();
    let action = &action;

    let fut = async move {
        let (execute_result, command_reports) = executor
            .execute(materialized_inputs, &action, cancellation)
            .await;

        let allow_omit_details = execute_result.is_ok();

        let commands = future::join_all(
            command_reports
                .iter()
                .map(|r| command_execution_report_to_proto(r, allow_omit_details)),
        )
        .await;

        let action_result;
        let execution_kind;
        let wall_time;
        let error;
        let output_size;

        let mut prefers_local = None;
        let mut requires_local = None;
        let mut allows_cache_upload = None;
        let mut did_cache_upload = None;
        let mut eligible_for_full_hybrid = None;

        let mut buck2_revision = None;
        let mut buck2_build_time = None;
        let mut hostname = None;

        match execute_result {
            Ok((outputs, meta)) => {
                output_size = outputs.calc_output_count_and_bytes().bytes;
                action_result = Ok(outputs);
                execution_kind = Some(meta.execution_kind.as_enum());
                wall_time = Some(meta.timing.wall_time);
                error = None;

                if let Some(command) = meta.execution_kind.command() {
                    prefers_local = Some(command.prefers_local);
                    requires_local = Some(command.requires_local);
                    allows_cache_upload = Some(command.allows_cache_upload);
                    did_cache_upload = Some(command.did_cache_upload);
                    eligible_for_full_hybrid = Some(command.eligible_for_full_hybrid);
                }
            }
            Err(e) => {
                // Because we already are sending the error message in the
                // ActionExecutionEnd event, we slim the error down in the result.
                // We can then unconditionally print the error message for compute(),
                // including ones near the beginning of this method, and also not
                // duplicate any error messages.
                action_result = Err(anyhow::anyhow!("Failed to build '{}'", action.owner()));
                // TODO (torozco): Remove (see protobuf file)?
                execution_kind = command_reports
                    .last()
                    .and_then(|r| r.status.execution_kind())
                    .map(|e| e.as_enum());
                wall_time = None;
                error = Some(e.as_proto());
                output_size = 0;
                // We define the below fields only in the instance of an action error
                // so as to reduce Scribe traffic and log it in buck2_action_errors
                buck2_revision = buck2_build_info::revision().map(|s| s.to_owned());
                buck2_build_time = buck2_build_info::time_iso8601().map(|s| s.to_owned());
                hostname = buck2_events::metadata::hostname();
            }
        };

        let outputs = action_result
            .as_ref()
            .map(|outputs| {
                outputs
                    .iter()
                    .filter_map(|(_artifact, value)| {
                        Some(buck2_data::ActionOutput {
                            tiny_digest: value.digest()?.tiny_digest().to_string(),
                        })
                    })
                    .collect()
            })
            .unwrap_or_default();

        (
            (action_result, wall_time),
            Box::new(buck2_data::ActionExecutionEnd {
                key: Some(action.key().as_proto()),
                kind: action.kind().into(),
                name: Some(buck2_data::ActionName {
                    category: action.category().as_str().to_owned(),
                    identifier: action.identifier().unwrap_or("").to_owned(),
                }),
                failed: error.is_some(),
                error,
                always_print_stderr: action.always_print_stderr(),
                wall_time: wall_time.and_then(|d| d.try_into().ok()),
                execution_kind: execution_kind.unwrap_or(buck2_data::ActionExecutionKind::NotSet)
                    as i32,
                output_size,
                commands,
                outputs,
                prefers_local: prefers_local.unwrap_or_default(),
                requires_local: requires_local.unwrap_or_default(),
                allows_cache_upload: allows_cache_upload.unwrap_or_default(),
                did_cache_upload: did_cache_upload.unwrap_or_default(),
                eligible_for_full_hybrid,
                buck2_revision,
                buck2_build_time,
                hostname,
            }),
        )
    };

    // boxed() the future so that we don't need to allocate space for it while waiting on input dependencies.
    let ((res, wall_time), spans) =
        async_record_root_spans(span_async(start_event, fut.boxed())).await;

    // TODO: This wall time is rather wrong. We should report a wall time on failures too.
    ctx.store_evaluation_data(BuildKeyActivationData {
        action: action.dupe(),
        duration: NodeDuration {
            user: wall_time.unwrap_or_default(),
            total: now.elapsed(),
        },
        spans,
    })?;

    res
}

pub struct BuildKeyActivationData {
    pub action: Arc<RegisteredAction>,
    pub duration: NodeDuration,
    pub spans: SmallVec<[SpanId; 1]>,
}

/// The cost of these calls are particularly critical. To control the cost (particularly size) of these calls
/// we drop the `async_trait` common in other `*Calculation` types and avoid `async fn` (for
/// build_action/build_artifact at least).
#[async_trait]
impl ActionCalculation for DiceComputations {
    async fn get_action(&self, action_key: &ActionKey) -> anyhow::Result<Arc<RegisteredAction>> {
        // TODO add async/deferred stuff
        self.compute_deferred_data(action_key.deferred_data())
            .await
            .map(|a| (*a).dupe())
            .with_context(|| format!("for action key `{}`", action_key))
    }

    async fn build_action(&self, action_key: &ActionKey) -> anyhow::Result<ActionOutputs> {
        // build_action is called for every action key. We don't use `async fn` to ensure that it has minimal cost.
        // We don't currently consume this in buck_e2e but it's good to log for debugging purposes.
        debug!("build_action {}", action_key);
        self.compute(BuildKey::ref_cast(action_key))
            .map(|v| v?.unshared_error())
            .await
    }

    async fn build_artifact(&self, artifact: &BuildArtifact) -> anyhow::Result<ActionOutputs> {
        self.build_action(artifact.key()).await
    }
}

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

#[async_trait]
impl Key for BuildKey {
    type Value = SharedResult<ActionOutputs>;

    async fn compute(
        &self,
        ctx: &DiceComputations,
        cancellation: &CancellationContext,
    ) -> Self::Value {
        build_action_impl(ctx, cancellation, &self.0)
            .await
            .shared_error()
    }

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

    fn validity(x: &Self::Value) -> bool {
        // we don't cache any kind of errors. Ideally, we could try to distinguish different
        // error types and try to cache non-transient error types, but practically there
        // are too many unknowns that may cause more harm than good if we cached errors.
        // So, don't cache it for now, until someday we decide to really need to.
        x.is_ok()
    }
}

async fn command_execution_report_to_proto(
    report: &CommandExecutionReport,
    allow_omit_details: bool,
) -> buck2_data::CommandExecution {
    let details = command_details(report, allow_omit_details).await;

    let status = match &report.status {
        CommandExecutionStatus::Success { .. } => buck2_data::command_execution::Success {}.into(),
        CommandExecutionStatus::Cancelled => buck2_data::command_execution::Cancelled {}.into(),
        CommandExecutionStatus::Failure { .. } => buck2_data::command_execution::Failure {}.into(),
        CommandExecutionStatus::TimedOut { duration, .. } => {
            buck2_data::command_execution::Timeout {
                duration: (*duration).try_into().ok(),
            }
            .into()
        }
        CommandExecutionStatus::Error { stage, error } => buck2_data::command_execution::Error {
            stage: (*stage).to_owned(),
            error: format!("{:#}", error),
        }
        .into(),
    };

    buck2_data::CommandExecution {
        details: Some(details),
        status: Some(status),
    }
}

pub async fn command_details(
    command: &CommandExecutionReport,
    allow_omit_details: bool,
) -> buck2_data::CommandExecutionDetails {
    // If the top-level command failed then we don't want to omit any details. If it succeeded and
    // so did this command (it could succeed while not having a success here if we have rejected
    // executions), then we'll strip non-relevant stuff.
    let omit_details =
        allow_omit_details && matches!(command.status, CommandExecutionStatus::Success { .. });

    let signed_exit_code = command.exit_code;

    let stdout;
    let stderr;

    if omit_details {
        stdout = Default::default();
        stderr = command.std_streams.to_lossy_stderr().await;
    } else {
        let pair = command.std_streams.to_lossy().await;
        stdout = pair.stdout;
        stderr = pair.stderr;
    };

    let command_data = command.status.execution_kind().map(|kind| match kind {
        CommandExecutionKind::Local {
            command,
            env,
            digest,
        } => {
            if omit_details {
                buck2_data::OmittedLocalCommand {
                    action_digest: digest.to_string(),
                }
                .into()
            } else {
                buck2_data::LocalCommand {
                    action_digest: digest.to_string(),
                    argv: command.to_owned(),
                    env: env
                        .iter()
                        .map(|(key, value)| buck2_data::EnvironmentEntry {
                            key: key.clone(),
                            value: value.clone(),
                        })
                        .collect(),
                }
                .into()
            }
        }
        CommandExecutionKind::Remote { digest } => buck2_data::RemoteCommand {
            action_digest: digest.to_string(),
            cache_hit: false,
            queue_time: command.timing.re_queue_time.and_then(|d| d.try_into().ok()),
        }
        .into(),
        CommandExecutionKind::ActionCache { digest } => buck2_data::RemoteCommand {
            action_digest: digest.to_string(),
            cache_hit: true,
            queue_time: command.timing.re_queue_time.and_then(|d| d.try_into().ok()),
        }
        .into(),
        CommandExecutionKind::LocalWorkerInit { command, env } => buck2_data::WorkerInitCommand {
            argv: command.to_owned(),
            env: env
                .iter()
                .map(|(key, value)| buck2_data::EnvironmentEntry {
                    key: key.clone(),
                    value: value.clone(),
                })
                .collect(),
        }
        .into(),
        CommandExecutionKind::LocalWorker {
            command,
            env,
            digest,
            fallback_exe,
        } => buck2_data::WorkerCommand {
            action_digest: digest.to_string(),
            argv: command.to_owned(),
            env: env
                .iter()
                .map(|(key, value)| buck2_data::EnvironmentEntry {
                    key: key.clone(),
                    value: value.clone(),
                })
                .collect(),
            fallback_exe: fallback_exe.to_owned(),
        }
        .into(),
    });

    buck2_data::CommandExecutionDetails {
        stdout,
        stderr,
        command: command_data,
        signed_exit_code,
        execution_stats: command.timing.execution_stats,
        input_materialization_duration: command
            .timing
            .input_materialization_duration
            .try_into()
            .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::sync::Arc;

use buck2_artifact::actions::key::ActionKey;
use buck2_artifact::deferred::data::DeferredData;

use crate::actions::RegisteredAction;

pub trait ActionKeyExt {
    #[allow(clippy::new_ret_no_self)]
    fn new(key: DeferredData<Arc<RegisteredAction>>) -> ActionKey;
    fn deferred_data(&self) -> &DeferredData<Arc<RegisteredAction>>;
}

impl ActionKeyExt for ActionKey {
    fn new(key: DeferredData<Arc<RegisteredAction>>) -> ActionKey {
        ActionKey::unchecked_new(key.into_deferred_key())
    }

    fn deferred_data(&self) -> &DeferredData<Arc<RegisteredAction>> {
        DeferredData::unchecked_new_ref(self.deferred_key())
    }
}