/*
 * 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::query::bxl::BxlCqueryFunctions;
use buck2_build_api::query::bxl::NEW_BXL_CQUERY_FUNCTIONS;
use buck2_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_build_api::query::oneshot::QUERY_FRONTEND;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_query::query::syntax::simple::eval::set::TargetSetExt;
use buck2_query::query::syntax::simple::functions::helpers::CapturedExpr;
use derivative::Derivative;
use derive_more::Display;
use dupe::Dupe;
use gazebo::prelude::*;
use starlark::any::ProvidesStaticType;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::eval::Evaluator;
use starlark::starlark_module;
use starlark::values::none::NoneOr;
use starlark::values::starlark_value;
use starlark::values::AllocValue;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueError;
use starlark::StarlarkDocs;

use crate::bxl::starlark_defs::context::BxlContext;
use crate::bxl::starlark_defs::file_set::FileSetExpr;
use crate::bxl::starlark_defs::file_set::StarlarkFileSet;
use crate::bxl::starlark_defs::query_util::parse_query_evaluation_result;
use crate::bxl::starlark_defs::target_expr::filter_incompatible;
use crate::bxl::starlark_defs::target_expr::TargetExpr;
use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;
use crate::bxl::starlark_defs::uquery::unpack_unconfigured_query_args;
use crate::bxl::value_as_starlark_target_label::ValueAsStarlarkTargetLabel;

#[derive(
    ProvidesStaticType,
    Derivative,
    Display,
    Trace,
    NoSerialize,
    StarlarkDocs,
    Allocative
)]
#[starlark_docs(directory = "bxl")]
#[derivative(Debug)]
#[display(fmt = "{:?}", self)]
#[allocative(skip)]
pub struct StarlarkCQueryCtx<'v> {
    #[trace(unsafe_ignore)]
    #[derivative(Debug = "ignore")]
    ctx: &'v BxlContext<'v>,
    #[derivative(Debug = "ignore")]
    target_platform: Option<TargetLabel>,
}

#[starlark_value(type = "cqueryctx", StarlarkTypeRepr, UnpackValue)]
impl<'v> StarlarkValue<'v> for StarlarkCQueryCtx<'v> {
    fn get_methods() -> Option<&'static Methods> {
        static RES: MethodsStatic = MethodsStatic::new();
        RES.methods(register_cquery)
    }
}

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

pub(crate) async fn get_cquery_env<'v>(
    ctx: &'v BxlContext<'v>,
    target_platform: Option<TargetLabel>,
) -> anyhow::Result<Box<dyn BxlCqueryFunctions<'v> + 'v>> {
    (NEW_BXL_CQUERY_FUNCTIONS.get()?)(
        ctx.async_ctx.0,
        target_platform,
        ctx.project_root().dupe(),
        ctx.cell_name,
    )
    .await
}

impl<'v> StarlarkCQueryCtx<'v> {
    pub async fn new(
        ctx: &'v BxlContext<'v>,
        global_target_platform: Value<'v>,
        default_target_platform: &Option<TargetLabel>,
    ) -> anyhow::Result<StarlarkCQueryCtx<'v>> {
        let target_platform = global_target_platform.parse_target_platforms(
            &ctx.target_alias_resolver,
            &ctx.cell_resolver,
            ctx.cell_name,
            default_target_platform,
        )?;

        Ok(Self {
            ctx,
            target_platform,
        })
    }
}

/// The context for performing `cquery` operations in bxl. The functions offered on this ctx are
/// the same behaviour as the query functions available within cquery command.
///
/// Query results are `[StarlarkTargetSet]`s of `[ConfiguredTargetNod]`s, which supports iteration,
/// indexing, `len()`, set addition/subtraction, and `equals()`.
#[starlark_module]
fn register_cquery(builder: &mut MethodsBuilder) {
    /// The `allpaths` query for computing all dependency paths.
    fn allpaths<'v>(
        this: &StarlarkCQueryCtx<'v>,
        from: Value<'v>,
        to: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx.async_ctx.via(|| async {
            get_cquery_env(this.ctx, this.target_platform.dupe())
                .await?
                .allpaths(
                    &filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            from,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?,
                    &filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            to,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?,
                )
                .await
                .map(StarlarkTargetSet::from)
        })
    }

    // The somepath query.
    fn somepath<'v>(
        this: &StarlarkCQueryCtx<'v>,
        from: Value<'v>,
        to: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx.async_ctx.via(|| async {
            get_cquery_env(this.ctx, this.target_platform.dupe())
                .await?
                .somepath(
                    &filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            from,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?,
                    &filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            to,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?,
                )
                .await
                .map(StarlarkTargetSet::from)
        })
    }

    /// The attrfilter query for rule attribute filtering.
    fn attrfilter<'v>(
        this: &StarlarkCQueryCtx<'v>,
        attr: &str,
        value: &str,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx.async_ctx.via(|| async {
            filter_incompatible(
                TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                    targets,
                    &this.target_platform,
                    this.ctx,
                    eval,
                )
                .await?
                .get(this.ctx.async_ctx.0)
                .await?
                .into_iter(),
                this.ctx,
            )?
            .attrfilter(attr, &|v| Ok(v == value))
            .map(StarlarkTargetSet::from)
        })
    }

    /// The kind query for filtering targets by rule type.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_kind(ctx):
    ///     kind = ctx.cquery().kind(".*1", "bin/kind/...")
    ///     ctx.output.print(kind)
    /// ```
    fn kind<'v>(
        this: &StarlarkCQueryCtx<'v>,
        regex: &str,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx.async_ctx.via(|| async {
            filter_incompatible(
                TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                    targets,
                    &this.target_platform,
                    this.ctx,
                    eval,
                )
                .await?
                .get(this.ctx.async_ctx.0)
                .await?
                .into_iter(),
                this.ctx,
            )?
            .kind(regex)
            .map(StarlarkTargetSet::from)
        })
    }

    /// The attrregexfilter query for rule attribute filtering with regex.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_attrregexfilter(ctx):
    ///     filtered = ctx.cquery().attrregexfilter("foo", "he.lo", "bin/kind/...")
    ///     ctx.output.print(filtered)
    /// ```
    fn attrregexfilter<'v>(
        this: &StarlarkCQueryCtx<'v>,
        attribute: &str,
        value: &str,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx.async_ctx.via(|| async {
            filter_incompatible(
                TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                    targets,
                    &this.target_platform,
                    this.ctx,
                    eval,
                )
                .await?
                .get(this.ctx.async_ctx.0)
                .await?
                .into_iter(),
                this.ctx,
            )?
            .attrregexfilter(attribute, value)
            .map(StarlarkTargetSet::from)
        })
    }

    /// The owner query for finding targets that own specified files.
    ///
    /// Sample usage:
    /// ```text
    /// def _owner_impl(ctx):
    ///     owner = ctx.cquery().owner("bin/TARGETS.fixture")
    ///     ctx.output.print(owner)
    /// ```
    fn owner<'v>(
        this: &StarlarkCQueryCtx,
        files: FileSetExpr,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                get_cquery_env(this.ctx, this.target_platform.dupe())
                    .await?
                    .owner(files.get(this.ctx).await?.as_ref())
                    .await
            })
            .map(StarlarkTargetSet::from)
    }

    /// The deps query for finding the transitive closure of dependencies.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_deps(ctx):
    ///     result = ctx.cquery().deps("root//bin:the_binary", 1)
    ///     ctx.output.print(result)
    /// ```
    fn deps<'v>(
        this: &StarlarkCQueryCtx<'v>,
        universe: Value<'v>,
        #[starlark(default = NoneOr::None)] depth: NoneOr<i32>,
        #[starlark(default = NoneOr::None)] filter: NoneOr<&'v str>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                let filter = filter
                    .into_option()
                    .try_map(buck2_query_parser::parse_expr)?;

                get_cquery_env(this.ctx, this.target_platform.dupe())
                    .await?
                    .deps(
                        &filter_incompatible(
                            TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                                universe,
                                &this.target_platform,
                                this.ctx,
                                eval,
                            )
                            .await?
                            .get(this.ctx.async_ctx.0)
                            .await?
                            .into_iter(),
                            this.ctx,
                        )?,
                        depth.into_option(),
                        filter
                            .as_ref()
                            .map(|span| CapturedExpr { expr: span })
                            .as_ref(),
                    )
                    .await
            })
            .map(StarlarkTargetSet::from)
    }

    /// The filter query for filtering targets by name.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_filter(ctx):
    ///     result = ctx.cquery().filter(".*the_binary", "root//...")
    ///     ctx.output.print(result)
    /// ```
    fn filter<'v>(
        this: &StarlarkCQueryCtx<'v>,
        regex: &str,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                filter_incompatible(
                    TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                        targets,
                        &this.target_platform,
                        this.ctx,
                        eval,
                    )
                    .await?
                    .get(this.ctx.async_ctx.0)
                    .await?
                    .into_iter(),
                    this.ctx,
                )?
                .filter_name(regex)
            })
            .map(StarlarkTargetSet::from)
    }

    /// The inputs query for finding input files.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_inputs(ctx):
    ///     result = ctx.cquery().inputs("root//bin:the_binary")
    ///     ctx.output.print(result)
    /// ```
    fn inputs<'v>(
        this: &StarlarkCQueryCtx<'v>,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkFileSet> {
        this.ctx
            .async_ctx
            .via(|| async {
                filter_incompatible(
                    TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                        targets,
                        &this.target_platform,
                        this.ctx,
                        eval,
                    )
                    .await?
                    .get(this.ctx.async_ctx.0)
                    .await?
                    .into_iter(),
                    this.ctx,
                )?
                .inputs()
            })
            .map(StarlarkFileSet::from)
    }

    /// The testsof query for listing the tests of the specified targets.
    fn testsof<'v>(
        this: &StarlarkCQueryCtx<'v>,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                get_cquery_env(this.ctx, this.target_platform.dupe())
                    .await?
                    .testsof(&filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            targets,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?)
                    .await
            })
            .map(StarlarkTargetSet::from)
    }

    /// The testsof query for listing the tests of the specified targets. Performs default target platform
    /// resolution under the hood for the tests found.
    fn testsof_with_default_target_platform<'v>(
        this: &StarlarkCQueryCtx<'v>,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                let maybe_compatibles = get_cquery_env(this.ctx, this.target_platform.dupe())
                    .await?
                    .testsof_with_default_target_platform(&filter_incompatible(
                        TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                            targets,
                            &this.target_platform,
                            this.ctx,
                            eval,
                        )
                        .await?
                        .get(this.ctx.async_ctx.0)
                        .await?
                        .into_iter(),
                        this.ctx,
                    )?)
                    .await?;

                filter_incompatible(maybe_compatibles.into_iter(), this.ctx)
            })
            .map(StarlarkTargetSet::from)
    }

    /// The rdeps query for finding the transitive closure of reverse dependencies.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_rdeps(ctx):
    ///     result = ctx.cquery().rdeps("root//bin:the_binary", "//lib:file1", 100)
    ///     ctx.output.print(result)
    /// ```
    fn rdeps<'v>(
        this: &StarlarkCQueryCtx<'v>,
        universe: Value<'v>,
        from: Value<'v>,
        depth: Option<i32>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        this.ctx
            .async_ctx
            .via(|| async {
                get_cquery_env(this.ctx, this.target_platform.dupe())
                    .await?
                    .rdeps(
                        &filter_incompatible(
                            TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                                universe,
                                &this.target_platform,
                                this.ctx,
                                eval,
                            )
                            .await?
                            .get(this.ctx.async_ctx.0)
                            .await?
                            .into_iter(),
                            this.ctx,
                        )?,
                        &filter_incompatible(
                            TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                                from,
                                &this.target_platform,
                                this.ctx,
                                eval,
                            )
                            .await?
                            .get(this.ctx.async_ctx.0)
                            .await?
                            .into_iter(),
                            this.ctx,
                        )?,
                        depth,
                    )
                    .await
            })
            .map(StarlarkTargetSet::from)
    }

    /// Evaluates some general query string. `query_args` can be a target_set of unconfigured nodes, or
    /// a list of strings.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_eval(ctx):
    ///     result1 = ctx.cquery().eval("inputs(root//bin:the_binary)")
    ///     ctx.output.print(result1)
    ///
    ///     result2 = ctx.cquery().eval("inputs(%s)", query_args = ["cell//path/to/file:target"])
    ///     ctx.output.print(result2)
    /// ```
    fn eval<'v>(
        this: &StarlarkCQueryCtx<'v>,
        query: &'v str,
        #[starlark(default = NoneOr::None)] query_args: NoneOr<Value<'v>>,
        #[starlark(default = NoneOr::None)] target_universe: NoneOr<Vec<String>>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Value<'v>> {
        let query_args = if query_args.is_none() {
            Vec::new()
        } else {
            let unwrapped_query_args = query_args.into_option().unwrap();
            if let Some(query_args) = unpack_unconfigured_query_args(unwrapped_query_args)? {
                query_args
            } else {
                let err = Err(ValueError::IncorrectParameterTypeWithExpected(
                    "list of strings, or a target_set of unconfigured nodes".to_owned(),
                    query_args.into_option().unwrap().get_type().to_owned(),
                )
                .into());

                if <&StarlarkTargetSet<ConfiguredTargetNode>>::unpack_value(unwrapped_query_args)
                    .is_some()
                {
                    return err
                        .context("target_set with configured nodes are currently not supported");
                }

                return err;
            }
        };

        this.ctx.async_ctx.via_dice(|ctx| async {
            parse_query_evaluation_result(
                QUERY_FRONTEND
                    .get()?
                    .eval_cquery(
                        ctx,
                        &this.ctx.working_dir()?,
                        CqueryOwnerBehavior::Deprecated,
                        query,
                        &query_args,
                        this.target_platform.dupe(),
                        target_universe.into_option().as_ref().map(|v| &v[..]),
                    )
                    .await?,
                eval,
            )
        })
    }

    /// Find the build file(s) that defines a target or a target set.
    ///
    /// Sample usage:
    /// ```text
    /// def _buildfile_impl(ctx):
    ///     owner = ctx.cquery().owner(["bin/TARGET", "bin/kind"])
    ///     result = ctx.cquery().buildfile(owner)
    ///     ctx.output.print(result)
    /// ```
    fn buildfile<'v>(
        this: &StarlarkCQueryCtx<'v>,
        targets: Value<'v>,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<StarlarkFileSet> {
        this.ctx
            .async_ctx
            .via(|| async {
                let targets = &filter_incompatible(
                    TargetExpr::<'v, ConfiguredTargetNode>::unpack(
                        targets,
                        &this.target_platform,
                        this.ctx,
                        eval,
                    )
                    .await?
                    .get(this.ctx.async_ctx.0)
                    .await?
                    .into_iter(),
                    this.ctx,
                )?;

                Ok(targets.buildfile())
            })
            .map(StarlarkFileSet::from)
    }
}


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

//! Parse some inputs to a `[`StarlarkUserEvent`].
//!

use std::collections::HashMap;

use buck2_data::StarlarkUserEvent;
use buck2_data::StarlarkUserMetadataValue;
use starlark::values::dict::DictRef;
use starlark::values::UnpackValue;
use starlark::values::Value;
use thiserror::Error;

#[derive(Error, Debug)]
enum StarlarkUserEventUnpack {
    #[error(
        "Metadata should be a dict where keys are strings, and values are strings, ints, or bools. Got type: `{0}`"
    )]
    InvalidMetadata(String),
    #[error("Metadata keys should be strings. Got type: `{0}`")]
    InvalidKey(String),
    #[error("Metadata values should be strings, ints, or bools. Key `{0}` had value type `{1}`")]
    InvalidValue(String, String),
}

pub(crate) fn to_starlark_user_event<'v>(
    id: &str,
    metadata: Value<'v>,
) -> anyhow::Result<StarlarkUserEvent> {
    let metadata_value = match DictRef::from_value(metadata) {
        Some(metadata_value) => metadata_value,
        None => {
            return Err(
                StarlarkUserEventUnpack::InvalidMetadata(metadata.get_type().to_owned()).into(),
            );
        }
    };

    let metadata_map = metadata_value
        .iter()
        .map(|(k, v)| {
            let k = match k.unpack_str() {
                Some(k) => k.to_owned(),
                None => {
                    return Err(StarlarkUserEventUnpack::InvalidKey(k.get_type().to_owned()).into());
                }
            };

            if let Some(v) = v.unpack_str() {
                Ok((
                    k,
                    StarlarkUserMetadataValue {
                        value: Some(
                            buck2_data::starlark_user_metadata_value::Value::StringValue(v.into()),
                        ),
                    },
                ))
            } else if let Some(v) = v.unpack_bool() {
                Ok((
                    k,
                    StarlarkUserMetadataValue {
                        value: Some(buck2_data::starlark_user_metadata_value::Value::BoolValue(
                            v,
                        )),
                    },
                ))
            } else if let Some(v) = v.unpack_i32() {
                Ok((
                    k,
                    StarlarkUserMetadataValue {
                        value: Some(buck2_data::starlark_user_metadata_value::Value::IntValue(v)),
                    },
                ))
            // Let's also accept floats since `instant()` methods return floats, but cast them to ints
            } else if let Some(v) = f64::unpack_value(v) {
                Ok((
                    k,
                    StarlarkUserMetadataValue {
                        value: Some(buck2_data::starlark_user_metadata_value::Value::IntValue(
                            v as i32,
                        )),
                    },
                ))
            } else {
                return Err(
                    StarlarkUserEventUnpack::InvalidValue(k, v.get_type().to_owned()).into(),
                );
            }
        })
        .collect::<anyhow::Result<HashMap<_, _>>>()?;

    Ok(StarlarkUserEvent {
        id: id.to_owned(),
        metadata: metadata_map,
    })
}


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

use buck2_artifact::artifact::source_artifact::SourceArtifact;
use buck2_build_api::interpreter::rule_defs::artifact::starlark_artifact_like::ValueAsArtifactLike;
use buck2_build_api::interpreter::rule_defs::artifact::StarlarkArtifact;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::data::HasIoProvider;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::instance::CellInstance;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::cells::CellAliasResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::pattern::maybe_split_cell_alias_and_relative_path;
use derive_more::Display;
use dupe::Dupe;
use either::Either;
use starlark::typing::Ty;
use starlark::values::type_repr::StarlarkTypeRepr;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueLike;

use crate::bxl::starlark_defs::context::starlark_async::BxlSafeDiceComputations;
use crate::bxl::starlark_defs::file_set::StarlarkFileNode;

/// FileExpr is just a simple type that can be used in starlark_module
/// functions for arguments that should be files. It will accept either a
/// literal (like `//path/to/some/file.txt`) or a `SourceArtifact` via `StarlarkArtifact`
/// or a `StarlarkFileNode`
#[derive(Debug, Display, Clone)]
pub enum FileExpr<'a> {
    Literal(&'a str),
    SourceArtifact(SourceArtifact),
    StarlarkFileNode(StarlarkFileNode),
}

fn parse_cell_path_as_file_expr_literal(
    val: &str,
    cell_alias_resolver: &CellAliasResolver,
) -> anyhow::Result<Option<CellPath>> {
    Ok(match maybe_split_cell_alias_and_relative_path(val)? {
        Some((alias, path)) => {
            let cell_name = cell_alias_resolver.resolve(alias.as_str())?;

            let cell_relative_path = CellRelativePath::new(path);

            Some(CellPath::new(cell_name, cell_relative_path.to_buf()))
        }
        None => None,
    })
}

impl<'a> FileExpr<'a> {
    pub fn get(
        self,
        dice: &BxlSafeDiceComputations<'_>,
        cell_instance: &CellInstance,
    ) -> anyhow::Result<CellPath> {
        match self {
            FileExpr::Literal(val) => {
                match parse_cell_path_as_file_expr_literal(
                    val,
                    cell_instance.cell_alias_resolver(),
                )? {
                    Some(cell_path) => Ok(cell_path),
                    None => {
                        let fs = dice.0.global_data().get_io_provider().project_root().dupe();
                        let path = Path::new(val);
                        let rel = if path.is_absolute() {
                            Cow::Owned(fs.relativize_any(AbsPath::new(path)?)?)
                        } else {
                            Cow::Borrowed(<&ProjectRelativePath>::try_from(val)?)
                        };
                        dice.via_dice(async move |ctx| {
                            ctx.get_cell_resolver().await?.get_cell_path(&rel)
                        })
                    }
                }
            }
            FileExpr::SourceArtifact(val) => Ok(val.get_path().to_cell_path()),
            FileExpr::StarlarkFileNode(val) => Ok(val.0),
        }
    }
}

impl<'v> UnpackValue<'v> for FileExpr<'v> {
    fn expected() -> String {
        "literal, or source artifact, or file node".to_owned()
    }

    fn unpack_value(value: Value<'v>) -> Option<Self> {
        if let Some(s) = value.unpack_str() {
            Some(FileExpr::Literal(s))
        } else if let Some(v) = value.downcast_ref::<StarlarkFileNode>() {
            Some(FileExpr::StarlarkFileNode(v.to_owned()))
        } else if let Some(v) = ValueAsArtifactLike::unpack_value(value) {
            v.0.get_bound_artifact()
                .map(|a| a.get_source())
                .ok()?
                .map(FileExpr::SourceArtifact)
        } else {
            None
        }
    }
}

impl<'v> StarlarkTypeRepr for FileExpr<'v> {
    fn starlark_type_repr() -> Ty {
        Either::<String, StarlarkArtifact>::starlark_type_repr()
    }
}

#[cfg(test)]
mod tests {

    use buck2_core::cells::alias::NonEmptyCellAlias;
    use buck2_core::cells::name::CellName;
    use buck2_core::cells::CellAliasResolver;
    use maplit::hashmap;

    use super::*;

    #[test]
    fn test_parse_cell_path_as_file_expr_literal() -> anyhow::Result<()> {
        let cell1 = CellName::testing_new("cell1");

        let map = hashmap![
            NonEmptyCellAlias::new("cell1".to_owned()).unwrap() => CellName::testing_new("cell1"),
            NonEmptyCellAlias::new("cell2".to_owned()).unwrap() => CellName::testing_new("cell2"),
            NonEmptyCellAlias::new("cell3".to_owned()).unwrap() => CellName::testing_new("cell3"),
        ];

        let cell_alias_resolver = CellAliasResolver::new(cell1, map)?;

        let actual_cell_path =
            parse_cell_path_as_file_expr_literal("//foo/bar", &cell_alias_resolver)?.unwrap();
        let expected_cell_path = CellPath::testing_new("cell1//foo/bar");
        assert_eq!(actual_cell_path, expected_cell_path);

        let actual_cell_path =
            parse_cell_path_as_file_expr_literal("@//foo/bar", &cell_alias_resolver)?.unwrap();
        let expected_cell_path = CellPath::testing_new("cell1//foo/bar");
        assert_eq!(actual_cell_path, expected_cell_path);

        let actual_cell_path =
            parse_cell_path_as_file_expr_literal("@cell1//foo/bar", &cell_alias_resolver)?.unwrap();
        let expected_cell_path = CellPath::testing_new("cell1//foo/bar");
        assert_eq!(actual_cell_path, expected_cell_path);

        let actual_cell_path =
            parse_cell_path_as_file_expr_literal("cell2//foo/bar", &cell_alias_resolver)?.unwrap();
        let expected_cell_path = CellPath::testing_new("cell2//foo/bar");
        assert_eq!(actual_cell_path, expected_cell_path);

        let actual_cell_path =
            parse_cell_path_as_file_expr_literal("@cell3//foo/bar", &cell_alias_resolver)?.unwrap();
        let expected_cell_path = CellPath::testing_new("cell3//foo/bar");
        assert_eq!(actual_cell_path, expected_cell_path);

        assert!(
            parse_cell_path_as_file_expr_literal("/abs/path/foo/bar", &cell_alias_resolver)?
                .is_none()
        );
        assert!(
            parse_cell_path_as_file_expr_literal("rel/path/foo/bar", &cell_alias_resolver)?
                .is_none()
        );

        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::borrow::Cow;
use std::fmt;
use std::ops::Deref;
use std::sync::Arc;

use allocative::Allocative;
use buck2_common::file_ops::SimpleDirEntry;
use buck2_core::cells::cell_path::CellPath;
use buck2_query::query::syntax::simple::eval::file_set::FileNode;
use buck2_query::query::syntax::simple::eval::file_set::FileSet;
use derive_more::Display;
use display_container::fmt_container;
use either::Either;
use gazebo::prelude::VecExt;
use indexmap::IndexSet;
use starlark::any::ProvidesStaticType;
use starlark::starlark_simple_value;
use starlark::typing::Ty;
use starlark::values::starlark_value;
use starlark::values::type_repr::StarlarkTypeRepr;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueError;
use starlark::values::ValueLike;
use starlark::StarlarkDocs;

use crate::bxl::starlark_defs::context::BxlContext;

/// FileSetExpr is just a simple type that can be used in starlark_module
/// functions for arguments that should be file sets. It will accept either a
/// literal (like `//path/to/some/file.txt`) or a FileSet Value (from one of the
/// bxl functions that return them).
pub enum FileSetExpr<'a> {
    Literal(&'a str),
    Literals(Vec<&'a str>),
    FileSet(&'a FileSet),
}

impl<'a> FileSetExpr<'a> {
    pub async fn get(self, bxl: &BxlContext<'_>) -> anyhow::Result<Cow<'a, FileSet>> {
        let set = match self {
            FileSetExpr::Literal(val) => Cow::Owned(FileSet::from_iter([FileNode(
                bxl.parse_query_file_literal(val)?,
            )])),
            FileSetExpr::Literals(val) => {
                let mut file_set = FileSet::new(IndexSet::new());
                for arg in val.iter() {
                    file_set.insert(FileNode(bxl.parse_query_file_literal(arg)?));
                }
                Cow::Owned(file_set)
            }
            FileSetExpr::FileSet(val) => Cow::Borrowed(val),
        };
        Ok(set)
    }

    // This will unpack a Value to a FileSetExpr, but doesn't accept as single string literal,
    // only a FileSetExpr or a list of string literals.
    fn unpack_set(value: Value<'a>) -> Option<Self> {
        value
            .downcast_ref::<StarlarkFileSet>()
            .map(|s| FileSetExpr::FileSet(s))
    }
}

impl<'v> StarlarkTypeRepr for FileSetExpr<'v> {
    fn starlark_type_repr() -> Ty {
        Either::<String, StarlarkFileSet>::starlark_type_repr()
    }
}

impl<'v> UnpackValue<'v> for FileSetExpr<'v> {
    fn expected() -> String {
        "literal or set of file names".to_owned()
    }

    fn unpack_value(value: Value<'v>) -> Option<Self> {
        if let Some(s) = value.unpack_str() {
            Some(FileSetExpr::Literal(s))
        } else if let Some(s) = <Vec<&'v str>>::unpack_value(value) {
            Some(FileSetExpr::Literals(s))
        } else {
            FileSetExpr::unpack_set(value)
        }
    }
}

#[derive(Debug, Display, ProvidesStaticType, Allocative, StarlarkDocs)]
#[derive(NoSerialize)] // TODO maybe this should be
#[starlark_docs(directory = "bxl")]
pub struct StarlarkFileSet(
    /// Set of files or directories.
    pub FileSet,
);

starlark_simple_value!(StarlarkFileSet);

#[starlark_value(type = "file_set")]
impl<'v> StarlarkValue<'v> for StarlarkFileSet {
    fn iterate_collect(&self, heap: &'v Heap) -> anyhow::Result<Vec<Value<'v>>> {
        Ok(self
            .0
            .iter()
            .map(|cell_path| heap.alloc(StarlarkFileNode(cell_path.clone())))
            .collect())
    }

    fn at(&self, index: Value<'v>, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        let i = index.unpack_i32().ok_or_else(|| {
            ValueError::IncorrectParameterTypeWithExpected(
                "int".to_owned(),
                index.get_type().to_owned(),
            )
        })?;
        if i < 0 {
            return Err(anyhow::anyhow!(ValueError::IndexOutOfBound(i)));
        }
        self.0
            .get_index(i as usize)
            .ok_or_else(|| anyhow::anyhow!(ValueError::IndexOutOfBound(i)))
            .map(|cell_path| heap.alloc(StarlarkFileNode(cell_path.clone())))
    }

    fn length(&self) -> anyhow::Result<i32> {
        match i32::try_from(self.0.len()) {
            Ok(l) => Ok(l),
            Err(e) => Err(e.into()),
        }
    }
}

impl From<FileSet> for StarlarkFileSet {
    fn from(v: FileSet) -> Self {
        Self(v)
    }
}

impl Deref for StarlarkFileSet {
    type Target = FileSet;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

#[derive(Debug, Display, ProvidesStaticType, Clone, Allocative, StarlarkDocs)]
#[derive(NoSerialize)]
#[starlark_docs(directory = "bxl")]
pub struct StarlarkFileNode(
    /// Cell path to the file or directory.
    pub CellPath,
);

starlark_simple_value!(StarlarkFileNode);

#[starlark_value(type = "file_node")]
impl<'v> StarlarkValue<'v> for StarlarkFileNode {}

#[derive(Debug, ProvidesStaticType, Clone, Allocative, StarlarkDocs)]
#[derive(NoSerialize)]
#[starlark_docs(directory = "bxl")]
pub struct StarlarkReadDirSet {
    /// Cell path to the directory/files.
    pub cell_path: CellPath,
    /// Files that are not ignored within the buckconfig.
    /// Sorted.
    pub included: Arc<[SimpleDirEntry]>,
    /// Only return directories when iterating or printing.
    pub dirs_only: bool,
}

starlark_simple_value!(StarlarkReadDirSet);

impl StarlarkReadDirSet {
    fn children(&self) -> anyhow::Result<Vec<CellPath>> {
        let mut result: Vec<CellPath> = Vec::with_capacity(self.included.len());
        result.extend(self.included.iter().filter_map(|e| {
            if !self.dirs_only || e.file_type.is_dir() {
                Some(self.cell_path.join(&e.file_name))
            } else {
                None
            }
        }));
        result.sort();
        Ok(result)
    }
}

impl fmt::Display for StarlarkReadDirSet {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.children() {
            Ok(children) => fmt_container(f, "[", "]", children),
            Err(e) => write!(f, "<Error: {}>", e),
        }
    }
}

#[starlark_value(type = "read_dir_set")]
impl<'v> StarlarkValue<'v> for StarlarkReadDirSet {
    fn iterate_collect(&self, heap: &'v Heap) -> anyhow::Result<Vec<Value<'v>>> {
        Ok(self
            .children()?
            .into_map(|cell_path| heap.alloc(StarlarkFileNode(cell_path))))
    }
}


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

use anyhow::Context;
use buck2_build_api::interpreter::rule_defs::artifact::StarlarkArtifact;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::NonDefaultProvidersName;
use buck2_core::provider::label::ProviderName;
use buck2_core::provider::label::ProvidersLabel;
use buck2_core::provider::label::ProvidersName;
use buck2_interpreter::types::label::Label;
use buck2_interpreter::types::label::StarlarkProvidersLabel;
use buck2_interpreter::types::target_label::StarlarkConfiguredTargetLabel;
use buck2_interpreter::types::target_label::StarlarkTargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use dupe::Dupe;
use starlark::environment::GlobalsBuilder;
use starlark::starlark_module;
use starlark::values::list::AllocList;
use starlark::values::list::ListRef;
use starlark::values::none::NoneType;
use starlark::values::Heap;
use starlark::values::StringValue;
use starlark::values::Value;
use starlark::values::ValueError;
use starlark::values::ValueLike;
use thiserror::Error;

use super::artifacts::visit_artifact_path_without_associated_deduped;
use super::context::output::get_artifact_path_display;
use super::context::output::get_cmd_line_inputs;
use crate::bxl::starlark_defs::context::BxlContext;
use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;
use crate::bxl::starlark_defs::time::StarlarkInstant;

/// Global methods on the target label.
#[starlark_module]
pub fn register_label_function(builder: &mut GlobalsBuilder) {
    /// Converts a `TargetLabel` into its corresponding `ProvidersLabel` given the subtarget names,
    /// which is a list for each layer of subtarget
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_sub_target(ctx):
    ///     owners = ctx.cquery().owner("bin/TARGETS.fixture")
    ///     for owner in owners:
    ///         configured_label = owner.label
    ///         unconfigured_label = configured_label.raw_target()
    ///         ctx.output.print(sub_target(unconfigured_label))
    ///         ctx.output.print(sub_target(unconfigured_label, "subtarget1"))
    ///         ctx.output.print(sub_target(unconfigured_label, ["subtarget1", "subtarget2"))
    /// ```
    fn sub_target<'v>(
        target: &StarlarkTargetLabel,
        #[starlark(default = AllocList::EMPTY)] subtarget_name: Value<'v>,
    ) -> anyhow::Result<StarlarkProvidersLabel> {
        let providers_name = value_to_providers_name(subtarget_name)?;

        Ok(StarlarkProvidersLabel::new(ProvidersLabel::new(
            target.label().dupe(),
            providers_name,
        )))
    }

    /// Converts a `TargetLabel` into its corresponding `ProvidersLabel` given the subtarget name
    /// which is a list for each layer of subtarget
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_sub_target(ctx):
    ///     owners = ctx.cquery().owner("bin/TARGETS.fixture")
    ///     for owner in owners:
    ///         configured_label = owner.label
    ///         ctx.output.print(configured_sub_target(configured_label))
    ///         ctx.output.print(configured_sub_target(configured_label, "subtarget1"))
    ///         ctx.output.print(configured_sub_target(configured_label, ["subtarget1", "subtarget2"))
    /// ```
    fn configured_sub_target<'v>(
        target: &StarlarkConfiguredTargetLabel,
        #[starlark(default = AllocList::EMPTY)] subtarget_name: Value<'v>,
    ) -> anyhow::Result<Label> {
        let providers_name = value_to_providers_name(subtarget_name)?;

        Ok(Label::new(ConfiguredProvidersLabel::new(
            target.label().dupe(),
            providers_name,
        )))
    }
}

/// Global methods on the target set.
#[starlark_module]
pub fn register_target_function(builder: &mut GlobalsBuilder) {
    /// Creates an empty target set for configured nodes.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_ctarget_set(ctx):
    ///     targets = ctarget_set()
    ///     ctx.output.print(type(targets))
    ///     ctx.output.print(len(targets))
    /// ```
    fn ctarget_set() -> anyhow::Result<StarlarkTargetSet<ConfiguredTargetNode>> {
        Ok(StarlarkTargetSet::from(TargetSet::new()))
    }

    /// Creates an empty target set for unconfigured nodes.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_utarget_set(ctx):
    ///     targets = utarget_set()
    ///     ctx.output.print(type(targets))
    ///     ctx.output.print(len(targets))
    /// ```
    fn utarget_set() -> anyhow::Result<StarlarkTargetSet<TargetNode>> {
        Ok(StarlarkTargetSet::from(TargetSet::new()))
    }
}

/// Global methods on artifacts.
#[starlark_module]
pub fn register_artifact_function(builder: &mut GlobalsBuilder) {
    /// The output path of a source or build artifact. Takes an optional boolean to print the absolute or relative path.
    /// Note that this method returns an artifact path without asking for the artifact to be materialized,
    /// (i.e. it may not actually exist on the disk yet).
    ///
    /// This is a risky function to call because you may accidentally pass this path to further BXL actions
    /// that expect the artifact to be materialized. If this happens, the BXL script will error out.
    /// If you want the path without materialization for other uses that donâ€™t involve passing them into
    /// further actions, then itâ€™s safe.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_get_path_without_materialization(ctx):
    ///     owner = ctx.cquery().owner("cell//path/to/file")[0]
    ///     artifact = owner.get_source("cell//path/to/file", ctx)
    ///     source_artifact_project_rel_path = get_path_without_materialization(artifact, ctx)
    ///     ctx.output.print(source_artifact_project_rel_path) # Note this artifact is NOT ensured or materialized
    /// ```
    fn get_path_without_materialization<'v>(
        this: &'v StarlarkArtifact,
        ctx: &'v BxlContext<'v>,
        #[starlark(require = named, default = false)] abs: bool,
        heap: &'v Heap,
    ) -> anyhow::Result<StringValue<'v>> {
        let path = get_artifact_path_display(
            this.artifact().get_path(),
            abs,
            &ctx.output_stream.project_fs,
            &ctx.output_stream.artifact_fs,
        )?;

        Ok(heap.alloc_str(&path))
    }

    /// The output paths of a `cmd_args()` inputs. The output paths will be returned as a list.
    /// Takes an optional boolean to print the absolute or relative path.
    /// Note that this method returns an artifact path without asking for the artifact to be materialized,
    /// (i.e. it may not actually exist on the disk yet).
    ///
    /// This is a risky function to call because you may accidentally pass this path to further BXL actions
    /// that expect the artifact to be materialized. If this happens, the BXL script will error out.
    /// If you want the path without materialization for other uses that donâ€™t involve passing them into
    /// further actions, then itâ€™s safe.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_get_paths_without_materialization(ctx):
    ///     node = ctx.configured_targets("root//bin:the_binary")
    ///     providers = ctx.analysis(node).providers()
    ///     path = get_paths_without_materialization(providers[RunInfo], abs=True) # Note this artifact is NOT ensured or materialized
    ///     ctx.output.print(path)
    /// ```
    fn get_paths_without_materialization<'v>(
        this: Value<'v>,
        ctx: &'v BxlContext<'v>,
        #[starlark(require = named, default = false)] abs: bool,
        heap: &'v Heap,
    ) -> anyhow::Result<Value<'v>> {
        match this.as_command_line() {
            Some(cmd_line) => {
                let inputs = get_cmd_line_inputs(cmd_line)?;
                let mut result = Vec::new();

                for artifact_group in &inputs.inputs {
                    result.push(artifact_group.dupe());
                }

                let mut paths = Vec::new();

                let _result = ctx.async_ctx.via_dice(|dice_ctx| {
                    visit_artifact_path_without_associated_deduped(
                        &result,
                        abs,
                        |artifact_path, abs| {
                            let path = get_artifact_path_display(
                                artifact_path,
                                abs,
                                &ctx.output_stream.project_fs,
                                &ctx.output_stream.artifact_fs,
                            )?;

                            paths.push(path);
                            Ok(())
                        },
                        dice_ctx,
                    )
                });
                Ok(heap.alloc(paths))
            }
            None => Err(anyhow::anyhow!("Expected a cmd_args()")),
        }
    }
}

/// Global methods for Instant.
#[starlark_module]
pub fn register_instant_function(builder: &mut GlobalsBuilder) {
    /// Creates an Instant at the current time.
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_elapsed_millis(ctx):
    ///     instant = now()
    ///     time_a = instant.elapsed_millis()
    ///     # do something that takes a long time
    ///     time_b = instant.elapsed_millis()
    ///
    ///     ctx.output.print(time_a)
    ///     ctx.output.print(time_b)
    /// ```
    fn now<'v>(heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        Ok(heap.alloc(StarlarkInstant(Instant::now())))
    }
}

/// This is used to mark the error returned by `fail_no_stacktrace()` (via context chaining).
/// We check if this marker is present after finishing BXL evaluation. If this marker is present,
/// then we hide the stacktrace. Otherwise, we emit the stacktrace to users.
#[derive(Debug, Error, Clone)]
#[error("fail:{0}")]
pub(crate) struct BxlErrorWithoutStacktrace(String);

/// Global method for error handling.
#[starlark_module]
pub fn register_error_handling_function(builder: &mut GlobalsBuilder) {
    fn fail_no_stacktrace(#[starlark(args)] args: Vec<Value>) -> anyhow::Result<NoneType> {
        let mut s = String::new();
        for x in args {
            s.push(' ');
            match x.unpack_str() {
                Some(x) => s.push_str(x),
                None => x.collect_repr(&mut s),
            }
        }
        Err(BxlErrorWithoutStacktrace(s).into())
    }
}

fn value_to_providers_name<'v>(subtarget_name: Value<'v>) -> anyhow::Result<ProvidersName> {
    let subtarget = if let Some(list) = ListRef::from_value(subtarget_name) {
        list.iter()
            .map(|name| {
                name.unpack_str()
                    .ok_or_else(|| {
                        anyhow::anyhow!(ValueError::IncorrectParameterTypeNamedWithExpected(
                            "subtarget_name".to_owned(),
                            "list of str or str".to_owned(),
                            name.get_type().to_owned(),
                        ))
                    })
                    .and_then(|name| {
                        ProviderName::new(name.to_owned())
                            .context("for parameter `subtarget_name`")
                            .map_err(|e| anyhow::anyhow!(e))
                    })
            })
            .collect::<anyhow::Result<Vec<_>>>()?
    } else if let Some(str) = subtarget_name.unpack_str() {
        vec![ProviderName::new(str.to_owned()).context("for parameter `subtarget_name`")?]
    } else {
        return Err(anyhow::anyhow!(
            ValueError::IncorrectParameterTypeNamedWithExpected(
                "subtarget_name".to_owned(),
                "list of str or str".to_owned(),
                subtarget_name.get_type().to_owned()
            )
        ));
    };

    Ok(if subtarget.is_empty() {
        ProvidersName::Default
    } else {
        ProvidersName::NonDefault(Box::new(NonDefaultProvidersName::Named(
            subtarget.into_boxed_slice(),
        )))
    })
}


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

//! Definitions of core functionality just for bxl functions to access

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

use allocative::Allocative;
use anyhow::Context;
use buck2_build_api::bxl::types::BxlFunctionLabel;
use buck2_interpreter::build_context::STARLARK_PATH_FROM_BUILD_CONTEXT;
use buck2_interpreter::bxl::BxlFunctions;
use buck2_interpreter::bxl::BXL_FUNCTIONS;
use buck2_interpreter::path::BxlFilePath;
use buck2_util::collections::ordered_map::OrderedMap;
use cli_args::CliArgs;
use derive_more::Display;
use starlark::any::ProvidesStaticType;
use starlark::collections::SmallMap;
use starlark::collections::SmallSet;
use starlark::environment::GlobalsBuilder;
use starlark::eval::Evaluator;
use starlark::starlark_simple_value;
use starlark::values::dict::DictOf;
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::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::Value;
use thiserror::Error;

use crate::bxl::starlark_defs::cli_args::ArgAccessor;
use crate::bxl::starlark_defs::cli_args::CliArgError;
use crate::bxl::starlark_defs::cli_args::CliArgValue;
use crate::bxl::starlark_defs::functions::register_artifact_function;
use crate::bxl::starlark_defs::functions::register_error_handling_function;
use crate::bxl::starlark_defs::functions::register_instant_function;
use crate::bxl::starlark_defs::functions::register_label_function;
use crate::bxl::starlark_defs::functions::register_target_function;
pub mod alloc_node;
pub mod analysis_result;
pub mod artifacts;
pub mod audit;
pub mod build_result;
pub mod cli_args;
pub mod context;
pub mod cquery;
pub mod event;
pub mod file_expr;
pub mod file_set;
pub mod functions;
pub mod nodes;
pub mod providers_expr;
mod query_util;
pub mod target_expr;
pub mod targetset;
pub mod time;
pub mod uquery;
use starlark::starlark_module;

use crate::bxl::eval::CliResolutionCtx;

#[starlark_module]
pub fn register_bxl_function(builder: &mut GlobalsBuilder) {
    fn bxl<'v>(
        #[starlark(require = named)] r#impl: Value<'v>,
        #[starlark(require = named)] cli_args: DictOf<'v, &'v str, &'v CliArgs>,
        #[starlark(require = named, default = "")] doc: &str,
        eval: &mut Evaluator<'v, '_>,
    ) -> anyhow::Result<Value<'v>> {
        let implementation = r#impl;

        let bxl_path = (*(STARLARK_PATH_FROM_BUILD_CONTEXT.get()?)(eval)?
            .unpack_bxl_file()
            .ok_or_else(|| anyhow::anyhow!("`bxl` can only be declared in bxl files"))?)
        .clone();

        let mut unresolved_cli_args = SmallMap::new();
        let mut short_args = SmallSet::new();

        for (arg, def) in cli_args.to_dict().into_iter() {
            if let Some(short) = def.short {
                if short_args.contains(&short) {
                    return Err(CliArgError::DuplicateShort(short.to_owned()).into());
                } else {
                    short_args.insert(short.to_owned());
                }
            }
            unresolved_cli_args.insert(arg.to_owned(), def.clone());
        }

        Ok(eval.heap().alloc(BxlFunction {
            bxl_path,
            id: RefCell::new(None),
            implementation,
            cli_args: unresolved_cli_args,
            docs: Some(doc.to_owned()),
        }))
    }
}

pub(crate) fn init_bxl_functions() {
    BXL_FUNCTIONS.init(BxlFunctions {
        register_cli_args_struct: |g| g.struct_("cli_args", cli_args::register_cli_args_module),
        register_bxl_function,
        register_artifact_function,
        register_label_function,
        register_target_function,
        register_instant_function,
        register_error_handling_function,
    });
}

/// Errors around rule declaration, instantiation, validation, etc
#[derive(Debug, Error)]
enum BxlError {
    #[error("Bxl defined in `{0}` must be assigned to a variable, e.g. `my_bxl = bxl(...)`")]
    BxlNotAssigned(String),
}

/// The callable created by `bxl()`
#[derive(Debug, Clone, ProvidesStaticType, Trace, NoSerialize, Allocative)]
pub(crate) struct BxlFunction<'v> {
    // The bxl path that contains the bxl() call; stored here so we can retrieve extra
    /// information during `export_as()`
    bxl_path: BxlFilePath,
    /// Once exported, the `import_path` and `name` of the callable. Used in DICE to retrieve bxl
    /// implementations
    id: RefCell<Option<BxlFunctionLabel>>,
    /// The implementation function for this bxl. Must be callable and take a ctx
    implementation: Value<'v>,
    /// the cli args to this bxl function
    cli_args: SmallMap<String, CliArgs>,
    docs: Option<String>,
}

impl<'v> Display for BxlFunction<'v> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match &*self.id.borrow() {
            Some(id) => write!(f, "{}()", id.name),
            None => write!(f, "<unbound bxl>"),
        }
    }
}

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

#[starlark_value(type = "bxl")]
impl<'v> StarlarkValue<'v> for BxlFunction<'v> {
    fn export_as(&self, variable_name: &str, _eval: &mut Evaluator<'v, '_>) {
        *self.id.borrow_mut() = Some(BxlFunctionLabel {
            bxl_path: self.bxl_path.clone(),
            name: variable_name.to_owned(),
        });
    }
}

impl<'v> Freeze for BxlFunction<'v> {
    type Frozen = FrozenBxlFunction;
    fn freeze(self, freezer: &Freezer) -> anyhow::Result<Self::Frozen> {
        let frozen_impl = self.implementation.freeze(freezer)?;
        let docs = self.docs;
        let id = match self.id.into_inner() {
            Some(x) => x,
            None => return Err(BxlError::BxlNotAssigned(self.bxl_path.to_string()).into()),
        };
        let bxl_id = Arc::new(id);

        Ok(FrozenBxlFunction {
            implementation: frozen_impl,
            cli_args: self.cli_args,
            bxl_id,
            docs,
        })
    }
}

#[derive(Debug, Display, ProvidesStaticType, NoSerialize, Allocative)]
#[display(fmt = "{}()", "bxl_id.name")]
pub struct FrozenBxlFunction {
    implementation: FrozenValue,
    cli_args: SmallMap<String, CliArgs>,
    bxl_id: Arc<BxlFunctionLabel>,
    docs: Option<String>,
}
starlark_simple_value!(FrozenBxlFunction);

#[starlark_value(type = "bxl")]
impl<'v> StarlarkValue<'v> for FrozenBxlFunction {}

impl FrozenBxlFunction {
    pub fn implementation(&self) -> FrozenValue {
        self.implementation
    }

    pub fn to_clap<'v>(&'v self, mut clap: clap::Command<'v>) -> clap::Command<'v> {
        if let Some(docs) = self.docs.as_ref() {
            clap = clap.about(docs.as_str())
        }

        for (arg, def) in self.cli_args.iter() {
            clap = clap.arg(def.to_clap(clap::Arg::new(arg.as_str()).long(arg.as_str())))
        }

        clap
    }

    /// Parses the cli args as defined by this bxl function. Automatically changes the CLI args
    /// to snakecase when accessed from the bxl context.
    pub(crate) async fn parse_clap<'a>(
        &self,
        clap: clap::ArgMatches,
        ctx: &CliResolutionCtx<'a>,
    ) -> anyhow::Result<OrderedMap<String, CliArgValue>> {
        let mut res = OrderedMap::with_capacity(self.cli_args.len());

        for (arg, cli) in self.cli_args.iter() {
            let snake_case_args = arg.replace('-', "_");
            if res.contains_key(&snake_case_args) {
                return Err(CliArgError::DefinedBothKebabAndSnakeCase(arg.clone()).into());
            }
            res.insert(
                snake_case_args,
                cli.parse_clap(ArgAccessor::Clap { clap: &clap, arg }, ctx)
                    .await
                    .with_context(|| {
                        format!("Error parsing cli flag `{}` for bxl function", arg)
                    })?,
            );
        }

        Ok(res)
    }
}


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

pub mod action;
pub mod configured;
pub mod unconfigured;


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

use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::pattern::pattern_type::ProvidersPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::ProvidersLabel;
use buck2_core::provider::label::ProvidersLabelMaybeConfigured;
use buck2_core::provider::label::ProvidersName;
use buck2_core::target::label::TargetLabel;
use buck2_interpreter::types::label::Label;
use buck2_interpreter::types::label::StarlarkProvidersLabel;
use buck2_interpreter::types::target_label::StarlarkConfiguredTargetLabel;
use buck2_interpreter::types::target_label::StarlarkTargetLabel;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use dupe::Dupe;
use itertools::Either;
use starlark::eval::Evaluator;
use starlark::values::list::ListRef;
use starlark::values::Value;
use starlark::values::ValueLike;
use thiserror::Error;

use crate::bxl::starlark_defs::context::BxlContext;
use crate::bxl::starlark_defs::nodes::configured::StarlarkConfiguredTargetNode;
use crate::bxl::starlark_defs::nodes::unconfigured::StarlarkTargetNode;
use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;
use crate::bxl::value_as_starlark_target_label::ValueAsStarlarkTargetLabel;

#[derive(Debug, Error)]
enum ProviderExprError {
    #[error("Expected a list of target like items, but was `{0}`")]
    NotAListOfTargets(String),
    #[error("Expected a single target like item, but was `{0}`")]
    NotATarget(String),
}

/// ProvidersExpr is just a simple type that can be used in starlark_module
/// functions for arguments that should be a set of provider labels. It will accept a
/// literal (like `//some:target[subtarget]`) or list of literals or a single provider label
pub(crate) enum ProvidersExpr<P: ProvidersLabelMaybeConfigured> {
    Literal(P),
    Iterable(Vec<P>),
}

impl ProvidersExpr<ConfiguredProvidersLabel> {
    pub fn unpack<'v>(
        value: Value<'v>,
        target_platform: Value<'v>,
        ctx: &BxlContext,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<Self> {
        let target_platform = target_platform.parse_target_platforms(
            &ctx.target_alias_resolver,
            &ctx.cell_resolver,
            ctx.cell_name,
            &ctx.global_target_platform,
        )?;

        Ok(
            if let Some(resolved) = Self::unpack_literal(value, &target_platform, ctx)? {
                resolved
            } else if let Some(resolved) = Self::unpack_iterable(value, ctx, eval, |v, ctx| {
                Self::unpack_literal(v, &target_platform, ctx)
            })? {
                resolved
            } else {
                return Err(anyhow::anyhow!(ProviderExprError::NotAListOfTargets(
                    value.to_repr()
                )));
            },
        )
    }

    fn unpack_literal<'v>(
        value: Value<'v>,
        target_platform: &Option<TargetLabel>,
        ctx: &BxlContext,
    ) -> anyhow::Result<Option<Self>> {
        if let Some(configured_target) = value.downcast_ref::<StarlarkConfiguredTargetNode>() {
            Ok(Some(Self::Literal(ConfiguredProvidersLabel::new(
                configured_target.0.label().dupe(),
                ProvidersName::Default,
            ))))
        } else if let Some(configured_target) =
            value.downcast_ref::<StarlarkConfiguredTargetLabel>()
        {
            Ok(Some(Self::Literal(ConfiguredProvidersLabel::new(
                configured_target.label().dupe(),
                ProvidersName::Default,
            ))))
        } else if let Some(configured_target) = value.downcast_ref::<Label>() {
            Ok(Some(Self::Literal(configured_target.label().clone())))
        } else {
            Self::unpack_providers_label(value, ctx)?.map_or(Ok(None), |label| {
                let result: anyhow::Result<_> = try {
                    Self::Literal(ctx.async_ctx.via_dice(|ctx| {
                        ctx.get_configured_provider_label(&label, target_platform.as_ref())
                    })?)
                };
                result.map(Some)
            })
        }
    }
}

impl ProvidersExpr<ProvidersLabel> {
    pub fn unpack<'v>(
        value: Value<'v>,
        ctx: &BxlContext<'_>,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<Self> {
        Ok(if let Some(resolved) = Self::unpack_literal(value, ctx)? {
            resolved
        } else if let Some(resolved) =
            Self::unpack_iterable(value, ctx, eval, Self::unpack_literal)?
        {
            resolved
        } else {
            return Err(anyhow::anyhow!(ProviderExprError::NotAListOfTargets(
                value.to_repr()
            )));
        })
    }

    fn unpack_literal<'v>(value: Value<'v>, ctx: &BxlContext<'_>) -> anyhow::Result<Option<Self>> {
        Self::unpack_providers_label(value, ctx)?
            .map_or(Ok(None), |label| Ok(Some(Self::Literal(label))))
    }
}

impl<P: ProvidersLabelMaybeConfigured> ProvidersExpr<P> {
    pub fn labels(&self) -> impl Iterator<Item = &P> {
        match &self {
            ProvidersExpr::Literal(item) => Either::Left(std::iter::once(item)),
            ProvidersExpr::Iterable(iter) => Either::Right(iter.iter()),
        }
    }

    fn unpack_providers_label<'v>(
        value: Value<'v>,
        ctx: &BxlContext<'_>,
    ) -> anyhow::Result<Option<ProvidersLabel>> {
        #[allow(clippy::manual_map)] // `if else if` looks better here
        Ok(if let Some(s) = value.unpack_str() {
            Some(
                ParsedPattern::<ProvidersPatternExtra>::parse_relaxed(
                    &ctx.target_alias_resolver,
                    // TODO(nga): Parse relaxed relative to cell root is incorrect.
                    CellPathRef::new(ctx.cell_name, CellRelativePath::empty()),
                    s,
                    &ctx.cell_resolver,
                )?
                .as_providers_label(s)?,
            )
        } else if let Some(target) = value.downcast_ref::<StarlarkTargetLabel>() {
            Some(ProvidersLabel::new(
                target.label().dupe(),
                ProvidersName::Default,
            ))
        } else if let Some(label) = value.downcast_ref::<StarlarkProvidersLabel>() {
            Some(label.label().clone())
        } else if let Some(node) = value.downcast_ref::<StarlarkTargetNode>() {
            Some(ProvidersLabel::new(
                node.0.label().dupe(),
                ProvidersName::Default,
            ))
        } else {
            None
        })
    }

    fn unpack_iterable<'v>(
        value: Value<'v>,
        ctx: &BxlContext<'_>,
        eval: &Evaluator<'v, '_>,
        unpack_literal: impl for<'a> Fn(
            Value<'v>,
            &BxlContext<'a>,
        ) -> anyhow::Result<Option<ProvidersExpr<P>>>,
    ) -> anyhow::Result<Option<ProvidersExpr<P>>> {
        Ok(Some(Self::Iterable(
            #[allow(clippy::manual_map)] // `if else if` looks better here
            if let Some(s) = value.downcast_ref::<StarlarkTargetSet<TargetNode>>() {
                Some(Either::Left(Either::Left(s.iter(eval.heap()))))
            } else if let Some(s) = value.downcast_ref::<StarlarkTargetSet<ConfiguredTargetNode>>()
            {
                Some(Either::Left(Either::Right(s.iter(eval.heap()))))
            } else if let Some(iterable) = ListRef::from_value(value) {
                Some(Either::Right(iterable.iter()))
            } else {
                None
            }
            .ok_or_else(|| ProviderExprError::NotATarget(value.to_repr()))?
            .map(|val| {
                if let Some(ProvidersExpr::Literal(resolved_val)) = unpack_literal(val, ctx)? {
                    Ok(resolved_val)
                } else {
                    Err(anyhow::anyhow!(ProviderExprError::NotATarget(
                        val.to_repr()
                    )))
                }
            })
            .collect::<anyhow::Result<_>>()?,
        )))
    }
}


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

use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationValue;
use starlark::eval::Evaluator;
use starlark::values::dict::Dict;
use starlark::values::Value;

use super::targetset::NodeLike;
use crate::bxl::starlark_defs::file_set::StarlarkFileSet;
use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;

pub(crate) fn parse_query_evaluation_result<'v, T: NodeLike>(
    result: QueryEvaluationResult<T>,
    eval: &mut Evaluator<'v, '_>,
) -> anyhow::Result<Value<'v>> {
    Ok(match result {
        QueryEvaluationResult::Single(result) => match result {
            QueryEvaluationValue::TargetSet(targets) => {
                eval.heap().alloc(StarlarkTargetSet::from(targets))
            }
            QueryEvaluationValue::FileSet(files) => eval.heap().alloc(StarlarkFileSet::from(files)),
        },
        QueryEvaluationResult::Multiple(multi) => eval.heap().alloc(Dict::new(
            multi
                .0
                .into_iter()
                .map(|(q, res)| {
                    Ok((
                        eval.heap().alloc(q).get_hashed()?,
                        match res? {
                            QueryEvaluationValue::TargetSet(targets) => {
                                eval.heap().alloc(StarlarkTargetSet::from(targets))
                            }
                            QueryEvaluationValue::FileSet(files) => {
                                eval.heap().alloc(StarlarkFileSet::from(files))
                            }
                        },
                    ))
                })
                .collect::<anyhow::Result<_>>()?,
        )),
    })
}


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

use std::ops::Deref;

use allocative::Allocative;
use buck2_query::query::environment::QueryTarget;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use buck2_query::query::syntax::simple::eval::set::TargetSetExt;
use derive_more::Display;
use dupe::Dupe;
use starlark::any::ProvidesStaticType;
use starlark::typing::Ty;
use starlark::values::starlark_value;
use starlark::values::type_repr::StarlarkTypeRepr;
use starlark::values::AllocValue;
use starlark::values::Freeze;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::StarlarkValue;
use starlark::values::UnpackValue;
use starlark::values::Value;
use starlark::values::ValueError;
use starlark::values::ValueLike;

use crate::bxl::starlark_defs::alloc_node::AllocNode;

pub trait NodeLike = QueryTarget + std::fmt::Debug + Eq + Dupe + AllocNode + Allocative;

#[derive(Debug, Display, Clone)]
#[derive(NoSerialize, Allocative)] // TODO maybe this should be
/// The StarlarkValue implementation for TargetSet to expose it to starlark.
pub struct StarlarkTargetSet<Node: QueryTarget>(pub TargetSet<Node>);

// TODO(nga): derive it.
unsafe impl<'a, Node: QueryTarget + 'static> ProvidesStaticType<'a> for StarlarkTargetSet<Node> {
    type StaticType = Self;
}

impl<Node: QueryTarget + AllocNode> StarlarkTargetSet<Node> {
    pub(crate) fn iter<'a, 'v>(&'a self, heap: &'v Heap) -> impl Iterator<Item = Value<'v>> + 'a
    where
        'v: 'a,
    {
        self.0
            .iter()
            .map(|target_node| target_node.dupe().alloc(heap))
    }
}

impl<Node: QueryTarget> Freeze for StarlarkTargetSet<Node> {
    type Frozen = StarlarkTargetSet<Node>;

    fn freeze(self, _freezer: &starlark::values::Freezer) -> anyhow::Result<Self::Frozen> {
        Ok(self)
    }
}

impl<'v, Node: NodeLike> StarlarkTypeRepr for &'v StarlarkTargetSet<Node> {
    fn starlark_type_repr() -> Ty {
        StarlarkTargetSet::<Node>::starlark_type_repr()
    }
}

impl<'v, Node: NodeLike> UnpackValue<'v> for &'v StarlarkTargetSet<Node> {
    fn unpack_value(x: Value<'v>) -> Option<Self> {
        StarlarkTargetSet::from_value(x)
    }
}

impl<'v, Node: NodeLike> AllocValue<'v> for StarlarkTargetSet<Node> {
    fn alloc_value(self, heap: &'v Heap) -> Value<'v> {
        heap.alloc_simple(self)
    }
}

#[starlark_value(type = "target_set")]
impl<'v, Node: NodeLike> StarlarkValue<'v> for StarlarkTargetSet<Node> {
    fn iterate_collect(&self, heap: &'v Heap) -> anyhow::Result<Vec<Value<'v>>> {
        Ok(self.iter(heap).collect())
    }

    fn at(&self, index: Value<'v>, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        let i = index.unpack_i32().ok_or_else(|| {
            ValueError::IncorrectParameterTypeWithExpected(
                "int".to_owned(),
                index.get_type().to_owned(),
            )
        })?;
        if i < 0 {
            return Err(anyhow::anyhow!(ValueError::IndexOutOfBound(i)));
        }
        self.0
            .get_index(i as usize)
            .ok_or_else(|| anyhow::anyhow!(ValueError::IndexOutOfBound(i)))
            .map(|node| node.dupe().alloc(heap))
    }

    fn length(&self) -> anyhow::Result<i32> {
        Ok(self.0.len().try_into()?)
    }

    fn add(&self, other: Value<'v>, heap: &'v Heap) -> Option<anyhow::Result<Value<'v>>> {
        let other = other.downcast_ref::<Self>()?;
        let union = self.0.union(&other.0);
        Some(Ok(heap.alloc(Self(union))))
    }

    fn sub(&self, other: Value<'v>, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        let other = other
            .downcast_ref::<Self>()
            .ok_or(ValueError::IncorrectParameterType)?;
        let union = self.0.difference(&other.0)?;
        Ok(heap.alloc(Self(union)))
    }

    fn equals(&self, other: Value<'v>) -> anyhow::Result<bool> {
        match other.downcast_ref::<StarlarkTargetSet<Node>>() {
            Some(other) => Ok(self.0 == other.0),
            None => Ok(false),
        }
    }
}

impl<Node: QueryTarget> From<TargetSet<Node>> for StarlarkTargetSet<Node> {
    fn from(v: TargetSet<Node>) -> Self {
        Self(v)
    }
}

impl<Node: QueryTarget> Deref for StarlarkTargetSet<Node> {
    type Target = TargetSet<Node>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<Node: NodeLike> StarlarkTargetSet<Node> {
    pub fn from_value<'v>(x: Value<'v>) -> Option<&'v Self> {
        ValueLike::downcast_ref::<Self>(x)
    }
}


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

use buck2_build_api::configure_targets::get_maybe_compatible_targets;
use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::paths::CellRelativePath;
use buck2_core::configuration::compatibility::IncompatiblePlatformReason;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::TargetLabel;
use buck2_interpreter::types::target_label::StarlarkConfiguredTargetLabel;
use buck2_interpreter::types::target_label::StarlarkTargetLabel;
use buck2_node::load_patterns::load_patterns;
use buck2_node::load_patterns::MissingTargetBehavior;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::configured_frontend::ConfiguredTargetNodeCalculation;
use buck2_node::nodes::frontend::TargetGraphCalculation;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use buck2_query::query::environment::QueryTarget;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use buck2_util::truncate::truncate;
use dice::DiceComputations;
use dupe::Dupe;
use dupe::IterDupedExt;
use either::Either;
use futures::TryFutureExt;
use starlark::collections::SmallSet;
use starlark::eval::Evaluator;
use starlark::values::list::ListRef;
use starlark::values::Value;
use starlark::values::ValueLike;
use thiserror::Error;

use crate::bxl::starlark_defs::context::BxlContext;
use crate::bxl::starlark_defs::nodes::configured::StarlarkConfiguredTargetNode;
use crate::bxl::starlark_defs::nodes::unconfigured::StarlarkTargetNode;
use crate::bxl::starlark_defs::targetset::StarlarkTargetSet;

/// TargetExpr is just a simple type that can be used in starlark_module
/// functions for arguments that should be target sets. It will accept a
/// literal (like `//some:target`) or list of literals or a TargetSet Value (from one of the
/// BXL functions that return them). It can be resolved to a `&TargetSet` with
/// the help of the `targets!()` macro.
pub enum TargetExpr<'v, Node: QueryTarget> {
    Node(Node),
    Label(Cow<'v, Node::NodeRef>),
    Iterable(Vec<Either<Node, Cow<'v, Node::NodeRef>>>),
    TargetSet(Cow<'v, TargetSet<Node>>),
}

impl<'v> TargetExpr<'v, ConfiguredTargetNode> {
    /// Get a vector of maybe compatible `ConfiguredTargetNode`s from the `TargetExpr`.
    /// Any callers of this function will need to call `filter_incompatible()` on the result
    /// in order to get the `TargetSet<ConfiguredTargetNode>`.
    pub async fn get(
        self,
        dice: &'v DiceComputations,
    ) -> anyhow::Result<Vec<MaybeCompatible<ConfiguredTargetNode>>> {
        match self {
            TargetExpr::Node(val) => Ok(vec![dice.get_configured_target_node(val.label()).await?]),
            TargetExpr::Label(label) => {
                Ok(vec![dice.get_configured_target_node(label.as_ref()).await?])
            }
            TargetExpr::Iterable(val) => {
                let futs = val.into_iter().map(|node_or_ref| async {
                    match node_or_ref {
                        Either::Left(node) => dice.get_configured_target_node(node.label()).await,
                        Either::Right(label) => {
                            dice.get_configured_target_node(label.as_ref()).await
                        }
                    }
                });

                futures::future::join_all(futs).await.into_iter().collect()
            }
            TargetExpr::TargetSet(val) => futures::future::join_all(val.iter().map(|node| {
                dice.get_configured_target_node(node.label())
                    .map_err(anyhow::Error::from)
            }))
            .await
            .into_iter()
            .collect(),
        }
    }
}

// Filters out incompatible targets and emits the error message
pub(crate) fn filter_incompatible(
    targets: impl Iterator<Item = MaybeCompatible<ConfiguredTargetNode>>,
    bxl_ctx: &BxlContext,
) -> anyhow::Result<TargetSet<ConfiguredTargetNode>> {
    let mut target_set = TargetSet::new();
    let mut incompatible_targets = SmallSet::new();

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

    if !incompatible_targets.is_empty() {
        bxl_ctx.print_to_error_stream(
            IncompatiblePlatformReason::skipping_message_for_multiple(incompatible_targets.iter()),
        )?;
    }

    Ok(target_set)
}

fn unpack_target_label<'v>(value: Value<'v>) -> Option<&'v TargetLabel> {
    if let Some(target) = value.downcast_ref::<StarlarkTargetLabel>() {
        Some(target.label())
    } else if let Some(node) = value.downcast_ref::<StarlarkTargetNode>() {
        Some(node.0.label())
    } else {
        None
    }
}

impl<'v> TargetExpr<'v, TargetNode> {
    /// Get a `TargetSet<TargetNode>` from the `TargetExpr`
    pub async fn get(
        self,
        ctx: &DiceComputations,
    ) -> anyhow::Result<Cow<'v, TargetSet<TargetNode>>> {
        match self {
            TargetExpr::Node(val) => {
                let mut set = TargetSet::new();
                set.insert(val);
                Ok(Cow::Owned(set))
            }
            TargetExpr::Label(label) => {
                let node = ctx.get_target_node(&label).await?;
                let mut set = TargetSet::new();
                set.insert(node);
                Ok(Cow::Owned(set))
            }
            TargetExpr::Iterable(val) => {
                let mut set = TargetSet::new();
                let futs = val.into_iter().map(|node_or_ref| async {
                    match node_or_ref {
                        Either::Left(node) => Ok(node),
                        Either::Right(node_ref) => ctx.get_target_node(&node_ref).await,
                    }
                });

                for node in futures::future::join_all(futs).await {
                    set.insert(node?);
                }

                Ok(Cow::Owned(set))
            }
            TargetExpr::TargetSet(val) => Ok(val),
        }
    }
}

#[derive(Debug, Error)]
pub enum TargetExprError {
    #[error(
        "Expected a list of target like items, but was `{0}`. If you have passed in a list of `label`s, make sure to call `configured_target()` to get the underlying configured target label."
    )]
    NotAListOfTargets(String),
    #[error(
        "Expected a single target like item, but was `{0}`. If you have passed in a `label`, make sure to call `configured_target()` to get the underlying configured target label."
    )]
    NotATarget(String),
}

impl<'v> TargetExpr<'v, ConfiguredTargetNode> {
    pub async fn unpack(
        value: Value<'v>,
        target_platform: &Option<TargetLabel>,
        ctx: &BxlContext<'v>,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<TargetExpr<'v, ConfiguredTargetNode>> {
        Ok(
            if let Some(resolved) = Self::unpack_literal(value, target_platform, ctx).await? {
                resolved
            } else if let Some(resolved) =
                Self::unpack_iterable(value, target_platform, ctx, eval).await?
            {
                resolved
            } else {
                return Err(anyhow::anyhow!(TargetExprError::NotAListOfTargets(
                    value.to_repr()
                )));
            },
        )
    }

    async fn unpack_literal(
        value: Value<'v>,
        target_platform: &Option<TargetLabel>,
        ctx: &BxlContext<'v>,
    ) -> anyhow::Result<Option<TargetExpr<'v, ConfiguredTargetNode>>> {
        if let Some(configured_target) = value.downcast_ref::<StarlarkConfiguredTargetNode>() {
            Ok(Some(Self::Node(configured_target.0.dupe())))
        } else if let Some(configured_target) =
            value.downcast_ref::<StarlarkConfiguredTargetLabel>()
        {
            Ok(Some(Self::Label(Cow::Borrowed(configured_target.label()))))
        } else if let Some(s) = value.unpack_str() {
            match ParsedPattern::<TargetPatternExtra>::parse_relaxed(
                &ctx.target_alias_resolver,
                // TODO(nga): Parse relaxed relative to cell root is incorrect.
                CellPathRef::new(ctx.cell_name, CellRelativePath::empty()),
                s,
                &ctx.cell_resolver,
            )? {
                ParsedPattern::Target(pkg, name, TargetPatternExtra) => {
                    Ok(Some(Self::Label(Cow::Owned(
                        ctx.async_ctx
                            .0
                            .get_configured_target(
                                &TargetLabel::new(pkg, name.as_ref()),
                                target_platform.as_ref(),
                            )
                            .await?,
                    ))))
                }
                pattern => {
                    let loaded_patterns =
                        load_patterns(ctx.async_ctx.0, vec![pattern], MissingTargetBehavior::Fail)
                            .await?;

                    let maybe_compatible = get_maybe_compatible_targets(
                        ctx.async_ctx.0,
                        loaded_patterns.iter_loaded_targets_by_package(),
                        target_platform.dupe(),
                    )
                    .await?
                    .collect::<Result<Vec<_>, _>>()?;

                    let result = filter_incompatible(maybe_compatible.into_iter(), ctx)?;
                    Ok(Some(Self::TargetSet(Cow::Owned(result))))
                }
            }
        } else {
            match unpack_target_label(value) {
                None => Ok(None),
                Some(label) => Ok(Some(Self::Label(Cow::Owned(
                    ctx.async_ctx
                        .0
                        .get_configured_target(label, target_platform.as_ref())
                        .await?,
                )))),
            }
        }
    }

    async fn unpack_iterable(
        value: Value<'v>,
        target_platform: &Option<TargetLabel>,
        ctx: &BxlContext<'v>,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<Option<TargetExpr<'v, ConfiguredTargetNode>>> {
        if let Some(s) = value.downcast_ref::<StarlarkTargetSet<ConfiguredTargetNode>>() {
            return Ok(Some(Self::TargetSet(Cow::Borrowed(s))));
        }

        #[allow(clippy::manual_map)] // `if else if` looks better here
        let items = if let Some(s) = value.downcast_ref::<StarlarkTargetSet<TargetNode>>() {
            Some(Either::Left(s.iter(eval.heap())))
        } else if let Some(iterable) = ListRef::from_value(value) {
            Some(Either::Right(iterable.iter()))
        } else {
            None
        }
        .ok_or_else(|| TargetExprError::NotAListOfTargets(value.to_repr()))?;

        let mut resolved = vec![];

        for item in items {
            let unpacked = Self::unpack_literal(item, target_platform, ctx).await?;

            match unpacked {
                Some(TargetExpr::Node(node)) => resolved.push(Either::Left(node)),
                Some(TargetExpr::Label(label)) => resolved.push(Either::Right(label)),
                Some(TargetExpr::TargetSet(set)) => match set {
                    Cow::Borrowed(s) => itertools::Either::Left(s.iter().duped()),
                    Cow::Owned(s) => itertools::Either::Right(s.into_iter()),
                }
                .for_each(|t| resolved.push(Either::Left(t))),
                _ => {
                    return Err(anyhow::anyhow!(TargetExprError::NotATarget(item.to_repr()))
                        .context(format!(
                            "Error resolving list `{}`",
                            truncate(&value.to_repr(), 150)
                        )));
                }
            }
        }

        Ok(Some(Self::Iterable(resolved)))
    }
}

impl<'v> TargetExpr<'v, TargetNode> {
    pub async fn unpack(
        value: Value<'v>,
        ctx: &BxlContext<'v>,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<TargetExpr<'v, TargetNode>> {
        Ok(
            if let Some(resolved) = Self::unpack_literal(value, ctx).await? {
                resolved
            } else if let Some(resolved) = Self::unpack_iterable(value, ctx, eval).await? {
                resolved
            } else {
                return Err(anyhow::anyhow!(TargetExprError::NotAListOfTargets(
                    value.to_repr()
                )));
            },
        )
    }

    async fn unpack_literal(
        value: Value<'v>,
        ctx: &BxlContext<'v>,
    ) -> anyhow::Result<Option<TargetExpr<'v, TargetNode>>> {
        if let Some(target) = value.downcast_ref::<StarlarkTargetNode>() {
            Ok(Some(Self::Node(target.0.dupe())))
        } else if let Some(label) = value.downcast_ref::<StarlarkTargetLabel>() {
            Ok(Some(Self::Label(Cow::Borrowed(label.label()))))
        } else if let Some(s) = value.unpack_str() {
            match ParsedPattern::<TargetPatternExtra>::parse_relaxed(
                &ctx.target_alias_resolver,
                // TODO(nga): Parse relaxed relative to cell root is incorrect.
                CellPathRef::new(ctx.cell_name, CellRelativePath::empty()),
                s,
                &ctx.cell_resolver,
            )? {
                ParsedPattern::Target(pkg, name, TargetPatternExtra) => Ok(Some(Self::Label(
                    Cow::Owned(TargetLabel::new(pkg, name.as_ref())),
                ))),
                pattern => {
                    let loaded_patterns =
                        load_patterns(ctx.async_ctx.0, vec![pattern], MissingTargetBehavior::Fail)
                            .await?;
                    let mut target_set = TargetSet::new();
                    for (_package, results) in loaded_patterns.into_iter() {
                        target_set.extend(results?.into_values());
                    }
                    Ok(Some(Self::TargetSet(Cow::Owned(target_set))))
                }
            }
        } else {
            match unpack_target_label(value) {
                None => Ok(None),
                Some(label) => Ok(Some(Self::Label(Cow::Borrowed(label)))),
            }
        }
    }

    async fn unpack_iterable(
        value: Value<'v>,
        ctx: &BxlContext<'v>,
        eval: &Evaluator<'v, '_>,
    ) -> anyhow::Result<Option<TargetExpr<'v, TargetNode>>> {
        if let Some(s) = value.downcast_ref::<StarlarkTargetSet<TargetNode>>() {
            return Ok(Some(Self::TargetSet(Cow::Borrowed(s))));
        }

        #[allow(clippy::manual_map)] // `if else if` looks better here
        let items = if let Some(s) = value.downcast_ref::<StarlarkTargetSet<TargetNode>>() {
            Some(Either::Left(s.iter(eval.heap())))
        } else if let Some(iterable) = ListRef::from_value(value) {
            Some(Either::Right(iterable.iter()))
        } else {
            None
        }
        .ok_or_else(|| TargetExprError::NotAListOfTargets(value.to_repr()))?;

        let mut resolved = vec![];

        for item in items {
            let unpacked = Self::unpack_literal(item, ctx).await?;

            match unpacked {
                Some(TargetExpr::Node(node)) => resolved.push(Either::Left(node)),
                Some(TargetExpr::Label(label)) => resolved.push(Either::Right(label)),
                Some(TargetExpr::TargetSet(set)) => match set {
                    Cow::Borrowed(s) => itertools::Either::Left(s.iter().duped()),
                    Cow::Owned(s) => itertools::Either::Right(s.into_iter()),
                }
                .for_each(|t| resolved.push(Either::Left(t))),
                _ => {
                    return Err(anyhow::anyhow!(TargetExprError::NotATarget(item.to_repr()))
                        .context(format!(
                            "Error resolving list `{}`",
                            truncate(&value.to_repr(), 150)
                        )));
                }
            }
        }
        Ok(Some(Self::Iterable(resolved)))
    }
}


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

use std::time::Instant;

use allocative::Allocative;
use starlark::environment::Methods;
use starlark::environment::MethodsBuilder;
use starlark::environment::MethodsStatic;
use starlark::starlark_module;
use starlark::starlark_simple_value;
use starlark::values::starlark_value;
use starlark::values::Heap;
use starlark::values::NoSerialize;
use starlark::values::ProvidesStaticType;
use starlark::values::StarlarkValue;
use starlark::values::Value;
use starlark::values::ValueLike;
use starlark::StarlarkDocs;

/// Starlark object for Instant.
#[derive(
    Clone,
    Debug,
    derive_more::Display,
    ProvidesStaticType,
    NoSerialize,
    StarlarkDocs,
    Allocative
)]
#[starlark_docs(directory = "bxl")]
#[display(fmt = "{:?}", _0)]
pub struct StarlarkInstant(pub(crate) Instant);

/// Instant methods, to aid in debugging/timing individual pieces of the bxl script.
#[starlark_module]
fn starlark_instant_methods(builder: &mut MethodsBuilder) {
    /// Elapsed time in secs as a float
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_elapsed_secs(ctx):
    ///     now = now()
    ///     time_a = now.elapsed_secs()
    ///     # do something that takes a long time
    ///     time_b = now.elapsed_secs()
    ///
    ///     ctx.output.print(time_a)
    ///     ctx.output.print(time_b)
    /// ```
    fn elapsed_secs<'v>(this: Value<'v>, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        let secs = this
            .downcast_ref::<StarlarkInstant>()
            .unwrap()
            .0
            .elapsed()
            .as_secs() as f64;

        Ok(heap.alloc(secs))
    }

    /// Elapsed time in millis as a float
    ///
    /// Sample usage:
    /// ```text
    /// def _impl_elapsed_millis(ctx):
    ///     now = now()
    ///     time_a = now.elapsed_millis()
    ///     # do something that takes a long time
    ///     time_b = now.elapsed_millis()
    ///
    ///     ctx.output.print(time_a)
    ///     ctx.output.print(time_b)
    /// ```
    fn elapsed_millis<'v>(this: Value<'v>, heap: &'v Heap) -> anyhow::Result<Value<'v>> {
        let millis = this
            .downcast_ref::<StarlarkInstant>()
            .unwrap()
            .0
            .elapsed()
            .as_millis() as f64;

        Ok(heap.alloc(millis))
    }
}

starlark_simple_value!(StarlarkInstant);

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