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

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::execute::action_executor::ActionExecutionKind;
use buck2_build_api::actions::execute::action_executor::ActionExecutionMetadata;
use buck2_build_api::actions::execute::action_executor::ActionOutputs;
use buck2_build_api::actions::Action;
use buck2_build_api::actions::ActionExecutable;
use buck2_build_api::actions::ActionExecutionCtx;
use buck2_build_api::actions::IncrementalActionExecutable;
use buck2_build_api::actions::UnregisteredAction;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::DefaultCommandLineContext;
use buck2_core::category::Category;
use buck2_execute::artifact::fs::ExecutorFs;
use buck2_execute::execute::command_executor::ActionExecutionTimingData;
use buck2_execute::materialize::materializer::WriteRequest;
use dupe::Dupe;
use indexmap::indexmap;
use indexmap::IndexMap;
use indexmap::IndexSet;
use once_cell::sync::Lazy;
use starlark::values::OwnedFrozenValue;
use thiserror::Error;

#[derive(Debug, Error)]
enum WriteActionValidationError {
    #[error("WriteAction received inputs")]
    TooManyInputs,
    #[error("WriteAction received no outputs")]
    NoOutputs,
    #[error("WriteAction received more than one output")]
    TooManyOutputs,
    #[error("Expected command line value, got {0}")]
    ContentsNotCommandLineValue(String),
}

#[derive(Allocative)]
pub(crate) struct UnregisteredWriteAction {
    is_executable: bool,
    macro_files: Option<IndexSet<Artifact>>,
}

impl UnregisteredWriteAction {
    pub(crate) fn new(is_executable: bool, macro_files: Option<IndexSet<Artifact>>) -> Self {
        Self {
            is_executable,
            macro_files,
        }
    }
}

impl UnregisteredAction for UnregisteredWriteAction {
    fn register(
        self: Box<Self>,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
        starlark_data: Option<OwnedFrozenValue>,
    ) -> anyhow::Result<Box<dyn Action>> {
        let contents = starlark_data.expect("module data to be present");

        let write_action = WriteAction::new(
            contents,
            self.is_executable,
            inputs,
            self.macro_files,
            outputs,
        )?;
        Ok(Box::new(write_action))
    }
}

#[derive(Debug, Allocative)]
struct WriteAction {
    contents: OwnedFrozenValue, // StarlarkCommandLine
    is_executable: bool,
    macro_files: Option<IndexSet<Artifact>>,
    output: BuildArtifact,
}

impl WriteAction {
    fn new(
        contents: OwnedFrozenValue,
        is_executable: bool,
        inputs: IndexSet<ArtifactGroup>,
        macro_files: Option<IndexSet<Artifact>>,
        outputs: IndexSet<BuildArtifact>,
    ) -> anyhow::Result<Self> {
        let mut outputs = outputs.into_iter();

        let output = match (outputs.next(), outputs.next()) {
            (Some(o), None) => o,
            (None, ..) => return Err(WriteActionValidationError::NoOutputs.into()),
            (Some(..), Some(..)) => return Err(WriteActionValidationError::TooManyOutputs.into()),
        };

        if !inputs.is_empty() {
            return Err(WriteActionValidationError::TooManyInputs.into());
        }

        if contents.value().as_command_line().is_none() {
            return Err(WriteActionValidationError::ContentsNotCommandLineValue(
                contents.value().to_repr(),
            )
            .into());
        }

        Ok(WriteAction {
            contents,
            is_executable,
            macro_files,
            output,
        })
    }

    fn get_contents(&self, fs: &ExecutorFs) -> anyhow::Result<String> {
        let mut cli = Vec::<String>::new();

        let mut ctx = if let Some(macro_files) = &self.macro_files {
            DefaultCommandLineContext::new_with_write_to_file_macros_support(fs, macro_files)
        } else {
            DefaultCommandLineContext::new(fs)
        };

        self.contents
            .value()
            .as_command_line()
            .unwrap()
            .add_to_command_line(&mut cli, &mut ctx)?;

        Ok(cli.join("\n"))
    }
}

#[async_trait]
impl Action for WriteAction {
    fn kind(&self) -> buck2_data::ActionKind {
        buck2_data::ActionKind::Write
    }

    fn inputs(&self) -> anyhow::Result<Cow<'_, [ArtifactGroup]>> {
        Ok(Cow::Borrowed(&[]))
    }

    fn outputs(&self) -> anyhow::Result<Cow<'_, [BuildArtifact]>> {
        Ok(Cow::Borrowed(slice::from_ref(&self.output)))
    }

    fn as_executable(&self) -> ActionExecutable<'_> {
        ActionExecutable::Incremental(self)
    }

    fn category(&self) -> &Category {
        static WRITE_CATEGORY: Lazy<Category> = Lazy::new(|| Category::try_from("write").unwrap());

        &WRITE_CATEGORY
    }

    fn identifier(&self) -> Option<&str> {
        Some(self.output.get_path().path().as_str())
    }

    fn aquery_attributes(&self, fs: &ExecutorFs) -> IndexMap<String, String> {
        // TODO(cjhopman): We should change this api to support returning a Result.
        indexmap! {
            "contents".to_owned() => match self.get_contents(fs) {
                Ok(v) => v,
                Err(e) => format!("ERROR: constructing contents ({})", e)
            }
        }
    }
}

#[async_trait]
impl IncrementalActionExecutable for WriteAction {
    async fn execute(
        &self,
        ctx: &mut dyn ActionExecutionCtx,
    ) -> anyhow::Result<(ActionOutputs, ActionExecutionMetadata)> {
        let fs = ctx.fs();

        let mut execution_start = None;

        let value = ctx
            .materializer()
            .declare_write(Box::new(|| {
                execution_start = Some(Instant::now());
                let content = self.get_contents(&ctx.executor_fs())?.into_bytes();
                Ok(vec![WriteRequest {
                    path: fs.resolve_build(self.output.get_path()),
                    content,
                    is_executable: self.is_executable,
                }])
            }))
            .await?
            .into_iter()
            .next()
            .context("Write did not execute")?;

        let wall_time = execution_start
            .context("Action did not set execution_start")?
            .elapsed();

        Ok((
            ActionOutputs::new(indexmap![self.output.get_path().dupe() => value]),
            ActionExecutionMetadata {
                execution_kind: ActionExecutionKind::Simple,
                timing: ActionExecutionTimingData { wall_time },
            },
        ))
    }
}

#[cfg(test)]
mod tests {
    // TODO: This needs proper tests, but right now it's kind of a pain to get the
    //       action framework up and running to test actions
    #[test]
    fn writes_file() {}
}


/*
 * 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::convert::TryFrom;
use std::fmt;
use std::io::Write;
use std::slice;
use std::time::Instant;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::execute::action_executor::ActionExecutionKind;
use buck2_build_api::actions::execute::action_executor::ActionExecutionMetadata;
use buck2_build_api::actions::execute::action_executor::ActionOutputs;
use buck2_build_api::actions::impls::json;
use buck2_build_api::actions::impls::json::validate_json;
use buck2_build_api::actions::Action;
use buck2_build_api::actions::ActionExecutable;
use buck2_build_api::actions::ActionExecutionCtx;
use buck2_build_api::actions::IncrementalActionExecutable;
use buck2_build_api::actions::UnregisteredAction;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineBuilder;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineContext;
use buck2_build_api::interpreter::rule_defs::cmd_args::WriteToFileMacroVisitor;
use buck2_core::category::Category;
use buck2_execute::artifact::fs::ExecutorFs;
use buck2_execute::execute::command_executor::ActionExecutionTimingData;
use buck2_execute::materialize::materializer::WriteRequest;
use dupe::Dupe;
use indexmap::indexmap;
use indexmap::IndexMap;
use indexmap::IndexSet;
use once_cell::sync::Lazy;
use starlark::any::ProvidesStaticType;
use starlark::coerce::Coerce;
use starlark::starlark_complex_value;
use starlark::values::starlark_value;
use starlark::values::Demand;
use starlark::values::Freeze;
use starlark::values::NoSerialize;
use starlark::values::OwnedFrozenValue;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::Value;
use starlark::values::ValueLike;
use thiserror::Error;

#[derive(Debug, Error)]
enum WriteJsonActionValidationError {
    #[error("WriteJsonAction received inputs")]
    TooManyInputs,
    #[error("WriteJsonAction received no outputs")]
    NoOutputs,
    #[error("WriteJsonAction received more than one output")]
    TooManyOutputs,
}

#[derive(Allocative)]
pub(crate) struct UnregisteredWriteJsonAction;

impl UnregisteredWriteJsonAction {
    pub(crate) fn new() -> Self {
        Self
    }

    pub(crate) fn cli<'v>(
        artifact: Value<'v>,
        content: Value<'v>,
    ) -> anyhow::Result<WriteJsonCommandLineArg<'v>> {
        Ok(WriteJsonCommandLineArg { artifact, content })
    }
}

impl UnregisteredAction for UnregisteredWriteJsonAction {
    fn register(
        self: Box<Self>,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
        starlark_data: Option<OwnedFrozenValue>,
    ) -> anyhow::Result<Box<dyn Action>> {
        let contents = starlark_data.expect("module data to be present");
        let action = WriteJsonAction::new(contents, inputs, outputs)?;
        Ok(Box::new(action))
    }
}

#[derive(Debug, Allocative)]
struct WriteJsonAction {
    contents: OwnedFrozenValue, // JSON value
    output: BuildArtifact,
}

impl WriteJsonAction {
    fn new(
        contents: OwnedFrozenValue,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
    ) -> anyhow::Result<Self> {
        validate_json(contents.value())?;

        let mut outputs = outputs.into_iter();

        let output = match (outputs.next(), outputs.next()) {
            (Some(o), None) => o,
            (None, ..) => return Err(WriteJsonActionValidationError::NoOutputs.into()),
            (Some(..), Some(..)) => {
                return Err(WriteJsonActionValidationError::TooManyOutputs.into());
            }
        };

        if !inputs.is_empty() {
            return Err(WriteJsonActionValidationError::TooManyInputs.into());
        }

        Ok(WriteJsonAction { contents, output })
    }

    fn write(&self, fs: &ExecutorFs, writer: impl Write) -> anyhow::Result<()> {
        json::write_json(self.contents.value(), Some(fs), writer)
    }

    fn get_contents(&self, fs: &ExecutorFs) -> anyhow::Result<Vec<u8>> {
        let mut contents = Vec::new();
        self.write(fs, &mut contents)?;
        Ok(contents)
    }
}

#[async_trait]
impl Action for WriteJsonAction {
    fn kind(&self) -> buck2_data::ActionKind {
        buck2_data::ActionKind::Write
    }

    fn inputs(&self) -> anyhow::Result<Cow<'_, [ArtifactGroup]>> {
        Ok(Cow::Borrowed(&[]))
    }

    fn outputs(&self) -> anyhow::Result<Cow<'_, [BuildArtifact]>> {
        Ok(Cow::Borrowed(slice::from_ref(&self.output)))
    }

    fn as_executable(&self) -> ActionExecutable<'_> {
        ActionExecutable::Incremental(self)
    }

    fn category(&self) -> &Category {
        static WRITE_CATEGORY: Lazy<Category> =
            Lazy::new(|| Category::try_from("write_json").unwrap());

        &WRITE_CATEGORY
    }

    fn identifier(&self) -> Option<&str> {
        Some(self.output.get_path().path().as_str())
    }

    fn aquery_attributes(&self, fs: &ExecutorFs) -> IndexMap<String, String> {
        let res: anyhow::Result<String> = try { String::from_utf8(self.get_contents(fs)?)? };
        // TODO(cjhopman): We should change this api to support returning a Result.
        indexmap! {
            "contents".to_owned() => match res {
                Ok(v) => v,
                Err(e) => format!("ERROR: constructing contents ({})", e)
            }
        }
    }
}

#[async_trait]
impl IncrementalActionExecutable for WriteJsonAction {
    async fn execute(
        &self,
        ctx: &mut dyn ActionExecutionCtx,
    ) -> anyhow::Result<(ActionOutputs, ActionExecutionMetadata)> {
        let fs = ctx.fs();

        let mut execution_start = None;

        let value = ctx
            .materializer()
            .declare_write(Box::new(|| {
                execution_start = Some(Instant::now());
                let content = self.get_contents(&ctx.executor_fs())?;
                Ok(vec![WriteRequest {
                    path: fs.resolve_build(self.output.get_path()),
                    content,
                    is_executable: false,
                }])
            }))
            .await?
            .into_iter()
            .next()
            .context("Write did not execute")?;

        let wall_time = execution_start
            .context("Action did not set execution_start")?
            .elapsed();

        Ok((
            ActionOutputs::new(indexmap![self.output.get_path().dupe() => value]),
            ActionExecutionMetadata {
                execution_kind: ActionExecutionKind::Simple,
                timing: ActionExecutionTimingData { wall_time },
            },
        ))
    }
}

/// WriteJsonCommandLineArgGen represents the artifact produced by write_json in a way that it can
/// be added to commandlines while including the artifacts referenced by cmdargs in the content that
/// was written.
#[derive(Debug, Clone, Trace, Coerce, Freeze, ProvidesStaticType, Allocative)]
#[derive(NoSerialize)] // TODO we should probably have a serialization for transitive set
#[repr(C)]
pub(crate) struct WriteJsonCommandLineArgGen<V> {
    artifact: V,
    // The list of artifacts here could be large and we don't want to hold those explicitly (due to
    // the memory cost) and so we hold the same content value that the write_json action itself will and
    // only traverse it when artifacts are requested.
    content: V,
}

impl<'v, V: ValueLike<'v>> fmt::Display for WriteJsonCommandLineArgGen<V> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str("<write_json_cli_args>")
    }
}

starlark_complex_value!(pub(crate) WriteJsonCommandLineArg);

#[starlark_value(type = "write_json_cli_args")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for WriteJsonCommandLineArgGen<V>
where
    Self: ProvidesStaticType<'v>,
{
    fn provide(&'v self, demand: &mut Demand<'_, 'v>) {
        demand.provide_value::<&dyn CommandLineArgLike>(self);
    }
}

impl<'v, V: ValueLike<'v>> CommandLineArgLike for WriteJsonCommandLineArgGen<V> {
    fn add_to_command_line(
        &self,
        builder: &mut dyn CommandLineBuilder,
        context: &mut dyn CommandLineContext,
    ) -> anyhow::Result<()> {
        self.artifact
            .to_value()
            .as_command_line_err()?
            .add_to_command_line(builder, context)
    }

    fn visit_artifacts(&self, visitor: &mut dyn CommandLineArtifactVisitor) -> anyhow::Result<()> {
        let artifact = self.artifact.to_value();
        let content = self.content.to_value();
        artifact.as_command_line_err()?.visit_artifacts(visitor)?;
        json::visit_json_artifacts(content, visitor)
    }

    fn contains_arg_attr(&self) -> bool {
        // In the write_json implementation, the commandlinebuilders we use don't support args.
        false
    }

    fn visit_write_to_file_macros(
        &self,
        _visitor: &mut dyn WriteToFileMacroVisitor,
    ) -> anyhow::Result<()> {
        // In the write_json implementation, the commandlinebuilders we use don't support args.
        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::time::Instant;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::execute::action_executor::ActionExecutionKind;
use buck2_build_api::actions::execute::action_executor::ActionExecutionMetadata;
use buck2_build_api::actions::execute::action_executor::ActionOutputs;
use buck2_build_api::actions::Action;
use buck2_build_api::actions::ActionExecutable;
use buck2_build_api::actions::ActionExecutionCtx;
use buck2_build_api::actions::IncrementalActionExecutable;
use buck2_build_api::actions::UnregisteredAction;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::interpreter::rule_defs::cmd_args::arg_builder::ArgBuilder;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineContext;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineLocation;
use buck2_build_api::interpreter::rule_defs::cmd_args::DefaultCommandLineContext;
use buck2_build_api::interpreter::rule_defs::cmd_args::WriteToFileMacroVisitor;
use buck2_build_api::interpreter::rule_defs::resolved_macro::ResolvedMacro;
use buck2_core::category::Category;
use buck2_core::fs::paths::RelativePathBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_execute::artifact::fs::ExecutorFs;
use buck2_execute::execute::command_executor::ActionExecutionTimingData;
use buck2_execute::materialize::materializer::WriteRequest;
use dupe::Dupe;
use indexmap::IndexSet;
use once_cell::sync::Lazy;
use starlark::values::OwnedFrozenValue;
use thiserror::Error;

#[derive(Allocative)]
pub(crate) struct UnregisteredWriteMacrosToFileAction {
    identifier: String,
}

impl UnregisteredWriteMacrosToFileAction {
    pub(crate) fn new(identifier: String) -> Self {
        Self { identifier }
    }
}

impl UnregisteredAction for UnregisteredWriteMacrosToFileAction {
    fn register(
        self: Box<Self>,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
        starlark_data: Option<OwnedFrozenValue>,
    ) -> anyhow::Result<Box<dyn Action>> {
        let contents = starlark_data.expect("Action data should be present");
        let action = WriteMacrosToFileAction::new(self.identifier, contents, inputs, outputs)?;
        Ok(Box::new(action))
    }
}

#[derive(Debug, Error)]
enum WriteMacrosActionValidationError {
    #[error("At least one output file must be specified for a write macros action")]
    NoOutputsSpecified,
    #[error("Expected command line value, got {0}")]
    ContentsNotCommandLineValue(String),
    #[error(
        "Number of write-to-file macro artifacts during analysis time should be consistent with when those artifacts are created"
    )]
    InconsistentNumberOfMacroArtifacts,
}

#[derive(Debug, Allocative)]
struct WriteMacrosToFileAction {
    identifier: String,
    contents: OwnedFrozenValue, // StarlarkCommandLine
    inputs: Box<[ArtifactGroup]>,
    outputs: Box<[BuildArtifact]>,
}

impl WriteMacrosToFileAction {
    fn new(
        identifier: String,
        contents: OwnedFrozenValue,
        inputs: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
    ) -> anyhow::Result<Self> {
        if outputs.is_empty() {
            Err(anyhow::anyhow!(
                WriteMacrosActionValidationError::NoOutputsSpecified
            ))
        } else if contents.value().as_command_line().is_none() {
            Err(anyhow::anyhow!(
                WriteMacrosActionValidationError::ContentsNotCommandLineValue(
                    contents.value().to_repr()
                )
            ))
        } else {
            Ok(Self {
                identifier,
                contents,
                inputs: inputs.into_iter().collect(),
                outputs: outputs.into_iter().collect(),
            })
        }
    }
}

#[async_trait]
impl Action for WriteMacrosToFileAction {
    fn kind(&self) -> buck2_data::ActionKind {
        buck2_data::ActionKind::WriteMacrosToFile
    }

    fn inputs(&self) -> anyhow::Result<Cow<'_, [ArtifactGroup]>> {
        Ok(Cow::Borrowed(&self.inputs))
    }

    fn outputs(&self) -> anyhow::Result<Cow<'_, [BuildArtifact]>> {
        Ok(Cow::Borrowed(&self.outputs))
    }

    fn as_executable(&self) -> ActionExecutable<'_> {
        ActionExecutable::Incremental(self)
    }

    fn category(&self) -> &Category {
        static WRITE_MACROS_CATEGORY: Lazy<Category> =
            Lazy::new(|| Category::try_from("write_macros_to_file").unwrap());

        &WRITE_MACROS_CATEGORY
    }

    fn identifier(&self) -> Option<&str> {
        Some(&self.identifier)
    }
}

#[async_trait]
impl IncrementalActionExecutable for WriteMacrosToFileAction {
    async fn execute(
        &self,
        ctx: &mut dyn ActionExecutionCtx,
    ) -> anyhow::Result<(ActionOutputs, ActionExecutionMetadata)> {
        let mut execution_start = None;

        let values = ctx
            .materializer()
            .declare_write(Box::new(|| {
                execution_start = Some(Instant::now());

                let fs = ctx.executor_fs();

                let mut output_contents = Vec::with_capacity(self.outputs.len());
                let mut macro_writer = MacroToFileWriter::new(&fs, &mut output_contents);

                self.contents
                    .value()
                    .as_command_line()
                    .unwrap()
                    .visit_write_to_file_macros(&mut macro_writer)?;

                if self.outputs.len() != output_contents.len() {
                    return Err(
                        WriteMacrosActionValidationError::InconsistentNumberOfMacroArtifacts.into(),
                    );
                }

                Ok(
                    std::iter::zip(self.outputs.iter(), output_contents.into_iter())
                        .map(|(output, content)| WriteRequest {
                            path: fs.fs().resolve_build(output.get_path()),
                            content: content.into_bytes(),
                            is_executable: false,
                        })
                        .collect(),
                )
            }))
            .await?;

        let wall_time = execution_start
            .context("Action did not set execution_start")?
            .elapsed();

        let output_values = std::iter::zip(self.outputs.iter(), values.into_iter())
            .map(|(output, value)| (output.get_path().dupe(), value))
            .collect();

        Ok((
            ActionOutputs::new(output_values),
            ActionExecutionMetadata {
                execution_kind: ActionExecutionKind::Simple,
                timing: ActionExecutionTimingData { wall_time },
            },
        ))
    }
}

struct MacroToFileWriter<'a> {
    fs: &'a ExecutorFs<'a>,
    outputs: &'a mut Vec<String>,
    relative_to_path: Option<RelativePathBuf>,
}

impl<'a> MacroToFileWriter<'a> {
    fn new(fs: &'a ExecutorFs<'a>, outputs: &'a mut Vec<String>) -> Self {
        Self {
            fs,
            outputs,
            relative_to_path: None,
        }
    }
}

impl WriteToFileMacroVisitor for MacroToFileWriter<'_> {
    fn visit_write_to_file_macro(&mut self, resolved_macro: &ResolvedMacro) -> anyhow::Result<()> {
        let content = {
            let mut builder = MacroOutput {
                result: String::new(),
            };
            let mut ctx = MacroContext::new(self.fs, &self.relative_to_path);
            resolved_macro.add_to_arg(&mut builder, &mut ctx)?;
            builder.result
        };

        self.outputs.push(content);
        Ok(())
    }

    fn set_current_relative_to_path(
        &mut self,
        gen: &dyn Fn(&dyn CommandLineContext) -> anyhow::Result<Option<RelativePathBuf>>,
    ) -> anyhow::Result<()> {
        self.relative_to_path = gen(&DefaultCommandLineContext::new(self.fs))?;
        Ok(())
    }
}

struct MacroContext<'a> {
    fs: &'a ExecutorFs<'a>,
    maybe_relative_to_path: &'a Option<RelativePathBuf>,
}

impl<'a> MacroContext<'a> {
    fn new(fs: &'a ExecutorFs, maybe_relative_to_path: &'a Option<RelativePathBuf>) -> Self {
        Self {
            fs,
            maybe_relative_to_path,
        }
    }

    fn relativize_path(&self, path: ProjectRelativePathBuf) -> RelativePathBuf {
        if let Some(relative_to_path) = self.maybe_relative_to_path {
            relative_to_path.relative(path)
        } else {
            path.into()
        }
    }
}

impl CommandLineContext for MacroContext<'_> {
    fn resolve_project_path(
        &self,
        path: ProjectRelativePathBuf,
    ) -> anyhow::Result<CommandLineLocation> {
        Ok(CommandLineLocation::from_relative_path(
            self.relativize_path(path),
            self.fs.path_separator(),
        ))
    }

    fn fs(&self) -> &ExecutorFs {
        self.fs
    }

    fn next_macro_file_path(&mut self) -> anyhow::Result<RelativePathBuf> {
        unreachable!("write-to-file macros could not be nested")
    }
}

// TODO(torozco): Just remove this, and ArgBuilder
struct MacroOutput {
    result: String,
}

impl ArgBuilder for MacroOutput {
    fn push_str(&mut self, s: &str) {
        self.result.push_str(s)
    }
}


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

use anyhow::Context;
use buck2_build_api::actions::artifact::get_artifact_fs::GetArtifactFs;
use buck2_build_api::audit_dep_files::AUDIT_DEP_FILES;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::category::Category;
use buck2_core::directory::Directory;
use buck2_core::directory::DirectoryIterator;
use buck2_core::target::label::ConfiguredTargetLabel;
use buck2_execute::digest_config::HasDigestConfig;
use buck2_execute::materialize::materializer::HasMaterializer;
use dice::DiceTransaction;

use crate::actions::impls::run::dep_files::get_dep_files;
use crate::actions::impls::run::dep_files::read_dep_files;
use crate::actions::impls::run::dep_files::DepFilesKey;
use crate::actions::impls::run::dep_files::StoredFingerprints;

pub(crate) fn init_audit_dep_files() {
    AUDIT_DEP_FILES.init(|ctx, label, category, identifier, stdout| {
        Box::pin(audit_dep_files(ctx, label, category, identifier, stdout))
    });
}

async fn audit_dep_files(
    ctx: &DiceTransaction,
    label: ConfiguredTargetLabel,
    category: Category,
    identifier: Option<String>,
    stdout: &mut (dyn Write + Send),
) -> anyhow::Result<()> {
    let key = DepFilesKey::new(BaseDeferredKey::TargetLabel(label), category, identifier);

    let state = get_dep_files(&key).context("Failed to find dep files")?;

    let dep_files = read_dep_files(
        state.has_signatures(),
        state.declared_dep_files(),
        &ctx.get_artifact_fs().await?,
        ctx.per_transaction_data().get_materializer().as_ref(),
    )
    .await
    .context("Failed to read dep files")?
    .context("Dep fils have expired")?;

    let fingerprints = state.locked_compute_fingerprints(
        Cow::Owned(dep_files),
        true,
        ctx.global_data().get_digest_config(),
    );

    let dirs = match &*fingerprints {
        StoredFingerprints::Digests(..) => {
            // This is bit awkward but this only for testing right now so that's OK
            return Err(anyhow::anyhow!("Fingerprints were stored as digests!"));
        }
        StoredFingerprints::Dirs(dirs) => dirs,
    };

    for (path, ..) in dirs
        .untagged
        .ordered_walk()
        .with_paths()
        .filter_map(|(p, e)| Some((p, e.into_leaf()?)))
    {
        writeln!(stdout, "untagged\t{}", path)?;
    }

    for (tag, dir) in dirs.tagged.iter() {
        for (path, ..) in dir
            .ordered_walk()
            .with_paths()
            .filter_map(|(p, e)| Some((p, e.into_leaf()?)))
        {
            writeln!(stdout, "{}\t{}", tag, path)?;
        }
    }

    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::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Display;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use buck2_artifact::artifact::artifact_type::Artifact;
use buck2_artifact::artifact::artifact_type::OutputArtifact;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::execute::action_execution_target::ActionExecutionTarget;
use buck2_build_api::actions::execute::action_executor::ActionExecutionKind;
use buck2_build_api::actions::execute::action_executor::ActionExecutionMetadata;
use buck2_build_api::actions::execute::action_executor::ActionOutputs;
use buck2_build_api::actions::impls::expanded_command_line::ExpandedCommandLineDigest;
use buck2_build_api::actions::ActionExecutionCtx;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::interpreter::rule_defs::artifact_tagging::ArtifactTag;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::category::Category;
use buck2_core::directory::DirectorySelector;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::fs_util;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::soft_error;
use buck2_events::dispatch::span_async;
use buck2_execute::artifact::artifact_dyn::ArtifactDyn;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::expand_selector_for_dependencies;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::directory::ActionImmutableDirectory;
use buck2_execute::directory::ActionSharedDirectory;
use buck2_execute::directory::INTERNER;
use buck2_execute::materialize::materializer::MaterializationError;
use buck2_execute::materialize::materializer::Materializer;
use buck2_file_watcher::dep_files::FLUSH_DEP_FILES;
use dashmap::DashMap;
use derive_more::Display;
use dupe::Dupe;
use futures::StreamExt;
use once_cell::sync::Lazy;
use parking_lot::MappedMutexGuard;
use parking_lot::Mutex;
use parking_lot::MutexGuard;
use thiserror::Error;
use tracing::instrument;

use crate::actions::impls::run::PreparedRunAction;

#[allocative::root]
static DEP_FILES: Lazy<DashMap<DepFilesKey, Arc<DepFileState>>> = Lazy::new(DashMap::new);

/// When this is set, we retain directories after fingerprintig, so that we can output them later
/// for debugging via `buck2 audit dep-files`.
static KEEP_DIRECTORIES: EnvHelper<bool> = EnvHelper::new("BUCK2_KEEP_DEP_FILE_DIRECTORIES");

/// Forget about all dep files. This isn't really meant to be commonly used, but if an invalid dep
/// file was produced and the user wants unblocking, this will provide it.
fn flush_dep_files() {
    tracing::info!("Flushing {} dep files", DEP_FILES.len());
    DEP_FILES.clear();
}

pub(crate) fn init_flush_dep_files() {
    FLUSH_DEP_FILES.init(flush_dep_files);
}

pub(crate) fn get_dep_files(key: &DepFilesKey) -> Option<Arc<DepFileState>> {
    DEP_FILES.get(key).map(|s| s.dupe())
}

/// A key used to associate a RunAction with a possible previous dep file.
#[derive(Eq, PartialEq, Hash, Display, Allocative)]
#[display(
    fmt = "{} {} {}",
    owner,
    category,
    "identifier.as_deref().unwrap_or(\"<no identifier>\")"
)]
pub struct DepFilesKey {
    owner: BaseDeferredKey,
    category: Category,
    identifier: Option<String>,
}

impl DepFilesKey {
    pub fn new(owner: BaseDeferredKey, category: Category, identifier: Option<String>) -> Self {
        Self {
            owner,
            category,
            identifier,
        }
    }

    pub fn from_action_execution_target(target: ActionExecutionTarget<'_>) -> Self {
        Self {
            owner: target.owner().dupe(),
            category: target.category().clone(),
            identifier: target.identifier().map(|t| t.to_owned()),
        }
    }
}

/// The input signatures for a DepFileState. We compute those lazily, so we either have the input
/// directories (no computation done), or the actual signatures (computation was done).
#[derive(Allocative)]
enum DepFileStateInputSignatures {
    /// Deferred(Some) means we have the input directories but haven't filtered them using the dep
    /// file yet (which in fact we haven't downloaded yet). Deferred(None) is a case that should
    /// never happen, since it would mean we panicked while producing the signatures, which is not
    /// fallible. It's the moral equivalent of a poisoned mutex.
    Deferred(Option<PartitionedInputs<ActionSharedDirectory>>),

    /// Computed represents the case where we have produced the input signatures. We only do this
    /// once at most.
    Computed(StoredFingerprints),
}

impl StoredFingerprints {
    fn to_dep_file_state(
        self,
        digests: CommandDigests,
        declared_dep_files: DeclaredDepFiles,
        result: &ActionOutputs,
    ) -> DepFileState {
        let input_signatures = Mutex::new(DepFileStateInputSignatures::Computed(self));
        DepFileState {
            digests,
            input_signatures,
            declared_dep_files,
            result: result.dupe(),
        }
    }
}

#[derive(Allocative)]
pub enum StoredFingerprints {
    /// Store only digests. This is what we use in prod because it is small.
    Digests(PartitionedInputs<TrackedFileDigest>),

    /// Store digests + dirs. We allow this via BUCK2_KEEP_DEP_FILE_DIRECTORIES because it gives
    /// more debuggability.
    Dirs(PartitionedInputs<ActionImmutableDirectory>),
}

impl PartialEq<PartitionedInputs<ActionImmutableDirectory>> for StoredFingerprints {
    fn eq(&self, other: &PartitionedInputs<ActionImmutableDirectory>) -> bool {
        let fingerprints = match self {
            Self::Digests(fingerprints) => Cow::Borrowed(fingerprints),
            Self::Dirs(dirs) => Cow::Owned(dirs.as_fingerprints()),
        };

        *fingerprints == other.as_fingerprints()
    }
}

/// The state that resulted from the previous evaluation of a command that produced dep files. This
/// contains everything we need to determine whether re-evaluation is necessary (and if it isn't,
/// to return the previous value).
#[derive(Allocative)]
pub struct DepFileState {
    digests: CommandDigests,
    input_signatures: Mutex<DepFileStateInputSignatures>,
    declared_dep_files: DeclaredDepFiles,
    result: ActionOutputs,
}

#[derive(Allocative)]
pub struct CommandDigests {
    pub cli: ExpandedCommandLineDigest,
    pub directory: FileDigest,
}

impl DepFileState {
    pub(crate) fn has_signatures(&self) -> bool {
        match *self.input_signatures.lock() {
            DepFileStateInputSignatures::Computed(..) => true,
            DepFileStateInputSignatures::Deferred(..) => false,
        }
    }

    pub(crate) fn declared_dep_files(&self) -> &DeclaredDepFiles {
        &self.declared_dep_files
    }

    /// Compute the signature for this DepFileState, having provided the dep files from
    /// read_dep_files.
    pub(crate) fn locked_compute_fingerprints<'a>(
        &'a self,
        dep_files: Cow<'_, ConcreteDepFiles>,
        keep_directories: bool,
        digest_config: DigestConfig,
    ) -> MappedMutexGuard<'a, StoredFingerprints> {
        // Now we need to know the signatures on the original action. Produce them if they're
        // missing. We're either storing input directories or outputs here.

        let mut guard = self.input_signatures.lock();

        if let DepFileStateInputSignatures::Deferred(ref mut directories) = *guard {
            let fingerprints = compute_fingerprints(
                directories
                    .take()
                    .expect("Poisoned DepFileStateInputSignatures")
                    .unshare(),
                dep_files.into_owned(),
                digest_config,
                keep_directories,
            );

            *guard = DepFileStateInputSignatures::Computed(fingerprints);
        }

        MutexGuard::map(guard, |v| match v {
            DepFileStateInputSignatures::Computed(signatures) => signatures,
            DepFileStateInputSignatures::Deferred(..) => unreachable!(),
        })
    }
}

/// The set of dep files declared by a RunAction, matching tags to their labels. We enforce at
/// creation time that tags and labels are both unique.
#[derive(Debug, Allocative)]
pub(crate) struct RunActionDepFiles {
    pub(crate) labels: HashMap<ArtifactTag, Arc<str>>,
}

impl Display for RunActionDepFiles {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let values: Vec<&str> = self
            .labels
            .values()
            .map(|x| x.as_ref())
            .collect::<Vec<&str>>();
        write!(f, "{:?}", values)
    }
}

impl RunActionDepFiles {
    pub fn new() -> Self {
        Self {
            labels: HashMap::new(),
        }
    }
}

pub(crate) struct LocalDepFileLookUpKey {
    pub dep_files_key: DepFilesKey,
    pub digests: CommandDigests,
    pub declared_inputs: PartitionedInputs<Vec<ArtifactGroup>>,
    pub declared_dep_files: DeclaredDepFiles,
}

pub(crate) fn make_local_dep_file_lookup_key(
    ctx: &mut dyn ActionExecutionCtx,
    visitor: DepFilesCommandLineVisitor<'_>,
    prepared_action: &PreparedRunAction,
) -> LocalDepFileLookUpKey {
    let digests = CommandDigests {
        cli: prepared_action.expanded.fingerprint(),
        directory: prepared_action
            .paths
            .input_directory()
            .fingerprint()
            .data()
            .dupe(),
    };
    let dep_files_key = DepFilesKey::from_action_execution_target(ctx.target());

    let DepFilesCommandLineVisitor {
        inputs: declared_inputs,
        outputs: declared_dep_files,
        ..
    } = visitor;
    LocalDepFileLookUpKey {
        dep_files_key,
        digests,
        declared_inputs,
        declared_dep_files,
    }
}

pub(crate) async fn check_local_dep_file_cache(
    ctx: &mut dyn ActionExecutionCtx,
    declared_outputs: &[BuildArtifact],
    dep_files: &LocalDepFileLookUpKey,
) -> anyhow::Result<Option<(ActionOutputs, ActionExecutionMetadata)>> {
    let LocalDepFileLookUpKey {
        dep_files_key,
        digests,
        declared_inputs,
        declared_dep_files,
    } = dep_files;

    let matching_result = span_async(buck2_data::MatchDepFilesStart {}, async {
        let res: anyhow::Result<_> = try {
            match_or_clear_dep_file(
                dep_files_key,
                digests,
                declared_inputs,
                declared_outputs,
                declared_dep_files,
                ctx,
            )
            .await?
        };

        (res, buck2_data::MatchDepFilesEnd {})
    })
    .await?;

    if let Some(matching_result) = matching_result {
        Ok(Some((
            matching_result,
            ActionExecutionMetadata {
                execution_kind: ActionExecutionKind::LocalDepFile,
                timing: Default::default(),
            },
        )))
    } else {
        Ok(None)
    }
}

/// Match the dep file recorded for key, or clear it from the map (if it exists).
#[instrument(level = "debug", skip(digests, declared_inputs, ctx), fields(key = %key))]
pub(crate) async fn match_or_clear_dep_file(
    key: &DepFilesKey,
    digests: &CommandDigests,
    declared_inputs: &PartitionedInputs<Vec<ArtifactGroup>>,
    declared_outputs: &[BuildArtifact],
    declared_dep_files: &DeclaredDepFiles,
    ctx: &dyn ActionExecutionCtx,
) -> anyhow::Result<Option<ActionOutputs>> {
    let previous_state = match get_dep_files(key) {
        Some(d) => d.dupe(),
        None => return Ok(None),
    };

    if dep_files_match(
        &previous_state,
        digests,
        declared_inputs,
        declared_outputs,
        declared_dep_files,
        ctx,
    )
    .await?
    {
        let fs = ctx.fs();

        // Finally, we need to make sure that the artifacts in the materializer actually
        // match. This is necessary in case a different action wrote to those artifacts and
        // didn't use the same cache key.
        let output_matches = previous_state
            .result
            .iter()
            .map(|(path, value)| (fs.buck_out_path_resolver().resolve_gen(path), value.dupe()))
            .collect();

        let materializer_accepts = ctx
            .materializer()
            .declare_match(output_matches)
            .await?
            .is_match();

        if materializer_accepts {
            tracing::trace!("Dep files are a hit");
            return Ok(Some(previous_state.result.dupe()));
        }
    }

    tracing::trace!("Dep files are a miss");
    DEP_FILES.remove(key);
    Ok(None)
}

async fn dep_files_match(
    previous_state: &DepFileState,
    digests: &CommandDigests,
    declared_inputs: &PartitionedInputs<Vec<ArtifactGroup>>,
    declared_outputs: &[BuildArtifact],
    declared_dep_files: &DeclaredDepFiles,
    ctx: &dyn ActionExecutionCtx,
) -> anyhow::Result<bool> {
    if !declared_dep_files.declares_same_dep_files(&previous_state.declared_dep_files) {
        // We first need to check if the same dep files existed before or not. If not, then we
        // can't assume they'll still be on disk, and we have to bail.
        tracing::trace!("Dep files miss: Dep files declaration has changed");
        return Ok(false);
    }

    if !outputs_are_reusable(declared_outputs, &previous_state.result) {
        tracing::trace!("Dep files miss: Output declaration has changed");
        return Ok(false);
    }

    if digests.cli != previous_state.digests.cli {
        tracing::trace!("Dep files miss: Command line has changed");
        return Ok(false);
    }

    if digests.directory == previous_state.digests.directory {
        tracing::trace!("Dep files hit: Command line and directory have not changed");
        return Ok(true);
    }

    let dep_files = read_dep_files(
        previous_state.has_signatures(),
        previous_state.declared_dep_files(),
        ctx.fs(),
        ctx.materializer(),
    )
    .await
    .context(
        "Error reading persisted dep files. \
            Fix the command that produced an invalid dep file. \
            You may also use `buck2 debug flush-dep-files` to drop all dep file state.",
    )?;

    let dep_files = match dep_files {
        Some(dep_files) => dep_files,
        None => {
            tracing::trace!("Dep files miss: Dep files cannot be materialized");
            return Ok(false);
        }
    };

    // Now we need to know the fingerprints on the original action. Produce them if they're
    // missing. We're either storing input directories or outputs here.

    let digest_config = ctx.digest_config();

    let fingerprints_match = {
        // NOTE: We don't bother releasing the guard here (we'd have to clone the fingerprints to do
        // so), because this Mutex won't be contended: only one action will look at its value.
        let previous_fingerprints = previous_state.locked_compute_fingerprints(
            Cow::Borrowed(&dep_files),
            KEEP_DIRECTORIES.get_copied()?.unwrap_or_default(),
            digest_config,
        );

        // NOTE: We use the new directory to e.g. resolve symlinks referenced in the dep file. This
        // makes sense: if a path in the depfile is still a symlink, then we'll compare the new
        // destination and the old. If it's not, then we can assume that the tool wouldn't traverse
        // the symlink anymore.
        let new_fingerprints = declared_inputs
            .to_directories(ctx)?
            .filter(dep_files)
            .fingerprint(digest_config);

        *previous_fingerprints == new_fingerprints
    };

    Ok(fingerprints_match)
}

/// If an action is unchanged but now requires a different set of outputs, that's not a cache hit
/// because we need to rehash the outputs. Having to re-run the action isn't the best, but it's
/// probably infrequent enough that we seem unlikely to care.
fn outputs_are_reusable(declared_outputs: &[BuildArtifact], outputs: &ActionOutputs) -> bool {
    for out in declared_outputs {
        if outputs.get(out.get_path()).is_none() {
            return false;
        }
    }

    true
}

/// Read the dep files for this DepFileState. This will return None if the dep files cannot be
/// materialized (because they e.g. expired).
pub(crate) async fn read_dep_files(
    has_signatures: bool,
    declared_dep_files: &DeclaredDepFiles,
    fs: &ArtifactFs,
    materializer: &dyn Materializer,
) -> anyhow::Result<Option<ConcreteDepFiles>> {
    // NOTE: We only materialize if we haven't computed our signatures yet, since we know we
    // can't have computed our signatures without having read the dep file already. In an ideal
    // world this wouldn't be necessary, but in practice contention on the materializer makes
    // this slower.
    if !has_signatures {
        match declared_dep_files.materialize(fs, materializer).await {
            Ok(()) => {}
            Err(MaterializeDepFilesError::NotFound) => return Ok(None),
            Err(e) => return Err(e.into()),
        };
    }

    let dep_files = declared_dep_files
        .read(fs)
        .context("Error reading dep files, verify that the action produced valid output")?;

    Ok(dep_files)
}

fn compute_fingerprints(
    directories: PartitionedInputs<ActionDirectoryBuilder>,
    dep_files: ConcreteDepFiles,
    digest_config: DigestConfig,
    keep_directories: bool,
) -> StoredFingerprints {
    let filtered_directories = directories.filter(dep_files).fingerprint(digest_config);
    if keep_directories {
        StoredFingerprints::Dirs(filtered_directories)
    } else {
        StoredFingerprints::Digests(filtered_directories.as_fingerprints())
    }
}

pub(crate) async fn eagerly_compute_fingerprints(
    declared_inputs: &PartitionedInputs<Vec<ArtifactGroup>>,
    declared_dep_files: &DeclaredDepFiles,
    ctx: &dyn ActionExecutionCtx,
) -> anyhow::Result<StoredFingerprints> {
    let directories = declared_inputs.to_directories(ctx)?;
    let dep_files = read_dep_files(false, declared_dep_files, ctx.fs(), ctx.materializer())
        .await?
        .context("Dep file not found")?;

    let fingerprints = compute_fingerprints(
        directories,
        dep_files,
        ctx.digest_config(),
        KEEP_DIRECTORIES.get_copied()?.unwrap_or_default(),
    );
    Ok(fingerprints)
}

/// Post-process the dep files produced by an action.
pub(crate) async fn populate_dep_files(
    key: DepFilesKey,
    digests: CommandDigests,
    declared_inputs: PartitionedInputs<Vec<ArtifactGroup>>,
    declared_dep_files: DeclaredDepFiles,
    result: &ActionOutputs,
    ctx: &dyn ActionExecutionCtx,
) -> anyhow::Result<()> {
    let state = if declared_dep_files.is_empty() || ctx.run_action_knobs().eager_dep_files {
        let fingerprint =
            eagerly_compute_fingerprints(&declared_inputs, &declared_dep_files, ctx).await?;
        fingerprint.to_dep_file_state(digests, declared_dep_files, result)
    } else {
        let directories = declared_inputs.to_directories(ctx)?;
        DepFileState {
            digests,
            input_signatures: Mutex::new(DepFileStateInputSignatures::Deferred(Some(
                directories.share(ctx.digest_config()),
            ))),
            declared_dep_files,
            result: result.dupe(),
        }
    };

    DEP_FILES.insert(key, Arc::new(state));
    Ok(())
}

/// Inputs partitioned by tag. `D` is the representation of the set of inputs.
#[derive(Clone, PartialEq, Eq, Allocative)]
pub struct PartitionedInputs<D> {
    pub untagged: D,
    pub tagged: HashMap<Arc<str>, D>,
}

impl<D> PartitionedInputs<D> {
    pub fn iter(&self) -> impl Iterator<Item = &D> {
        std::iter::once(&self.untagged).chain(self.tagged.values())
    }
}

impl Default for PartitionedInputs<Vec<ArtifactGroup>> {
    fn default() -> Self {
        Self {
            tagged: Default::default(),
            untagged: Default::default(),
        }
    }
}

impl PartitionedInputs<Vec<ArtifactGroup>> {
    /// Add an input to this set of PartitionedInputs.
    fn visit_input(
        &mut self,
        input: ArtifactGroup,
        tag: Option<&ArtifactTag>,
        dep_files: &RunActionDepFiles,
    ) {
        let input_group = match tag {
            None => &mut self.untagged,
            Some(tag) => {
                // NOTE: If an input has a tag that doesn't match a dep file, we don't care about
                // it.
                match dep_files.labels.get(tag) {
                    Some(label) => self.tagged.entry(label.dupe()).or_default(),
                    None => &mut self.untagged,
                }
            }
        };

        input_group.push(input);
    }

    /// Produce Directories from this set of PartitionedInputs. One directory will be produced for
    /// each tag (and one for untagged). This will actually allocate directories (whereas until now
    /// we only held references to artifacts).
    fn to_directories(
        &self,
        ctx: &dyn ActionExecutionCtx,
    ) -> anyhow::Result<PartitionedInputs<ActionDirectoryBuilder>> {
        let reduce = |inputs: &[ArtifactGroup]| {
            let mut builder = ActionDirectoryBuilder::empty();

            for input in inputs {
                let input = ctx.artifact_values(input);
                input.add_to_directory(&mut builder, ctx.fs())?;
            }

            anyhow::Ok(builder)
        };

        Ok(PartitionedInputs {
            untagged: reduce(&self.untagged)?,
            tagged: self
                .tagged
                .iter()
                .map(|(tag, inputs)| anyhow::Ok((tag.dupe(), reduce(inputs)?)))
                .collect::<Result<_, _>>()?,
        })
    }
}

impl PartitionedInputs<ActionSharedDirectory> {
    fn unshare(self) -> PartitionedInputs<ActionDirectoryBuilder> {
        PartitionedInputs {
            untagged: self.untagged.into_builder(),
            tagged: self
                .tagged
                .into_iter()
                .map(|(k, v)| (k, v.into_builder()))
                .collect(),
        }
    }
}

impl PartitionedInputs<ActionDirectoryBuilder> {
    fn share(self, digest_config: DigestConfig) -> PartitionedInputs<ActionSharedDirectory> {
        PartitionedInputs {
            untagged: self
                .untagged
                .fingerprint(digest_config.as_directory_serializer())
                .shared(&*INTERNER),
            tagged: self
                .tagged
                .into_iter()
                .map(|(k, v)| {
                    (
                        k,
                        v.fingerprint(digest_config.as_directory_serializer())
                            .shared(&*INTERNER),
                    )
                })
                .collect(),
        }
    }

    /// Filter this set of PartitionedInputs using dep files that were produced by the command (or
    /// a previous invocation thereof). For each partition of inputs (i.e. untagged, or specific
    /// tags), only the inputs listed in the dep file will be retained.
    fn filter(mut self, mut concrete_dep_files: ConcreteDepFiles) -> Self {
        fn filter(
            label: &Arc<str>,
            builder: &mut ActionDirectoryBuilder,
            concrete_dep_files: &mut ConcreteDepFiles,
        ) {
            let matching_selector = match concrete_dep_files.contents.get_mut(label) {
                Some(s) => s,
                None => return,
            };

            expand_selector_for_dependencies(builder, matching_selector);

            // NOTE: We ignore the filtering if it produces an invalid directory. If we can't
            // filter, that means we're selecting into a leaf, which means one of two things:
            //
            // - We hit an external symlink, and we can't traverse through those.
            // - The directory structure changed and was a dir is now a leaf, in which case we'll
            // select the leaf but when comparing this to the dir it'll conflict.
            let _ignored = matching_selector.filter(builder);
        }

        for (label, dir) in self.tagged.iter_mut() {
            filter(label, dir, &mut concrete_dep_files);
        }

        self
    }

    /// Compute fingerprints for all the PartitionedInputs in here. This lets us do comparisons.
    fn fingerprint(
        self,
        digest_config: DigestConfig,
    ) -> PartitionedInputs<ActionImmutableDirectory> {
        PartitionedInputs {
            untagged: self
                .untagged
                .fingerprint(digest_config.as_directory_serializer()),
            tagged: self
                .tagged
                .into_iter()
                .map(|(k, v)| (k, v.fingerprint(digest_config.as_directory_serializer())))
                .collect(),
        }
    }
}

impl PartitionedInputs<ActionImmutableDirectory> {
    fn as_fingerprints(&self) -> PartitionedInputs<TrackedFileDigest> {
        PartitionedInputs {
            untagged: self.untagged.fingerprint().dupe(),
            tagged: self
                .tagged
                .iter()
                .map(|(k, v)| (k.dupe(), v.fingerprint().dupe()))
                .collect(),
        }
    }
}

/// A dep file declared by a command. This includes its label (which we use for stability across
/// command executions), as well as the artifact that the command will write the dep file to.
#[derive(Hash, PartialEq, Eq, Debug, Dupe, Clone, Allocative)]
struct DeclaredDepFile {
    label: Arc<str>,
    output: Artifact,
}

/// All the dep files declared by a command;
#[derive(Default, Debug, Allocative)]
pub(crate) struct DeclaredDepFiles {
    tagged: HashMap<ArtifactTag, DeclaredDepFile>,
}

impl DeclaredDepFiles {
    fn is_empty(&self) -> bool {
        self.tagged.is_empty()
    }

    /// Add dep file to this set.
    fn visit_output(
        &mut self,
        artifact: OutputArtifact,
        tag: Option<&ArtifactTag>,
        dep_files: &RunActionDepFiles,
    ) {
        match tag {
            None => {}
            Some(tag) => {
                // NOTE: We have validated tags earlier, so we know that if a tag does not point to
                // a dep file here, it's safe to ignore it. We also know that we'll have exactly 1
                // dep file per tag.
                if let Some(label) = dep_files.labels.get(tag) {
                    // NOTE: analysis has been done so we know inputs are bound now.
                    let output = (*artifact).dupe().ensure_bound().unwrap().into_artifact();

                    self.tagged.insert(
                        tag.dupe(),
                        DeclaredDepFile {
                            label: label.dupe(),
                            output,
                        },
                    );
                }
            }
        }
    }

    /// Given an ActionOutputs, materialize this set of dep files, so that we may read them later.
    async fn materialize(
        &self,
        fs: &ArtifactFs,
        materializer: &dyn Materializer,
    ) -> Result<(), MaterializeDepFilesError> {
        let mut paths = Vec::with_capacity(self.tagged.len());

        for declared_dep_file in self.tagged.values() {
            let dep_file = &declared_dep_file.output;
            let path = dep_file
                .resolve_path(fs)
                .map_err(|e| MaterializeDepFilesError::MaterializationFailed { source: e })?;
            paths.push(path);
        }

        if paths.is_empty() {
            return Ok(());
        }

        let mut has_not_found = false;
        let mut stream = materializer
            .materialize_many(paths)
            .await
            .map_err(|e| MaterializeDepFilesError::MaterializationFailed { source: e })?;

        while let Some(dep_file) = stream.next().await {
            match dep_file {
                Ok(()) => {}
                Err(MaterializationError::NotFound { .. }) => {
                    has_not_found = true;
                }
                Err(e) => {
                    return Err(MaterializeDepFilesError::MaterializationFailed {
                        source: e.into(),
                    });
                }
            }
        }

        if has_not_found {
            Err(MaterializeDepFilesError::NotFound)
        } else {
            Ok(())
        }
    }

    /// Read this set of dep files, producing ConcreteDepFiles. This can then be used to compute
    /// signatures for the input set that was used, and for future input sets.
    fn read(&self, fs: &ArtifactFs) -> anyhow::Result<Option<ConcreteDepFiles>> {
        let mut contents = HashMap::with_capacity(self.tagged.len());

        for declared_dep_file in self.tagged.values() {
            let mut selector = DirectorySelector::empty();

            let dep_file = declared_dep_file.output.resolve_path(fs)?;

            let read_dep_file: anyhow::Result<()> = try {
                let dep_file_path = fs.fs().resolve(&dep_file);
                let dep_file = fs_util::read_to_string_opt(&dep_file_path)?;

                let dep_file = match dep_file {
                    Some(dep_file) => dep_file,
                    None => {
                        soft_error!(
                            "missing_dep_file",
                            anyhow::anyhow!("Dep file is missing at {}", dep_file_path)
                        )?;
                        return Ok(None);
                    }
                };

                for line in dep_file.split('\n') {
                    let line = line.trim();
                    if line.is_empty() {
                        continue;
                    }
                    let path = ProjectRelativePath::new(line)
                        .context("Invalid line encountered in dep file")?;

                    selector.select(path);
                }
            };

            read_dep_file.with_context(|| {
                format!(
                    "Action execution produced an invalid `{}` dep file at `{}`",
                    declared_dep_file.label, dep_file,
                )
            })?;

            contents.insert(declared_dep_file.label.dupe(), selector);
        }

        Ok(Some(ConcreteDepFiles { contents }))
    }

    /// Returns whether two DeclaredDepFile instances have the same dep files. This ignores the tag
    /// identity, but it requires the same paths declared using the same name. This is a
    /// pre-requisite for being able to reuse dep files from a previous invocation.
    fn declares_same_dep_files(&self, other: &Self) -> bool {
        let this = self.tagged.values().collect::<HashSet<_>>();
        let other = other.tagged.values().collect::<HashSet<_>>();
        this == other
    }
}

#[derive(Error, Debug)]
enum MaterializeDepFilesError {
    #[error("Error materializing dep file")]
    MaterializationFailed {
        #[source]
        source: anyhow::Error,
    },

    #[error("A dep file was not found")]
    NotFound,
}

/// A set of concrete dep files. That is, given a label, a selector that represents the subset of
/// files whose tags matches this label that should be considered relevant.
#[derive(Clone)]
pub struct ConcreteDepFiles {
    contents: HashMap<Arc<str>, DirectorySelector>,
}

/// A command line visitor to collect inputs and outputs in a form relevant for dep files
/// computations.
pub(crate) struct DepFilesCommandLineVisitor<'a> {
    pub inputs: PartitionedInputs<Vec<ArtifactGroup>>,
    pub outputs: DeclaredDepFiles,
    dep_files: &'a RunActionDepFiles,
}

impl<'a> DepFilesCommandLineVisitor<'a> {
    pub(crate) fn new(dep_files: &'a RunActionDepFiles) -> Self {
        Self {
            inputs: Default::default(),
            outputs: Default::default(),
            dep_files,
        }
    }
}

impl CommandLineArtifactVisitor for DepFilesCommandLineVisitor<'_> {
    fn visit_input(&mut self, input: ArtifactGroup, tag: Option<&ArtifactTag>) {
        self.inputs.visit_input(input, tag, self.dep_files);
    }

    fn visit_output(&mut self, artifact: OutputArtifact, tag: Option<&ArtifactTag>) {
        self.outputs.visit_output(artifact, tag, self.dep_files);
    }
}

#[cfg(test)]
mod test {

    use buck2_artifact::artifact::artifact_type::testing::BuildArtifactTestingExt;
    use buck2_artifact::deferred::id::DeferredId;
    use buck2_core::configuration::data::ConfigurationData;
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
    use buck2_core::target::label::ConfiguredTargetLabel;
    use maplit::hashmap;

    use super::*;

    #[test]
    fn test_declares_same_dep_files() {
        let target =
            ConfiguredTargetLabel::testing_parse("cell//pkg:foo", ConfigurationData::testing_new());

        let artifact1 = Artifact::from(BuildArtifact::testing_new(
            target.dupe(),
            ForwardRelativePathBuf::unchecked_new("foo/bar1.h".to_owned()),
            DeferredId::testing_new(0),
        ));

        let artifact2 = Artifact::from(BuildArtifact::testing_new(
            target.dupe(),
            ForwardRelativePathBuf::unchecked_new("foo/bar2.h".to_owned()),
            DeferredId::testing_new(0),
        ));

        let depfile1 = DeclaredDepFile {
            label: Arc::from("foo"),
            output: artifact1,
        };

        let depfile2 = DeclaredDepFile {
            label: Arc::from("foo"),
            output: artifact2.dupe(),
        };

        let depfile3 = DeclaredDepFile {
            label: Arc::from("bar"),
            output: artifact2,
        };

        let tag1 = ArtifactTag::new();

        let tag2 = ArtifactTag::new();

        let decl1 = DeclaredDepFiles {
            tagged: hashmap! { tag1.dupe() => depfile1.dupe() },
        };

        let decl2 = DeclaredDepFiles {
            tagged: hashmap! { tag2.dupe() => depfile1.dupe() },
        };

        let decl3 = DeclaredDepFiles {
            tagged: hashmap! { tag2.dupe() => depfile2.dupe() },
        };

        let decl4 = DeclaredDepFiles {
            tagged: hashmap! { tag2.dupe() => depfile3.dupe() },
        };

        assert!(decl1.declares_same_dep_files(&decl1));
        assert!(decl1.declares_same_dep_files(&decl2));
        assert!(!decl2.declares_same_dep_files(&decl3));
        assert!(!decl3.declares_same_dep_files(&decl4));
    }
}


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

use std::fmt::Display;

use buck2_build_api::artifact_groups::ArtifactGroupValues;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::Directory;
use buck2_core::directory::DirectoryEntry;
use buck2_core::directory::DirectoryIterator;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::directory::ActionDirectoryMember;
use serde::Serialize;
use serde::Serializer;

fn stringify<T, S>(value: &T, serializer: S) -> Result<S::Ok, S::Error>
where
    T: Display,
    S: Serializer,
{
    serializer.collect_str(value)
}

pub(crate) fn metadata_content(
    fs: &ArtifactFs,
    inputs: &[&ArtifactGroupValues],
    digest_config: DigestConfig,
) -> anyhow::Result<(Vec<u8>, TrackedFileDigest)> {
    let mut builder = ActionDirectoryBuilder::empty();
    for &group in inputs {
        group.add_to_directory(&mut builder, fs)?;
    }

    #[derive(Serialize)]
    struct PathWithDigest<'a> {
        path: ForwardRelativePathBuf,
        #[serde(serialize_with = "stringify")]
        digest: &'a FileDigest,
    }

    #[derive(Serialize)]
    struct MetadataJson<'a> {
        version: i32,
        digests: Vec<PathWithDigest<'a>>,
    }

    let mut digests = Vec::new();
    let mut walk = builder.ordered_walk();
    while let Some((path, item)) = walk.next() {
        match item {
            DirectoryEntry::Leaf(ActionDirectoryMember::File(metadata)) => {
                digests.push(PathWithDigest {
                    path: path.get(),
                    digest: metadata.digest.data(),
                });
            }
            // Omit symlinks and let user script detect and handle symlinks in inputs.
            // Metadata will contain artifacts which are symlinked, meaning the user
            // can resolve the symlink and get the digest of the symlinked artifact.
            DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(_))
            | DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(_)) => {}
            // Only interested in actual content, skip.
            DirectoryEntry::Dir(_) => {}
        }
    }

    let json = MetadataJson {
        digests,
        // Increment this version if format changes
        version: 1,
    };
    let json_string = serde_json::to_string(&json)?;
    let digest =
        TrackedFileDigest::from_content(json_string.as_bytes(), digest_config.cas_digest_config());
    Ok((json_string.into(), digest))
}


/*
 * 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::Display;
use std::ops::ControlFlow;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_artifact::artifact::build_artifact::BuildArtifact;
use buck2_build_api::actions::box_slice_set::BoxSliceSet;
use buck2_build_api::actions::execute::action_executor::ActionExecutionMetadata;
use buck2_build_api::actions::execute::action_executor::ActionOutputs;
use buck2_build_api::actions::impls::expanded_command_line::ExpandedCommandLine;
use buck2_build_api::actions::Action;
use buck2_build_api::actions::ActionExecutable;
use buck2_build_api::actions::ActionExecutionCtx;
use buck2_build_api::actions::IncrementalActionExecutable;
use buck2_build_api::actions::UnregisteredAction;
use buck2_build_api::artifact_groups::ArtifactGroup;
use buck2_build_api::artifact_groups::ArtifactGroupValues;
use buck2_build_api::interpreter::rule_defs::cmd_args::value_as::ValueAsCommandLineLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArgLike;
use buck2_build_api::interpreter::rule_defs::cmd_args::CommandLineArtifactVisitor;
use buck2_build_api::interpreter::rule_defs::cmd_args::DefaultCommandLineContext;
use buck2_build_api::interpreter::rule_defs::cmd_args::SimpleCommandLineArtifactVisitor;
use buck2_build_api::interpreter::rule_defs::provider::builtin::worker_info::WorkerInfo;
use buck2_core::category::Category;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_execute::artifact::fs::ExecutorFs;
use buck2_execute::execute::environment_inheritance::EnvironmentInheritance;
use buck2_execute::execute::request::ActionMetadataBlob;
use buck2_execute::execute::request::CommandExecutionInput;
use buck2_execute::execute::request::CommandExecutionOutput;
use buck2_execute::execute::request::CommandExecutionPaths;
use buck2_execute::execute::request::CommandExecutionRequest;
use buck2_execute::execute::request::ExecutorPreference;
use buck2_execute::execute::request::WorkerId;
use buck2_execute::execute::request::WorkerSpec;
use derive_more::Display;
use dupe::Dupe;
use gazebo::prelude::*;
use host_sharing::HostSharingRequirements;
use host_sharing::WeightClass;
use indexmap::indexmap;
use indexmap::IndexSet;
use itertools::Itertools;
use serde_json::json;
use sorted_vector_map::SortedVectorMap;
use starlark::coerce::Coerce;
use starlark::starlark_complex_value;
use starlark::values::dict::DictRef;
use starlark::values::none::NoneOr;
use starlark::values::starlark_value;
use starlark::values::Freeze;
use starlark::values::NoSerialize;
use starlark::values::OwnedFrozenValue;
use starlark::values::OwnedFrozenValueTyped;
use starlark::values::ProvidesStaticType;
use starlark::values::StarlarkValue;
use starlark::values::Trace;
use starlark::values::UnpackValue;
use starlark::values::ValueLike;
use starlark::values::ValueOf;
use thiserror::Error;

use crate::actions::impls::run::dep_files::check_local_dep_file_cache;
use crate::actions::impls::run::dep_files::make_local_dep_file_lookup_key;
use crate::actions::impls::run::dep_files::populate_dep_files;
use crate::actions::impls::run::dep_files::DepFilesCommandLineVisitor;
use crate::actions::impls::run::dep_files::LocalDepFileLookUpKey;
use crate::actions::impls::run::dep_files::RunActionDepFiles;
use crate::actions::impls::run::metadata::metadata_content;

pub(crate) mod audit_dep_files;
pub mod dep_files;
mod metadata;

#[derive(Debug, Error)]
enum RunActionValidationError {
    #[error("Expected command line value, got {0}")]
    ContentsNotCommandLineValue(String),
}

#[derive(Debug, Allocative)]
pub(crate) struct MetadataParameter {
    /// Name of the environment variable which is set to contain
    /// resolved path of the metadata file when requested by user.
    pub(crate) env_var: String,
    /// User-defined path in the output directory of the metadata file.
    pub(crate) path: ForwardRelativePathBuf,
}

impl Display for MetadataParameter {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let json = json!({
            "env_var": self.env_var,
            "path": self.path,
        });
        write!(f, "{}", json)
    }
}

#[derive(Debug, Error)]
enum LocalPreferenceError {
    #[error("cannot have `local_only = True` and `prefer_local = True` at the same time")]
    LocalOnlyAndPreferLocal,
    #[error("cannot have `local_only = True` and `prefer_remote = True` at the same time")]
    LocalOnlyAndPreferRemote,
    #[error(
        "cannot have `local_only = True`, `prefer_local = True` and `prefer_remote = True` at the same time"
    )]
    LocalOnlyAndPreferLocalAndPreferRemote,
    #[error("cannot have `prefer_local = True` and `prefer_remote = True` at the same time")]
    PreferLocalAndPreferRemote,
}

pub(crate) fn new_executor_preference(
    local_only: bool,
    prefer_local: bool,
    prefer_remote: bool,
) -> anyhow::Result<ExecutorPreference> {
    match (local_only, prefer_local, prefer_remote) {
        (true, false, false) => Ok(ExecutorPreference::LocalRequired),
        (true, false, true) => Err(anyhow::anyhow!(
            LocalPreferenceError::LocalOnlyAndPreferRemote
        )),
        (false, true, false) => Ok(ExecutorPreference::LocalPreferred),
        (false, true, true) => Err(anyhow::anyhow!(
            LocalPreferenceError::PreferLocalAndPreferRemote
        )),
        (false, false, false) => Ok(ExecutorPreference::Default),
        (false, false, true) => Ok(ExecutorPreference::RemotePreferred),
        (true, true, false) => Err(anyhow::anyhow!(
            LocalPreferenceError::LocalOnlyAndPreferLocal
        )),
        (true, true, true) => Err(anyhow::anyhow!(
            LocalPreferenceError::LocalOnlyAndPreferLocalAndPreferRemote
        )),
    }
}

#[derive(Debug, Allocative)]
pub(crate) struct UnregisteredRunAction {
    pub(crate) category: Category,
    pub(crate) identifier: Option<String>,
    pub(crate) executor_preference: ExecutorPreference,
    pub(crate) always_print_stderr: bool,
    pub(crate) weight: WeightClass,
    pub(crate) low_pass_filter: bool,
    pub(crate) dep_files: RunActionDepFiles,
    pub(crate) metadata_param: Option<MetadataParameter>,
    pub(crate) no_outputs_cleanup: bool,
    pub(crate) allow_cache_upload: bool,
    pub(crate) force_full_hybrid_if_capable: bool,
}

impl UnregisteredAction for UnregisteredRunAction {
    fn register(
        self: Box<Self>,
        _: IndexSet<ArtifactGroup>,
        outputs: IndexSet<BuildArtifact>,
        starlark_data: Option<OwnedFrozenValue>,
    ) -> anyhow::Result<Box<dyn Action>> {
        let starlark_values =
            starlark_data.context("module data to be present (internal error)")?;
        let run_action = RunAction::new(*self, starlark_values, outputs)?;
        Ok(Box::new(run_action))
    }
}

#[derive(
    Debug,
    Display,
    Trace,
    ProvidesStaticType,
    NoSerialize,
    Allocative,
    Freeze,
    Coerce
)]
#[display(fmt = "run_action_values")]
#[repr(C)]
pub(crate) struct StarlarkRunActionValuesGen<V> {
    pub(crate) exe: V,
    pub(crate) args: V,
    pub(crate) env: V,
    /// `WorkerInfo` or `None`.
    pub(crate) worker: V,
}

#[starlark_value(type = "run_action_values")]
impl<'v, V: ValueLike<'v> + 'v> StarlarkValue<'v> for StarlarkRunActionValuesGen<V> where
    Self: ProvidesStaticType<'v>
{
}

starlark_complex_value!(pub(crate) StarlarkRunActionValues);

impl<'v, V: ValueLike<'v>> StarlarkRunActionValuesGen<V> {
    pub(crate) fn worker(&self) -> anyhow::Result<ValueOf<'v, NoneOr<&'v WorkerInfo<'v>>>> {
        ValueOf::unpack_value_err(self.worker.to_value())
    }
}

struct UnpackedRunActionValues<'v> {
    exe: &'v dyn CommandLineArgLike,
    args: &'v dyn CommandLineArgLike,
    env: Vec<(&'v str, &'v dyn CommandLineArgLike)>,
    worker: Option<(&'v dyn CommandLineArgLike, WorkerId)>,
}

#[derive(Debug, Allocative)]
pub(crate) struct RunAction {
    inner: UnregisteredRunAction,
    starlark_values: OwnedFrozenValueTyped<FrozenStarlarkRunActionValues>,
    outputs: BoxSliceSet<BuildArtifact>,
}

impl RunAction {
    fn unpack(
        values: &OwnedFrozenValueTyped<FrozenStarlarkRunActionValues>,
    ) -> anyhow::Result<UnpackedRunActionValues> {
        let exe = values.exe.to_value().as_command_line_err()?;
        let args = values.args.to_value().as_command_line_err()?;
        let env = if values.env.is_none() {
            Vec::new()
        } else {
            let d = DictRef::from_value(values.env.to_value()).context("expecting dict")?;
            let mut res = Vec::with_capacity(d.len());
            for (k, v) in d.iter() {
                res.push((
                    k.unpack_str().context("expecting string")?,
                    v.as_command_line_err()?,
                ));
            }
            res
        };
        let worker: NoneOr<&WorkerInfo> = values.worker()?.typed;

        let worker = if let Some(worker) = worker.into_option() {
            let worker_exe = worker.exe_command_line();
            let worker_id = WorkerId(worker.id);
            Some((worker_exe, worker_id))
        } else {
            None
        };

        Ok(UnpackedRunActionValues {
            exe,
            args,
            env,
            worker,
        })
    }

    /// Get the command line expansion for this RunAction.
    fn expand_command_line_and_worker(
        &self,
        fs: &ExecutorFs,
        artifact_visitor: &mut impl CommandLineArtifactVisitor,
    ) -> anyhow::Result<(ExpandedCommandLine, Option<WorkerSpec>)> {
        let mut ctx = DefaultCommandLineContext::new(fs);
        let values = Self::unpack(&self.starlark_values)?;

        let mut exe_rendered = Vec::<String>::new();
        values
            .exe
            .add_to_command_line(&mut exe_rendered, &mut ctx)?;
        values.exe.visit_artifacts(artifact_visitor)?;

        let worker = if let Some((worker_exe, worker_id)) = values.worker {
            let mut worker_rendered = Vec::<String>::new();
            worker_exe.add_to_command_line(&mut worker_rendered, &mut ctx)?;
            worker_exe.visit_artifacts(artifact_visitor)?;
            Some(WorkerSpec {
                id: worker_id,
                exe: worker_rendered,
            })
        } else {
            None
        };

        let mut args_rendered = Vec::<String>::new();
        values
            .args
            .add_to_command_line(&mut args_rendered, &mut ctx)?;
        values.args.visit_artifacts(artifact_visitor)?;

        let cli_env: anyhow::Result<SortedVectorMap<_, _>> = values
            .env
            .into_iter()
            .map(|(k, v)| {
                let mut env = Vec::<String>::new(); // TODO (torozco): Use a String.
                let mut ctx = DefaultCommandLineContext::new(fs);
                v.add_to_command_line(&mut env, &mut ctx)?;
                v.visit_artifacts(artifact_visitor)?;
                let var = env.join(" ");
                Ok((k.to_owned(), var))
            })
            .collect();

        Ok((
            ExpandedCommandLine {
                exe: exe_rendered,
                args: args_rendered,
                env: cli_env?,
            },
            worker,
        ))
    }

    pub(crate) fn new(
        inner: UnregisteredRunAction,
        starlark_values: OwnedFrozenValue,
        outputs: IndexSet<BuildArtifact>,
    ) -> anyhow::Result<Self> {
        let starlark_values = match starlark_values.downcast() {
            Ok(starlark_values) => starlark_values,
            Err(starlark_values) => {
                return Err(RunActionValidationError::ContentsNotCommandLineValue(
                    starlark_values.value().to_repr(),
                )
                .into());
            }
        };

        Self::unpack(&starlark_values)?;

        Ok(RunAction {
            inner,
            starlark_values,
            outputs: BoxSliceSet::from(outputs),
        })
    }

    fn prepare(
        &self,
        visitor: &mut impl RunActionVisitor,
        ctx: &mut dyn ActionExecutionCtx,
    ) -> anyhow::Result<PreparedRunAction> {
        let fs = ctx.fs();

        let (expanded, worker) =
            self.expand_command_line_and_worker(&ctx.executor_fs(), visitor)?;

        // TODO (@torozco): At this point, might as well just receive the list already. Finding
        // those things in a HashMap is just not very useful.
        let artifact_inputs: Vec<&ArtifactGroupValues> = visitor
            .inputs()
            .map(|group| ctx.artifact_values(group))
            .collect();

        let mut inputs: Vec<CommandExecutionInput> =
            artifact_inputs[..].map(|&i| CommandExecutionInput::Artifact(Box::new(i.dupe())));

        // Handle case when user requested file with action metadata to be generated.
        // Generate content and output path for the file. It will be either passed
        // to RE as a blob or written to disk in local executor.
        // Path to this file is passed to user in environment variable which is selected by user.
        let extra_env = if let Some(metadata_param) = &self.inner.metadata_param {
            let path = BuckOutPath::new(ctx.target().owner().dupe(), metadata_param.path.clone());
            let resolved_path = fs.buck_out_path_resolver().resolve_gen(&path);
            let extra = (metadata_param.env_var.to_owned(), resolved_path.to_string());
            let (data, digest) = metadata_content(fs, &artifact_inputs, ctx.digest_config())?;
            inputs.push(CommandExecutionInput::ActionMetadata(ActionMetadataBlob {
                data,
                digest,
                path,
            }));
            Some(extra)
        } else {
            None
        };

        let paths = CommandExecutionPaths::new(
            inputs,
            self.outputs
                .iter()
                .map(|b| CommandExecutionOutput::BuildArtifact {
                    path: b.get_path().dupe(),
                    output_type: b.output_type(),
                })
                .collect(),
            ctx.fs(),
            ctx.digest_config(),
        )?;

        Ok(PreparedRunAction {
            expanded,
            extra_env,
            paths,
            worker,
        })
    }
}

pub(crate) struct PreparedRunAction {
    expanded: ExpandedCommandLine,
    extra_env: Option<(String, String)>,
    paths: CommandExecutionPaths,
    worker: Option<WorkerSpec>,
}

impl PreparedRunAction {
    fn into_command_execution_request(self) -> CommandExecutionRequest {
        let Self {
            expanded: ExpandedCommandLine { exe, args, mut env },
            extra_env,
            paths,
            worker,
        } = self;

        for (k, v) in extra_env.into_iter() {
            env.insert(k, v);
        }

        CommandExecutionRequest::new(exe, args, paths, env).with_worker(worker)
    }
}

trait RunActionVisitor: CommandLineArtifactVisitor {
    type Iter<'a>: Iterator<Item = &'a ArtifactGroup>
    where
        Self: 'a;

    fn inputs<'a>(&'a self) -> Self::Iter<'a>;
}

impl RunActionVisitor for SimpleCommandLineArtifactVisitor {
    type Iter<'a> = impl Iterator<Item = &'a ArtifactGroup> where Self: 'a;

    fn inputs<'a>(&'a self) -> Self::Iter<'a> {
        self.inputs.iter()
    }
}

impl RunActionVisitor for DepFilesCommandLineVisitor<'_> {
    type Iter<'a> = impl Iterator<Item = &'a ArtifactGroup> where Self: 'a;

    fn inputs<'a>(&'a self) -> Self::Iter<'a> {
        self.inputs.iter().flat_map(|g| g.iter())
    }
}

#[async_trait]
impl Action for RunAction {
    fn kind(&self) -> buck2_data::ActionKind {
        buck2_data::ActionKind::Run
    }

    fn inputs(&self) -> anyhow::Result<Cow<'_, [ArtifactGroup]>> {
        let values = Self::unpack(&self.starlark_values)?;
        let mut artifact_visitor = SimpleCommandLineArtifactVisitor::new();
        values.args.visit_artifacts(&mut artifact_visitor)?;
        values.exe.visit_artifacts(&mut artifact_visitor)?;
        if let Some((worker_exe, _)) = values.worker {
            worker_exe.visit_artifacts(&mut artifact_visitor)?;
        }
        for (_, v) in values.env.iter() {
            v.visit_artifacts(&mut artifact_visitor)?;
        }
        Ok(Cow::Owned(artifact_visitor.inputs.into_iter().collect()))
    }

    fn outputs(&self) -> anyhow::Result<Cow<'_, [BuildArtifact]>> {
        Ok(Cow::Borrowed(self.outputs.as_slice()))
    }

    fn as_executable(&self) -> ActionExecutable<'_> {
        ActionExecutable::Incremental(self)
    }

    fn category(&self) -> &Category {
        &self.inner.category
    }

    fn identifier(&self) -> Option<&str> {
        self.inner.identifier.as_deref()
    }

    fn always_print_stderr(&self) -> bool {
        self.inner.always_print_stderr
    }

    fn aquery_attributes(&self, fs: &ExecutorFs) -> indexmap::IndexMap<String, String> {
        let mut cli_rendered = Vec::<String>::new();
        let mut ctx = DefaultCommandLineContext::new(fs);
        let values = Self::unpack(&self.starlark_values).unwrap();
        values
            .exe
            .add_to_command_line(&mut cli_rendered, &mut ctx)
            .unwrap();
        values
            .args
            .add_to_command_line(&mut cli_rendered, &mut ctx)
            .unwrap();
        let cmd = format!("[{}]", cli_rendered.iter().join(", "));
        indexmap! {
            "cmd".to_owned() => cmd,
            "executor_preference".to_owned() => self.inner.executor_preference.to_string(),
            "always_print_stderr".to_owned() => self.inner.always_print_stderr.to_string(),
            "weight".to_owned() => self.inner.weight.to_string(),
            "dep_files".to_owned() => self.inner.dep_files.to_string(),
            "metadata_param".to_owned() => match &self.inner.metadata_param {
                None => "None".to_owned(),
                Some(x) => x.to_string(),
            },
            "no_outputs_cleanup".to_owned() => self.inner.no_outputs_cleanup.to_string(),
        }
    }
}

#[async_trait]
impl IncrementalActionExecutable for RunAction {
    async fn execute(
        &self,
        ctx: &mut dyn ActionExecutionCtx,
    ) -> anyhow::Result<(ActionOutputs, ActionExecutionMetadata)> {
        let knobs = ctx.run_action_knobs();
        let process_dep_files = !self.inner.dep_files.labels.is_empty() || knobs.hash_all_commands;
        let (prepared_run_action, dep_files) = if !process_dep_files {
            (
                self.prepare(&mut SimpleCommandLineArtifactVisitor::new(), ctx)?,
                None,
            )
        } else {
            let mut visitor = DepFilesCommandLineVisitor::new(&self.inner.dep_files);
            let prepared = self.prepare(&mut visitor, ctx)?;
            let dep_files = make_local_dep_file_lookup_key(ctx, visitor, &prepared);
            (prepared, Some(dep_files))
        };

        // Run actions are assumed to be shared
        let host_sharing_requirements = HostSharingRequirements::Shared(self.inner.weight);

        let req = prepared_run_action
            .into_command_execution_request()
            .with_prefetch_lossy_stderr(true)
            .with_executor_preference(self.inner.executor_preference)
            .with_host_sharing_requirements(host_sharing_requirements)
            .with_low_pass_filter(self.inner.low_pass_filter)
            .with_outputs_cleanup(!self.inner.no_outputs_cleanup)
            .with_local_environment_inheritance(EnvironmentInheritance::local_command_exclusions())
            .with_force_full_hybrid_if_capable(self.inner.force_full_hybrid_if_capable)
            .with_custom_tmpdir(ctx.target().custom_tmpdir());

        // First prepare the action, check the action cache, check dep_files if needed, and execute the command
        let prepared_action = ctx.prepare_action(&req)?;
        let manager = ctx.command_execution_manager();
        let (mut result, dep_files) = match ctx.action_cache(manager, &req, &prepared_action).await
        {
            ControlFlow::Break(res) => (res, dep_files),
            ControlFlow::Continue(manager) => {
                let dep_files = if let Some(dep_files) = dep_files {
                    match check_local_dep_file_cache(ctx, self.outputs.as_slice(), &dep_files)
                        .await?
                    {
                        Some(m) => {
                            // We have a dep_file based match, return early
                            return Ok(m);
                        }
                        None => Some(dep_files),
                    }
                } else {
                    None
                };

                (
                    ctx.exec_cmd(manager, &req, &prepared_action).await,
                    dep_files,
                )
            }
        };

        if self.inner.allow_cache_upload {
            result.did_cache_upload = ctx
                .cache_upload(prepared_action.action.dupe(), &result)
                .await?;
        }

        let (outputs, metadata) =
            ctx.unpack_command_execution_result(&req, result, self.inner.allow_cache_upload)?;

        if let Some(dep_files) = dep_files {
            let LocalDepFileLookUpKey {
                dep_files_key,
                digests,
                declared_inputs,
                declared_dep_files,
            } = dep_files;

            populate_dep_files(
                dep_files_key,
                digests,
                declared_inputs,
                declared_dep_files,
                &outputs,
                ctx,
            )
            .await?;
        }

        Ok((outputs, metadata))
    }
}


/*
 * 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_build_api::analysis::registry::AnalysisRegistry;
use buck2_build_api::interpreter::rule_defs::context::AnalysisContext;
use buck2_build_api::interpreter::rule_defs::register_rule_defs;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::configuration::data::ConfigurationData;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::ProvidersName;
use buck2_core::target::label::TargetLabel;
use buck2_execute::digest_config::DigestConfig;
use buck2_interpreter::types::label::Label;
use buck2_node::configuration::execution::ExecutionPlatformResolution;
use dupe::Dupe;
use indoc::indoc;
use maplit::hashmap;
use starlark::environment::GlobalsBuilder;
use starlark::environment::Module;
use starlark::eval::Evaluator;
use starlark::eval::ReturnFileLoader;
use starlark::syntax::AstModule;
use starlark::syntax::Dialect;
use starlark::values::structs::AllocStruct;
use starlark::values::UnpackValue;
use starlark::values::Value;

fn run_ctx_test(
    content: &str,
    result_handler: impl FnOnce(anyhow::Result<Value>) -> anyhow::Result<()>,
) -> anyhow::Result<()> {
    let func_mod = Module::new();
    let globals = GlobalsBuilder::extended().with(register_rule_defs).build();
    let prelude = indoc!(
        r#"
         def assert_eq(a, b):
             if a != b:
                 fail("Expected {}, got {}".format(a, b))
         "#
    );
    let full_content = format!("{}\n{}", prelude, content);

    {
        let mut eval = Evaluator::new(&func_mod);
        let ast = AstModule::parse("foo.bzl", full_content, &Dialect::Extended).unwrap();
        eval.eval_module(ast, &globals).unwrap();
    };
    let frozen_func_mod = func_mod.freeze()?;
    let test_function = frozen_func_mod.get("test").unwrap();

    let modules = hashmap!["func_mod" => &frozen_func_mod];

    let env = Module::new();
    let file_loader = ReturnFileLoader { modules: &modules };
    let test_function = test_function.owned_value(env.frozen_heap());
    let mut eval = Evaluator::new(&env);
    eval.set_loader(&file_loader);
    let label = TargetLabel::testing_parse("root//foo/bar:some_name")
        .configure(ConfigurationData::testing_new());
    let registry = AnalysisRegistry::new_from_owner(
        BaseDeferredKey::TargetLabel(label.dupe()),
        ExecutionPlatformResolution::unspecified(),
    )?;
    let attributes = eval.heap().alloc(AllocStruct([("name", "some_name")]));

    let ctx = eval.heap().alloc(AnalysisContext::new(
        eval.heap(),
        attributes,
        Some(
            eval.heap()
                .alloc_typed(Label::new(ConfiguredProvidersLabel::new(
                    label,
                    ProvidersName::Default,
                ))),
        ),
        registry,
        DigestConfig::testing_default(),
    ));

    let returned = eval.eval_function(test_function, &[ctx], &[]);
    result_handler(returned)
}

#[test]
fn ctx_instantiates() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(ctx):
             assert_eq("foo/bar", ctx.label.package)
             assert_eq("some_name", ctx.label.name)
             assert_eq(None, ctx.label.sub_target)
             return ctx.attrs.name
         "#
    );
    run_ctx_test(content, |ret| {
        assert_eq!("some_name", ret.unwrap().unpack_str().unwrap());
        Ok(())
    })
}

#[test]
fn declare_output_declares_outputs() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             out = c.actions.declare_output("foo/bar.cpp")
             return (out.basename, out.short_path)
         "#
    );

    run_ctx_test(content, |ret| {
        let a = <(&str, &str)>::unpack_value(ret.unwrap()).unwrap();
        assert_eq!("bar.cpp", a.0);
        assert_eq!("foo/bar.cpp", a.1);
        Ok(())
    })
}

#[test]
fn declare_output_with_prefix() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             out = c.actions.declare_output("out/test", "foo/bar.cpp")
             return (out.basename, out.short_path)
         "#
    );

    run_ctx_test(content, |ret| {
        let a = <(&str, &str)>::unpack_value(ret.unwrap()).unwrap();
        assert_eq!("bar.cpp", a.0);
        assert_eq!("foo/bar.cpp", a.1);
        Ok(())
    })
}

#[test]
fn declare_output_dot() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             return c.actions.declare_output("magic", ".")
         "#
    );

    let expect = "artifact with an empty filename component";
    run_ctx_test(content, |ret| match ret {
        Err(e) if e.to_string().contains(expect) => Ok(()),
        _ => panic!(
            "Expected a specific failure containing `{}`, got {:?}",
            expect, ret
        ),
    })
}

#[test]
fn declare_output_dot_bad() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             return c.actions.declare_output("..")
         "#
    );

    let expect = "expected a normalized path";
    run_ctx_test(content, |ret| match ret {
        Err(e) if e.to_string().contains(expect) => Ok(()),
        _ => panic!(
            "Expected a specific failure containing `{}`, got {:?}",
            expect, ret
        ),
    })
}

#[test]
fn declare_output_dotdot() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             return c.actions.declare_output("foo/..")
         "#
    );

    let expect = "expected a normalized path";
    run_ctx_test(content, |ret| match ret {
        Err(e) if e.to_string().contains(expect) => Ok(()),
        _ => panic!(
            "Expected a specific failure containing `{}`, got {:?}",
            expect, ret
        ),
    })
}

#[test]
fn declare_output_require_bound() -> anyhow::Result<()> {
    let content = indoc!(
        r#"
         def test(c):
             a = c.actions.declare_output("a")
             b = c.actions.declare_output("b")
             c.actions.run([a, b.as_output()], category = "test_category")
         "#
    );

    let expect = "must be bound by now";
    run_ctx_test(content, |ret| match ret {
        Err(e) if e.to_string().contains(expect) => Ok(()),
        _ => panic!(
            "Expected a specific failure containing `{}`, got {:?}",
            expect, ret
        ),
    })
}


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

#![cfg(test)]

mod context;

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


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

#![feature(async_closure)]
#![feature(box_patterns)]
#![feature(try_blocks)]

pub mod analysis;
pub mod attrs;

pub fn init_late_bindings() {
    analysis::calculation::init_rule_analysis_calculation();
}