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

use std::time::Duration;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_cli_proto::profile_request::ProfileOpts;
use buck2_cli_proto::profile_request::Profiler;
use buck2_cli_proto::target_profile::Action;
use buck2_cli_proto::BxlProfile;
use buck2_cli_proto::ProfileRequest;
use buck2_cli_proto::ProfileResponse;
use buck2_cli_proto::TargetProfile;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonCommandOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::streaming::BuckSubcommand;
use buck2_client_ctx::streaming::StreamingCommand;
use dupe::Dupe;

use super::bxl::BxlCommandOptions;

#[derive(Debug, clap::Parser)]
#[clap(about = "Profiling mechanisms")]
pub enum ProfileCommand {
    #[clap(about = "Profile analysis")]
    Analysis(BuckProfileOptions),

    #[clap(about = "Profile loading")]
    Loading(BuckProfileOptions),

    #[clap(about = "Profile BXL script")]
    Bxl(BxlProfileOptions),
}

pub enum ProfileOptionsType {
    BuckProfileOptions {
        opts: AnalysisLoadProfileOptions,
        action: Action,
    },
    BxlProfileOptions {
        opts: BxlCommandOptions,
    },
}

impl ProfileCommand {
    pub fn exec(self, matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let submatches = matches.subcommand().expect("subcommand not found").1;
        match self {
            Self::Analysis(opts) => ProfileSubcommand {
                opts: ProfileOptionsType::BuckProfileOptions {
                    opts: opts.buck_opts,
                    action: Action::Analysis,
                },
                profile_common_opts: opts.profile_common_opts,
            },
            Self::Loading(opts) => ProfileSubcommand {
                opts: ProfileOptionsType::BuckProfileOptions {
                    opts: opts.buck_opts,
                    action: Action::Loading,
                },
                profile_common_opts: opts.profile_common_opts,
            },
            Self::Bxl(opts) => ProfileSubcommand {
                opts: ProfileOptionsType::BxlProfileOptions {
                    opts: opts.bxl_opts,
                },
                profile_common_opts: opts.profile_common_opts,
            },
        }
        .exec(submatches, ctx)
    }

    pub fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        argv.no_need_to_sanitize()
    }
}

#[derive(clap::ValueEnum, Dupe, Clone, Debug)]
enum BuckProfileMode {
    TimeFlame,
    HeapFlameAllocated,
    HeapFlameRetained,
    HeapSummaryAllocated,
    HeapSummaryRetained,
    Statement,
    Bytecode,
    BytecodePairs,
    Typecheck,
}

#[derive(Debug, clap::Parser)]
pub struct BxlProfileOptions {
    #[clap(flatten)]
    bxl_opts: BxlCommandOptions,

    #[clap(flatten)]
    profile_common_opts: ProfileCommonOptions,
}

#[derive(Debug, clap::Parser)]
pub struct BuckProfileOptions {
    #[clap(flatten)]
    buck_opts: AnalysisLoadProfileOptions,

    #[clap(flatten)]
    profile_common_opts: ProfileCommonOptions,
}

#[derive(Debug, clap::Parser)]
pub struct AnalysisLoadProfileOptions {
    #[clap(value_name = "TARGET")]
    target_pattern: String,

    /// In analysis profiling, capture the profile of the target and its dependencies,
    /// and output the merged profile.
    #[clap(long, short = 'r')]
    recursive: bool,
}

#[derive(Debug, clap::Parser)]
pub struct ProfileCommonOptions {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    /// Output file path for profile data.
    ///
    /// File will be created if it does not exist, and overwritten if it does.
    #[clap(long, short = 'o', value_name = "PATH")]
    output: PathArg,

    /// Profile mode.
    ///
    /// Memory profiling modes have suffixes either `-allocated` or `retained`.
    ///
    /// `-retained` means memory kept in frozen starlark heap after analysis complete.
    /// `-retained` does not work when profiling loading,
    /// because no memory is retained after loading and frozen heap is not even created.
    /// This is probably what you want when profiling analysis.
    ///
    /// `-allocated` means allocated memory, including memory which is later garbage collected.
    #[clap(long, short = 'm', value_enum)]
    mode: BuckProfileMode,
}

pub struct ProfileSubcommand {
    opts: ProfileOptionsType,
    profile_common_opts: ProfileCommonOptions,
}

fn profile_mode_to_profile(mode: &BuckProfileMode) -> Profiler {
    match mode {
        BuckProfileMode::TimeFlame => Profiler::TimeFlame,
        BuckProfileMode::HeapFlameAllocated => Profiler::HeapFlameAllocated,
        BuckProfileMode::HeapFlameRetained => Profiler::HeapFlameRetained,
        BuckProfileMode::HeapSummaryAllocated => Profiler::HeapSummaryAllocated,
        BuckProfileMode::HeapSummaryRetained => Profiler::HeapSummaryRetained,
        BuckProfileMode::Statement => Profiler::Statement,
        BuckProfileMode::Bytecode => Profiler::Bytecode,
        BuckProfileMode::BytecodePairs => Profiler::BytecodePairs,
        BuckProfileMode::Typecheck => Profiler::Typecheck,
    }
}

#[async_trait]
impl StreamingCommand for ProfileSubcommand {
    const COMMAND_NAME: &'static str = "profile";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.client_context(
            &self.profile_common_opts.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;

        let destination_path = self.profile_common_opts.output.resolve(&ctx.working_dir);

        let profile_mode = &self.profile_common_opts.mode;

        let destination_path = destination_path.into_string()?;

        let console_opts = ctx.stdin().console_interaction_stream(self.console_opts());

        let response = match self.opts {
            ProfileOptionsType::BuckProfileOptions { opts, action } => {
                let target_opts = TargetProfile {
                    target_pattern: Some(buck2_data::TargetPattern {
                        value: opts.target_pattern,
                    }),
                    recursive: opts.recursive,
                    action: action.into(),
                };

                buckd
                    .with_flushing()
                    .profile(
                        ProfileRequest {
                            context: Some(context),
                            profile_opts: Some(ProfileOpts::TargetProfile(target_opts)),
                            destination_path,
                            profiler: profile_mode_to_profile(profile_mode).into(),
                        },
                        console_opts,
                        &mut NoPartialResultHandler,
                    )
                    .await??
            }
            ProfileOptionsType::BxlProfileOptions { opts } => {
                let bxl_opts = BxlProfile {
                    bxl_label: opts.bxl_label,
                    bxl_args: opts.bxl_args,
                };

                buckd
                    .with_flushing()
                    .profile(
                        ProfileRequest {
                            context: Some(context),
                            profile_opts: Some(ProfileOpts::BxlProfile(bxl_opts)),
                            destination_path,
                            profiler: profile_mode_to_profile(profile_mode).into(),
                        },
                        console_opts,
                        &mut NoPartialResultHandler,
                    )
                    .await??
            }
        };

        let ProfileResponse {
            elapsed,
            total_retained_bytes,
        } = response;

        let elapsed = elapsed
            .context("Missing duration")
            .and_then(|d| {
                Duration::try_from(d).map_err(|_| anyhow::anyhow!("Duration is negative"))
            })
            .context("Elapsed is invalid")?;

        buck2_client_ctx::println!(
            "Starlark {:?} profile has been written to {}",
            profile_mode,
            self.profile_common_opts.output.display(),
        )?;
        buck2_client_ctx::println!("Elapsed: {:.3}s", elapsed.as_secs_f64())?;
        buck2_client_ctx::println!("Total retained bytes: {}", total_retained_bytes)?;

        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        &self.profile_common_opts.common_opts.console_opts
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.profile_common_opts.common_opts.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.profile_common_opts.common_opts.config_opts
    }
}


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

use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::path_arg::PathArg;
use buck2_common::invocation_roots::find_invocation_roots;
use thiserror::Error;

#[derive(Debug)]
enum RootKind {
    Package,
    Cell,
    Project,
    Daemon,
}

impl FromStr for RootKind {
    type Err = String;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "package" => Ok(Self::Package),
            "cell" => Ok(Self::Cell),
            "project" => Ok(Self::Project),
            "daemon" => Ok(Self::Daemon),
            _ => Err("expected one of `package`, `cell`, `project`, or `daemon`".to_owned()),
        }
    }
}

#[derive(Debug, clap::Parser)]
#[clap(about = "Find buck cell, project or package root")]
pub struct RootCommand {
    #[clap(short, long, help("which root to print"), default_value("cell"), possible_values(&["package", "cell", "project", "daemon"]))]
    kind: RootKind,
    #[clap(
        help(
            "determine the root for a specific directory (if not provided, finds the root for the current directory)"
        ),
        value_name = "PATH",
        long
    )]
    dir: Option<PathArg>,
}

#[derive(Debug, Error)]
enum RootError {
    #[error("Finding package root isn't yet implemented.")]
    PackageRootUnimplemented,
}

impl RootCommand {
    pub fn exec(
        self,
        _matches: &clap::ArgMatches,
        ctx: ClientCommandContext<'_>,
    ) -> anyhow::Result<()> {
        let root = match self.kind {
            RootKind::Package => return Err(RootError::PackageRootUnimplemented.into()),
            RootKind::Cell => match self.dir {
                Some(dir) => find_invocation_roots(&dir.resolve(&ctx.working_dir))?.cell_root,
                None => ctx.paths()?.cell_root().to_owned(),
            },
            RootKind::Project => match self.dir {
                Some(dir) => find_invocation_roots(&dir.resolve(&ctx.working_dir))?
                    .project_root
                    .root()
                    .to_owned(),
                None => ctx.paths()?.project_root().root().to_owned(),
            },
            RootKind::Daemon => ctx.paths()?.daemon_dir()?.path,
        };

        buck2_client_ctx::println!("{}", root.to_string_lossy())?;
        Ok(())
    }

    pub fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        argv.no_need_to_sanitize()
    }
}


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

use std::collections::HashMap;
use std::fs::File;
use std::io::Write;

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::build_request::build_providers;
use buck2_cli_proto::build_request::BuildProviders;
use buck2_cli_proto::build_request::Materializations;
use buck2_cli_proto::BuildRequest;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::command_outcome::CommandOutcome;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonBuildOptions;
use buck2_client_ctx::common::CommonCommandOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_core::soft_error;
use buck2_wrapper_common::BUCK2_WRAPPER_ENV_VAR;
use buck2_wrapper_common::BUCK_WRAPPER_UUID_ENV_VAR;
use serde::Serialize;
use thiserror::Error;

use crate::commands::build::print_build_result;

/// Build and run the selected target.
///
/// The Build ID for the underlying build execution is made available to the target in
/// the `BUCK_RUN_BUILD_ID` environment variable.
#[derive(Debug, clap::Parser)]
#[clap(
    name = "run",
    setting = clap::AppSettings::TrailingVarArg
)]
pub struct RunCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    #[clap(flatten)]
    build_opts: CommonBuildOptions,

    #[clap(long = "providers", help = "Print the providers of each target")]
    print_providers: bool,

    #[clap(
        long = "command-args-file",
        help = "Write the command to a file instead of executing it.",
        group = "exec_options"
    )]
    command_args_file: Option<String>,

    #[clap(
        long = "chdir",
        help = "Set the current working directory of the executable being run",
        group = "exec_options"
    )]
    chdir: Option<String>,

    /// Instead of running the command, print out the command
    /// formatted for shell interpolation, use as: $(buck2 run --emit-shell ...)
    #[clap(long, group = "exec_options")]
    emit_shell: bool,

    #[clap(name = "TARGET", help = "Target to build and run")]
    target: String,

    #[clap(
        name = "TARGET_ARGS",
        help = "Additional arguments passed to the target when running it"
    )]
    extra_run_args: Vec<String>,
}

#[async_trait]
impl StreamingCommand for RunCommand {
    const COMMAND_NAME: &'static str = "run";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;
        // TODO(rafaelc): fail fast on the daemon if the target doesn't have RunInfo
        let response = buckd
            .with_flushing()
            .build(
                BuildRequest {
                    context: Some(context),
                    // TODO(wendyy): glob patterns should be prohibited, and command should fail before the build event happens.
                    target_patterns: vec![buck2_data::TargetPattern {
                        value: self.target.clone(),
                    }],
                    unstable_print_providers: self.print_providers,
                    build_providers: Some(BuildProviders {
                        default_info: build_providers::Action::Skip as i32,
                        run_info: build_providers::Action::Build as i32,
                        test_info: build_providers::Action::Skip as i32,
                    }),
                    response_options: None,
                    build_opts: Some(self.build_opts.to_proto()),
                    final_artifact_materializations: Materializations::Materialize as i32,
                    target_universe: Vec::new(),
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut NoPartialResultHandler,
            )
            .await;

        let console = self.common_opts.console_opts.final_console();
        let success = match &response {
            Ok(CommandOutcome::Success(response)) => response.error_messages.is_empty(),
            Ok(CommandOutcome::Failure(_)) => false,
            Err(_) => false,
        };
        if !success {
            console.print_error("BUILD FAILED")?;
        }
        let response = response??;
        print_build_result(&console, &response.error_messages)?;

        if !success {
            return ExitResult::failure();
        }

        if response.build_targets.len() > 1 {
            soft_error!(
                "buck_run_multiple_targets",
                RunCommandError::MultipleTargets.into()
            )?;
        }

        // TODO(rafaelc): use absolute paths for artifacts in the cli
        //      we should run the command from the current dir, not the project root
        if response.build_targets.is_empty() || response.build_targets[0].run_args.is_empty() {
            return ExitResult::err(RunCommandError::NonBinaryRule(self.target).into());
        }
        let mut run_args = response.build_targets[0].run_args.clone();
        run_args.extend(self.extra_run_args);

        // Special case for recursive invocations of buck; `BUCK2_WRAPPER` is set by wrapper scripts that execute
        // Buck2. We're not a wrapper script, so we unset it to prevent `run` from inheriting it.
        std::env::remove_var(BUCK2_WRAPPER_ENV_VAR);
        std::env::remove_var(BUCK_WRAPPER_UUID_ENV_VAR);

        if let Some(file_path) = self.command_args_file {
            let mut output = File::create(&file_path).with_context(|| {
                format!("Failed to create/open `{}` to print command", file_path)
            })?;

            let command = CommandArgsFile {
                path: run_args[0].clone(),
                argv: run_args,
                envp: std::env::vars().collect(),
                is_fix_script: false,
                print_command: false,
            };
            let serialized =
                serde_json::to_string(&command).context("Failed to serialize command")?;
            output
                .write_all(serialized.as_bytes())
                .context("Failed to write command")?;

            return ExitResult::success();
        }

        if self.emit_shell {
            if cfg!(unix) {
                buck2_client_ctx::println!("{}", shlex::join(run_args.iter().map(|a| a.as_str())))?;
                return ExitResult::success();
            } else {
                return ExitResult::err(RunCommandError::EmitShellNotSupportedOnWindows.into());
            }
        }

        ExitResult::exec(
            run_args[0].clone(),
            run_args,
            self.chdir,
            vec![("BUCK_RUN_BUILD_ID".to_owned(), ctx.trace_id.to_string())],
        )
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        &self.common_opts.console_opts
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.common_opts.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.common_opts.config_opts
    }

    fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        let Argv {
            argv,
            expanded_argv,
        } = argv;
        let to_redact: std::collections::HashSet<_> = self.extra_run_args.iter().collect();
        SanitizedArgv {
            argv: argv
                .into_iter()
                .filter(|arg| !to_redact.contains(arg))
                .collect(),
            expanded_argv: expanded_argv
                .into_iter()
                .filter(|arg| !to_redact.contains(arg))
                .collect(),
        }
    }
}

#[derive(Serialize)]
struct CommandArgsFile {
    path: String,
    argv: Vec<String>,
    envp: HashMap<String, String>,
    // Not used. For buck_v1 back compatibility only.
    is_fix_script: bool,
    // Not used. For buck_v1 back compatibility only.
    print_command: bool,
}

#[derive(Error, Debug)]
pub enum RunCommandError {
    #[error("Target `{0}` is not a binary rule (only binary rules can be `run`)")]
    NonBinaryRule(String),
    #[error("`--emit-shell` is not supported on Windows")]
    EmitShellNotSupportedOnWindows,
    #[error(
        "`buck2 run` only supports a single target, but multiple targets were requested. Only executing the first one built."
    )]
    MultipleTargets,
}


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

use async_trait::async_trait;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
#[clap(about = "Start, query, and control the http server")]
pub struct ServerCommand {}

#[async_trait]
impl StreamingCommand for ServerCommand {
    const COMMAND_NAME: &'static str = "server";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let status = buckd.with_flushing().status(false).await?;
        buck2_client_ctx::println!("buckd.endpoint={}", status.process_info.unwrap().endpoint)?;
        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        CommonConsoleOptions::simple_ref()
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        CommonDaemonCommandOptions::default_ref()
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        CommonBuildConfigurationOptions::default_ref()
    }
}


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

use std::time::Duration;

use anyhow::Context;
use buck2_cli_proto::StatusResponse;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::daemon::client::connect::establish_connection_existing;
use buck2_client_ctx::daemon::client::connect::BuckdConnectOptions;
use buck2_client_ctx::subscribers::stdout_stderr_forwarder::StdoutStderrForwarder;
use buck2_common::daemon_dir::DaemonDir;
use buck2_common::invocation_paths::InvocationPaths;
use chrono::NaiveDateTime;
use humantime::format_duration;
use walkdir::WalkDir;

#[derive(Debug, thiserror::Error)]
enum StatusError {
    #[error("Incorrect seconds/nanos argument")]
    NativeDateTime,
}

#[derive(Debug, clap::Parser)]
#[clap(about = "Buckd status")]
pub struct StatusCommand {
    #[clap(long, help = "Whether to include a state snapshot in the output.")]
    snapshot: bool,
    #[clap(long, help = "Enable printing status for all running buckd")]
    all: bool,
}

impl StatusCommand {
    pub fn exec(
        self,
        _matches: &clap::ArgMatches,
        ctx: ClientCommandContext<'_>,
    ) -> anyhow::Result<()> {
        ctx.with_runtime(async move |ctx| {
            if self.all {
                let mut daemon_dirs = Vec::new();
                let root = InvocationPaths::common_buckd_dir()?;
                let walker = WalkDir::new(&root).follow_links(false).into_iter();
                for entry in walker {
                    let entry = entry?;
                    if entry.file_type().is_dir() {
                        let dir = DaemonDir {
                            path: entry.into_path().try_into()?,
                        };

                        if dir.buckd_info().exists() {
                            daemon_dirs.push(dir);
                        }
                    }
                }

                let mut statuses = Vec::new();
                for dir in daemon_dirs {
                    if let Ok(bootstrap_client) = establish_connection_existing(&dir).await {
                        statuses.push(process_status(
                            bootstrap_client
                                .with_subscribers(vec![Box::new(StdoutStderrForwarder)])
                                .with_flushing()
                                .status(self.snapshot)
                                .await?,
                        )?);
                    }
                }

                buck2_client_ctx::println!("{}", serde_json::to_string_pretty(&statuses)?)?;
            } else {
                match ctx
                    .connect_buckd(BuckdConnectOptions::existing_only_no_console())
                    .await
                {
                    Err(_) => {
                        buck2_client_ctx::eprintln!("no buckd running")?;
                        // Should this be an error?
                    }
                    Ok(mut client) => {
                        let json_status =
                            process_status(client.with_flushing().status(self.snapshot).await?)?;
                        buck2_client_ctx::println!(
                            "{}",
                            serde_json::to_string_pretty(&json_status)?
                        )?;
                    }
                }
            }

            Ok(())
        })
    }

    pub fn sanitize_argv(&self, argv: Argv) -> SanitizedArgv {
        argv.no_need_to_sanitize()
    }
}

fn timestamp_to_string(seconds: u64, nanos: u32) -> anyhow::Result<String> {
    Ok(NaiveDateTime::from_timestamp_opt(seconds as i64, nanos)
        .context(StatusError::NativeDateTime)?
        .format("%Y-%m-%dT%H:%M:%SZ")
        .to_string())
}

fn duration_to_string(duration: Duration) -> String {
    let duration = Duration::from_secs(duration.as_secs());
    format_duration(duration).to_string()
}

fn process_status(status: StatusResponse) -> anyhow::Result<serde_json::Value> {
    let timestamp = match status.start_time {
        None => "unknown".to_owned(),
        Some(timestamp) => timestamp_to_string(timestamp.seconds as u64, timestamp.nanos as u32)?,
    };
    let uptime = match status.uptime {
        None => "unknown".to_owned(),
        Some(uptime) => {
            let uptime = Duration::new(uptime.seconds as u64, uptime.nanos as u32);
            duration_to_string(uptime)
        }
    };

    Ok(serde_json::json!({
        "start_time": timestamp,
        "uptime": uptime,
        "process_info": serde_json::to_value(status.process_info)?,
        "daemon_constraints": serde_json::to_value(status.daemon_constraints)?,
        "snapshot": serde_json::to_value(status.snapshot)?,
        "project_root": status.project_root,
        "isolation_dir": status.isolation_dir,
        "forkserver_pid": serde_json::to_value(status.forkserver_pid)?,
        "supports_vpnless": status.supports_vpnless.unwrap_or_default(),
    }))
}

#[cfg(test)]
mod tests {
    use std::time::Duration;

    use crate::commands::status::duration_to_string;
    use crate::commands::status::timestamp_to_string;

    #[test]
    fn test_timestamp_to_string() {
        // Check with `TZ=UTC date -r 1662516832 -Iseconds`.
        assert_eq!(
            "2022-09-07T02:13:52Z",
            timestamp_to_string(1662516832, 123).unwrap(),
        );
    }

    #[test]
    fn test_duration_to_string() {
        assert_eq!(
            "1h 2m 3s",
            duration_to_string(Duration::new(3600 + 120 + 3, 123456789))
        );
    }
}


/*
 * 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 anyhow::Context as _;
use async_trait::async_trait;
use buck2_cli_proto::protobuf_util::ProtobufSplitter;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::common::ConsoleType;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::events_ctx::PartialResultCtx;
use buck2_client_ctx::events_ctx::PartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stream_util::reborrow_stream_for_static;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_subscription_proto::SubscriptionRequest;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use once_cell::sync::Lazy;
use prost::Message;
use tokio_util::codec::FramedRead;

/// Open a subscription channel to the Buck2 daemon. This allows you to interact with the Buck2
/// daemon via the `stdin` and `stdout` of this command: you send requests to the daemon by writing
/// to `stdin`, and you get responses via `stdout`.
///
/// The protocol used by this command is length-prefixed protobuf. This format is a repeated series
/// of a varint followed by a record of the length indicated by said varint.
///
/// The protobuf spec for those records is described in
/// `buck2_subscription_proto/subscription.proto`. The client writes `SubscriptionRequest` and
/// reads `SubscriptionResponse`. See the documentation in `subscription.proto` to discover
/// available APIs.
///
/// This API does not (currently) allow invalid requests and will error out when one is sent.
#[derive(Debug, clap::Parser)]
#[clap(about = "Subscribe to updates from the Buck2 daemon")]
pub struct SubscribeCommand {
    #[clap(flatten)]
    config_opts: CommonBuildConfigurationOptions,

    #[clap(flatten)]
    event_log_opts: CommonDaemonCommandOptions,

    /// Whether to request command snapshots.
    #[clap(long)]
    active_commands: bool,

    /// Whether to get output as JSON. The JSON format is deemed unstable so this should only be
    /// used for debugging.
    #[clap(long)]
    unstable_json: bool,
}

#[async_trait]
impl StreamingCommand for SubscribeCommand {
    const COMMAND_NAME: &'static str = "subscribe";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let client_context =
            ctx.client_context(&self.config_opts, matches, ctx.sanitized_argv.argv.clone())?;

        let stream = FramedRead::new(ctx.stdin(), ProtobufSplitter)
            .and_then(|bytes| {
                futures::future::ready(
                    SubscriptionRequest::decode_length_delimited(bytes)
                        .context("Error decoding SubscriptionRequest"),
                )
            })
            .map(|res| match res {
                Ok(r) => r,
                Err(e) => {
                    // NOTE: if stderr is gone there is not much we can do besides not write to
                    // stderr.
                    let reason = format!("Error parsing request: {:#}", e);
                    let _ignored = buck2_client_ctx::eprintln!("{}", reason);
                    SubscriptionRequest {
                        request: Some(
                            buck2_subscription_proto::Disconnect { reason, ok: false }.into(),
                        ),
                    }
                }
            });

        let mut partial_result_handler = SubscriptionPartialResultHandler {
            buffer: Vec::new(),
            json: self.unstable_json,
            ok: true,
        };

        let stream = if self.active_commands {
            futures::stream::once(futures::future::ready(SubscriptionRequest {
                request: Some(buck2_subscription_proto::SubscribeToActiveCommands {}.into()),
            }))
            .chain(stream)
            .left_stream()
        } else {
            stream.right_stream()
        };

        let stream = stream.map(|request| buck2_cli_proto::SubscriptionRequestWrapper {
            request: Some(request),
        });

        {
            let partial_result_handler = &mut partial_result_handler;

            reborrow_stream_for_static(
                stream,
                |stream| async move {
                    buckd
                        .with_flushing()
                        .subscription(client_context, stream, partial_result_handler)
                        .await
                },
                || {
                    Some(buck2_cli_proto::SubscriptionRequestWrapper {
                        request: Some(SubscriptionRequest {
                            request: Some(
                                buck2_subscription_proto::Disconnect {
                                    reason: "EOF on stdin".to_owned(),
                                    ok: true,
                                }
                                .into(),
                            ),
                        }),
                    })
                },
            )
            .await??;
        };

        if partial_result_handler.ok {
            ExitResult::success()
        } else {
            ExitResult::failure()
        }
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        // This should only be communicated with by an IDE, so disable anything other
        // than the simple console
        static SIMPLE_CONSOLE: Lazy<CommonConsoleOptions> = Lazy::new(|| CommonConsoleOptions {
            console_type: ConsoleType::Simple,
            ui: vec![],
            no_interactive_console: true,
        });
        &SIMPLE_CONSOLE
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.config_opts
    }

    fn should_show_waiting_message(&self) -> bool {
        // It's normal to get no open spans for this command
        false
    }
}

/// Outputs subscription messages.
struct SubscriptionPartialResultHandler {
    /// We reuse our output buffer here.
    buffer: Vec<u8>,
    json: bool,
    ok: bool,
}

#[async_trait]
impl PartialResultHandler for SubscriptionPartialResultHandler {
    type PartialResult = buck2_cli_proto::SubscriptionResponseWrapper;

    async fn handle_partial_result(
        &mut self,
        mut ctx: PartialResultCtx<'_, '_>,
        partial_res: Self::PartialResult,
    ) -> anyhow::Result<()> {
        let response = partial_res
            .response
            .context("Empty `SubscriptionResponseWrapper`")?;

        if let Some(buck2_subscription_proto::subscription_response::Response::Goodbye(goodbye)) =
            &response.response
        {
            self.ok = self.ok && goodbye.ok;
        }

        self.buffer.clear();

        if self.json {
            serde_json::to_writer(&mut self.buffer, &response).context("JSON encoding failed")?;
            self.buffer.push(b'\n');
        } else {
            response
                .encode_length_delimited(&mut self.buffer)
                .context("Encoding failed")?;
        }

        ctx.stdout(&self.buffer).await
    }
}


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

use async_trait::async_trait;
use buck2_cli_proto::targets_request;
use buck2_cli_proto::targets_request::OutputFormat;
use buck2_cli_proto::TargetsRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonCommandOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::daemon::client::StdoutPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::query_args::CommonAttributeArgs;
use buck2_client_ctx::stdin::Stdin;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use dupe::Dupe;
use gazebo::prelude::*;

#[derive(thiserror::Error, Debug)]
enum TargetsError {
    /// Clap should report it, but if we missed something, this is a fallback.
    #[error("Flags are mutually exclusive")]
    IncompatibleArguments,
}

// Use non-camel case so the possible values match buck1's
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Dupe, clap::ArgEnum)]
#[clap(rename_all = "snake_case")]
enum TargetHashFileMode {
    PathsOnly,
    PathsAndContents,
    None,
}

#[derive(Debug, clap::ArgEnum, Clone, Dupe)]
enum TargetHashGraphType {
    None,
    Unconfigured,
    Configured,
}

// Use non-camel case so the possible values match buck1's
/// Possible values for the --target-hash-function arg. We don't actually
/// honor the specific algorithms, we use them as a hint to pick "fast" or "strong".
#[allow(non_camel_case_types)]
#[derive(Debug, clap::ArgEnum, Clone, Dupe)]
enum TargetHashFunction {
    Sha1,
    Sha256,
    Murmur_Hash3,
    Fast,
    Strong,
}

#[derive(Debug, clap::Parser)]
#[clap(name = "targets", about = "Show details about the specified targets")]
pub struct TargetsCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    /// Print targets as JSON
    #[clap(long)]
    json: bool,

    /// Print targets as JSON-lines
    #[clap(long, conflicts_with = "json")]
    json_lines: bool,

    /// Print statistics of how many entries were processed
    #[clap(long)]
    stats: bool,

    /// Print the fully-qualified build target for the specified aliases
    #[clap(long, alias = "resolvealias")]
    resolve_alias: bool,

    /// Print a stable hash of each target after the target name. Incompatible with '--show-rulekey'.
    #[clap(long, conflicts_with = "streaming")]
    show_target_hash: bool,

    /// Print a stable unconfigured hash of each target after the target name.
    #[clap(long, conflicts_with = "show-target-hash")]
    show_unconfigured_target_hash: bool,

    /// Modifies computation of target hashes. If set to `PATHS_AND_CONTENTS` (the default), the contents
    /// of all files referenced from the targets will be used to compute the target hash. If set to
    /// `PATHS_ONLY`, only files' paths contribute to the hash. If set to `NONE` no files will be used.
    /// See also --target-hash-modified-paths.
    #[clap(
        long,
        arg_enum,
        ignore_case = true,
        default_value = "paths_and_contents",
        conflicts_with = "streaming"
    )]
    target_hash_file_mode: TargetHashFileMode,

    /// Modifies computation of target hashes. Only effective when --target-hash-file-mode is set to
    /// `PATHS_ONLY`. If a target or its dependencies reference a file from this set, the target's hash
    /// will be different than if this option was omitted. Otherwise, the target's hash will be the same
    /// as if this option was omitted.
    #[clap(long, multiple_values = true, conflicts_with = "streaming")]
    target_hash_modified_paths: Vec<PathArg>,

    /// Selects either the "fast" or the "strong" target hash function to be used for computing target hashes.
    /// While we don't specify the exact algorithm, the "strong" algorithm should be a reasonable cryptographic
    /// hash (ex. blake3) while the "fast" function will likely be a non-crypto hash. Both functions are
    /// guaranteed to be deterministic and to have the same value across different platforms/architectures.
    #[clap(long, ignore_case = true, default_value = "fast", arg_enum)]
    target_hash_function: TargetHashFunction,

    /// When true, emit the hash or target node and all dependencies recursively.
    /// When false, hash only the target node.
    #[clap(long, action = clap::ArgAction::Set, default_value = "true", conflicts_with = "streaming")]
    target_hash_recursive: bool,

    #[clap(flatten)]
    attributes: CommonAttributeArgs,

    /// Enables printing of default attributes. This would be attributes in a target that aren't
    /// explicitly set in the target but instead use the default set in the rule declaration.
    #[clap(long)]
    include_defaults: bool,

    /// Print the path to the output for each of the rules relative to the cell
    #[clap(long)]
    show_output: bool,

    /// Print the absolute path to the output for each of the rules relative to the cell
    #[clap(long)]
    show_full_output: bool,

    /// On loading errors, put buck.error in the output stream and continue
    #[clap(long)]
    keep_going: bool,

    /// Write output as soon as it is available. The order of the output items
    /// is non-deterministic and if multiple patterns cover the same target, may
    /// have duplicates.
    #[clap(long)]
    streaming: bool,

    /// Don't cache the target information on the build graph
    #[clap(long, requires = "streaming")]
    no_cache: bool,

    /// Show the imports of each package/import. Shows an additional output per package/import
    /// (not per target), including implicit dependencies (e.g. the prelude) but only direct
    /// dependencies (not the transitive closure).
    #[clap(long, requires = "streaming")]
    imports: bool,

    /// File to put the output in, rather than sending to stdout.
    ///
    /// File will be created if it does not exist, and overwritten if it does.
    #[clap(long, short = 'o', value_name = "PATH")]
    output: Option<PathArg>,

    /// Patterns to interpret
    #[clap(name = "TARGET_PATTERNS")]
    patterns: Vec<String>,

    /// Number of threads to use during execution (default is # cores)
    #[clap(short = 'j', long = "num-threads", value_name = "THREADS")]
    pub num_threads: Option<u32>,
}

impl TargetsCommand {
    fn output_format(&self) -> anyhow::Result<OutputFormat> {
        if self.json {
            if self.json_lines || self.stats {
                return Err(TargetsError::IncompatibleArguments.into());
            }
            Ok(OutputFormat::Json)
        } else if self.json_lines {
            if self.json || self.stats {
                return Err(TargetsError::IncompatibleArguments.into());
            }
            Ok(OutputFormat::JsonLines)
        } else if !self.attributes.get()?.is_empty() {
            Ok(OutputFormat::Json)
        } else if self.stats {
            if self.json || self.json_lines {
                return Err(TargetsError::IncompatibleArguments.into());
            }
            Ok(OutputFormat::Stats)
        } else {
            Ok(OutputFormat::Text)
        }
    }
}

#[async_trait]
impl StreamingCommand for TargetsCommand {
    const COMMAND_NAME: &'static str = "targets";

    async fn exec_impl(
        mut self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let target_hash_use_fast_hash = match self.target_hash_function {
            TargetHashFunction::Sha1 | TargetHashFunction::Sha256 => {
                buck2_client_ctx::eprintln!(
                    "buck2 only supports \"fast\" and \"strong\" target hash functions. Using the \"strong\" hash."
                )?;
                false
            }
            TargetHashFunction::Murmur_Hash3 => {
                buck2_client_ctx::eprintln!(
                    "buck2 only supports \"fast\" and \"strong\" target hash functions. Using the \"fast\" hash."
                )?;
                true
            }
            TargetHashFunction::Fast => true,
            TargetHashFunction::Strong => false,
        };

        let output_attributes = self.attributes.get()?;
        let target_hash_graph_type =
            match (self.show_target_hash, self.show_unconfigured_target_hash) {
                (true, true) => {
                    return ExitResult::err(anyhow::Error::new(
                        TargetsError::IncompatibleArguments,
                    ));
                }
                (true, false) => targets_request::TargetHashGraphType::Configured as i32,
                (false, true) => targets_request::TargetHashGraphType::Unconfigured as i32,
                (false, false) => targets_request::TargetHashGraphType::None as i32,
            };

        let output_format = self.output_format()?;

        let context = Some(ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?);

        let target_hash_modified_paths = self
            .target_hash_modified_paths
            .into_try_map(|path| path.resolve(&ctx.working_dir).into_string())?;

        let target_request = TargetsRequest {
            context,
            target_patterns: self.patterns.map(|pat| buck2_data::TargetPattern {
                value: pat.to_owned(),
            }),
            output_format: output_format as i32,
            targets: Some(if self.resolve_alias {
                targets_request::Targets::ResolveAlias(targets_request::ResolveAlias {})
            } else {
                targets_request::Targets::Other(targets_request::Other {
                    output_attributes,
                    target_hash_file_mode: match self.target_hash_file_mode {
                        TargetHashFileMode::PathsOnly => {
                            targets_request::TargetHashFileMode::PathsOnly as i32
                        }
                        TargetHashFileMode::PathsAndContents => {
                            targets_request::TargetHashFileMode::PathsAndContents as i32
                        }
                        TargetHashFileMode::None => {
                            targets_request::TargetHashFileMode::NoFiles as i32
                        }
                    },
                    target_hash_modified_paths,
                    target_hash_use_fast_hash,
                    target_hash_graph_type,
                    include_default_attributes: self.include_defaults,
                    target_hash_recursive: self.target_hash_recursive,
                    keep_going: self.keep_going,
                    streaming: self.streaming,
                    cached: !self.no_cache,
                    imports: self.imports,
                })
            }),
            output: self
                .output
                .try_map(|x| x.resolve(&ctx.working_dir).into_string())?,
            concurrency: self
                .num_threads
                .map(|num| buck2_cli_proto::Concurrency { concurrency: num }),
        };

        if self.show_output {
            targets_show_outputs(
                ctx.stdin(),
                buckd,
                target_request,
                None,
                &self.common_opts.console_opts,
            )
            .await
        } else if self.show_full_output {
            let project_root = ctx.paths()?.roots.project_root.clone();
            targets_show_outputs(
                ctx.stdin(),
                buckd,
                target_request,
                Some(project_root.root()),
                &self.common_opts.console_opts,
            )
            .await
        } else {
            targets(
                ctx.stdin(),
                buckd,
                target_request,
                &self.common_opts.console_opts,
            )
            .await
        }
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        &self.common_opts.console_opts
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.common_opts.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.common_opts.config_opts
    }
}

async fn targets_show_outputs(
    stdin: &mut Stdin,
    buckd: &mut BuckdClientConnector<'_>,
    target_request: TargetsRequest,
    root_path: Option<&AbsNormPath>,
    console_opts: &CommonConsoleOptions,
) -> ExitResult {
    let response = buckd
        .with_flushing()
        .targets_show_outputs(
            target_request,
            stdin.console_interaction_stream(console_opts),
            &mut NoPartialResultHandler,
        )
        .await??;
    for target_paths in response.targets_paths {
        for path in target_paths.paths {
            let path = if cfg!(windows) {
                path.replace('/', "\\")
            } else {
                path
            };
            match root_path {
                Some(root) => {
                    buck2_client_ctx::println!(
                        "{} {}",
                        target_paths.target,
                        root.as_path().join(path).display()
                    )
                }
                None => buck2_client_ctx::println!("{} {}", target_paths.target, path),
            }?;
        }
    }
    ExitResult::success()
}

async fn targets(
    stdin: &mut Stdin,
    buckd: &mut BuckdClientConnector<'_>,
    target_request: TargetsRequest,
    console_opts: &CommonConsoleOptions,
) -> ExitResult {
    let response = buckd
        .with_flushing()
        .targets(
            target_request,
            stdin.console_interaction_stream(console_opts),
            &mut StdoutPartialResultHandler,
        )
        .await??;
    if !response.serialized_targets_output.is_empty() {
        buck2_client_ctx::print!("{}", response.serialized_targets_output)?;
    }
    ExitResult::success()
}


/*
 * 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 anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::CounterWithExamples;
use buck2_cli_proto::TestRequest;
use buck2_cli_proto::TestSessionOptions;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonBuildOptions;
use buck2_client_ctx::common::CommonCommandOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::final_console::FinalConsole;
use buck2_client_ctx::output_destination_arg::OutputDestinationArg;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::stdio::eprint_line;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_client_ctx::subscribers::superconsole::test::span_from_build_failure_count;
use buck2_client_ctx::subscribers::superconsole::test::TestCounterColumn;
use buck2_core::fs::fs_util;
use buck2_core::fs::working_dir::WorkingDir;
use gazebo::prelude::*;
use superconsole::Line;
use superconsole::Span;

use crate::commands::build::print_build_result;

fn forward_output_to_path(
    output: &str,
    path_arg: &PathArg,
    working_dir: &WorkingDir,
) -> anyhow::Result<()> {
    fs_util::write(path_arg.resolve(working_dir), output)
        .context("Failed to write test executor output to path")
}

fn print_error_counter(
    console: &FinalConsole,
    counter: &CounterWithExamples,
    error_type: &str,
    symbol: &str,
) -> anyhow::Result<()> {
    if counter.count > 0 {
        console.print_error(&format!("{} {}", counter.count, error_type))?;
        for test_name in &counter.example_tests {
            console.print_error(&format!("  {} {}", symbol, test_name))?;
        }
        if counter.count > counter.max {
            console.print_error(&format!(
                "  ...and {} more not shown...",
                counter.count - counter.max
            ))?;
        }
    }
    Ok(())
}
#[derive(Debug, clap::Parser)]
#[clap(name = "test", about = "Build and test the specified targets")]
pub struct TestCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    #[clap(flatten)]
    build_opts: CommonBuildOptions,

    #[clap(
        long = "exclude",
        multiple_values = true,
        help = "Labels on targets to exclude from tests"
    )]
    exclude: Vec<String>,

    #[clap(
        long = "include",
        alias = "labels",
        help = "Labels on targets to include from tests. Prefixing with `!` means to exclude. First match wins unless overridden by `always-exclude` flag.\n\
If include patterns are present, regardless of whether exclude patterns are present, then all targets are by default excluded unless explicitly included.",
        multiple_values = true
    )]
    include: Vec<String>,

    #[clap(
        long = "always-exclude",
        alias = "always_exclude",
        help = "Whether to always exclude if the label appears in `exclude`, regardless of which appears first"
    )]
    always_exclude: bool,

    #[clap(
        long = "build-filtered",
        help = "Whether to build tests that are excluded via labels."
    )]
    build_filtered_targets: bool, // TODO(bobyf) this flag should always override the buckconfig option when we use it

    /// This option does nothing. It is here to keep compatibility with Buck1 and ci
    #[allow(unused)] // for v1 compat
    #[clap(long = "deep")]
    deep: bool,

    // ignored. only for e2e tests. compatibility with v1.
    #[clap(long = "xml")]
    #[allow(unused)] // for v1 compat
    xml: Option<String>,

    /// Will allow tests that are compatible with RE (setup to run from the repo root and
    /// use relative paths) to run from RE.
    #[clap(long, group = "re_options", alias = "unstable-allow-tests-on-re")]
    unstable_allow_compatible_tests_on_re: bool,

    /// Will run tests to on RE even if they are missing required settings (running from the root +
    /// relative paths). Those required settings just get overridden.
    #[clap(long, group = "re_options", alias = "unstable-force-tests-on-re")]
    unstable_allow_all_tests_on_re: bool,

    #[clap(name = "TARGET_PATTERNS", help = "Patterns to test")]
    patterns: Vec<String>,

    /// Writes the test executor stdout to the provided path
    ///
    /// --test-executor-stdout=- will write to stdout
    ///
    /// --test-executor-stdout=FILEPATH will write to the provided filepath, overwriting the current
    /// file if it exists
    ///
    /// By default the test executor's stdout stream is captured
    #[clap(long)]
    test_executor_stdout: Option<OutputDestinationArg>,

    /// Writes the test executor stderr to the provided path
    ///
    /// --test-executor-stderr=- will write to stderr
    ///
    /// --test-executor-stderr=FILEPATH will write to the provided filepath, overwriting the current
    /// file if it exists
    ///
    /// By default test executor's stderr stream is captured
    #[clap(long)]
    test_executor_stderr: Option<OutputDestinationArg>,

    /// Additional arguments passed to the test executor.
    ///
    /// Test executor is expected to have `--env` flag to pass environment variables.
    /// Can be used like this:
    ///
    /// buck2 test //foo:bar -- --env PRIVATE_KEY=123
    #[clap(name = "TEST_EXECUTOR_ARGS", raw = true)]
    test_executor_args: Vec<String>,
}

#[async_trait]
impl StreamingCommand for TestCommand {
    const COMMAND_NAME: &'static str = "test";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;
        let response = buckd
            .with_flushing()
            .test(
                TestRequest {
                    context: Some(context),
                    target_patterns: self
                        .patterns
                        .map(|pat| buck2_data::TargetPattern { value: pat.clone() }),
                    test_executor_args: self.test_executor_args,
                    excluded_labels: self.exclude,
                    included_labels: self.include,
                    always_exclude: self.always_exclude,
                    build_filtered_targets: self.build_filtered_targets,
                    // we don't currently have a different flag for this, so just use the build one.
                    concurrency: self.build_opts.num_threads.unwrap_or(0),
                    build_opts: Some(self.build_opts.to_proto()),
                    session_options: Some(TestSessionOptions {
                        allow_re: self.unstable_allow_compatible_tests_on_re
                            || self.unstable_allow_all_tests_on_re,
                        force_use_project_relative_paths: self.unstable_allow_all_tests_on_re,
                        force_run_from_project_root: self.unstable_allow_all_tests_on_re,
                    }),
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut NoPartialResultHandler,
            )
            .await??;

        let statuses = response
            .test_statuses
            .as_ref()
            .expect("Daemon to not return empty statuses");

        let listing_failed = statuses
            .listing_failed
            .as_ref()
            .context("Missing `listing_failed`")?;
        let passed = statuses.passed.as_ref().context("Missing `passed`")?;
        let failed = statuses.failed.as_ref().context("Missing `failed`")?;
        let fatals = statuses.fatals.as_ref().context("Missing `fatals`")?;
        let skipped = statuses.skipped.as_ref().context("Missing `skipped`")?;

        let console = self.common_opts.console_opts.final_console();
        print_build_result(&console, &response.error_messages)?;
        if !response.error_messages.is_empty() {
            console.print_error(&format!("{} BUILDS FAILED", response.error_messages.len()))?;
        }

        // TODO(nmj): Might make sense for us to expose the event ctx, and use its
        //            handle_stdout method, instead of raw buck2_client::println!s here.
        // TODO: also remove the duplicate information when the above is done.

        let mut line = Line::default();
        line.push(Span::new_unstyled_lossy("Tests finished: "));
        if listing_failed.count > 0 {
            line.push(TestCounterColumn::LISTING_FAIL.to_span_from_test_statuses(statuses)?);
            line.push(Span::new_unstyled_lossy(". "));
        }
        let columns = [
            TestCounterColumn::PASS,
            TestCounterColumn::FAIL,
            TestCounterColumn::FATAL,
            TestCounterColumn::SKIP,
        ];
        for column in columns {
            line.push(column.to_span_from_test_statuses(statuses)?);
            line.push(Span::new_unstyled_lossy(". "));
        }
        line.push(span_from_build_failure_count(
            response.error_messages.len(),
        )?);
        eprint_line(&line)?;

        print_error_counter(&console, listing_failed, "LISTINGS FAILED", "âš ")?;
        print_error_counter(&console, failed, "TESTS FAILED", "âœ—")?;
        print_error_counter(&console, fatals, "TESTS FATALS", "âš ")?;
        if passed.count + failed.count + fatals.count + skipped.count == 0 {
            console.print_warning("NO TESTS RAN")?;
        }

        let info_messages = response.executor_info_messages;
        for message in info_messages {
            console.print_stderr(message.as_str())?;
        }

        match self.test_executor_stderr {
            Some(OutputDestinationArg::Path(path)) => {
                forward_output_to_path(&response.executor_stderr, &path, &ctx.working_dir)?;
            }
            Some(OutputDestinationArg::Stream) => {
                console.print_error(&response.executor_stderr)?;
            }
            _ => {}
        }

        let exit_result = if let Some(exit_code) = response.exit_code {
            ExitResult::status_extended(exit_code)
        } else {
            ExitResult::failure()
        };

        match self.test_executor_stdout {
            Some(OutputDestinationArg::Path(path)) => {
                forward_output_to_path(&response.executor_stdout, &path, &ctx.working_dir)?;
                exit_result
            }
            Some(OutputDestinationArg::Stream) => {
                exit_result.with_stdout(response.executor_stdout.into_bytes())
            }
            _ => exit_result,
        }
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        &self.common_opts.console_opts
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        &self.common_opts.event_log_opts
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        &self.common_opts.config_opts
    }
}


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

use async_trait::async_trait;
use buck2_cli_proto::AllocativeRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct AllocativeCommand {
    /// Output directory path for profile data.
    ///
    /// Directory will be created if it does not exist.
    #[clap(
        long,
        short = 'o',
        value_name = "PATH",
        default_value = "allocative-out"
    )]
    output: PathArg,
}

#[async_trait]
impl StreamingCommand for AllocativeCommand {
    const COMMAND_NAME: &'static str = "allocative";

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.empty_client_context()?;
        buckd
            .with_flushing()
            .allocative(
                AllocativeRequest {
                    context: Some(context),
                    output_path: self.output.resolve(&ctx.working_dir).into_string()?,
                },
                ctx.stdin().console_interaction_stream(self.console_opts()),
                &mut NoPartialResultHandler,
            )
            .await??;
        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        CommonConsoleOptions::default_ref()
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        CommonDaemonCommandOptions::default_ref()
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        CommonBuildConfigurationOptions::default_ref()
    }
}


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

use async_trait::async_trait;
use buck2_cli_proto::UnstableAllocatorStatsRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct AllocatorStatsCommand {
    /// Options to pass to allocator stats. We use JEMalloc, so the docs for `malloc_stats_print`
    /// indicate what is available (<https://jemalloc.net/jemalloc.3.html>). The default
    /// configuration prints minimal output, formatted as JSON.
    #[clap(short, long, default_value = "Jmdablxg", value_name = "OPTION")]
    options: String,
}

#[async_trait]
impl StreamingCommand for AllocatorStatsCommand {
    const COMMAND_NAME: &'static str = "allocator_stats";

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let res = buckd
            .with_flushing()
            .unstable_allocator_stats(UnstableAllocatorStatsRequest {
                options: self.options,
            })
            .await?;

        buck2_client_ctx::println!("{}", res.response)?;

        ExitResult::success()
    }

    fn console_opts(&self) -> &CommonConsoleOptions {
        CommonConsoleOptions::none_ref()
    }

    fn event_log_opts(&self) -> &CommonDaemonCommandOptions {
        CommonDaemonCommandOptions::default_ref()
    }

    fn common_opts(&self) -> &CommonBuildConfigurationOptions {
        CommonBuildConfigurationOptions::default_ref()
    }
}