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

use allocative::Allocative;
use anyhow::Context as _;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_build_api_derive::internal_provider;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_interpreter::types::label::Label;
use starlark::any::ProvidesStaticType;
use starlark::collections::SmallMap;
use starlark::environment::GlobalsBuilder;
use starlark::values::dict::DictRef;
use starlark::values::type_repr::DictType;
use starlark::values::Coerce;
use starlark::values::Freeze;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOf;
use thiserror::Error;

use crate::interpreter::rule_defs::artifact::StarlarkArtifact;
use crate::interpreter::rule_defs::artifact::ValueAsArtifactLike;
// Provider that signals a rule is installable (ex. android_binary)

#[derive(Debug, Error)]
enum InstallInfoProviderErrors {
    #[error("expected a label, got `{0}` (type `{1}`)")]
    ExpectedLabel(String, String),
    #[error("Expected a dictionary of artifacts but key `{key}` contained `{got}`")]
    ExpectedArtifact { key: String, got: String },
    #[error("Expected a dictionary with string keys, but got key `{0}`")]
    ExpectedStringKey(String),
    #[error("File with key `{key}`: `{artifact}` should not have any associated artifacts")]
    AssociatedArtifacts { key: String, artifact: String },
}

#[internal_provider(install_info_creator)]
#[derive(Clone, Coerce, Debug, Freeze, Trace, ProvidesStaticType, Allocative)]
#[repr(C)]
#[freeze(validator = validate_install_info, bounds = "V: ValueLike<'freeze>")]
pub struct InstallInfoGen<V> {
    // Label for the installer
    #[provider(field_type = "Label")]
    installer: V,
    // list of files that need to be installed
    #[provider(field_type = "DictType<String, StarlarkArtifact>")]
    files: V,
}

impl<'v, V: ValueLike<'v>> InstallInfoGen<V> {
    pub fn get_installer(&self) -> anyhow::Result<ConfiguredProvidersLabel> {
        let label = Label::from_value(self.installer.to_value())
            .ok_or_else(|| {
                InstallInfoProviderErrors::ExpectedLabel(
                    self.installer.to_value().to_repr(),
                    self.installer.to_value().get_type().to_owned(),
                )
            })?
            .label()
            .to_owned();
        Ok(label)
    }

    fn get_files_dict(&self) -> DictRef<'v> {
        DictRef::from_value(self.files.to_value()).expect("Value is a Dict")
    }

    fn get_files_iter<'a>(
        files: &'a DictRef<'v>,
    ) -> impl Iterator<Item = anyhow::Result<(&'v str, ValueAsArtifactLike<'v>)>> + 'a {
        files.iter().map(|(k, v)| {
            let k = k
                .unpack_str()
                .ok_or_else(|| InstallInfoProviderErrors::ExpectedStringKey(k.to_string()))?;
            Ok((
                k,
                ValueAsArtifactLike::unpack_value(v).ok_or_else(|| {
                    InstallInfoProviderErrors::ExpectedArtifact {
                        key: k.to_owned(),
                        got: v.get_type().to_owned(),
                    }
                })?,
            ))
        })
    }

    pub fn get_files(&self) -> anyhow::Result<SmallMap<&'v str, Artifact>> {
        Self::get_files_iter(&self.get_files_dict())
            .map(|x| {
                let (k, v) = x?;
                Ok((
                    k,
                    v.0.get_bound_artifact()
                        .with_context(|| format!("For key `{k}`"))?,
                ))
            })
            .collect()
    }
}

#[starlark_module]
fn install_info_creator(globals: &mut GlobalsBuilder) {
    fn InstallInfo<'v>(
        installer: ValueOf<'v, &'v Label>,
        files: ValueOf<'v, SmallMap<&'v str, Value<'v>>>,
    ) -> anyhow::Result<InstallInfo<'v>> {
        let info = InstallInfo {
            installer: *installer,
            files: files.value,
        };
        validate_install_info(&info)?;
        Ok(info)
    }
}

fn validate_install_info<'v, V>(info: &InstallInfoGen<V>) -> anyhow::Result<()>
where
    V: ValueLike<'v>,
{
    for x in InstallInfoGen::<V>::get_files_iter(&info.get_files_dict()) {
        let (k, v) = x?;
        if let Some(other_artifacts) = v.0.get_associated_artifacts() {
            if !other_artifacts.is_empty() {
                return Err(InstallInfoProviderErrors::AssociatedArtifacts {
                    key: k.to_owned(),
                    artifact: v.0.to_string(),
                }
                .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 allocative::Allocative;
use anyhow::Context;
use buck2_build_api_derive::internal_provider;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_interpreter::types::label::Label;
use dupe::Dupe;
use indexmap::IndexMap;
use starlark::any::ProvidesStaticType;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::values::dict::DictRef;
use starlark::values::type_repr::DictType;
use starlark::values::Coerce;
use starlark::values::Freeze;
use starlark::values::Trace;
use starlark::values::Value;

use crate::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use crate::interpreter::rule_defs::cmd_args::StarlarkCommandLine;
use crate::starlark::values::ValueLike;

#[internal_provider(local_resource_info_creator)]
#[derive(Clone, Debug, Freeze, Coerce, Trace, ProvidesStaticType, Allocative)]
#[freeze(validator = validate_local_resource_info, bounds = "V: ValueLike<'freeze>")]
#[repr(C)]
pub struct LocalResourceInfoGen<V> {
    /// Configured target that is providing this local resource.
    #[provider(field_type = "Label")]
    source_target: V,
    /// Command to run to initialize a local resource.
    /// Running this command writes a JSON to stdout.
    /// This JSON represents a pool of local resources which are ready to be used.
    /// Example JSON would be:
    /// {
    ///   "pid": 42,
    ///   "resources": [
    ///     {"socket_address": "foo:1"},
    ///     {"socket_address": "bar:2"}
    ///   ]
    /// }
    /// Where '"pid"` maps to a PID of a process which should be sent SIGTERM to release the pool of resources
    /// when they are no longer needed. `"resources"` maps to the pool of resources.
    /// When a local resource from this particular pool is needed for an execution command, single entity
    /// will be reserved from the pool, for example `{"socket_address": "bar:2"}` and environment variable with
    /// name resolved using mapping in `resource_env_vars` field and `"socket_address"` key will be added to
    /// execution command.
    #[provider(field_type = "StarlarkCommandLine")]
    setup: V,
    /// Mapping from environment variable (appended to an execution command which is dependent on this local resource)
    /// to keys in setup command JSON output.
    #[provider(field_type = "DictType<String, String>")]
    resource_env_vars: V,
}

fn validate_local_resource_info<'v, V>(info: &LocalResourceInfoGen<V>) -> anyhow::Result<()>
where
    V: ValueLike<'v>,
{
    _ = Label::from_value(info.source_target.to_value()).with_context(|| {
        format!(
            "Value for `source_target` field is not a label: `{}`",
            info.source_target
        )
    })?;
    let setup = StarlarkCommandLine::try_from_value(info.setup.to_value()).with_context(|| {
        format!(
            "Value for `setup` field is not a command line: `{}`",
            info.setup
        )
    })?;
    if setup.is_empty() {
        return Err(anyhow::anyhow!(
            "Value for `setup` field is an empty command line: `{}`",
            info.setup
        ));
    }

    let env_vars = DictRef::from_value(info.resource_env_vars.to_value()).with_context(|| {
        format!(
            "Value for `resource_env_vars` field is not a dictionary: `{}`",
            info.resource_env_vars
        )
    })?;

    if env_vars.iter().count() == 0 {
        return Err(anyhow::anyhow!(
            "Value for `resource_env_vars` field is an empty dictionary: `{}`",
            info.resource_env_vars
        ));
    }

    let validation_iter = env_vars.iter().map(|(key, value)| {
        _ = key.unpack_str().with_context(|| {
            format!(
                "Invalid key in `resource_env_vars`: Expected a str, got: `{}`",
                key
            )
        })?;

        _ = value.unpack_str().with_context(|| {
            format!(
                "Invalid value in `resource_env_vars`: Expected a str, got: `{}`",
                value
            )
        })?;

        Ok::<(), anyhow::Error>(())
    });

    for validation_item in validation_iter {
        validation_item?;
    }

    Ok(())
}

#[starlark_module]
fn local_resource_info_creator(globals: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenLocalResourceInfo)]
    fn LocalResourceInfo<'v>(
        #[starlark(require = named)] source_target: Value<'v>,
        #[starlark(require = named)] setup: Value<'v>,
        #[starlark(require = named)] resource_env_vars: Value<'v>,
        _eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<LocalResourceInfo<'v>> {
        let result = LocalResourceInfo {
            source_target,
            setup,
            resource_env_vars,
        };
        validate_local_resource_info(&result)?;
        Ok(result)
    }
}

impl FrozenLocalResourceInfo {
    /// Configured target that is providing this local resource.
    pub fn source_target_label(&self) -> ConfiguredTargetLabel {
        self.source_target
            .downcast_ref::<Label>()
            .unwrap()
            .label()
            .target()
            .dupe()
    }

    /// Mapping from keys in setup command JSON output to environment variables keys which
    /// should be appended to execution commands dependent on this local resource.
    pub fn env_var_mapping(&self) -> IndexMap<String, String> {
        let env_vars = DictRef::from_value(self.resource_env_vars.to_value()).unwrap();
        env_vars
            .iter()
            .map(|(k, v)| {
                (
                    k.unpack_str().unwrap().to_owned(),
                    v.unpack_str().unwrap().to_owned(),
                )
            })
            .collect()
    }

    pub fn setup_command_line(&self) -> &dyn CommandLineArgLike {
        self.setup.to_value().as_command_line().unwrap()
    }
}


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

//! Builtin providers.

pub mod configuration_info;
pub mod constraint_setting_info;
pub mod constraint_value_info;
pub mod default_info;
pub mod execution_platform_info;
pub mod execution_platform_registration_info;
pub mod external_runner_test_info;
pub mod install_info;
pub mod local_resource_info;
pub mod platform_info;
pub mod run_info;
pub mod template_placeholder_info;
pub mod worker_info;
pub mod worker_run_info;


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

use std::fmt::Debug;

use allocative::Allocative;
use buck2_build_api_derive::internal_provider;
use buck2_core::configuration::data::ConfigurationData;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::GlobalsBuilder;
use starlark::values::Freeze;
use starlark::values::Heap;
use starlark::values::StringValue;
use starlark::values::Trace;
use starlark::values::ValueLike;
use starlark::values::ValueOf;

use crate::interpreter::rule_defs::provider::builtin::configuration_info::ConfigurationInfo;

#[internal_provider(platform_info_creator)]
#[derive(Clone, Debug, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[repr(C)]
pub struct PlatformInfoGen<V> {
    #[provider(field_type = "String")]
    label: V,
    #[provider(field_type = "ConfigurationInfo")]
    configuration: V,
}

impl<'v, V: ValueLike<'v>> PlatformInfoGen<V> {
    pub fn to_configuration(&self) -> anyhow::Result<ConfigurationData> {
        ConfigurationData::from_platform(
            self.label
                .to_value()
                .unpack_str()
                .expect("type checked during construction")
                .to_owned(),
            ConfigurationInfo::from_value(self.configuration.to_value())
                .expect("type checked during construction")
                .to_configuration_data()?,
        )
    }
}

impl<'v> PlatformInfo<'v> {
    pub fn from_configuration(
        cfg: &ConfigurationData,
        heap: &'v Heap,
    ) -> anyhow::Result<PlatformInfo<'v>> {
        let label = heap.alloc_str(cfg.label()?).to_value();
        let configuration = heap.alloc(ConfigurationInfo::from_configuration_data(
            cfg.data()?,
            heap,
        ));
        Ok(PlatformInfoGen {
            label,
            configuration,
        })
    }
}

#[starlark_module]
fn platform_info_creator(globals: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenPlatformInfo)]
    fn PlatformInfo<'v>(
        #[starlark(require = named)] label: StringValue<'v>,
        #[starlark(require = named)] configuration: ValueOf<'v, &'v ConfigurationInfo<'v>>,
    ) -> anyhow::Result<PlatformInfo<'v>> {
        Ok(PlatformInfo {
            label: label.to_value(),
            configuration: *configuration,
        })
    }
}

#[cfg(test)]
mod tests {
    use buck2_node::attrs::internal::internal_attrs_platform_info_provider_id;

    use crate::interpreter::rule_defs::provider::builtin::platform_info::PlatformInfoCallable;

    #[test]
    fn test_platform_info_provider_id_in_internal_attrs_correct() {
        assert_eq!(
            internal_attrs_platform_info_provider_id(),
            PlatformInfoCallable::provider_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::fmt::Debug;

use allocative::Allocative;
use buck2_build_api_derive::internal_provider;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::values::list::AllocList;
use starlark::values::Freeze;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;

use crate::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use crate::interpreter::rule_defs::cmd_args::CommandLineBuilder;
use crate::interpreter::rule_defs::cmd_args::CommandLineContext;
use crate::interpreter::rule_defs::cmd_args::StarlarkCommandLine;
use crate::interpreter::rule_defs::cmd_args::WriteToFileMacroVisitor;

/// Provider that signals that a rule is runnable
#[internal_provider(run_info_creator)]
#[derive(Clone, Debug, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[repr(transparent)]
pub struct RunInfoGen<V> {
    /// The command to run, stored as CommandLine
    #[provider(field_type = "StarlarkCommandLine")]
    args: V,
}

#[starlark_module]
fn run_info_creator(globals: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenRunInfo)]
    fn RunInfo<'v>(
        #[starlark(default = AllocList::EMPTY)] args: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<RunInfo<'v>> {
        let heap = eval.heap();
        let valid_args = StarlarkCommandLine::try_from_value(args)?;
        Ok(RunInfo {
            args: heap.alloc(valid_args),
        })
    }
}

impl<'v, V: ValueLike<'v>> CommandLineArgLike for RunInfoGen<V> {
    fn add_to_command_line(
        &self,
        cli: &mut dyn CommandLineBuilder,
        context: &mut dyn CommandLineContext,
    ) -> anyhow::Result<()> {
        self.args
            .to_value()
            .as_command_line()
            .expect("a command line from construction")
            .add_to_command_line(cli, context)?;
        Ok(())
    }

    fn visit_artifacts(&self, visitor: &mut dyn CommandLineArtifactVisitor) -> anyhow::Result<()> {
        self.args
            .to_value()
            .as_command_line()
            .expect("a command line from construction")
            .visit_artifacts(visitor)?;
        Ok(())
    }

    fn contains_arg_attr(&self) -> bool {
        self.args
            .to_value()
            .as_command_line()
            .expect("a command line from construction")
            .contains_arg_attr()
    }

    fn visit_write_to_file_macros(
        &self,
        _visitor: &mut dyn WriteToFileMacroVisitor,
    ) -> anyhow::Result<()> {
        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::fmt::Debug;

use allocative::Allocative;
use buck2_build_api_derive::internal_provider;
use either::Either;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::collections::SmallMap;
use starlark::environment::GlobalsBuilder;
use starlark::values::dict::AllocDict;
use starlark::values::dict::DictRef;
use starlark::values::dict::FrozenDictRef;
use starlark::values::Freeze;
use starlark::values::FrozenRef;
use starlark::values::FrozenValue;
use starlark::values::Trace;
use starlark::values::Value;
use thiserror::Error;

use crate::interpreter::rule_defs::cmd_args::value::FrozenCommandLineArg;
use crate::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use crate::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;

#[derive(Debug, Error)]
enum TemplatePlaceholderInfoError {
    #[error(
        "Expected TemplatePlaceholderInfo.{field_key} to be a dict of String->arg-like Value, got `{value_repr}`."
    )]
    VariablesNotADict {
        field_key: String,
        value_repr: String,
    },
    #[error(
        "Expected TemplatePlaceholderInfo.{field_key} to be a dict of String->Value, where each Value is ArgLike or is itself a dict of String->ArgLike but {key} isn't either of those."
    )]
    ValueNotValidKeyedPlaceholder { field_key: String, key: String },
    #[error(
        "Expected TemplatePlaceholderInfo.{field_key}[{key}] to be a dict of String->Arglike but {inner_key} isn't an ArgLike Value."
    )]
    InnerValueNotCommandLineLike {
        field_key: String,
        key: String,
        inner_key: String,
    },
}
/// A provider that is used for expansions in string attribute templates
///
/// String attribute templates allow two types of user-defined placeholders, "unkeyed placeholders"
/// like `$(CXX)` or `$(aapt)` and "keyed placeholders"  that include a target key like
/// `$(cxxppflags //some:target)`. The expansion of each of these types is based on the
/// `TemplatePlaceholderInfo` providers.
///
/// "keyed placeholders" are used for the form `$(<key> <target>)` or `$(<key> <target> <arg>)`. In both cases
/// the lookup will expect a `TemplatePlaceholderInfo` in the providers of `<target>`. It will then lookup
/// `<key>` in the keyed_variables (call this the `value`). There are then four valid possibilities:
///  1. no-arg placeholder, an arg-like `value`: resolve to `value`
///  2. no-arg placeholder, a dictionary `value`: resolve to `value["DEFAULT"]`
///  3. arg placeholder, a non-dictionary `value`: this is an error
///  4. arg placeholder, a dictionary `value`: resolve to `value[<arg>]`
///
/// "unkeyed placeholders" are resolved by matching to any of the deps of the target. `$(CXX)` will resolve
/// to the "CXX" value in any dep's `TemplateProviderInfo.unkeyed_variables`
///
/// Fields:
///  - unkeyed_variables: A mapping of names to arg-like values. These are used for "unkeyed placeholder" expansion.
///  - keyed_variables: A mapping of names to arg-like values or dictionary of string to
///        arg-like values. These are used for "keyed placeholder" expansion.
#[internal_provider(template_placeholder_info_creator)]
#[derive(Clone, Debug, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[repr(C)]
pub struct TemplatePlaceholderInfoGen<V> {
    // Dict[String, CommandLineArg]
    unkeyed_variables: V,
    // Dict[String, Either<CommandLineArg, Dict[String, CommandLineArg]>]
    keyed_variables: V,
}

impl FrozenTemplatePlaceholderInfo {
    pub fn unkeyed_variables(&self) -> SmallMap<FrozenRef<'static, str>, FrozenCommandLineArg> {
        FrozenDictRef::from_frozen_value(self.unkeyed_variables)
            .expect("should be a dict-like object")
            .iter()
            .map(|(k, v)| {
                (
                    k.downcast_frozen_str().expect("should have string keys"),
                    FrozenCommandLineArg::new(v).unwrap(),
                )
            })
            .collect()
    }

    pub fn keyed_variables(
        &self,
    ) -> SmallMap<
        FrozenRef<'static, str>,
        Either<FrozenCommandLineArg, SmallMap<FrozenRef<'static, str>, FrozenCommandLineArg>>,
    > {
        FrozenDictRef::from_frozen_value(self.keyed_variables)
            .expect("should be a dict-like object")
            .iter()
            .map(|(k, v)| {
                (k.downcast_frozen_str().expect("should have string keys"), {
                    if let Some(dict) = FrozenDictRef::from_frozen_value(v) {
                        Either::Right(
                            dict.iter()
                                .map(|(k, v)| {
                                    (
                                        k.downcast_frozen_str().expect("should have string keys"),
                                        FrozenCommandLineArg::new(v).unwrap(),
                                    )
                                })
                                .collect(),
                        )
                    } else if let Ok(cmd) = FrozenCommandLineArg::new(v) {
                        Either::Left(cmd)
                    } else {
                        unreachable!("should be dict or command line")
                    }
                })
            })
            .collect()
    }

    pub fn _lookup_provider_raw(providers: &FrozenProviderCollectionValue) -> Option<&FrozenValue> {
        let provider_id = TemplatePlaceholderInfoCallable::provider_id();
        providers
            .provider_collection()
            .get_provider_raw(provider_id)
    }
}

fn verify_variables_type(field_key: &str, variables: Value) -> anyhow::Result<()> {
    match DictRef::from_value(variables) {
        None => Err(TemplatePlaceholderInfoError::VariablesNotADict {
            field_key: field_key.to_owned(),
            value_repr: variables.to_repr(),
        }
        .into()),
        Some(dict) => {
            for (key, value) in dict.iter() {
                if value.as_command_line().is_some() {
                    continue;
                }

                if let Some(dict) = DictRef::from_value(value) {
                    for (inner_key, value) in dict.iter() {
                        if value.as_command_line().is_none() {
                            return Err(
                                TemplatePlaceholderInfoError::InnerValueNotCommandLineLike {
                                    field_key: field_key.to_owned(),
                                    key: key.to_string(),
                                    inner_key: inner_key.to_string(),
                                }
                                .into(),
                            );
                        }
                    }
                    continue;
                }

                return Err(
                    TemplatePlaceholderInfoError::ValueNotValidKeyedPlaceholder {
                        field_key: field_key.to_owned(),
                        key: key.to_string(),
                    }
                    .into(),
                );
            }

            Ok(())
        }
    }
}

impl<'v> TemplatePlaceholderInfo<'v> {
    fn new(unkeyed_variables: Value<'v>, keyed_variables: Value<'v>) -> anyhow::Result<Self> {
        verify_variables_type("unkeyed_variables", unkeyed_variables)?;
        verify_variables_type("keyed_variables", keyed_variables)?;
        Ok(Self {
            unkeyed_variables,
            keyed_variables,
        })
    }
}

#[starlark_module]
fn template_placeholder_info_creator(builder: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenTemplatePlaceholderInfo)]
    fn TemplatePlaceholderInfo<'v>(
        #[starlark(default = AllocDict::EMPTY)] unkeyed_variables: Value<'v>,
        #[starlark(default = AllocDict::EMPTY)] keyed_variables: Value<'v>,
    ) -> anyhow::Result<TemplatePlaceholderInfo<'v>> {
        TemplatePlaceholderInfo::new(unkeyed_variables, keyed_variables)
    }
}


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

use std::fmt::Debug;
use std::sync::atomic;
use std::sync::atomic::AtomicU64;

use allocative::Allocative;
use anyhow::Context;
use buck2_build_api_derive::internal_provider;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::values::list::AllocList;
use starlark::values::Freeze;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;

use crate::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use crate::interpreter::rule_defs::cmd_args::StarlarkCommandLine;

/// Provider that signals that a rule is a worker tool
#[internal_provider(worker_info_creator)]
#[derive(Clone, Debug, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[freeze(validator = validate_worker_info, bounds = "V: ValueLike<'freeze>")]
#[repr(C)]
pub struct WorkerInfoGen<V> {
    // Command to spawn a new worker
    #[provider(field_type = "StarlarkCommandLine")]
    pub exe: V,

    pub id: u64,
}

fn next_id() -> u64 {
    static LAST_ID: AtomicU64 = AtomicU64::new(0);
    LAST_ID.fetch_add(1, atomic::Ordering::Relaxed) + 1
}

#[starlark_module]
fn worker_info_creator(globals: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenWorkerInfo)]
    fn WorkerInfo<'v>(
        #[starlark(default = AllocList::EMPTY)] exe: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<WorkerInfo<'v>> {
        let heap = eval.heap();
        let valid_exe = StarlarkCommandLine::try_from_value(exe)?;
        let exe = heap.alloc(valid_exe);
        let id = next_id();
        Ok(WorkerInfo { exe, id })
    }
}

impl<'v, V: ValueLike<'v>> WorkerInfoGen<V> {
    pub fn exe_command_line(&self) -> &'v dyn CommandLineArgLike {
        self.exe
            .to_value()
            .as_command_line()
            .expect("validated at construction")
    }
}

fn validate_worker_info<'v, V>(info: &WorkerInfoGen<V>) -> anyhow::Result<()>
where
    V: ValueLike<'v>,
{
    let exe = StarlarkCommandLine::try_from_value(info.exe.to_value()).with_context(|| {
        format!(
            "Value for `exe` field is not a command line: `{}`",
            info.exe
        )
    })?;
    if exe.is_empty() {
        return Err(anyhow::anyhow!(
            "Value for `exe` field is an empty command line: `{}`",
            info.exe
        ));
    }

    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::fmt::Debug;

use allocative::Allocative;
use buck2_build_api_derive::internal_provider;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::values::list::AllocList;
use starlark::values::Freeze;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOf;
use starlark::values::ValueTyped;

use crate::interpreter::rule_defs::cmd_args::StarlarkCommandLine;
use crate::interpreter::rule_defs::provider::builtin::worker_info::WorkerInfo;

/// Provider that signals that a rule can run using a worker
#[internal_provider(worker_run_info_creator)]
#[derive(Clone, Debug, Coerce, Trace, Freeze, ProvidesStaticType, Allocative)]
#[repr(C)]
pub struct WorkerRunInfoGen<V> {
    // Configuration needed to spawn a new worker
    #[provider(field_type = "WorkerInfo")]
    worker: V,

    // Command to execute without spawning a worker, when the build environment or configuration does not support workers
    #[provider(field_type = "StarlarkCommandLine")]
    exe: V,
}

#[starlark_module]
fn worker_run_info_creator(globals: &mut GlobalsBuilder) {
    #[starlark(as_type = FrozenWorkerRunInfo)]
    fn WorkerRunInfo<'v>(
        #[starlark(require = named)] worker: ValueOf<'v, &'v WorkerInfo<'v>>,
        #[starlark(require = named, default = AllocList::EMPTY)] exe: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<WorkerRunInfo<'v>> {
        let heap = eval.heap();
        let valid_exe = StarlarkCommandLine::try_from_value(exe)?;
        Ok(WorkerRunInfo {
            worker: *worker,
            exe: heap.alloc(valid_exe),
        })
    }
}

impl<'v, V: ValueLike<'v>> WorkerRunInfoGen<V> {
    pub fn worker(&self) -> ValueOf<'v, &'v WorkerInfo<'v>> {
        ValueOf::unpack_value(self.worker.to_value()).expect("validated at construction")
    }

    pub fn exe(&self) -> ValueTyped<'v, StarlarkCommandLine<'v>> {
        ValueTyped::new(self.exe.to_value()).expect("validated at construction")
    }
}


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

#![allow(clippy::module_inception)]

mod transitive_set;
mod transitive_set_args_projection;
mod transitive_set_definition;
mod transitive_set_iterator;
mod transitive_set_json_projection;
mod traversal;

use thiserror::Error;

pub use self::transitive_set::FrozenTransitiveSet;
pub use self::transitive_set::TransitiveSet;
use self::transitive_set::TransitiveSetGen;
pub use self::transitive_set::TransitiveSetLike;
pub use self::transitive_set_args_projection::FrozenTransitiveSetArgsProjection;
pub use self::transitive_set_args_projection::TransitiveSetArgsProjection;
pub use self::transitive_set_definition::transitive_set_definition_from_value;
pub use self::transitive_set_definition::FrozenTransitiveSetDefinition;
pub use self::transitive_set_definition::TransitiveSetDefinition;
pub use self::transitive_set_definition::TransitiveSetOperations;
pub use self::transitive_set_definition::TransitiveSetProjectionKind;
pub use self::transitive_set_definition::TransitiveSetProjectionSpec;
use self::transitive_set_iterator::BfsTransitiveSetIteratorGen;
use self::transitive_set_iterator::PostorderTransitiveSetIteratorGen;
use self::transitive_set_iterator::PreorderTransitiveSetIteratorGen;
use self::transitive_set_iterator::TopologicalTransitiveSetIteratorGen;
use self::transitive_set_iterator::TransitiveSetIteratorLike;
pub use self::transitive_set_json_projection::FrozenTransitiveSetJsonProjection;
pub use self::transitive_set_json_projection::TransitiveSetJsonProjection;
pub use self::traversal::TransitiveSetOrdering;

#[derive(Debug, Error)]
pub(crate) enum TransitiveSetError {
    #[error(
        "Transitive set type must be assigned to a top-level variable, e.g. `MySet = transitive_set()`"
    )]
    TransitiveSetNotAssigned,

    #[error(
        "Transitive set was provided with a definition that is not the output of transitive_set()"
    )]
    TransitiveSetDefinitionWasInvalid,

    #[error(
        "Transitive set was used before being assigned to a top-level variable, e.g. `MySet = transitive_set()`"
    )]
    TransitiveSetUsedBeforeAssignment,

    #[error("Transitive set transitive values must be transitive sets, got `{}`", .got)]
    TransitiveValueIsNotTransitiveSet { got: String },

    #[error(
        "Transitive set transitive values must be of the same transitive set type (expected: `{}`, got: `{}`)",
        .expected,
        .got
    )]
    TransitiveValueIsOfWrongType { expected: String, got: String },

    #[error(
        "The requested projection `{}` does not exist. Valid projections: {}",
        .projection,
        .valid_projections.join(", "),
    )]
    ProjectionDoesNotExist {
        projection: String,
        valid_projections: Vec<String>,
    },

    #[error(
        "Requested a {} projection, but `{}` is a `{}` projection (and should use `{}` instead)",
        .expected_kind.short_name(),
        .projection,
        .actual_kind.short_name(),
        .actual_kind.function_name(),
    )]
    ProjectionKindMismatch {
        projection: String,
        expected_kind: TransitiveSetProjectionKind,
        actual_kind: TransitiveSetProjectionKind,
    },

    #[error("Error evaluating transitive set projection {}", .name)]
    ProjectionError {
        name: String,

        #[source]
        error: anyhow::Error,
    },

    #[error("Expected args_projection `{}` function to take a single argument", .name)]
    ProjectionSignatureError { name: String },

    #[error("Error evaluating transitive set reduction {}", .name)]
    ReductionError {
        name: String,

        #[source]
        error: anyhow::Error,
    },

    #[error(
        "The requested reduction `{}` does not exist. Valid reduction: {}",
        .reduction,
        .valid_reductions.join(", "),
    )]
    ReductionDoesNotExist {
        reduction: String,
        valid_reductions: Vec<String>,
    },

    #[error("Expected ordering to be one of `preorder`, `postorder`, `topological`, or `bfs`,  but got `{0}`", .ordering)]
    OrderingUnexpectedValue { ordering: String },
}


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

use std::fmt;
use std::iter;

use allocative::Allocative;
use anyhow::Context as _;
use display_container::display_pair;
use display_container::fmt_container;
use display_container::iter_display_chain;
use dupe::Dupe;
use gazebo::prelude::*;
use serde::ser::SerializeMap;
use serde::Serialize;
use serde::Serializer;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::eval::Evaluator;
use starlark::values::list::AllocList;
use starlark::values::starlark_value;
use starlark::values::Freeze;
use starlark::values::Freezer;
use starlark::values::FrozenValue;
use starlark::values::Heap;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOf;

use crate::actions::impls::json::validate_json;
use crate::actions::impls::json::visit_json_artifacts;
use crate::artifact_groups::deferred::TransitiveSetKey;
use crate::artifact_groups::ArtifactGroup;
use crate::artifact_groups::TransitiveSetProjectionKey;
use crate::interpreter::rule_defs::cmd_args::SimpleCommandLineArtifactVisitor;
use crate::interpreter::rule_defs::transitive_set::transitive_set_definition::TransitiveSetProjectionKind;
use crate::interpreter::rule_defs::transitive_set::transitive_set_definition_from_value;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetOrdering;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetTraversal;
use crate::interpreter::rule_defs::transitive_set::BfsTransitiveSetIteratorGen;
use crate::interpreter::rule_defs::transitive_set::PostorderTransitiveSetIteratorGen;
use crate::interpreter::rule_defs::transitive_set::PreorderTransitiveSetIteratorGen;
use crate::interpreter::rule_defs::transitive_set::TopologicalTransitiveSetIteratorGen;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetArgsProjection;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetError;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetIteratorLike;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetJsonProjection;

#[derive(Debug, Clone, Trace, ProvidesStaticType, Allocative)]
#[repr(C)]
pub struct TransitiveSetGen<V> {
    /// A Deferred key that maps back to this set. This is used to compute its inputs.
    pub key: TransitiveSetKey,

    /// The TransitiveSetCallable that this set uses.
    pub(crate) definition: V,

    /// The immediate value of this node. If None, then this node will not yield anything when
    /// iterated over (but we'll still traverse to its children).
    pub(crate) node: Option<NodeGen<V>>,

    /// Pre-computed reductions. Those are arbitrary values based on the set's definition.
    pub(crate) reductions: Box<[V]>,

    /// Further transitive sets.
    pub children: Box<[V]>,
}

#[derive(Debug, Clone, Trace, Allocative)]
#[repr(C)]
pub struct NodeGen<V> {
    /// The value
    pub value: V,

    /// Pre-computed projections.
    pub projections: Box<[V]>,
}

unsafe impl<'v> Coerce<TransitiveSetGen<Value<'v>>> for TransitiveSetGen<FrozenValue> {}

impl<V: fmt::Display> fmt::Display for TransitiveSetGen<V> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        fmt_container(
            f,
            &format!("{}(", self.definition),
            ")",
            iter_display_chain(
                self.node
                    .as_ref()
                    .map(|node| display_pair("value", "=", &node.value)),
                iter::once(format!("{} children", self.children.len())),
            ),
        )
    }
}

impl<'v, V: ValueLike<'v>> Serialize for TransitiveSetGen<V> {
    fn serialize<S>(&self, s: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut s = s.serialize_map(Some(3))?;
        s.serialize_entry("definition", &self.definition)?;
        if let Some(node) = self.node.as_ref() {
            s.serialize_entry("value", &node.value)?;
        }

        s.serialize_entry("children", &self.children.len())?;
        s.end()
    }
}

impl<V> TransitiveSetGen<V> {
    pub fn key(&self) -> &TransitiveSetKey {
        &self.key
    }
}

impl<'v> NodeGen<Value<'v>> {
    fn freeze(self, freezer: &Freezer) -> anyhow::Result<NodeGen<FrozenValue>> {
        let Self { value, projections } = self;

        let value = value.freeze(freezer)?;
        let projections = projections.freeze(freezer)?;

        Ok(NodeGen { value, projections })
    }
}

impl<'v, V: ValueLike<'v>> TransitiveSetGen<V> {
    fn matches_definition(&self, definition: Value<'v>) -> bool {
        definition.ptr_eq(self.definition.to_value())
    }

    pub fn projection_name(&'v self, projection: usize) -> anyhow::Result<&'v str> {
        let def = transitive_set_definition_from_value(self.definition.to_value())
            .context("Invalid definition")?;

        Ok(def
            .operations()
            .projections
            .get_index(projection)
            .context("Invalid projection id")?
            .0
            .as_str())
    }

    pub fn get_projection_value(&self, projection: usize) -> anyhow::Result<Option<V>> {
        match &self.node {
            None => Ok(None),
            Some(node) => Ok(Some(
                *node
                    .projections
                    .get(projection)
                    .context("Invalid projection id")?,
            )),
        }
    }

    pub fn get_projection_key(&self, projection: usize) -> TransitiveSetProjectionKey {
        TransitiveSetProjectionKey {
            key: self.key.dupe(),
            projection,
        }
    }
}

impl FrozenTransitiveSet {
    pub fn get_projection_sub_inputs(
        &self,
        projection: usize,
    ) -> anyhow::Result<Vec<ArtifactGroup>> {
        let mut sub_inputs = Vec::new();

        if let Some(projection) = self.get_projection_value(projection)? {
            let mut visitor = SimpleCommandLineArtifactVisitor::new();
            // It's either an args-like or a json projection. visit_json_artifacts handles both the way we want.
            visit_json_artifacts(projection.to_value(), &mut visitor)?;
            sub_inputs.extend(visitor.inputs);
        }

        // Reuse the same projection for children sets.
        for v in self.children.iter() {
            let v = TransitiveSet::from_value(v.to_value()).context("Invalid deferred")?;
            sub_inputs.push(ArtifactGroup::TransitiveSetProjection(
                TransitiveSetProjectionKey {
                    key: v.key().dupe(),
                    projection,
                },
            ));
        }
        Ok(sub_inputs)
    }
}

impl<'v, V> TransitiveSetGen<V>
where
    V: ValueLike<'v>,
    TransitiveSetGen<V>: StarlarkValue<'v> + TransitiveSetLike<'v>,
{
    pub fn iter<'a>(
        &'a self,
        ordering: TransitiveSetOrdering,
    ) -> Box<dyn TransitiveSetIteratorLike<'a, 'v, V> + 'a>
    where
        'v: 'a,
    {
        match ordering {
            TransitiveSetOrdering::Preorder => {
                Box::new(PreorderTransitiveSetIteratorGen::new(self))
            }
            TransitiveSetOrdering::Postorder => {
                Box::new(PostorderTransitiveSetIteratorGen::new(self))
            }
            TransitiveSetOrdering::Topological => {
                Box::new(TopologicalTransitiveSetIteratorGen::new(self))
            }
            TransitiveSetOrdering::Bfs => Box::new(BfsTransitiveSetIteratorGen::new(self)),
        }
    }

    pub fn iter_values<'a>(
        &'a self,
        ordering: TransitiveSetOrdering,
    ) -> anyhow::Result<Box<dyn Iterator<Item = Value<'v>> + 'a>>
    where
        'v: 'a,
    {
        Ok(Box::new(
            self.iter(ordering)
                .values()
                .map(|node| node.value.to_value()),
        ))
    }

    pub(super) fn iter_projection_values<'a>(
        &'a self,
        ordering: TransitiveSetOrdering,
        projection: usize,
    ) -> anyhow::Result<Box<dyn Iterator<Item = Value<'v>> + 'a>>
    where
        'v: 'a,
    {
        let mut iter = self.iter(ordering).values().peekable();

        // Defensively, check the projection is valid. We know the set has the same definition
        // throughout so it'll be safe (enough) to unwrap if it is valid on the first one.
        if let Some(v) = iter.peek() {
            v.projections
                .get(projection)
                .context("Invalid projection")?;
        }

        Ok(Box::new(iter.map(move |node| {
            node.projections.get(projection).unwrap().to_value()
        })))
    }
}

pub trait TransitiveSetLike<'v> {
    fn from_value(v: Value<'v>) -> Option<&'v Self>;
}

impl<'v> TransitiveSetLike<'v> for TransitiveSet<'v> {
    fn from_value(v: Value<'v>) -> Option<&'v Self> {
        TransitiveSet::from_value(v)
    }
}

impl<'v> TransitiveSetLike<'v> for FrozenTransitiveSet {
    fn from_value(v: Value<'v>) -> Option<&'v Self> {
        // FrozenTransitiveSet::from_value(v)

        v.downcast_ref::<FrozenTransitiveSet>()
    }
}

starlark_complex_value!(pub TransitiveSet);

#[starlark_value(type = "transitive_set")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for TransitiveSetGen<V>
where
    Self: ProvidesStaticType<'v> + TransitiveSetLike<'v>,
{
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(transitive_set_methods)
    }

    fn matches_type(&self, ty: &str) -> bool {
        if ty == "transitive_set" {
            return true;
        }

        transitive_set_definition_from_value(self.definition.to_value())
            .map_or(false, |d| d.matches_type(ty))
    }
}

impl<'v> Freeze for TransitiveSet<'v> {
    type Frozen = FrozenTransitiveSet;

    fn freeze(self, freezer: &Freezer) -> anyhow::Result<Self::Frozen> {
        let Self {
            key,
            definition,
            node,
            reductions,
            children,
        } = self;
        let definition = definition.freeze(freezer)?;
        let node = node.try_map(|node| node.freeze(freezer))?;
        let children = children.freeze(freezer)?;
        let reductions = reductions.freeze(freezer)?;
        Ok(TransitiveSetGen {
            key,
            definition,
            node,
            reductions,
            children,
        })
    }
}

impl<'v> TransitiveSet<'v> {
    pub fn new(
        key: TransitiveSetKey,
        definition: Value<'v>,
        value: Option<Value<'v>>,
        children: impl IntoIterator<Item = Value<'v>>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Self> {
        let def = match transitive_set_definition_from_value(definition.to_value()) {
            Some(def) if def.has_id() => def,
            Some(..) => {
                return Err(TransitiveSetError::TransitiveSetUsedBeforeAssignment.into());
            }
            None => {
                return Err(TransitiveSetError::TransitiveSetDefinitionWasInvalid.into());
            }
        };

        let children = children.into_iter().collect::<Box<[_]>>();
        let children_sets = children.try_map(|v| match TransitiveSet::from_value(*v) {
            Some(set) if set.matches_definition(definition) => Ok(set),
            Some(set) => {
                fn format_def(def: Value<'_>) -> String {
                    match transitive_set_definition_from_value(def) {
                        Some(def) => format!("{:?}", def.as_debug()),
                        None => "<invalid>".to_owned(),
                    }
                }
                Err(TransitiveSetError::TransitiveValueIsOfWrongType {
                    expected: format_def(definition),
                    got: format_def(set.definition),
                })
            }
            None => {
                Err(TransitiveSetError::TransitiveValueIsNotTransitiveSet { got: v.to_string() })
            }
        })?;

        let node = value.try_map(|value| {
            let projections = def
                .operations()
                .projections
                .iter()
                .map(|(name, spec)| {
                    let projected_value = eval
                        .eval_function(spec.projection, &[value], &[])
                        .map_err(|error| TransitiveSetError::ProjectionError {
                            error,
                            name: name.clone(),
                        })?;
                    match spec.kind {
                        TransitiveSetProjectionKind::Args => {
                            TransitiveSetArgsProjection::as_command_line(projected_value)?;
                        }
                        TransitiveSetProjectionKind::Json => {
                            validate_json(projected_value)?;
                        }
                    }
                    anyhow::Ok(projected_value)
                })
                .collect::<Result<Box<[_]>, _>>()?;

            anyhow::Ok(NodeGen { value, projections })
        })?;

        let reductions = def
            .operations()
            .reductions
            .iter()
            .enumerate()
            .map(|(idx, (name, reduce))| {
                let children_values = children_sets.try_map(|c| {
                    c.reductions
                        .get(idx)
                        .copied()
                        .with_context(|| format!("Child {} is missing reduction {}", c, idx))
                })?;
                let children_values = eval.heap().alloc(AllocList(children_values));

                let value = value.unwrap_or_else(Value::new_none);

                let reduced = eval
                    .eval_function(*reduce, &[children_values, value], &[])
                    .map_err(|error| TransitiveSetError::ReductionError {
                        error,
                        name: name.clone(),
                    })?;

                anyhow::Ok(reduced)
            })
            .collect::<Result<Box<[_]>, _>>()?;

        Ok(Self {
            key,
            definition,
            node,
            reductions,
            children,
        })
    }

    pub fn new_from_values(
        key: TransitiveSetKey,
        definition: Value<'v>,
        value: Option<Value<'v>>,
        children: Option<Value<'v>>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Self> {
        let children = children
            .map(|v| v.iterate(eval.heap()))
            .transpose()?
            .into_iter()
            .flatten();

        Self::new(key, definition, value, children, eval)
    }
}

#[starlark_module]
fn transitive_set_methods(builder: &mut MethodsBuilder) {
    fn project_as_json<'v>(
        this: ValueOf<'v, &'v TransitiveSet<'v>>,
        projection: &str,
        #[starlark(require = named, default = "preorder")] ordering: &str,
        heap: &'v Heap,
    ) -> anyhow::Result<Value<'v>> {
        let def = transitive_set_definition_from_value(this.typed.definition)
            .context("Invalid this.definition")?;

        let index = def
            .operations()
            .get_index_of_projection(TransitiveSetProjectionKind::Json, projection)?;

        Ok(heap.alloc(TransitiveSetJsonProjection {
            transitive_set: this.value,
            projection: index,
            ordering: TransitiveSetOrdering::parse(ordering)?,
        }))
    }

    fn project_as_args<'v>(
        this: ValueOf<'v, &'v TransitiveSet<'v>>,
        projection: &str,
        #[starlark(require = named, default = "preorder")] ordering: &str,
        heap: &'v Heap,
    ) -> anyhow::Result<Value<'v>> {
        let def = transitive_set_definition_from_value(this.typed.definition)
            .context("Invalid this.definition")?;

        let index = def
            .operations()
            .get_index_of_projection(TransitiveSetProjectionKind::Args, projection)?;

        Ok(heap.alloc(TransitiveSetArgsProjection {
            transitive_set: this.value,
            projection: index,
            ordering: TransitiveSetOrdering::parse(ordering)?,
        }))
    }

    fn reduce<'v>(
        this: ValueOf<'v, &'v TransitiveSet<'v>>,
        reduction: &str,
    ) -> anyhow::Result<Value<'v>> {
        let def = transitive_set_definition_from_value(this.typed.definition)
            .context("Invalid this.definition")?;

        let index = match def.operations().reductions.get_index_of(reduction) {
            Some(index) => index,
            None => {
                return Err(TransitiveSetError::ReductionDoesNotExist {
                    reduction: reduction.into(),
                    valid_reductions: def
                        .operations()
                        .reductions
                        .keys()
                        .map(String::from)
                        .collect::<Vec<_>>(),
                }
                .into());
            }
        };

        this.typed
            .reductions
            .get(index)
            .copied()
            .with_context(|| format!("Missing reduction {}", index))
    }

    fn traverse<'v>(
        this: ValueOf<'v, &'v TransitiveSet<'v>>,
        heap: &'v Heap,
        #[starlark(require = named, default = "preorder")] ordering: &str,
    ) -> anyhow::Result<Value<'v>> {
        Ok(heap.alloc(TransitiveSetTraversal {
            inner: this.value,
            ordering: TransitiveSetOrdering::parse(ordering)?,
        }))
    }

    #[starlark(attribute)]
    fn definition<'v>(this: ValueOf<'v, &'v TransitiveSet<'v>>) -> anyhow::Result<Value<'v>> {
        Ok(this.typed.definition)
    }

    #[starlark(attribute)]
    fn value<'v>(this: ValueOf<'v, &'v TransitiveSet<'v>>) -> anyhow::Result<Value<'v>> {
        Ok(match this.typed.node.as_ref() {
            Some(node) => node.value,
            None => Value::new_none(),
        })
    }
}


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

use std::fmt;
use std::fmt::Display;
use std::iter;

use allocative::Allocative;
use anyhow::Context as _;
use display_container::display_pair;
use display_container::fmt_container;
use display_container::iter_display_chain;
use dupe::Dupe;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::values::list::ListRef;
use starlark::values::starlark_value;
use starlark::values::Demand;
use starlark::values::Freeze;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::StringValue;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOf;

use crate::artifact_groups::ArtifactGroup;
use crate::artifact_groups::TransitiveSetProjectionKey;
use crate::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use crate::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use crate::interpreter::rule_defs::cmd_args::CommandLineBuilder;
use crate::interpreter::rule_defs::cmd_args::CommandLineContext;
use crate::interpreter::rule_defs::cmd_args::WriteToFileMacroVisitor;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetOrdering;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetProjectionTraversal;
use crate::interpreter::rule_defs::transitive_set::TransitiveSet;

/// TransitiveSetArgsProjection is the starlark value returned from the starlark method `transitive_set.project_as_args()`
///
/// The projected values are all stored on the TransitiveSet itself and this value will reference back to that. The main
/// point of this object is to provide the implementation of CommandLineArgLike so that the args projection
/// can be used in places that accept command lines.
#[derive(Debug, Clone, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[derive(NoSerialize)] // TODO we should probably have a serialization for transitive set
#[repr(C)]
pub struct TransitiveSetArgsProjectionGen<V> {
    pub(super) transitive_set: V,

    /// The index of the projection. Once transitive sets are defined, their projections never
    /// change, so we can afford to just store the index here.
    pub projection: usize,

    /// The ordering to use when traversing the projection.
    pub ordering: TransitiveSetOrdering,
}

impl<'v, V: ValueLike<'v>> Display for TransitiveSetArgsProjectionGen<V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let projection_name = self.projection_name().unwrap_or("<invalid projection>");
        fmt_container(
            f,
            "TransitiveSetProjection(",
            ")",
            iter_display_chain(
                iter::once(projection_name),
                iter::once(display_pair("transitive_set", "=", &self.transitive_set)),
            ),
        )
    }
}

impl<'v, V: ValueLike<'v>> TransitiveSetArgsProjectionGen<V> {
    fn projection_name(&self) -> anyhow::Result<&'v str> {
        TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?
            .projection_name(self.projection)
    }
}

impl<'v, V: ValueLike<'v>> TransitiveSetArgsProjectionGen<V> {
    /// For args projections, we allow either a CommandLineArgLike or a list of them (just like `cmd_args.add()`).
    /// This function allows us to treat those two as the same.
    /// TODO(cjhopman): It may be better to wrap the list case in a new CommandLineArgLike impl when returned from
    /// the projection. Then we'd only have to verify the contents type once and it might be a bit simpler to use.
    pub(super) fn as_command_line(v: V) -> anyhow::Result<impl CommandLineArgLike + 'v> {
        enum Impl<'v> {
            Item(&'v dyn CommandLineArgLike),
            List(&'v [Value<'v>]),
        }
        impl<'v> CommandLineArgLike for Impl<'v> {
            fn add_to_command_line(
                &self,
                cli: &mut dyn CommandLineBuilder,
                context: &mut dyn CommandLineContext,
            ) -> anyhow::Result<()> {
                match self {
                    Impl::Item(v) => v.add_to_command_line(cli, context),
                    Impl::List(items) => {
                        for v in *items {
                            v.as_command_line()
                                .unwrap()
                                .add_to_command_line(cli, context)?;
                        }
                        Ok(())
                    }
                }
            }

            fn contains_arg_attr(&self) -> bool {
                match self {
                    Impl::Item(v) => v.contains_arg_attr(),
                    Impl::List(items) => {
                        for v in *items {
                            if v.as_command_line().unwrap().contains_arg_attr() {
                                return true;
                            }
                        }
                        false
                    }
                }
            }

            fn visit_write_to_file_macros(
                &self,
                visitor: &mut dyn WriteToFileMacroVisitor,
            ) -> anyhow::Result<()> {
                match self {
                    Impl::Item(v) => v.visit_write_to_file_macros(visitor),
                    Impl::List(items) => {
                        for v in *items {
                            v.as_command_line()
                                .unwrap()
                                .visit_write_to_file_macros(visitor)?;
                        }
                        Ok(())
                    }
                }
            }

            fn visit_artifacts(
                &self,
                visitor: &mut dyn CommandLineArtifactVisitor,
            ) -> anyhow::Result<()> {
                match self {
                    Impl::Item(v) => v.visit_artifacts(visitor),
                    Impl::List(items) => {
                        for v in *items {
                            v.as_command_line().unwrap().visit_artifacts(visitor)?;
                        }
                        Ok(())
                    }
                }
            }
        }

        let value = v.to_value();
        if let Some(values) = ListRef::from_value(value) {
            for v in values.content() {
                v.as_command_line_err()?;
            }
            Ok(Impl::List(values.content()))
        } else {
            Ok(Impl::Item(value.as_command_line_err()?))
        }
    }
}

starlark_complex_value!(pub TransitiveSetArgsProjection);

#[starlark_value(type = "transitive_set_args_projection")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for TransitiveSetArgsProjectionGen<V>
where
    Self: ProvidesStaticType<'v>,
{
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(transitive_set_args_projection_methods)
    }

    fn provide(&'v self, demand: &mut Demand<'_, 'v>) {
        demand.provide_value::<&dyn CommandLineArgLike>(self);
    }
}

impl<'v, V: ValueLike<'v>> CommandLineArgLike for TransitiveSetArgsProjectionGen<V> {
    fn add_to_command_line(
        &self,
        builder: &mut dyn CommandLineBuilder,
        context: &mut dyn CommandLineContext,
    ) -> anyhow::Result<()> {
        let set = TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?;

        for node in set.iter(self.ordering).values() {
            let projection = node
                .projections
                .get(self.projection)
                .context("Invalid projection id")?;

            TransitiveSetArgsProjection::as_command_line(*projection)?
                .add_to_command_line(builder, context)?;
        }

        Ok(())
    }

    fn visit_artifacts(&self, visitor: &mut dyn CommandLineArtifactVisitor) -> anyhow::Result<()> {
        let set = TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?;

        visitor.visit_input(
            ArtifactGroup::TransitiveSetProjection(TransitiveSetProjectionKey {
                key: set.key().dupe(),
                projection: self.projection,
            }),
            None,
        );

        Ok(())
    }

    fn contains_arg_attr(&self) -> bool {
        // TODO(cjhopman): This seems wrong, there's no enforcement that the projected values don't have arg attrs (and
        // in fact several of our common uses of in fact do have arg attrs). There doesn't seem to be any reason to
        // not allow them, either.
        false
    }

    fn visit_write_to_file_macros(
        &self,
        _visitor: &mut dyn WriteToFileMacroVisitor,
    ) -> anyhow::Result<()> {
        // TODO(cjhopman): This seems wrong, there's no verification that the projected
        // values don't have write_to_file_macros in them.
        Ok(())
    }
}

#[starlark_module]
fn transitive_set_args_projection_methods(builder: &mut MethodsBuilder) {
    fn traverse<'v>(
        this: ValueOf<'v, &'v TransitiveSetArgsProjection<'v>>,
        heap: &'v Heap,
    ) -> anyhow::Result<Value<'v>> {
        Ok(heap.alloc(TransitiveSetProjectionTraversal {
            transitive_set: this.typed.transitive_set,
            projection: this.typed.projection,
            ordering: this.typed.ordering,
        }))
    }

    #[starlark(attribute)]
    fn projection_name<'v>(
        this: ValueOf<'v, &'v TransitiveSetArgsProjection<'v>>,
        heap: &'v Heap,
    ) -> anyhow::Result<StringValue<'v>> {
        Ok(heap.alloc_str(this.typed.projection_name()?))
    }

    #[starlark(attribute)]
    fn transitive_set<'v>(
        this: ValueOf<'v, &'v TransitiveSetArgsProjection<'v>>,
    ) -> anyhow::Result<Value<'v>> {
        Ok(this.typed.transitive_set)
    }
}


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

use std::cell::RefCell;
use std::fmt;
use std::fmt::Display;
use std::hash::Hash;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_core::bzl::ImportPath;
use derive_more::Display;
use dupe::Dupe;
use serde::Serialize;
use serde::Serializer;
use starlark::any::ProvidesStaticType;
use starlark::coerce::coerce;
use starlark::coerce::Coerce;
use starlark::collections::SmallMap;
use starlark::collections::StarlarkHasher;
use starlark::eval::Evaluator;
use starlark::values::starlark_value;
use starlark::values::AllocValue;
use starlark::values::Freeze;
use starlark::values::Freezer;
use starlark::values::FrozenValue;
use starlark::values::Heap;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::Tracer;
use starlark::values::Value;
use starlark::values::ValueLike;

use crate::interpreter::rule_defs::transitive_set::TransitiveSetError;

#[derive(Debug, Clone, Dupe, Copy, Trace, Freeze, PartialEq, Allocative)]
pub enum TransitiveSetProjectionKind {
    Args,
    Json,
}

impl TransitiveSetProjectionKind {
    pub fn short_name(&self) -> &'static str {
        match self {
            TransitiveSetProjectionKind::Args => "args",
            TransitiveSetProjectionKind::Json => "json",
        }
    }

    pub fn function_name(&self) -> &'static str {
        match self {
            TransitiveSetProjectionKind::Args => "project_as_args",
            TransitiveSetProjectionKind::Json => "project_as_json",
        }
    }
}

// The Coerce derivation doesn't work if this is just a tuple in the SmallMap value.
#[derive(Debug, Clone, Trace, Coerce, Freeze, Allocative)]
#[repr(C)]
pub struct TransitiveSetProjectionSpec<V> {
    pub kind: TransitiveSetProjectionKind,
    pub projection: V,
}

/// A unique identity for a given [`TransitiveSetDefinition`].
#[derive(Debug, Clone, Display, Allocative, Hash)]
#[display(fmt = "{}", "name")]
struct TransitiveSetId {
    module_id: ImportPath,
    name: String,
}

#[derive(Debug, ProvidesStaticType, Allocative)]
pub struct TransitiveSetDefinition<'v> {
    /// The name of this transitive set. This is filed in by `export_as` when it's assigned to a
    /// top-level variable. This must be set before this is used.
    id: RefCell<Option<Arc<TransitiveSetId>>>,

    /// The module id where this `TransitiveSetDefinition` is created and assigned
    module_id: ImportPath,

    operations: TransitiveSetOperationsGen<Value<'v>>,
}

#[derive(Debug, Clone, Trace, Coerce, Freeze, Allocative)]
#[repr(C)]
pub struct TransitiveSetOperationsGen<V> {
    /// Callables that will project the values contained in transitive sets of this type to
    /// cmd_args or json. This can be used to include a transitive set into a command or json file.
    pub(crate) projections: SmallMap<String, TransitiveSetProjectionSpec<V>>,

    /// Callables that will reduce the values contained in transitive sets to a single value per
    /// node. This can be used to e.g. aggregate flags throughout a transitive set;
    pub(crate) reductions: SmallMap<String, V>,
}

pub type TransitiveSetOperations<'v> = TransitiveSetOperationsGen<Value<'v>>;

impl<V> TransitiveSetOperationsGen<V> {
    pub fn valid_projections(&self, kind: TransitiveSetProjectionKind) -> Vec<String> {
        self.projections
            .iter()
            .filter_map(|(k, spec)| {
                if kind == spec.kind {
                    Some(k.to_string())
                } else {
                    None
                }
            })
            .collect::<Vec<_>>()
    }

    pub fn get_index_of_projection(
        &self,
        kind: TransitiveSetProjectionKind,
        proj: &str,
    ) -> anyhow::Result<usize> {
        let index = match self.projections.get_index_of(proj) {
            Some(index) => index,
            None => {
                return Err(TransitiveSetError::ProjectionDoesNotExist {
                    projection: proj.to_owned(),
                    valid_projections: self.valid_projections(TransitiveSetProjectionKind::Args),
                }
                .into());
            }
        };

        let (_, spec) = self.projections.get_index(index).unwrap();
        if spec.kind != kind {
            return Err(TransitiveSetError::ProjectionKindMismatch {
                projection: proj.to_owned(),
                expected_kind: kind,
                actual_kind: spec.kind,
            }
            .into());
        }

        Ok(index)
    }
}

unsafe impl<'v> Trace<'v> for TransitiveSetDefinition<'v> {
    fn trace(&mut self, tracer: &Tracer<'v>) {
        self.operations.trace(tracer)
    }
}

impl<'v> Display for TransitiveSetDefinition<'v> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.id.try_borrow() {
            Ok(val) => match val.as_deref() {
                Some(id) => write!(f, "{}", id),
                None => write!(f, "unnamed transitive set"),
            },
            Err(..) => write!(f, "borrowed transitive set"),
        }
    }
}

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

impl<'v> TransitiveSetDefinition<'v> {
    pub fn new(module_id: ImportPath, operations: TransitiveSetOperations<'v>) -> Self {
        Self {
            id: RefCell::new(None),
            module_id,
            operations,
        }
    }

    pub fn has_id(&self) -> bool {
        self.id.borrow().is_some()
    }
}

impl<'v> AllocValue<'v> for TransitiveSetDefinition<'v> {
    fn alloc_value(self, heap: &'v Heap) -> Value<'v> {
        heap.alloc_complex(self)
    }
}

#[starlark_value(type = "transitive_set_definition")]
impl<'v> StarlarkValue<'v> for TransitiveSetDefinition<'v> {
    fn export_as(&self, variable_name: &str, _: &mut Evaluator<'v, '_>) {
        // First export wins
        let mut id = self.id.borrow_mut();
        if id.is_none() {
            let new_id = Arc::new(TransitiveSetId {
                module_id: self.module_id.clone(),
                name: variable_name.to_owned(),
            });
            *id = Some(new_id.dupe());
        }
    }

    fn dir_attr(&self) -> Vec<String> {
        vec!["type".to_owned()]
    }

    fn has_attr(&self, attribute: &str, _heap: &'v Heap) -> bool {
        attribute == "type"
    }

    fn get_attr(&self, attribute: &str, heap: &'v Heap) -> Option<Value<'v>> {
        if attribute == "type" {
            let id = self.id.borrow();
            let typ = id
                .as_ref()
                .map_or("transitive_set_definition", |id| id.name.as_str());
            Some(heap.alloc(typ))
        } else {
            None
        }
    }

    fn write_hash(&self, hasher: &mut StarlarkHasher) -> anyhow::Result<()> {
        let id = self.id.borrow();
        let id = id
            .as_deref()
            .context("cannot hash a transitive_set_definition without id")?;
        id.hash(hasher);
        Ok(())
    }

    // TODO (torozco): extra_memory()?
}

#[derive(Display, ProvidesStaticType, Allocative)]
#[display(fmt = "{}", id)]
pub struct FrozenTransitiveSetDefinition {
    /// The name of this transitive set. This is filed in by `export_as` when it's assigned to a
    /// top-level variable. This must be set before this is used.
    id: Arc<TransitiveSetId>,

    operations: TransitiveSetOperationsGen<FrozenValue>,
}

impl fmt::Debug for FrozenTransitiveSetDefinition {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "TransitiveSetDefinition({} declared in {})",
            self.id.name, self.id.module_id
        )
    }
}

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

#[starlark_value(type = "transitive_set_definition")]
impl<'v> StarlarkValue<'v> for FrozenTransitiveSetDefinition {
    fn dir_attr(&self) -> Vec<String> {
        vec!["type".to_owned()]
    }

    fn has_attr(&self, attribute: &str, _heap: &'v Heap) -> bool {
        attribute == "type"
    }

    fn get_attr(&self, attribute: &str, heap: &'v Heap) -> Option<Value<'v>> {
        if attribute == "type" {
            let typ = self.id.name.as_str();
            Some(heap.alloc(typ))
        } else {
            None
        }
    }

    fn write_hash(&self, hasher: &mut StarlarkHasher) -> anyhow::Result<()> {
        self.id.hash(hasher);
        Ok(())
    }
}

starlark_simple_value!(FrozenTransitiveSetDefinition);

impl<'v> Freeze for TransitiveSetDefinition<'v> {
    type Frozen = FrozenTransitiveSetDefinition;

    fn freeze(self, freezer: &Freezer) -> anyhow::Result<Self::Frozen> {
        let Self {
            id,
            module_id: _,
            operations,
        } = self;

        let id = match id.into_inner() {
            Some(x) => x,
            None => {
                // Unfortunately we have no name or location for the definition at this point.
                return Err(TransitiveSetError::TransitiveSetNotAssigned.into());
            }
        };

        let operations = operations.freeze(freezer)?;

        Ok(FrozenTransitiveSetDefinition { id, operations })
    }
}

pub fn transitive_set_definition_from_value<'v>(
    x: Value<'v>,
) -> Option<&dyn TransitiveSetDefinitionLike<'v>> {
    if let Some(x) = x.downcast_ref::<TransitiveSetDefinition>() {
        Some(x as &dyn TransitiveSetDefinitionLike<'v>)
    } else if let Some(x) = x.downcast_ref::<FrozenTransitiveSetDefinition>() {
        Some(x as &dyn TransitiveSetDefinitionLike<'v>)
    } else {
        None
    }
}

pub trait TransitiveSetDefinitionLike<'v> {
    fn has_id(&self) -> bool;

    fn as_debug(&self) -> &dyn fmt::Debug;

    fn matches_type(&self, ty: &str) -> bool;

    fn operations(&self) -> &TransitiveSetOperations<'v>;
}

impl<'v> TransitiveSetDefinitionLike<'v> for TransitiveSetDefinition<'v> {
    fn has_id(&self) -> bool {
        Self::has_id(self)
    }

    fn as_debug(&self) -> &dyn fmt::Debug {
        self
    }

    fn matches_type(&self, ty: &str) -> bool {
        self.id.borrow().as_ref().map_or(false, |id| id.name == ty)
    }

    fn operations(&self) -> &TransitiveSetOperations<'v> {
        &self.operations
    }
}

impl<'v> TransitiveSetDefinitionLike<'v> for FrozenTransitiveSetDefinition {
    fn has_id(&self) -> bool {
        true
    }

    fn as_debug(&self) -> &dyn fmt::Debug {
        self
    }

    fn matches_type(&self, ty: &str) -> bool {
        self.id.name == ty
    }

    fn operations(&self) -> &TransitiveSetOperations<'v> {
        coerce(&self.operations)
    }
}


/*
 * 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::hash_map::Entry;
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;

use anyhow::Context as _;
use fnv::FnvBuildHasher;
use starlark::values::Value;
use starlark::values::ValueIdentity;
use starlark::values::ValueLike;

use crate::interpreter::rule_defs::transitive_set::transitive_set::NodeGen;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetGen;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetLike;

pub trait TransitiveSetIteratorLike<'a, 'v, V>: Iterator<Item = &'a TransitiveSetGen<V>>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    fn values(self: Box<Self>) -> TransitiveSetValuesIteratorGen<'a, 'v, V>;
}

fn assert_transitive_set<'v, V>(child: Value<'v>) -> &'v TransitiveSetGen<V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
{
    TransitiveSetLike::from_value(child)
        .with_context(|| {
            format!(
                "Invalid set: expected {:?}, got: {:?}",
                std::any::type_name::<V>(),
                child
            )
        })
        .unwrap()
}

/// A DFS, left-to-right iterator over a TransitiveSet.
pub struct PreorderTransitiveSetIteratorGen<'a, 'v, V> {
    stack: Vec<&'a TransitiveSetGen<V>>,
    seen: HashSet<ValueIdentity<'v>, FnvBuildHasher>,
}

impl<'a, 'v, V> PreorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    pub fn new(set: &'a TransitiveSetGen<V>) -> Self {
        Self {
            stack: vec![set],
            seen: Default::default(),
        }
    }

    fn enqueue_children(&mut self, children: &'a [V]) {
        for child in children.iter().rev() {
            let child = child.to_value();

            if self.seen.insert(child.identity()) {
                self.stack.push(assert_transitive_set(child));
            }
        }
    }
}

impl<'a, 'v, V> TransitiveSetIteratorLike<'a, 'v, V> for PreorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    fn values(self: Box<Self>) -> TransitiveSetValuesIteratorGen<'a, 'v, V> {
        TransitiveSetValuesIteratorGen { inner: self }
    }
}

/// An iterator over values of a TransitiveSet. Notionally a FilterMap, but defined as its own type
/// since there are a few too many lifetimes involved to make a nice `impl Iterator<...>` work
/// here.
impl<'a, 'v, V> Iterator for PreorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    type Item = &'a TransitiveSetGen<V>;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.stack.pop()?;
        self.enqueue_children(&next.children);
        Some(next)
    }
}

/// A postorder traversal iterator over a TransitiveSet.
/// Traverses by children left-to-right, and then visits the current node.
pub struct PostorderTransitiveSetIteratorGen<'a, 'v, V> {
    stack: Vec<Option<&'a TransitiveSetGen<V>>>,
    parent_stack: Vec<&'a TransitiveSetGen<V>>,
    seen: HashSet<ValueIdentity<'v>, FnvBuildHasher>,
}

impl<'a, 'v, V> PostorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    pub fn new(set: &'a TransitiveSetGen<V>) -> Self {
        Self {
            stack: vec![Some(set)],
            parent_stack: vec![],
            seen: Default::default(),
        }
    }

    fn enqueue_children(&mut self, children: &'a [V]) {
        for child in children.iter().rev() {
            let child = child.to_value();

            if self.seen.insert(child.identity()) {
                self.stack.push(Some(assert_transitive_set(child)));
            }
        }
    }
}

impl<'a, 'v, V> TransitiveSetIteratorLike<'a, 'v, V>
    for PostorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    fn values(self: Box<Self>) -> TransitiveSetValuesIteratorGen<'a, 'v, V> {
        TransitiveSetValuesIteratorGen { inner: self }
    }
}

/// An iterator over values of a TransitiveSet. Notionally a FilterMap, but defined as its own type
/// since there are a few too many lifetimes involved to make a nice `impl Iterator<...>` work
/// here.
impl<'a, 'v, V> Iterator for PostorderTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    type Item = &'a TransitiveSetGen<V>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(next) = self.stack.pop()? {
            if next.children.is_empty() {
                return Some(next);
            }

            self.stack.push(None); // Add sentinel value to indicate end of the parent list.
            self.parent_stack.push(next);
            self.enqueue_children(&next.children);
        }

        // Found a sentinel value indicating children are traversed, return parent.
        self.parent_stack.pop()
    }
}

/// A topological traversal iterator over a TransitiveSet, such that nodes are listed after all
/// nodes that have them as descendants.
///
/// This is equivalent to a pre-order traversal, except that when nodes are shared with more than
/// one parent it is returned in the order of its last occurrence.
pub struct TopologicalTransitiveSetIteratorGen<'a, 'v, V> {
    output_stack: Vec<&'a TransitiveSetGen<V>>,
    instance_counts: HashMap<ValueIdentity<'v>, u32, FnvBuildHasher>,
}

impl<'a, 'v, V> TopologicalTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    pub fn new(set: &'a TransitiveSetGen<V>) -> Self {
        Self {
            output_stack: vec![set],
            instance_counts: TopologicalTransitiveSetIteratorGen::count_instances(set),
        }
    }

    fn count_instances(
        set: &'a TransitiveSetGen<V>,
    ) -> HashMap<ValueIdentity<'v>, u32, FnvBuildHasher> {
        let mut stack = vec![set];
        let mut instance_counts = HashMap::<ValueIdentity<'v>, u32, FnvBuildHasher>::default();

        while let Some(next) = stack.pop() {
            for child in next.children.iter().rev() {
                let child = child.to_value();

                match instance_counts.entry(child.identity()) {
                    Entry::Occupied(mut o) => {
                        *o.get_mut() += 1;
                    }
                    Entry::Vacant(v) => {
                        v.insert(1);
                        stack.push(assert_transitive_set(child));
                    }
                }
            }
        }

        instance_counts
    }

    fn enqueue_children(&mut self, children: &'a [V]) {
        for child in children.iter().rev() {
            let child = child.to_value();

            // It's safe to unwrap since instance_counts is populated during construction and contains
            // all nodes in the tree. `unwrap()` would only fail if the tree was modified.
            let count: &mut u32 = self.instance_counts.get_mut(&child.identity()).unwrap();

            // If this fails, the tree either contains cycles or was modified after construction.
            assert!(*count > 0, "Unexpected node when traversing tree");

            if *count == 1 {
                // Push the last occurrence of the node onto the output stack.
                self.output_stack.push(assert_transitive_set(child));
            }

            *count -= 1;
        }
    }
}

impl<'a, 'v, V> TransitiveSetIteratorLike<'a, 'v, V>
    for TopologicalTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    fn values(self: Box<Self>) -> TransitiveSetValuesIteratorGen<'a, 'v, V> {
        TransitiveSetValuesIteratorGen { inner: self }
    }
}

/// An iterator over values of a TransitiveSet. Notionally a FilterMap, but defined as its own type
/// since there are a few too many lifetimes involved to make a nice `impl Iterator<...>` work
/// here.
impl<'a, 'v, V> Iterator for TopologicalTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    type Item = &'a TransitiveSetGen<V>;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.output_stack.pop()?;
        self.enqueue_children(&next.children);
        Some(next)
    }
}

/// A breadth-first-search (BFS), left-to-right iterator over a TransitiveSet.
pub struct BfsTransitiveSetIteratorGen<'a, 'v, V> {
    queue: VecDeque<&'a TransitiveSetGen<V>>,
    seen: HashSet<ValueIdentity<'v>, FnvBuildHasher>,
}

impl<'a, 'v, V> BfsTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    pub fn new(set: &'a TransitiveSetGen<V>) -> Self {
        Self {
            queue: VecDeque::from(vec![set]),
            seen: Default::default(),
        }
    }

    fn enqueue_children(&mut self, children: &'a [V]) {
        for child in children.iter() {
            let child = child.to_value();

            if self.seen.insert(child.identity()) {
                self.queue.push_back(assert_transitive_set(child));
            }
        }
    }
}

impl<'a, 'v, V> TransitiveSetIteratorLike<'a, 'v, V> for BfsTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    fn values(self: Box<Self>) -> TransitiveSetValuesIteratorGen<'a, 'v, V> {
        TransitiveSetValuesIteratorGen { inner: self }
    }
}

/// An iterator over values of a TransitiveSet. Notionally a FilterMap, but defined as its own type
/// since there are a few too many lifetimes involved to make a nice `impl Iterator<...>` work
/// here.
impl<'a, 'v, V> Iterator for BfsTransitiveSetIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    type Item = &'a TransitiveSetGen<V>;

    fn next(&mut self) -> Option<Self::Item> {
        let next = self.queue.pop_front()?;
        self.enqueue_children(&next.children);
        Some(next)
    }
}

pub struct TransitiveSetValuesIteratorGen<'a, 'v, V> {
    inner: Box<dyn TransitiveSetIteratorLike<'a, 'v, V> + 'a>,
}

impl<'a, 'v, V> Iterator for TransitiveSetValuesIteratorGen<'a, 'v, V>
where
    V: 'v + Copy + ValueLike<'v>,
    TransitiveSetGen<V>: TransitiveSetLike<'v>,
    'v: 'a,
{
    type Item = &'a NodeGen<V>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            let next = self.inner.next()?;
            if let Some(node) = next.node.as_ref() {
                return Some(node);
            }
        }
    }
}


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

use std::fmt;
use std::fmt::Display;
use std::iter;

use allocative::Allocative;
use anyhow::Context as _;
use display_container::display_pair;
use display_container::fmt_container;
use display_container::iter_display_chain;
use dupe::Dupe;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::values::starlark_value;
use starlark::values::Freeze;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::StringValue;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::values::ValueOf;

use crate::artifact_groups::TransitiveSetProjectionKey;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetOrdering;
use crate::interpreter::rule_defs::transitive_set::traversal::TransitiveSetProjectionTraversal;
use crate::interpreter::rule_defs::transitive_set::TransitiveSet;

/// TransitiveSetJsonProjection is the starlark value returned from the starlark method `transitive_set.project_as_json()`
///
/// The projected values are all stored on the TransitiveSet itself and this value will reference back to that. The main
/// point of this object is to provide a distinct value that the write_json implementation understands so that the value
/// can be passed to that.
#[derive(Debug, Clone, Coerce, Trace, Freeze, ProvidesStaticType, Allocative)]
#[derive(NoSerialize)] // TODO we should probably have a serialization for transitive set
#[repr(C)]
pub struct TransitiveSetJsonProjectionGen<V> {
    pub(super) transitive_set: V,

    /// The index of the projection. Once transitive sets are defined, their projections never
    /// change, so we can afford to just store the index here.
    pub projection: usize,

    /// The ordering to use when traversing the projection.
    pub ordering: TransitiveSetOrdering,
}

impl<'v, V: ValueLike<'v>> Display for TransitiveSetJsonProjectionGen<V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let projection_name = self.projection_name().unwrap_or("<invalid projection>");
        fmt_container(
            f,
            "TransitiveSetProjection(",
            ")",
            iter_display_chain(
                iter::once(projection_name),
                iter::once(display_pair("transitive_set", "=", &self.transitive_set)),
            ),
        )
    }
}

impl<'v, V: ValueLike<'v>> TransitiveSetJsonProjectionGen<V> {
    fn projection_name(&self) -> anyhow::Result<&'v str> {
        TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?
            .projection_name(self.projection)
    }

    pub(crate) fn to_projection_key(&self) -> anyhow::Result<TransitiveSetProjectionKey> {
        let set = TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?;

        Ok(TransitiveSetProjectionKey {
            key: set.key().dupe(),
            projection: self.projection,
        })
    }
}

impl<'v, V: ValueLike<'v>> TransitiveSetJsonProjectionGen<V> {
    pub fn iter_values<'a>(&'a self) -> anyhow::Result<Box<dyn Iterator<Item = Value<'v>> + 'a>>
    where
        'v: 'a,
    {
        let set = TransitiveSet::from_value(self.transitive_set.to_value())
            .context("Invalid transitive_set")?;
        set.iter_projection_values(self.ordering, self.projection)
    }
}

starlark_complex_value!(pub TransitiveSetJsonProjection);

#[starlark_value(type = "transitive_set_json_projection")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for TransitiveSetJsonProjectionGen<V>
where
    Self: ProvidesStaticType<'v>,
{
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(transitive_set_json_projection_methods)
    }
}

#[starlark_module]
fn transitive_set_json_projection_methods(builder: &mut MethodsBuilder) {
    fn traverse<'v>(
        this: ValueOf<'v, &'v TransitiveSetJsonProjection<'v>>,
        heap: &'v Heap,
    ) -> anyhow::Result<Value<'v>> {
        Ok(heap.alloc(TransitiveSetProjectionTraversal {
            transitive_set: this.typed.transitive_set,
            projection: this.typed.projection,
            ordering: this.typed.ordering,
        }))
    }

    #[starlark(attribute)]
    fn projection_name<'v>(
        this: ValueOf<'v, &'v TransitiveSetJsonProjection<'v>>,
        heap: &'v Heap,
    ) -> anyhow::Result<StringValue<'v>> {
        Ok(heap.alloc_str(this.typed.projection_name()?))
    }

    #[starlark(attribute)]
    fn transitive_set<'v>(
        this: ValueOf<'v, &'v TransitiveSetJsonProjection<'v>>,
    ) -> anyhow::Result<Value<'v>> {
        Ok(this.typed.transitive_set)
    }
}


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

use allocative::Allocative;
use anyhow::Context as _;
use derive_more::Display;
use dupe::Dupe;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::values::starlark_value;
use starlark::values::Freeze;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;

use crate::interpreter::rule_defs::transitive_set::TransitiveSet;
use crate::interpreter::rule_defs::transitive_set::TransitiveSetError;

#[derive(Debug, Clone, Dupe, Copy, Trace, Freeze, PartialEq, Allocative)]
pub enum TransitiveSetOrdering {
    /// Preorder traversal, a good default behavior which traverses depth-first returning the current node, and then its children left-to-right.
    Preorder,
    /// Postorder traversal, which traverses depth-first, first for the children left-to-right, recursively, then returning the node itself.
    Postorder,
    /// Topological sort, such that nodes are listed after all nodes that have them as descendants.
    Topological,
    /// Breadth-first search traversal.
    Bfs,
}

impl TransitiveSetOrdering {
    pub fn parse(s: &str) -> anyhow::Result<TransitiveSetOrdering> {
        // NOTE: If this list is updated, update the OrderingUnexpectedValue error text.
        match s {
            "preorder" => Ok(Self::Preorder),
            "postorder" => Ok(Self::Postorder),
            "topological" => Ok(Self::Topological),
            "bfs" => Ok(Self::Bfs),
            _ => Err(anyhow::anyhow!(
                TransitiveSetError::OrderingUnexpectedValue {
                    ordering: s.to_owned()
                }
            )),
        }
    }
}

#[derive(
    Debug,
    Clone,
    Trace,
    Coerce,
    Freeze,
    Display,
    ProvidesStaticType,
    NoSerialize,
    Allocative
)]
#[display(fmt = "Traversal({})", inner)]
#[repr(C)]
pub struct TransitiveSetTraversalGen<V> {
    pub(super) inner: V,
    pub ordering: TransitiveSetOrdering,
}

starlark_complex_value!(pub TransitiveSetTraversal);

#[starlark_value(type = "transitive_set_iterator")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for TransitiveSetTraversalGen<V>
where
    Self: ProvidesStaticType<'v>,
{
    fn iterate_collect(&self, _heap: &'v Heap) -> anyhow::Result<Vec<Value<'v>>> {
        let tset = TransitiveSet::from_value(self.inner.to_value()).context("Invalid inner")?;
        Ok(tset.iter_values(self.ordering)?.collect())
    }
}

/// The type returned from .traverse() on a tset projection. This is shared by multiple types of
/// projection (as the traversals all just iterate over the projected values).
#[derive(
    Debug,
    Clone,
    Trace,
    Coerce,
    Freeze,
    Display,
    ProvidesStaticType,
    NoSerialize,
    Allocative
)]
#[display(fmt = "Traversal({}[\"{}\"])", transitive_set, projection)]
#[repr(C)]
pub struct TransitiveSetProjectionTraversalGen<V> {
    pub(super) transitive_set: V,
    pub projection: usize,
    pub ordering: TransitiveSetOrdering,
}

starlark_complex_value!(pub TransitiveSetProjectionTraversal);

#[starlark_value(type = "transitive_set_args_projection_iterator")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for TransitiveSetProjectionTraversalGen<V>
where
    Self: ProvidesStaticType<'v>,
{
    fn iterate_collect(&self, _heap: &'v Heap) -> anyhow::Result<Vec<Value<'v>>> {
        let set =
            TransitiveSet::from_value(self.transitive_set.to_value()).context("Invalid inner")?;
        Ok(set
            .iter_projection_values(self.ordering, self.projection)?
            .collect())
    }
}


/*
 * 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_artifact::artifact::artifact_type::Artifact;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_util::late_binding::LateBinding;
use dice::DiceComputations;
use indexmap::IndexMap;

/// Used by `audit classpath`.
pub static CLASSPATH_FOR_TARGETS: LateBinding<
    for<'c> fn(
        &'c DiceComputations,
        Vec<ConfiguredTargetLabel>,
    ) -> Pin<
        Box<
            dyn Future<Output = anyhow::Result<IndexMap<ConfiguredTargetLabel, Artifact>>>
                + Send
                + 'c,
        >,
    >,
> = LateBinding::new("CLASSPATH_FOR_TARGETS");