/*
 * 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 buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::events_ctx::EventsCtx;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::replayer::Replayer;
use buck2_client_ctx::subscribers::get::get_console_with_root;

use crate::commands::log::options::EventLogOptions;

/// Replay an event log.
///
/// This command allows visualizing an existing event log in a Superconsole.
#[derive(Debug, clap::Parser)]
#[clap(
    setting = clap::AppSettings::TrailingVarArg
)]
pub struct ReplayCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,

    #[clap(
        long,
        help = "Control the playback speed using a float (i.e. 0.5, 2, etc)",
        value_name = "NUMBER"
    )]
    pub speed: Option<f64>,

    /// Preload the event log. This is typically only useful for benchmarking.
    #[clap(long)]
    preload: bool,

    #[clap(flatten)]
    console_opts: CommonConsoleOptions,

    #[clap(help = "Override the arguments")]
    pub override_args: Vec<String>,
}

impl ReplayCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self {
            event_log,
            speed,
            preload,
            console_opts,
            override_args: _,
        } = self;

        ctx.with_runtime(async move |mut ctx| {
            let (replayer, invocation) =
                Replayer::new(event_log.get(&ctx).await?, speed, preload).await?;

            let console = get_console_with_root(
                invocation.trace_id,
                console_opts.console_type,
                ctx.verbosity,
                true,
                speed,
                "(replay)", // Could be better
                console_opts.superconsole_config(),
                ctx.paths()?.isolation.clone(),
            )?
            .context("You must request a console for replay")?;

            let res = EventsCtx::new(vec![console])
                .unpack_stream::<_, ReplayResult, _>(
                    &mut NoPartialResultHandler,
                    Box::pin(replayer),
                    None,
                    ctx.stdin().console_interaction_stream(&console_opts),
                )
                .await??;

            for e in &res.errors {
                buck2_client_ctx::eprintln!("{}", e)?;
            }

            ExitResult::success()
        })
    }
}

struct ReplayResult {
    errors: Vec<String>,
}

impl TryFrom<buck2_cli_proto::command_result::Result> for ReplayResult {
    type Error = buck2_cli_proto::command_result::Result;

    fn try_from(v: buck2_cli_proto::command_result::Result) -> Result<Self, Self::Error> {
        use buck2_cli_proto::command_result::Result;

        // It would be good to declare this as a extension trait on our types, but for now to
        // support Replay this is fine;
        let errors = match v {
            Result::Error(v) => v.messages,
            Result::BuildResponse(v) => v.error_messages,
            Result::TestResponse(v) => v.error_messages,
            Result::CqueryResponse(v) => v.error_messages,
            Result::UqueryResponse(v) => v.error_messages,
            Result::AqueryResponse(v) => v.error_messages,
            Result::BxlResponse(v) => v.error_messages,
            _ => Vec::new(),
        };

        Ok(Self { errors })
    }
}


/*
 * 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_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stdio;
use tokio_stream::StreamExt;

use crate::commands::log::options::EventLogOptions;

/// Outputs the log in JSON format from selected invocation.
#[derive(Debug, clap::Parser)]
pub struct ShowLogCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,
}

impl ShowLogCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self { event_log } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            let (invocation, mut events) = log_path.unpack_stream().await?;

            let mut buf = Vec::new();

            serde_json::to_writer(&mut buf, &invocation)?;
            stdio::print_bytes(&buf)?;
            stdio::print_bytes(b"\n")?;

            while let Some(event) = events.try_next().await? {
                buf.clear();
                serde_json::to_writer(&mut buf, &event)?;
                stdio::print_bytes(&buf)?;
                stdio::print_bytes(b"\n")?;
            }

            anyhow::Ok(())
        })?;
        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 buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stdio;
use buck2_client_ctx::subscribers::event_log::user_event_types::try_get_user_event_for_read;
use tokio_stream::StreamExt;

use crate::commands::log::options::EventLogOptions;

/// Converts the event log from a selected invocation into a user event log, in JSONL format.
#[derive(Debug, clap::Parser)]
pub struct ShowUserLogCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,
}

impl ShowUserLogCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self { event_log } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            let (invocation, mut events) = log_path.unpack_stream().await?;

            let mut buf = Vec::new();

            serde_json::to_writer(&mut buf, &invocation)?;
            stdio::print_bytes(&buf)?;
            stdio::print_bytes(b"\n")?;

            while let Some(event) = events.try_next().await? {
                buf.clear();

                if let Some(user_event) = try_get_user_event_for_read(&event)? {
                    serde_json::to_writer(&mut buf, &user_event)?;
                    stdio::print_bytes(&buf)?;
                    stdio::print_bytes(b"\n")?;
                }
            }

            anyhow::Ok(())
        })?;
        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 buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;

use crate::commands::log::options::EventLogOptions;

/// Show buck command line arguments from selected invocation.
///
/// This command output is not machine readable.
/// Robots, please use `buck2 log show`.
#[derive(Debug, clap::Parser)]
pub struct WhatCmdCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,

    /// Show @-expanded command line arguments instead of the original command line.
    #[clap(long)]
    expand: bool,
}

impl WhatCmdCommand {
    pub(crate) fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext) -> ExitResult {
        let WhatCmdCommand { event_log, expand } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;
            let (invocation, _events) = log_path.unpack_stream().await?;

            buck2_client_ctx::println!("# cd {}", invocation.working_dir)?;
            if expand {
                buck2_client_ctx::println!("{}", invocation.display_expanded_command_line())?;
            } else {
                buck2_client_ctx::println!("{}", invocation.display_command_line())?;
            }

            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 buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;

use crate::commands::log::what_ran::WhatRanCommand;
use crate::commands::log::what_ran::WhatRanCommandCommon;

/// Outputs every command that failed in the selected invocation.
///
/// Look at the help for what-ran to understand the output format.
#[derive(Debug, clap::Parser)]
pub struct WhatFailedCommand {
    #[clap(flatten)]
    pub common: WhatRanCommandCommon,
}

impl WhatFailedCommand {
    pub fn exec(self, matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        WhatRanCommand {
            common: self.common,
            failed: true,
        }
        .exec(matches, ctx)
    }
}


/*
 * 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_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stream_value::StreamValue;
use tokio_stream::StreamExt;

use crate::commands::log::options::EventLogOptions;
use crate::commands::log::LogCommandOutputFormat;

/// Outputs materializations from selected invocation.
///
/// The output is a tab-separated list containing the path,
/// the materialization method, the file count, and the total size (after decompression).
#[derive(Debug, clap::Parser)]
pub struct WhatMaterializedCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,
    #[clap(
        long = "format",
        help = "Which output format to use for this command",
        default_value = "tabulated",
        ignore_case = true,
        arg_enum
    )]
    pub output: LogCommandOutputFormat,
}

fn write_output(
    output: &LogCommandOutputFormat,
    materialization: &buck2_data::MaterializationEnd,
    method: &str,
) -> anyhow::Result<()> {
    #[derive(serde::Serialize)]
    struct Record<'a> {
        path: &'a str,
        method: &'a str,
        file_count: u64,
        total_bytes: u64,
    }

    match output {
        LogCommandOutputFormat::Tabulated => {
            buck2_client_ctx::println!(
                "{}\t{}\t{}\t{}",
                materialization.path,
                method,
                materialization.file_count,
                materialization.total_bytes
            )
        }
        LogCommandOutputFormat::Csv => buck2_client_ctx::stdio::print_with_writer(|w| {
            let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(w);
            writer.serialize(Record {
                path: &materialization.path,
                method,
                file_count: materialization.file_count,
                total_bytes: materialization.total_bytes,
            })
        }),
        LogCommandOutputFormat::Json => buck2_client_ctx::stdio::print_with_writer(|mut w| {
            let record = Record {
                path: &materialization.path,
                method,
                file_count: materialization.file_count,
                total_bytes: materialization.total_bytes,
            };
            serde_json::to_writer(&mut w, &record)?;
            w.write(b"\n").map(|_| ())
        }),
    }
}

impl WhatMaterializedCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self { event_log, output } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            let (invocation, mut events) = log_path.unpack_stream().await?;

            buck2_client_ctx::eprintln!(
                "Showing materializations from: {}",
                invocation.display_command_line()
            )?;

            while let Some(event) = events.try_next().await? {
                match event {
                    StreamValue::Event(event) => {
                        match &event.data {
                            Some(buck2_data::buck_event::Data::SpanEnd(ref end)) => match end
                                .data
                                .as_ref()
                            {
                                Some(buck2_data::span_end_event::Data::Materialization(m)) => {
                                    // Only log what has been materialized.
                                    if !m.success {
                                        continue;
                                    }

                                    let method = match m
                                        .method
                                        .and_then(buck2_data::MaterializationMethod::from_i32)
                                    {
                                        Some(buck2_data::MaterializationMethod::CasDownload) => {
                                            "cas"
                                        }
                                        Some(buck2_data::MaterializationMethod::LocalCopy) => {
                                            "copy"
                                        }
                                        Some(buck2_data::MaterializationMethod::HttpDownload) => {
                                            "http"
                                        }
                                        Some(buck2_data::MaterializationMethod::Write) => "write",
                                        _ => "<unknown>",
                                    };

                                    write_output(&output, m, method)?;
                                }
                                _ => {}
                            },
                            _ => {}
                        }
                    }
                    StreamValue::Result(..) | StreamValue::PartialResult(..) => {}
                }
            }

            anyhow::Ok(())
        })?;
        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 std::borrow::Cow;
use std::collections::HashMap;

use async_trait::async_trait;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stream_value::StreamValue;
use buck2_data::re_platform::Property;
use buck2_event_observer::what_ran;
use buck2_event_observer::what_ran::CommandReproducer;
use buck2_event_observer::what_ran::WhatRanOptions;
use buck2_event_observer::what_ran::WhatRanOutputCommand;
use buck2_event_observer::what_ran::WhatRanOutputCommandExtra;
use buck2_event_observer::what_ran::WhatRanOutputWriter;
use buck2_event_observer::what_ran::WhatRanRelevantAction;
use buck2_event_observer::what_ran::WhatRanState;
use futures::stream::Stream;
use futures::TryStreamExt;
use indexmap::IndexMap;

use crate::commands::log::options::EventLogOptions;
use crate::commands::log::LogCommandOutputFormat;

/// Output everything Buck2 ran from selected invocation.
///
/// The output is presented as a series of tab-delimited records with the following structure:
///
/// The reason for executing a given command. That's either to build or to test.
///
/// The identity of this command. This will include the target that ran required it.
///
/// The executor for this command. This will either be RE or local.
///
/// Details to reproduce it. For RE, that's the action digest. For local, the command.
///
///
/// To reproduce an action that ran on RE, use the following command then follow the instructions.
/// The DIGEST is of the form `hash:size`.
///
/// frecli cas download-action DIGEST
///
///
/// To reproduce an action that ran locally, make sure your working directory is the project root
/// (if unsure, use `buck2 root --kind project` to find it), then run the command. The command is
/// already shell-quoted.
#[derive(Debug, clap::Parser)]
pub struct WhatRanCommand {
    #[clap(flatten)]
    pub common: WhatRanCommandCommon,

    /// Show only commands that failed
    #[clap(long)]
    pub failed: bool,
}

#[derive(Debug, clap::Parser)]
pub struct WhatRanCommandCommon {
    #[clap(flatten)]
    event_log: EventLogOptions,

    #[clap(
        long = "format",
        help = "Which output format to use for this command",
        default_value = "tabulated",
        ignore_case = true,
        arg_enum
    )]
    pub output: LogCommandOutputFormat,

    #[clap(flatten)]
    pub options: WhatRanOptions,
}

impl WhatRanCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self {
            common:
                WhatRanCommandCommon {
                    event_log,
                    mut output,
                    options,
                },
            failed,
        } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            let (invocation, events) = log_path.unpack_stream().await?;

            buck2_client_ctx::eprintln!(
                "Showing commands from: {}",
                invocation.display_command_line()
            )?;

            if failed {
                WhatFailedImpl::execute(events, &mut output, &options).await?;
            } else {
                WhatRanImpl::execute(events, &mut output, &options).await?;
            };

            anyhow::Ok(())
        })?;

        ExitResult::success()
    }
}

#[async_trait]
trait WhatRanCommandImplementation: Default {
    fn event(
        &mut self,
        event: Box<buck2_data::BuckEvent>,
        output: &mut impl WhatRanOutputWriter,
        options: &WhatRanOptions,
    ) -> anyhow::Result<()>;

    async fn execute(
        mut events: impl Stream<Item = anyhow::Result<StreamValue>> + Unpin + Send,
        output: &mut (impl WhatRanOutputWriter + Send),
        options: &WhatRanOptions,
    ) -> anyhow::Result<()> {
        let mut cmd = Self::default();

        while let Some(event) = events.try_next().await? {
            match event {
                StreamValue::Event(event) => cmd.event(event, output, options)?,
                _ => {}
            }
        }

        Ok(())
    }
}

/// The state for a WhatRan command. This is all the events we have seen that are
/// WhatRanRelevantActions. This emits the actions immediately.
#[derive(Default)]
pub struct WhatRanImpl {
    /// Maps action spans to their details.
    known_actions: HashMap<u64, Box<buck2_data::BuckEvent>>,
}

impl WhatRanState<u64> for WhatRanImpl {
    fn get(&self, span_id: u64) -> Option<WhatRanRelevantAction<'_>> {
        self.known_actions
            .get(&span_id)
            .and_then(|e| e.data.as_ref())
            .and_then(WhatRanRelevantAction::from_buck_data)
    }
}

impl WhatRanCommandImplementation for WhatRanImpl {
    /// Receive a new event. We start by emitting it if it's relevant (since that only takes a
    /// borrow), and then if it's relevant as a parent, we store it for latter use. Note that in
    /// practice we don't expect the event to be *both* relevant to emit *and* a
    /// WhatRanRelevantAction, but it doesn't hurt to always check both.
    fn event(
        &mut self,
        event: Box<buck2_data::BuckEvent>,
        output: &mut impl WhatRanOutputWriter,
        options: &WhatRanOptions,
    ) -> anyhow::Result<()> {
        if let Some(data) = &event.data {
            what_ran::emit_event_if_relevant(event.parent_id, data, &*self, output, options)?;

            if WhatRanRelevantAction::from_buck_data(data).is_some() {
                self.known_actions.insert(event.span_id, event);
            }
        }

        Ok(())
    }
}

/// The state for a WhatRan command when only showing actions that failed. This stores all the events
/// we have seen that are WhatRanRelevantActions, and the CommandReproducer associated with them.
#[derive(Default)]
pub struct WhatFailedImpl {
    /// Maps action spans to their details.
    known_actions: HashMap<u64, WhatFailedEntry>,
}

#[allow(clippy::vec_box)]
struct WhatFailedEntry {
    /// Known to be a WhatRanRelevantAction.
    event: Box<buck2_data::BuckEvent>,

    /// Known to be a CommandReproducer.
    reproducers: Vec<Box<buck2_data::BuckEvent>>,
}

impl WhatRanState<u64> for WhatFailedImpl {
    fn get(&self, span_id: u64) -> Option<WhatRanRelevantAction<'_>> {
        self.known_actions
            .get(&span_id)
            .and_then(|e| e.event.data.as_ref())
            .and_then(WhatRanRelevantAction::from_buck_data)
    }
}

impl WhatRanCommandImplementation for WhatFailedImpl {
    /// Receive a new event. We start by emitting it if it's relevant (since that only takes a
    /// borrow), and then if it's relevant as a parent, we store it for latter use. Note that in
    /// practice we don't expect the event to be *both* relevant to emit *and* a
    /// WhatRanRelevantAction, but it doesn't hurt to always check both.
    fn event(
        &mut self,
        event: Box<buck2_data::BuckEvent>,
        output: &mut impl WhatRanOutputWriter,
        options: &WhatRanOptions,
    ) -> anyhow::Result<()> {
        if let Some(data) = &event.data {
            if WhatRanRelevantAction::from_buck_data(data).is_some() {
                self.known_actions.insert(
                    event.span_id,
                    WhatFailedEntry {
                        event,
                        reproducers: Default::default(),
                    },
                );
                return Ok(());
            }

            if CommandReproducer::from_buck_data(data, options).is_some() {
                if let Some(entry) = self.known_actions.get_mut(&event.parent_id) {
                    entry.reproducers.push(event);
                }
                return Ok(());
            }

            match data {
                buck2_data::buck_event::Data::SpanEnd(span) => match &span.data {
                    Some(buck2_data::span_end_event::Data::ActionExecution(action))
                        if action.failed =>
                    {
                        if let Some(entry) = self.known_actions.remove(&event.span_id) {
                            let action = WhatRanRelevantAction::from_buck_data(
                                entry.event.data.as_ref().expect("Checked above"),
                            );

                            for repro in entry.reproducers.iter() {
                                what_ran::emit_reproducer(
                                    action,
                                    CommandReproducer::from_buck_data(
                                        repro.data.as_ref().expect("Checked above"),
                                        options,
                                    )
                                    .expect("Checked above"),
                                    output,
                                )?;
                            }
                        }
                    }
                    _ => {}
                },
                _ => {}
            }
        }

        Ok(())
    }
}

/// An output that writes to stdout in a tabulated format.
impl WhatRanOutputWriter for LogCommandOutputFormat {
    fn emit_command(&mut self, command: WhatRanOutputCommand<'_>) -> anyhow::Result<()> {
        match self {
            Self::Tabulated => {
                buck2_client_ctx::println!(
                    "{}\t{}\t{}\t{}",
                    command.reason(),
                    command.identity(),
                    command.repro().executor(),
                    command.repro().as_human_readable()
                )
            }
            Self::Json => {
                let reproducer = match command.repro() {
                    CommandReproducer::CacheQuery(cache_hit) => JsonReproducer::CacheQuery {
                        digest: &cache_hit.action_digest,
                    },
                    CommandReproducer::CacheHit(cache_hit) => JsonReproducer::Cache {
                        digest: &cache_hit.action_digest,
                        action_key: cache_hit.action_key.as_deref(),
                    },
                    CommandReproducer::ReExecute(re_execute) => JsonReproducer::Re {
                        digest: &re_execute.action_digest,
                        platform_properties: into_index_map(&re_execute.platform),
                        action_key: re_execute.action_key.as_deref(),
                    },
                    CommandReproducer::LocalExecute(local_execute) => JsonReproducer::Local {
                        command: local_execute.command.as_ref().map_or_else(
                            || Cow::Owned(Vec::new()),
                            |command| Cow::Borrowed(command.argv.as_ref()),
                        ),
                        env: local_execute
                            .command
                            .as_ref()
                            .into_iter()
                            .flat_map(|command| command.env.iter())
                            .map(|entry| (entry.key.as_ref(), entry.value.as_ref()))
                            .collect(),
                    },
                    CommandReproducer::WorkerExecute(worker_execute) => JsonReproducer::Local {
                        command: worker_execute.command.as_ref().map_or_else(
                            || Cow::Owned(Vec::new()),
                            |command| Cow::Borrowed(command.argv.as_ref()),
                        ),
                        env: worker_execute
                            .command
                            .as_ref()
                            .into_iter()
                            .flat_map(|command| command.env.iter())
                            .map(|entry| (entry.key.as_ref(), entry.value.as_ref()))
                            .collect(),
                    },
                };

                let command = JsonCommand {
                    reason: command.reason(),
                    identity: command.identity(),
                    reproducer,
                    extra: command.extra().map(Into::into),
                };

                buck2_client_ctx::stdio::print_with_writer(|mut w| {
                    serde_json::to_writer(&mut w, &command)?;
                    w.write(b"\n").map(|_| ())
                })
            }
            Self::Csv => {
                #[derive(serde::Serialize)]
                struct Record<'a> {
                    reason: &'a str,
                    identity: &'a str,
                    executor: String,
                    reproducer: String,
                }

                buck2_client_ctx::stdio::print_with_writer(|w| {
                    let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(w);
                    writer.serialize(Record {
                        reason: command.reason(),
                        identity: command.identity(),
                        executor: command.repro().executor(),
                        reproducer: command.repro().as_human_readable().to_string(),
                    })
                })
            }
        }
    }
}

fn into_index_map(platform: &Option<buck2_data::RePlatform>) -> IndexMap<&str, &str> {
    platform.as_ref().map_or_else(IndexMap::new, |p| {
        p.properties
            .iter()
            .map(|Property { name, value }| (name.as_ref(), value.as_ref()))
            .collect()
    })
}

#[derive(serde::Serialize)]
struct JsonCommand<'a> {
    reason: &'a str,
    identity: &'a str,
    reproducer: JsonReproducer<'a>,
    #[serde(skip_serializing_if = "Option::is_none")]
    extra: Option<JsonExtra<'a>>,
}

mod json_reproducer {
    #![allow(clippy::ref_option_ref)] // within Serialize

    use super::*;

    #[derive(serde::Serialize)]
    #[serde(tag = "executor", content = "details")]
    pub enum JsonReproducer<'a> {
        CacheQuery {
            digest: &'a str,
        },
        Cache {
            digest: &'a str,
            #[serde(skip_serializing_if = "Option::is_none")]
            action_key: Option<&'a str>,
        },
        Re {
            digest: &'a str,
            platform_properties: IndexMap<&'a str, &'a str>,
            #[serde(skip_serializing_if = "Option::is_none")]
            action_key: Option<&'a str>,
        },
        Local {
            command: Cow<'a, [String]>,
            env: IndexMap<&'a str, &'a str>,
        },
    }
}

use json_reproducer::JsonReproducer;

#[derive(serde::Serialize)]
#[serde(rename_all = "lowercase")]
enum JsonExtra<'a> {
    TestCases(&'a [String]),
}

impl<'a> From<WhatRanOutputCommandExtra<'a>> for JsonExtra<'a> {
    fn from(extra: WhatRanOutputCommandExtra<'a>) -> JsonExtra<'a> {
        match extra {
            WhatRanOutputCommandExtra::TestCases(cases) => JsonExtra::TestCases(cases),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_base_command() -> JsonCommand<'static> {
        let command = Cow::Owned(vec!["some".to_owned(), "command".to_owned()]);
        let mut env = IndexMap::new();
        env.insert("KEY", "val");

        JsonCommand {
            reason: "test.run",
            identity: "some/target",
            reproducer: JsonReproducer::Local { command, env },
            extra: None,
        }
    }

    fn make_base_command_in_re() -> JsonCommand<'static> {
        JsonCommand {
            reason: "test.run",
            identity: "some/target",
            reproducer: JsonReproducer::Re {
                digest: "placeholder",
                platform_properties: indexmap::indexmap! {
                    "platform" => "linux-remote-execution"
                },
                action_key: None,
            },
            extra: None,
        }
    }

    #[test]
    fn serialize_what_ran_command_no_extr() -> anyhow::Result<()> {
        let command = make_base_command();

        let expected = r#"{
  "reason": "test.run",
  "identity": "some/target",
  "reproducer": {
    "executor": "Local",
    "details": {
      "command": [
        "some",
        "command"
      ],
      "env": {
        "KEY": "val"
      }
    }
  }
}"#;
        assert_eq!(expected, serde_json::to_string_pretty(&command)?);
        Ok(())
    }

    #[test]
    fn serialize_what_ran_command_with_extra() -> anyhow::Result<()> {
        let mut command = make_base_command();
        let cases = &["case".to_owned()];
        command.extra = Some(JsonExtra::TestCases(cases));

        let expected = r#"{
  "reason": "test.run",
  "identity": "some/target",
  "reproducer": {
    "executor": "Local",
    "details": {
      "command": [
        "some",
        "command"
      ],
      "env": {
        "KEY": "val"
      }
    }
  },
  "extra": {
    "testcases": [
      "case"
    ]
  }
}"#;
        assert_eq!(expected, serde_json::to_string_pretty(&command)?);
        Ok(())
    }

    #[test]
    fn serialize_what_ran_command_in_re() -> anyhow::Result<()> {
        let command = make_base_command_in_re();

        let expected = r#"{
  "reason": "test.run",
  "identity": "some/target",
  "reproducer": {
    "executor": "Re",
    "details": {
      "digest": "placeholder",
      "platform_properties": {
        "platform": "linux-remote-execution"
      }
    }
  }
}"#;
        assert_eq!(expected, serde_json::to_string_pretty(&command)?);
        Ok(())
    }
}


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

use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;
use std::time::SystemTimeError;

use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stream_value::StreamValue;
use buck2_client_ctx::subscribers::superconsole::session_info::SessionInfoComponent;
use buck2_client_ctx::subscribers::superconsole::timed_list::TimedList;
use buck2_client_ctx::subscribers::superconsole::StatefulSuperConsole;
use buck2_client_ctx::subscribers::superconsole::SuperConsoleConfig;
use buck2_client_ctx::subscribers::superconsole::SuperConsoleState;
use buck2_client_ctx::subscribers::superconsole::CUTOFFS;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_event_observer::verbosity::Verbosity;
use buck2_events::BuckEvent;
use superconsole::components::DrawVertical;
use superconsole::Component;
use superconsole::Dimensions;
use superconsole::DrawMode;
use superconsole::Lines;
use tokio_stream::StreamExt;

use crate::commands::log::options::EventLogOptions;

/// Show the spans that were open when the log ended.
#[derive(Debug, clap::Parser)]
pub struct WhatUpCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,

    /// Show spans after X amount of milliseconds
    #[clap(
        long,
        help = "Print the actions that where open after certain amount of milliseconds",
        value_name = "NUMBER"
    )]
    pub after: Option<u64>,
}

impl WhatUpCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self { event_log, after } = self;
        let cutoff_time = after.map(Duration::from_millis);

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            // Get events
            let (invocation, mut events) = log_path.unpack_stream().await?;

            let mut super_console = StatefulSuperConsole::console_builder()
                .build_forced(StatefulSuperConsole::FALLBACK_SIZE)?;

            let mut super_console_state = SuperConsoleState::new(
                None,
                invocation.trace_id,
                FileNameBuf::unchecked_new("placeholder"),
                Verbosity::Default,
                true,
                SuperConsoleConfig {
                    max_lines: 1000000,
                    ..Default::default()
                },
            )?;
            let mut first_timestamp = None;
            // Ignore any events that are truncated, hence unreadable
            while let Ok(Some(event)) = events.try_next().await {
                match event {
                    StreamValue::Event(event) => {
                        let e = BuckEvent::try_from(event)?;
                        match cutoff_time {
                            Some(cutoff_time) => {
                                if should_stop_reading(
                                    cutoff_time,
                                    e.timestamp(),
                                    *first_timestamp.get_or_insert(e.timestamp()),
                                )? {
                                    break;
                                }
                            }
                            _ => (),
                        }

                        super_console_state.update_event_observer(&Arc::new(e))?;
                    }
                    StreamValue::PartialResult(..) => {}
                    StreamValue::Result(result) => {
                        let result = StatefulSuperConsole::render_result_errors(&result);
                        super_console.emit(result);
                        super_console.finalize(&Self::component(&super_console_state))?;
                        buck2_client_ctx::eprintln!("No open spans to render when log ended")?;
                        return Ok(());
                    }
                }
            }

            super_console
                .finalize_with_mode(&Self::component(&super_console_state), DrawMode::Normal)?;
            anyhow::Ok(())
        })?;

        ExitResult::success()
    }

    fn component<'a>(state: &'a SuperConsoleState) -> impl Component + 'a {
        struct ComponentImpl<'a> {
            state: &'a SuperConsoleState,
        }

        impl<'a> Component for ComponentImpl<'a> {
            fn draw_unchecked(
                &self,
                dimensions: Dimensions,
                mode: DrawMode,
            ) -> anyhow::Result<Lines> {
                let mut draw = DrawVertical::new(dimensions);
                draw.draw(
                    &SessionInfoComponent {
                        session_info: self.state.session_info(),
                    },
                    mode,
                )?;
                draw.draw(&TimedList::new(&CUTOFFS, "", self.state), mode)?;
                Ok(draw.finish())
            }
        }

        ComponentImpl { state }
    }
}

fn should_stop_reading(
    after: Duration,
    event: SystemTime,
    first: SystemTime,
) -> Result<bool, SystemTimeError> {
    let elapsed = event.duration_since(first)?;
    if elapsed > after {
        return Ok(true);
    }
    Ok(false)
}


/*
 * 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 buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::stream_value::StreamValue;
use buck2_event_observer::display;
use buck2_event_observer::display::TargetDisplayOptions;
use tokio_stream::StreamExt;

use crate::commands::log::options::EventLogOptions;
use crate::commands::log::LogCommandOutputFormat;

/// Outputs stats about uploads to RE from the selected invocation.
#[derive(Debug, clap::Parser)]
pub struct WhatUploadedCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,
    #[clap(
        long = "format",
        help = "Which output format to use for this command",
        default_value = "tabulated",
        ignore_case = true,
        arg_enum
    )]
    pub output: LogCommandOutputFormat,
}

fn print_uploads(
    format: &LogCommandOutputFormat,
    upload: ReUploadEvent,
    state: &HashMap<u64, buck2_data::ActionExecutionStart>,
    total_digests: &mut u64,
    total_bytes: &mut u64,
) -> anyhow::Result<()> {
    let digests_uploaded = upload.inner.digests_uploaded.unwrap_or_default();
    let bytes_uploaded = upload.inner.bytes_uploaded.unwrap_or_default();
    *total_digests += digests_uploaded;
    *total_bytes += bytes_uploaded;

    let action_str = if let Some(action) = state.get(&upload.parent_span_id) {
        display::display_action_identity(
            action.key.as_ref(),
            action.name.as_ref(),
            TargetDisplayOptions::for_log(),
        )?
    } else {
        "unknown action".to_owned()
    };

    #[derive(serde::Serialize)]
    struct Record<'a> {
        action: &'a str,
        digests_uploaded: u64,
        bytes_uploaded: u64,
    }

    match format {
        LogCommandOutputFormat::Tabulated => {
            buck2_client_ctx::println!("{}\t{}\t{}", action_str, digests_uploaded, bytes_uploaded)
        }
        LogCommandOutputFormat::Csv => buck2_client_ctx::stdio::print_with_writer(|w| {
            let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(w);
            writer.serialize(Record {
                action: &action_str,
                digests_uploaded,
                bytes_uploaded,
            })
        }),
        LogCommandOutputFormat::Json => {
            buck2_client_ctx::stdio::print_with_writer(|w| {
                let record = Record {
                    action: &action_str,
                    digests_uploaded,
                    bytes_uploaded,
                };
                serde_json::to_writer(w, &record)
            })?;
            buck2_client_ctx::println!("")
        }
    }
}

struct ReUploadEvent<'a> {
    pub parent_span_id: u64,
    pub inner: &'a buck2_data::ReUploadEnd,
}

impl WhatUploadedCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let Self { event_log, output } = self;

        ctx.with_runtime(async move |ctx| {
            let log_path = event_log.get(&ctx).await?;

            let (invocation, mut events) = log_path.unpack_stream().await?;
            buck2_client_ctx::eprintln!(
                "Showing uploads from: {}",
                invocation.display_command_line()
            )?;

            let mut total_digests_uploaded = 0;
            let mut total_bytes_uploaded = 0;
            let mut state = HashMap::new();
            while let Some(event) = events.try_next().await? {
                match event {
                    // Insert parent span information so we can refer back to it later.
                    StreamValue::Event(event) => match event.data {
                        Some(buck2_data::buck_event::Data::SpanStart(start)) => match start.data {
                            Some(buck2_data::span_start_event::Data::ActionExecution(action)) => {
                                state.insert(event.span_id, action);
                            }
                            _ => {}
                        },
                        Some(buck2_data::buck_event::Data::SpanEnd(end)) => {
                            match end.data.as_ref() {
                                Some(buck2_data::span_end_event::Data::ReUpload(ref u)) => {
                                    print_uploads(
                                        &output,
                                        ReUploadEvent {
                                            parent_span_id: event.parent_id,
                                            inner: u,
                                        },
                                        &state,
                                        &mut total_digests_uploaded,
                                        &mut total_bytes_uploaded,
                                    )?;
                                }
                                _ => {}
                            }
                        }
                        _ => {}
                    },
                    StreamValue::Result(..) | StreamValue::PartialResult(..) => {}
                }
            }
            buck2_client_ctx::eprintln!(
                "total: digests: {}, bytes: {}",
                total_digests_uploaded,
                total_bytes_uploaded
            )?;

            anyhow::Ok(())
        })?;

        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 async_trait::async_trait;
use buck2_cli_proto::AqueryRequest;
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::StdoutPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

use crate::commands::query::common::CommonQueryOptions;

/// Perform queries on the action graph (experimental).
///
/// The action graph consists of all the declared actions for a build, with dependencies
/// when one action consumes the outputs of another action.
///
/// Examples:
///
/// Print the action producing a target's default output
///
/// `buck2 aquery //java/com/example/app:amazing`
///
/// List all the commands for run actions for building a target
///
/// `buck2 aquery 'kind(run, deps("//java/com/example/app:amazing+more"))' --output-attribute=cmd`
///
/// Dynamic outputs (`ctx.actions.dynamic_output`):
///
/// Currently, aquery interacts poorly with dynamic outputs. It may return incorrect results or otherwise
/// behave unexpectedly.
#[derive(Debug, clap::Parser)]
#[clap(name = "aquery")]
pub struct AqueryCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    #[clap(flatten)]
    query_common: CommonQueryOptions,
}

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

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let (query, query_args) = self.query_common.get_query();
        let unstable_output_format = self.query_common.output_format() as i32;
        let output_attributes = self.query_common.attributes.get()?;
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;

        let response = buckd
            .with_flushing()
            .aquery(
                AqueryRequest {
                    query,
                    query_args,
                    context: Some(context),
                    output_attributes,
                    unstable_output_format,
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut StdoutPartialResultHandler,
            )
            .await??;

        for message in &response.error_messages {
            buck2_client_ctx::eprintln!("{}", message)?;
        }

        if !response.error_messages.is_empty() {
            ExitResult::failure()
        } else {
            ExitResult::success()
        }
    }

    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 buck2_cli_proto::QueryOutputFormat;
use buck2_client_ctx::query_args::CommonAttributeArgs;
use buck2_query_parser::placeholder::QUERY_PERCENT_SS_PLACEHOLDER;
use dupe::Dupe;

#[derive(
    Debug,
    Clone,
    Dupe,
    clap::ArgEnum,
    serde::Serialize,
    serde::Deserialize
)]
#[clap(rename_all = "snake_case")]
enum QueryOutputFormatArg {
    Dot,
    Json,
    DotCompact,
}

/// Args common to all the query commands
#[derive(Debug, clap::Parser, serde::Serialize, serde::Deserialize)]
#[clap(group = clap::ArgGroup::new("output_attribute_flags").multiple(false))]
pub(crate) struct CommonQueryOptions {
    #[clap(name = "QUERY", help = "the query to evaluate")]
    query: String,

    #[clap(flatten)]
    pub attributes: CommonAttributeArgs,

    #[clap(long, help = "Output in JSON format")]
    json: bool,

    #[clap(long, help = "Output in Graphviz Dot format")]
    dot: bool,

    #[clap(long, help = "Output in a more compact format than Graphviz Dot")]
    dot_compact: bool,

    #[clap(
        long,
        ignore_case = true,
        help = "Output format (default: list).",
        long_help = "Output format (default: list). \n
           dot -  dot graph format. \n
           dot_compact - compact alternative to dot format. \n
           json - JSON format.
         ",
        value_name = "dot|dot_compact|json",
        arg_enum
    )]
    output_format: Option<QueryOutputFormatArg>,

    #[clap(
        name = "QUERY_ARGS",
        help = "list of literals for a multi-query (one containing `%s` or `%Ss`)"
    )]
    query_args: Vec<String>,
}

impl CommonQueryOptions {
    fn args_as_set(args: &[String]) -> String {
        let mut s = "set(".to_owned();
        for (i, v) in args.iter().enumerate() {
            if i != 0 {
                s += " ";
            }
            s += "'";
            s += v;
            s += "'";
        }
        s += ")";
        s
    }

    pub fn output_format(&self) -> QueryOutputFormat {
        match self.output_format {
            Some(QueryOutputFormatArg::Json) => QueryOutputFormat::Json,
            Some(QueryOutputFormatArg::Dot) => QueryOutputFormat::Dot,
            Some(QueryOutputFormatArg::DotCompact) => QueryOutputFormat::DotCompact,
            None => {
                if self.json {
                    QueryOutputFormat::Json
                } else if self.dot {
                    QueryOutputFormat::Dot
                } else if self.dot_compact {
                    QueryOutputFormat::DotCompact
                } else {
                    QueryOutputFormat::Default
                }
            }
        }
    }

    pub fn get_query(&self) -> (String, Vec<String>) {
        if self.query.contains(QUERY_PERCENT_SS_PLACEHOLDER) {
            let replacement = Self::args_as_set(&self.query_args);
            (
                self.query
                    .replace(QUERY_PERCENT_SS_PLACEHOLDER, &replacement),
                vec![],
            )
        } else {
            (self.query.clone(), self.query_args.clone())
        }
    }
}


/*
 * 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::CqueryRequest;
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::StdoutPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

use crate::commands::query::common::CommonQueryOptions;

/// Perform queries on the configured target graph.
///
/// The configured target graph includes information about the configuration (platforms) and
/// transitions involved in building targets. In the configured graph, `selects` are fully
/// resolved. The same target may appear in multiple different configurations (when printed,
/// the configuration is after the target in parentheses).
///
/// A user can specify a `--target-universe` flag to control how literals are resolved. When
/// provided, any literals will resolve to all matching targets within the universe (which
/// includes the targets passed as the universe and all transitive deps of them).
/// When not provided, we implicitly set the universe to be rooted at every target literal
/// in the `cquery`.
///
/// Run `buck2 docs cquery` for more documentation about the functions available in cquery
/// expressions.
///
/// Examples:
///
/// Print all the attributes of a target
///
/// `buck2 cquery //java/com/example/app:amazing --output-all-attributes`
///
/// List the deps of a target (special characters in a target will require quotes):
///
/// `buck2 cquery 'deps("//java/com/example/app:amazing+more")'`
#[derive(Debug, clap::Parser)]
#[clap(name = "cquery")]
pub struct CqueryCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    #[clap(flatten)]
    query_common: CommonQueryOptions,

    #[clap(
        long,
        use_delimiter = true,
        help = "Comma separated list of targets at which to root the queryable universe.
                This is useful since targets can exist in multiple configurations. While
                this argument isn't required, it's recommended for most non-trivial queries.
                In the absence of this argument, buck2 will use the target literals
                in your cquery expression as the argument to this."
    )]
    target_universe: Vec<String>,

    #[clap(
        long,
        help = "Show the providers of the query result instead of the attributes and labels"
    )]
    show_providers: bool,

    #[allow(rustdoc::bare_urls)]
    /// Enable deprecated `owner()` function behavior.
    ///
    /// See this post https://fburl.com/1mf2d2xj for details.
    #[clap(long)]
    deprecated_owner: bool,

    #[allow(rustdoc::bare_urls)]
    /// Enable correct `owner()` function behavior.
    ///
    /// See this post https://fburl.com/1mf2d2xj for details.
    #[clap(long)]
    correct_owner: bool,
}

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

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let (query, query_args) = self.query_common.get_query();
        let unstable_output_format = self.query_common.output_format() as i32;
        let output_attributes = self.query_common.attributes.get()?;
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;

        let correct_owner = match (self.correct_owner, self.deprecated_owner) {
            (true, false) => true,
            (false, true) => false,
            (false, false) => true,
            (true, true) => {
                return ExitResult::bail(
                    "Cannot specify both --correct-owner and --deprecated-owner",
                );
            }
        };

        let response = buckd
            .with_flushing()
            .cquery(
                CqueryRequest {
                    query,
                    query_args,
                    context: Some(context),
                    output_attributes,
                    target_universe: self.target_universe,
                    show_providers: self.show_providers,
                    unstable_output_format,
                    correct_owner,
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut StdoutPartialResultHandler,
            )
            .await??;

        for message in &response.error_messages {
            buck2_client_ctx::eprintln!("{}", message)?;
        }

        if !response.error_messages.is_empty() {
            ExitResult::failure()
        } else {
            ExitResult::success()
        }
    }

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

pub mod aquery;
pub(crate) mod common;
pub mod cquery;
pub mod uquery;


/*
 * 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::UqueryRequest;
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::StdoutPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

use crate::commands::query::common::CommonQueryOptions;

/// Perform queries on the unconfigured target graph.
///
/// The unconfigured target graph consists of the targets as they are defined in the build
/// files. In this graph, each target appears exactly once and `select()`s are in the unresolved
/// form. For large queries, the unconfigured graph may be much smaller than the configured
/// graph and queries can be much more efficiently performed there.
///
/// When querying the unconfigured graph, dependencies appearing in all branches of `select()`
/// dictionaries will be treated as dependencies.
///
/// Run `buck2 docs uquery` for more documentation about the functions available in cquery
/// expressions.
///
/// Examples:
///
/// Print all the attributes of a target
///
/// `buck2 uquery //java/com/example/app:amazing --output-all-attributes
///
/// List the deps of a target (special characters in a target will require quotes):
/// `buck2 uquery 'deps("//java/com/example/app:amazing+more")'`
///
/// select() encoding:
///
/// When printed, values with `select()`s use a special json encoding.
///
/// `1 + select({"//:a": 1, "DEFAULT": 2})` will be encoded as:
///
/// `{"__type": "concat", "items": [1, {"__type": "selector", "entries": {"//:a": 1, "DEFAULT": 2}}]}`
#[derive(Debug, clap::Parser)]
#[clap(name = "uquery")]
pub struct UqueryCommand {
    #[clap(flatten)]
    common_opts: CommonCommandOptions,

    #[clap(flatten)]
    query_common: CommonQueryOptions,
}

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

    async fn exec_impl(
        mut self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let (query, query_args) = self.query_common.get_query();
        let unstable_output_format = self.query_common.output_format() as i32;
        let output_attributes = self.query_common.attributes.get()?;
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;

        let response = buckd
            .with_flushing()
            .uquery(
                UqueryRequest {
                    query,
                    query_args,
                    context: Some(context),
                    output_attributes,
                    unstable_output_format,
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut StdoutPartialResultHandler,
            )
            .await??;

        for message in &response.error_messages {
            buck2_client_ctx::eprintln!("{}", message)?;
        }

        if !response.error_messages.is_empty() {
            ExitResult::failure()
        } else {
            ExitResult::success()
        }
    }

    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 std::fmt;
use std::time::Duration;
use std::time::SystemTime;

use buck2_client_ctx::stream_value::StreamValue;
use buck2_client_ctx::subscribers::event_log::read::EventLogPathBuf;
use buck2_events::BuckEvent;
use buck2_wrapper_common::invocation_id::TraceId;
use chrono::DateTime;
use chrono::Local;
use futures::TryStreamExt;
use humantime::format_duration;
use thiserror::Error;

#[derive(Debug, Error)]
enum BuildInfoError {
    #[error("Failed to read event log")]
    EventLogReadFail,
}

struct LogInfo {
    revision: Option<String>,
    daemon_uptime_s: Option<u64>,
    timestamp_end: Option<SystemTime>,
}

pub(crate) struct BuildInfo {
    uuid: TraceId,
    pub timestamp: DateTime<Local>,
    pub command: String,
    working_dir: String,
    pub buck2_revision: String,
    pub command_duration: Option<Duration>,
    pub daemon_uptime_s: Option<u64>,
}

impl fmt::Display for BuildInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(
            f,
            "buck2 UI: https://www.internalfb.com/buck2/{}
timestamp: {}
command: {}
working dir: {}
buck2_revision: {}
command duration: {}
daemon uptime: {}
        ",
            self.uuid,
            self.timestamp.format("%c %Z"),
            self.command,
            self.working_dir,
            self.buck2_revision,
            seconds_to_string(self.command_duration.map(|d| d.as_secs())),
            seconds_to_string(self.daemon_uptime_s),
        )
    }
}

pub(crate) async fn get(log: &EventLogPathBuf) -> anyhow::Result<BuildInfo> {
    let (invocation, events) = log.unpack_stream().await?;
    let mut filtered_events = events.try_filter_map(|log| {
        let maybe_buck_event = match log {
            StreamValue::Result(_) | StreamValue::PartialResult(_) => None,
            StreamValue::Event(buck_event) => Some(buck_event),
        };
        futures::future::ready(Ok(maybe_buck_event))
    });

    let first_event: BuckEvent = filtered_events
        .try_next()
        .await?
        .ok_or(BuildInfoError::EventLogReadFail)?
        .try_into()?;

    let mut info = LogInfo {
        revision: None,
        daemon_uptime_s: None,
        timestamp_end: None,
    };
    loop {
        let res = match filtered_events.try_next().await {
            Ok(Some(event)) => extract_info(&mut info, event),
            Ok(None) => break,
            Err(e) => Err(e),
        };
        if let Err(e) = res {
            buck2_client_ctx::eprintln!("Error found when iterating through logs: {:#}", e)?;
            break;
        }
    }

    let timestamp_start = first_event.timestamp();
    let duration = {
        if let Some(end) = info.timestamp_end {
            Some(end.duration_since(timestamp_start)?)
        } else {
            None
        }
    };

    let t_start: DateTime<Local> = timestamp_start.into();

    let output = BuildInfo {
        uuid: first_event.trace_id()?,
        timestamp: t_start,
        command: format_cmd(&invocation.command_line_args),
        working_dir: invocation.working_dir,
        buck2_revision: info.revision.unwrap_or_else(|| "".to_owned()),
        command_duration: duration,
        daemon_uptime_s: info.daemon_uptime_s,
    };

    Ok(output)
}

fn extract_info(info: &mut LogInfo, event: Box<buck2_data::BuckEvent>) -> anyhow::Result<()> {
    match event.data {
        Some(buck2_data::buck_event::Data::SpanStart(span)) => match &span.data {
            Some(buck2_data::span_start_event::Data::Command(action)) => {
                if info.revision.is_none() && action.metadata.contains_key("buck2_revision") {
                    if let Some(buck2_revision) = action.metadata.get("buck2_revision") {
                        info.revision.get_or_insert(buck2_revision.clone());
                    }
                }
            }
            _ => (),
        },
        Some(buck2_data::buck_event::Data::Instant(span)) => match &span.data {
            Some(buck2_data::instant_event::Data::Snapshot(snapshot)) => {
                info.daemon_uptime_s.get_or_insert(snapshot.daemon_uptime_s);
            }
            _ => (),
        },

        _ => (),
    }
    if let Some(timestamp) = event.timestamp {
        info.timestamp_end = Some(SystemTime::try_from(timestamp)?)
    };
    Ok(())
}

fn seconds_to_string(seconds: Option<u64>) -> String {
    if let Some(seconds) = seconds {
        let duration = Duration::from_secs(seconds);
        format_duration(duration).to_string()
    } else {
        "".to_owned()
    }
}

pub fn format_cmd(cmd_args: &[String]) -> String {
    if cmd_args.is_empty() {
        "???".to_owned()
    } else {
        let mut program_name: &str = &cmd_args[0];
        let program_args = &cmd_args[1..];
        if program_name.ends_with("fbcode/buck2/.buck2") {
            program_name = "buck2";
        }
        format!("{} {}", program_name, program_args.join(" "))
    }
}


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

use anyhow::Context;
use buck2_cli_proto::unstable_dice_dump_request::DiceDumpFormat;
use buck2_cli_proto::UnstableDiceDumpRequest;
use buck2_client_ctx::daemon::client::connect::BootstrapBuckdClient;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::manifold;
use buck2_core::fs::fs_util::create_dir_all;
use buck2_core::fs::fs_util::remove_all;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_util::process::background_command;

pub async fn upload_dice_dump(
    buckd: BootstrapBuckdClient,
    buck_out_dice: AbsNormPathBuf,
    manifold_id: &String,
) -> anyhow::Result<String> {
    let buckd = buckd.with_subscribers(Default::default());
    let manifold_filename = &format!("{}_dice-dump.tar", manifold_id);
    let this_dump_folder_name = chrono::Utc::now().format("%Y-%m-%dT%H:%M:%S").to_string();
    DiceDump::new(buck_out_dice, &this_dump_folder_name)
        .upload(buckd, manifold_filename)
        .await?;

    Ok(format!("buck2_rage_dumps/flat/{}", manifold_filename))
}

struct DiceDump {
    buck_out_dice: AbsNormPathBuf,
    dump_folder: AbsPathBuf,
}

impl DiceDump {
    fn new(buck_out_dice: AbsNormPathBuf, dump_folder_name: &str) -> Self {
        let dump_folder = buck_out_dice.as_abs_path().join(dump_folder_name);

        Self {
            buck_out_dice,
            dump_folder,
        }
    }

    async fn upload(
        &self,
        mut buckd: BuckdClientConnector<'_>,
        manifold_filename: &str,
    ) -> anyhow::Result<()> {
        buck2_client_ctx::eprintln!("Generating Buck2 DICE dump...")?;
        create_dir_all(&self.buck_out_dice).with_context(|| {
            format!(
                "Failed to create directory `{}`, no DICE dump will be created",
                self.buck_out_dice.display()
            )
        })?;

        buckd
            .with_flushing()
            .unstable_dice_dump(UnstableDiceDumpRequest {
                destination_path: self.dump_folder.to_str().unwrap().to_owned(),
                format: DiceDumpFormat::Tsv.into(),
            })
            .await
            .with_context(|| {
                format!(
                    "DICE dump at `{}` failed to complete",
                    self.dump_folder.display()
                )
            })?;

        // create DICE dump name using the old command being rage on and the trace id of this rage command.

        buck2_client_ctx::eprintln!(
            "Compressed DICE dump being uploaded to manifold as {}...",
            &manifold_filename
        )?;
        upload_to_manifold(&self.dump_folder, manifold_filename)
            .await
            .with_context(|| "Failed during manifold upload!")?;

        Ok(())
    }
}

async fn upload_to_manifold(dump_folder: &Path, manifold_filename: &str) -> anyhow::Result<()> {
    if !cfg!(target_os = "windows") {
        buck2_core::facebook_only();

        let tar = background_command("tar")
            .arg("-c")
            .arg(dump_folder)
            .stdin(std::process::Stdio::null())
            .stdout(std::process::Stdio::piped())
            .stderr(std::process::Stdio::null())
            .spawn()?;

        manifold::Upload::new(manifold::Bucket::RAGE_DUMPS, manifold_filename)
            .with_default_ttl()
            .from_stdio(tar.stdout.unwrap().into())?
            .spawn()
            .await?;
    }
    Ok(())
}

impl Drop for DiceDump {
    fn drop(&mut self) {
        if let Err(e) = remove_all(&self.dump_folder).with_context(|| {
            format!(
                "Failed to remove Buck2 DICE dump folder at `{}`. Please remove this manually as it could be quite large.",
                self.dump_folder.display()
            )
        }) { tracing::warn!("{:#}", e); };
    }
}


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

use std::io::Cursor;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_audit::deferred_materializer::DeferredMaterializerCommand;
use buck2_audit::deferred_materializer::DeferredMaterializerSubcommand;
use buck2_audit::AuditCommand;
use buck2_cli_proto::ClientContext;
use buck2_client_ctx::command_outcome::CommandOutcome;
use buck2_client_ctx::daemon::client::connect::BootstrapBuckdClient;
use buck2_client_ctx::events_ctx::PartialResultCtx;
use buck2_client_ctx::events_ctx::PartialResultHandler;
use buck2_client_ctx::manifold;
use buck2_client_ctx::subscribers::subscriber::EventSubscriber;
use buck2_common::result::SharedResult;

pub async fn upload_materializer_state(
    buckd: &SharedResult<BootstrapBuckdClient>,
    client_context: &ClientContext,
    manifold_id: &String,
) -> anyhow::Result<String> {
    let mut buckd = buckd
        .clone()?
        .with_subscribers(vec![Box::new(TracingSubscriber) as _]);

    let mut capture = CaptureStdout::new();

    let outcome = buckd
        .with_flushing()
        .audit(
            buck2_cli_proto::GenericRequest {
                context: Some(client_context.clone()),
                serialized_opts: serde_json::to_string(&AuditCommand::DeferredMaterializer(
                    DeferredMaterializerCommand {
                        common_opts: Default::default(),
                        subcommand: DeferredMaterializerSubcommand::List,
                    },
                ))?,
            },
            None,
            &mut capture,
        )
        .await?;

    match outcome {
        CommandOutcome::Success(..) => {}
        CommandOutcome::Failure(..) => return Err(anyhow::anyhow!("Command failed")),
    }

    let manifold_bucket = manifold::Bucket::RAGE_DUMPS;
    let manifold_filename = format!("{}_materializer_state", manifold_id);
    manifold::Upload::new(manifold_bucket, &manifold_filename)
        .with_default_ttl()
        .from_async_read(&mut Cursor::new(&capture.buf))?
        .spawn()
        .await
        .context("Error uploading to Manifold")?;

    Ok(format!(
        "{}/flat/{}",
        manifold_bucket.name, manifold_filename
    ))
}

/// Receive StdoutBytes, just capture them.
struct CaptureStdout {
    buf: Vec<u8>,
}

impl CaptureStdout {
    fn new() -> Self {
        Self { buf: Vec::new() }
    }
}

#[async_trait]
impl PartialResultHandler for CaptureStdout {
    type PartialResult = buck2_cli_proto::StdoutBytes;

    async fn handle_partial_result(
        &mut self,
        _ctx: PartialResultCtx<'_, '_>,
        partial_res: Self::PartialResult,
    ) -> anyhow::Result<()> {
        self.buf.extend(partial_res.data);
        Ok(())
    }
}

struct TracingSubscriber;

#[async_trait]
impl EventSubscriber for TracingSubscriber {
    async fn handle_tailer_stderr(&mut self, stderr: &str) -> anyhow::Result<()> {
        tracing::info!("{}", stderr);
        Ok(())
    }

    async fn handle_error(&mut self, error: &anyhow::Error) -> anyhow::Result<()> {
        tracing::info!("{:#}", error);
        Ok(())
    }

    async fn handle_command_result(
        &mut self,
        result: &buck2_cli_proto::CommandResult,
    ) -> anyhow::Result<()> {
        tracing::info!("{:?}", result);
        Ok(())
    }
}


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

mod build_info;
mod dice;
mod materializer;
mod source_control;
mod system_info;
mod thread_dump;

use std::collections::HashMap;
use std::fmt;
use std::future::Future;
use std::process::Stdio;
use std::time::Duration;
use std::time::SystemTime;

use anyhow::Context;
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::BootstrapBuckdClient;
use buck2_client_ctx::daemon::client::connect::BuckdConnectConstraints;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::manifold;
use buck2_client_ctx::stdin::Stdin;
use buck2_client_ctx::subscribers::event_log::file_names::do_find_log_by_trace_id;
use buck2_client_ctx::subscribers::event_log::file_names::get_local_logs;
use buck2_client_ctx::subscribers::event_log::read::EventLogPathBuf;
use buck2_client_ctx::subscribers::event_log::read::EventLogSummary;
use buck2_common::result::ToSharedResultExt;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_data::instant_event::Data;
use buck2_data::InstantEvent;
use buck2_data::RageResult;
use buck2_events::sink::scribe::new_thrift_scribe_sink_if_enabled;
use buck2_events::sink::scribe::ThriftScribeSink;
use buck2_events::BuckEvent;
use buck2_wrapper_common::invocation_id::TraceId;
use chrono::offset::Local;
use chrono::DateTime;
use dupe::Dupe;
use futures::future::FutureExt;
use futures::future::LocalBoxFuture;
use maplit::convert_args;
use maplit::hashmap;
use serde::Serialize;
use thiserror::Error;
use tokio::io::AsyncBufRead;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::process::Command;

#[derive(Debug, Error)]
enum RageError {
    #[error("Failed to get a valid user selection")]
    InvalidSelectionError,
    #[error("Failed to find the logs for command")]
    LogNotFoundError,
    #[error("Pastry command timeout, make sure you are on Lighthouse/VPN")]
    PastryTimeout,
    #[error("Failed to spawn pastry")]
    PastrySpawnError,
    #[error("Error writing to pastry")]
    PastryWriteError,
    #[error("Error reading pastry output")]
    PastryOutputError,
    #[error("Pastry command failed with code '{0}' and error '{1}' ")]
    PastryCommandError(i32, String),
    #[error("Failed to open file `{0}`")]
    OpenFileError(String),
}

#[derive(Debug, PartialEq, Serialize)]
struct RageSection<T> {
    title: String,
    status: CommandStatus<T>,
}

#[derive(Debug, PartialEq, Serialize)]
enum CommandStatus<T> {
    Success { output: T },
    Failure { error: String },
    Timeout,
    Skipped,
}

impl<'a, T> RageSection<T>
where
    T: std::fmt::Display + 'a,
{
    fn get<Fut>(
        title: String,
        timeout: Duration,
        command: impl FnOnce() -> Fut,
    ) -> LocalBoxFuture<'a, Self>
    where
        Fut: Future<Output = anyhow::Result<T>> + 'a,
    {
        let fut = command();
        async move {
            let status = match tokio::time::timeout(timeout, fut).await {
                Err(_) => CommandStatus::Timeout,
                Ok(Ok(output)) => CommandStatus::Success { output },
                Ok(Err(e)) => CommandStatus::Failure {
                    error: format!("Error: {:?}", e),
                },
            };
            RageSection { title, status }
        }
        .boxed_local()
    }

    fn get_skipped(title: String) -> LocalBoxFuture<'a, Self> {
        let status = CommandStatus::Skipped;
        async { RageSection { title, status } }.boxed_local()
    }

    fn output(&self) -> String {
        match &self.status {
            CommandStatus::Success { output } => output.to_string(),
            CommandStatus::Failure { error } => error.to_owned(),
            CommandStatus::Timeout {} => "Timeout".to_owned(),
            CommandStatus::Skipped {} => "Skipped".to_owned(),
        }
    }

    fn get_field<D>(&self, extract_field: impl FnOnce(&T) -> Option<D>) -> Option<D> {
        match &self.status {
            CommandStatus::Success { output } => extract_field(output),
            _ => None,
        }
    }

    fn pretty_print_section(
        &self,
        f: &mut fmt::Formatter,
        content: String,
    ) -> Result<(), std::fmt::Error> {
        let content_divider = "-".repeat(30);
        write!(
            f,
            "{title}\n{content_divider}\n{content}\n\n\n",
            title = self.title
        )
    }
}

impl<T> fmt::Display for RageSection<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.pretty_print_section(f, self.output())
    }
}

#[derive(Copy, Clone, PartialEq, Debug)]
#[derive(clap::ArgEnum)]
enum Origin {
    HangDetector,
    Unspecified,
}

impl fmt::Display for Origin {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

#[derive(Debug, clap::Parser)]
#[clap(
    name = "rage",
    about = "Record information about the previous failed buck2 command",
    group = clap::ArgGroup::new("invocation").multiple(false)
)]
pub struct RageCommand {
    /// Stop collecting information after `<timeout>` seconds
    #[clap(long, default_value = "60")]
    timeout: u64,
    /// Use value 0 to select last invocation, 1 to select second to last and so on
    #[clap(long, group = "invocation")]
    invocation_offset: Option<usize>,
    /// Select invocation directly using the invocation's UUID
    #[clap(long, group = "invocation")]
    invocation_id: Option<TraceId>,
    /// Collect rage report about buck2 in general, not about specific invocation
    #[clap(long, group = "invocation")]
    no_invocation: bool,
    /// We may want to omit paste if this is not a user
    /// or is called in a machine with no pastry command
    #[clap(long)]
    no_paste: bool,
    /// Where buck2 rage is being called from
    #[clap(long, arg_enum, default_value_t = Origin::Unspecified)]
    origin: Origin,
    /// Obtain thread dump for debugging
    #[clap(long)]
    thread_dump: bool,
}

impl RageCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        buck2_core::facebook_only();

        ctx.with_runtime(async move |mut ctx| {
            let timeout = Duration::from_secs(self.timeout);
            let paths = ctx.paths.as_ref().map_err(|e| e.dupe())?;
            let stderr_path = paths.daemon_dir()?.buckd_stderr();
            let logdir = paths.log_dir();
            let dice_dump_dir = paths.dice_dump_dir();

            let client_ctx = ctx.empty_client_context()?;

            let rage_id = TraceId::new();
            let mut manifold_id = format!("{}", rage_id);
            let sink = create_scribe_sink(&ctx)?;

            buck2_client_ctx::eprintln!(
                "Collection will terminate after {} seconds (override with --timeout param)",
                self.timeout
            )?;
            buck2_client_ctx::eprintln!("Collecting debug info...\n\n")?;

            // If there is a daemon, connect.
            let buckd = BootstrapBuckdClient::connect(paths, BuckdConnectConstraints::ExistingOnly)
                .await
                .shared_error();

            let selected_invocation = maybe_select_invocation(ctx.stdin(), &logdir, &self).await?;
            let invocation_id = get_trace_id(&selected_invocation).await?;
            if let Some(ref invocation_id) = invocation_id {
                manifold_id = format!("{}_{}", invocation_id, manifold_id);
            }

            let system_info_command =
                RageSection::get("System info".to_owned(), timeout, system_info::get);
            let daemon_stderr_command =
                RageSection::get("Daemon stderr Manifold path".to_owned(), timeout, || {
                    upload_daemon_stderr(stderr_path, &manifold_id)
                });
            let hg_snapshot_id_command = RageSection::get(
                "Source control".to_owned(),
                timeout,
                source_control::get_info,
            );
            let dice_dump_command =
                RageSection::get("Dice dump Manifold path".to_owned(), timeout, || async {
                    dice::upload_dice_dump(buckd.clone()?, dice_dump_dir, &manifold_id).await
                });
            let materializer_state = RageSection::get(
                "Materializer state Manifold path".to_owned(),
                timeout,
                || materializer::upload_materializer_state(&buckd, &client_ctx, &manifold_id),
            );
            let thread_dump = {
                let title = "Thread dump Manifold path".to_owned();
                if self.thread_dump {
                    RageSection::get(title, timeout, || {
                        thread_dump::upload_thread_dump(&buckd, &manifold_id)
                    })
                } else {
                    RageSection::get_skipped(title)
                }
            };
            let build_info_command = {
                let title = "Associated invocation info".to_owned();
                match selected_invocation.as_ref() {
                    None => RageSection::get_skipped(title),
                    Some(invocation) => {
                        RageSection::get(title, timeout, || build_info::get(invocation))
                    }
                }
            };

            let event_log_command = {
                let title = "Event log upload Manifold path".to_owned();
                match selected_invocation.as_ref() {
                    None => RageSection::get_skipped(title),
                    Some(path) => {
                        RageSection::get(title, timeout, || upload_event_logs(path, &manifold_id))
                    }
                }
            };

            let (
                system_info,
                daemon_stderr_dump,
                hg_snapshot_id,
                dice_dump,
                materializer_state,
                thread_dump,
                build_info,
                event_log_dump,
            ) = tokio::join!(
                system_info_command,
                daemon_stderr_command,
                hg_snapshot_id_command,
                dice_dump_command,
                materializer_state,
                thread_dump,
                build_info_command,
                event_log_command,
            );
            let sections = vec![
                system_info.to_string(),
                daemon_stderr_dump.to_string(),
                hg_snapshot_id.to_string(),
                dice_dump.to_string(),
                materializer_state.to_string(),
                thread_dump.to_string(),
                build_info.to_string(),
                event_log_dump.to_string(),
            ];
            output_rage(self.no_paste, &sections.join("")).await?;

            self.send_to_scuba(
                &rage_id,
                sink,
                invocation_id,
                system_info,
                daemon_stderr_dump,
                hg_snapshot_id,
                dice_dump,
                materializer_state,
                thread_dump,
                event_log_dump,
                build_info,
            )
            .await?;
            ExitResult::success()
        })
    }

    async fn send_to_scuba(
        &self,
        rage_id: &TraceId,
        sink: Option<ThriftScribeSink>,
        invocation_id: Option<TraceId>,
        system_info: RageSection<system_info::SystemInfo>,
        daemon_stderr_dump: RageSection<String>,
        hg_snapshot_id: RageSection<String>,
        dice_dump: RageSection<String>,
        materializer_state: RageSection<String>,
        thread_dump: RageSection<String>,
        event_log_dump: RageSection<String>,
        build_info: RageSection<build_info::BuildInfo>,
    ) -> anyhow::Result<()> {
        let mut string_data = convert_args!(
            keys = String::from,
            hashmap! (
                "dice_dump" => dice_dump.output(),
                "materializer_state" => materializer_state.output(),
                "thread_dump" => thread_dump.output(),
                "daemon_stderr_dump" => daemon_stderr_dump.output(),
                "hg_snapshot_id" => hg_snapshot_id.output(),
                "invocation_id" => invocation_id.map(|inv| inv.to_string()).unwrap_or_default(),
                "origin" => self.origin.to_string(),
                "event_log_dump" => event_log_dump.output(),
            )
        );

        let command = build_info.get_field(|o| Some(o.command.to_owned()));
        let buck2_revision = build_info.get_field(|o| Some(o.buck2_revision.to_owned()));
        let username = system_info.get_field(|o| o.username.to_owned());
        let hostname = system_info.get_field(|o| o.hostname.to_owned());
        let os = system_info.get_field(|o| Some(o.os.to_owned()));
        let os_version = system_info.get_field(|o| o.os_version.to_owned());

        insert_if_some(&mut string_data, "command", command);
        insert_if_some(&mut string_data, "buck2_revision", buck2_revision);
        insert_if_some(&mut string_data, "username", username);
        insert_if_some(&mut string_data, "hostname", hostname);
        insert_if_some(&mut string_data, "os", os);
        insert_if_some(&mut string_data, "os_version", os_version);

        let mut int_data = HashMap::new();
        let daemon_uptime_s = build_info.get_field(|o| o.daemon_uptime_s);
        insert_if_some(&mut int_data, "daemon_uptime_s", daemon_uptime_s);

        let timestamp = build_info.get_field(|o| Some(SystemTime::from(o.timestamp).into()));
        let command_duration = build_info.get_field(|o| {
            Some(prost_types::Duration {
                seconds: o.command_duration?.as_secs() as i64,
                nanos: o.command_duration?.subsec_nanos() as i32,
            })
        });

        dispatch_result_event(
            sink.as_ref(),
            rage_id,
            RageResult {
                string_data,
                int_data,
                timestamp,
                command_duration,
            },
        )
        .await
    }

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

fn insert_if_some<D>(data: &mut HashMap<String, D>, key: &str, value: Option<D>) {
    if let Some(value) = value {
        data.insert(key.to_owned(), value);
    }
}

async fn upload_daemon_stderr(
    path: AbsNormPathBuf,
    manifold_id: &String,
) -> anyhow::Result<String> {
    // can't use async_fs_util
    // the trait to convert from tokio::fs::File is not implemented for Stdio
    let upload_log_file: Stdio = std::fs::File::open(&path)
        .context(RageError::OpenFileError(path.display().to_string()))?
        .into();
    let filename = format!("{}.stderr", manifold_id);
    manifold::Upload::new(manifold::Bucket::RAGE_DUMPS, &filename)
        .with_default_ttl()
        .from_stdio(upload_log_file)?
        .spawn()
        .await?;
    Ok(format!("buck2_rage_dumps/flat/{}", filename))
}

async fn upload_event_logs(path: &EventLogPathBuf, manifold_id: &str) -> anyhow::Result<String> {
    let filename = format!("{}-event_log{}", manifold_id, path.extension());
    let bucket = manifold::Bucket::RAGE_DUMPS;
    manifold::Upload::new(bucket, &filename)
        .with_default_ttl()
        .from_file(path.path())?
        .spawn()
        .await?;
    Ok(format!("{}/flat/{}", bucket.name, filename))
}

async fn dispatch_result_event(
    sink: Option<&ThriftScribeSink>,
    rage_id: &TraceId,
    result: RageResult,
) -> anyhow::Result<()> {
    let data = Some(Data::RageResult(result));
    dispatch_event_to_scribe(sink, rage_id, InstantEvent { data }).await?;
    Ok(())
}

async fn dispatch_event_to_scribe(
    sink: Option<&ThriftScribeSink>,
    trace_id: &TraceId,
    event: InstantEvent,
) -> anyhow::Result<()> {
    if let Some(sink) = sink {
        sink.send_now(BuckEvent::new(
            SystemTime::now(),
            trace_id.to_owned(),
            None,
            None,
            event.into(),
        ))
        .await;
    } else {
        tracing::warn!(
            "Couldn't send rage results to scribe, rage ID `{}`",
            trace_id
        )
    };
    Ok(())
}

#[allow(unused_variables)] // Conditional compilation
fn create_scribe_sink(ctx: &ClientCommandContext) -> anyhow::Result<Option<ThriftScribeSink>> {
    // TODO(swgiillespie) scribe_logging is likely the right feature for this, but we should be able to inject a sink
    // without using configurations at the call site
    new_thrift_scribe_sink_if_enabled(
        ctx.fbinit(),
        /* buffer size */ 100,
        /* retry_backoff */ Duration::from_millis(500),
        /* retry_attempts */ 5,
        /* message_batch_size */ None,
    )
}

async fn maybe_select_invocation(
    stdin: &mut Stdin,
    logdir: &AbsNormPathBuf,
    command: &RageCommand,
) -> anyhow::Result<Option<EventLogPathBuf>> {
    if command.no_invocation {
        return Ok(None);
    };

    if let Some(trace_id) = &command.invocation_id {
        return Ok(Some(do_find_log_by_trace_id(logdir, trace_id)?));
    }

    let logs = get_local_logs(logdir)?;
    let mut logs = logs
        .into_iter()
        .rev() // newest first
        .collect::<Vec<_>>();
    if logs.is_empty() {
        return Ok(None);
    }
    let index = log_index(stdin, &logs, command.invocation_offset).await?;
    if index >= logs.len() {
        return Err(RageError::LogNotFoundError.into());
    }
    Ok(Some(logs.swap_remove(index)))
}

async fn log_index(
    stdin: &mut Stdin,
    logs: &[EventLogPathBuf],
    invocation_offset: Option<usize>,
) -> Result<usize, anyhow::Error> {
    let index = match invocation_offset {
        Some(i) => i,
        None => {
            let mut stdin = BufReader::new(stdin);
            user_prompt_select_log(&mut stdin, logs).await?
        }
    };
    Ok(index)
}

async fn user_prompt_select_log<'a>(
    stdin: impl AsyncBufRead + Unpin,
    logs: &'a [EventLogPathBuf],
) -> anyhow::Result<usize> {
    buck2_client_ctx::eprintln!("Which buck invocation would you like to report?\n")?;
    let logs_summary =
        futures::future::try_join_all(logs.iter().map(|log_path| log_path.get_summary())).await?;
    for (index, log_summary) in logs_summary.iter().enumerate() {
        print_log_summary(index, log_summary)?;
    }
    buck2_client_ctx::eprintln!()?;
    let prompt = format!(
        "Invocation: (type a number between 0 and {}) ",
        logs_summary.len() - 1
    );
    let selection = get_user_selection(stdin, &prompt, |i| i < logs_summary.len()).await?;

    let chosen_log = logs_summary.get(selection).expect("Selection out of range");

    let timestamp: DateTime<Local> = chosen_log.timestamp.into();
    buck2_client_ctx::eprintln!("Selected invocation at {}\n", timestamp.format("%c %Z"))?;
    Ok(selection)
}

async fn get_user_selection<P>(
    mut stdin: impl AsyncBufRead + Unpin,
    prompt: &str,
    predicate: P,
) -> anyhow::Result<usize>
where
    P: Fn(usize) -> bool,
{
    buck2_client_ctx::eprint!("{}", prompt)?;

    let mut input = String::new();
    stdin.read_line(&mut input).await?;

    match input.trim().parse() {
        Ok(selection) if predicate(selection) => Ok(selection),
        _ => Err(RageError::InvalidSelectionError.into()),
    }
}

fn print_log_summary(index: usize, log_summary: &EventLogSummary) -> anyhow::Result<()> {
    let cmd = build_info::format_cmd(&log_summary.invocation.command_line_args);

    let timestamp: DateTime<Local> = log_summary.timestamp.into();
    buck2_client_ctx::eprintln!(
        "{:<7} {}    {}",
        format!("[{}].", index),
        timestamp.format("%c %Z"),
        cmd
    )
}

async fn output_rage(no_paste: bool, output: &str) -> anyhow::Result<()> {
    if no_paste {
        buck2_client_ctx::println!("{}", output)?;
    } else {
        let paste = generate_paste("Buck2 Rage", output).await?;
        buck2_client_ctx::eprintln!(
            "\nPlease post in https://fb.workplace.com/groups/buck2users with the following link:\n\n{}\n",
            paste
        )?;
    };
    Ok(())
}

async fn generate_paste(title: &str, content: &str) -> anyhow::Result<String> {
    let mut pastry = Command::new("pastry")
        .args(["--title", title])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .kill_on_drop(true)
        .spawn()
        .context(RageError::PastrySpawnError)?;
    let mut stdin = pastry.stdin.take().expect("Stdin should open");

    let writer = async move {
        stdin
            .write_all(content.as_bytes())
            .await
            .context(RageError::PastryWriteError)
    };

    let reader = async move {
        let output = tokio::time::timeout(Duration::from_secs(10), pastry.wait_with_output())
            .await
            .context(RageError::PastryTimeout)?
            .context(RageError::PastryOutputError)?;
        if !output.status.success() {
            let error = String::from_utf8_lossy(&output.stderr).to_string();
            let code = output
                .status
                .code()
                .ok_or_else(|| RageError::PastryCommandError(1, error.clone()))?;
            return Err(RageError::PastryCommandError(code, error).into());
        }
        let output = String::from_utf8(output.stdout).context(RageError::PastryOutputError)?;
        Ok(output)
    };

    let ((), paste) = futures::future::try_join(writer, reader).await?;

    Ok(paste)
}

async fn get_trace_id(invocation: &Option<EventLogPathBuf>) -> anyhow::Result<Option<TraceId>> {
    let invocation_id = match invocation {
        None => None,
        Some(invocation) => Some(invocation.get_summary().await?.trace_id),
    };
    Ok(invocation_id)
}


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

use anyhow::Context;
use thiserror::Error;
use tokio::process::Command;

#[derive(Debug, Error)]
enum SourceControlError {
    #[error("HG command failed with code '{0}' and error '{1}' ")]
    HgCommand(i32, String),
    #[error("Git command failed with code '{0}' and error '{1}' ")]
    GitCommand(i32, String),
    #[error("`{0}` was not UTF-8")]
    Utf8(String),
}

enum CommandResult {
    Ok(String),
    RepoNotFound,
}

pub async fn get_info() -> anyhow::Result<String> {
    let hg_info = get_hg_info().await;
    if let Ok(CommandResult::Ok(output)) = hg_info {
        return Ok(output);
    }
    let git_info = get_git_info().await;
    if let Ok(CommandResult::Ok(output)) = git_info {
        return Ok(output);
    }
    hg_info?;
    git_info?;
    Ok("Current directory is not inside a repository (tried hg and git)".to_owned())
}

async fn get_hg_info() -> anyhow::Result<CommandResult> {
    let result = Command::new("hg")
        .args(["snapshot", "create"])
        .env("HGPLAIN", "1")
        .output()
        .await?;
    if result.status.success() {
        let output = from_utf8(result.stdout, "hg snapshot stdout")?;
        return Ok(CommandResult::Ok(format!("HG snapshot ID: {}", output)));
    };
    let error = from_utf8(result.stderr, "hg snapshot stderr")?;
    if error.contains("is not inside a repository") {
        return Ok(CommandResult::RepoNotFound);
    };
    let code = result.status.code().unwrap_or(1);
    Err(SourceControlError::HgCommand(code, error).into())
}

async fn get_git_info() -> anyhow::Result<CommandResult> {
    let commit_hash = Command::new("git")
        .args(["log", "-1", "--format=%H"])
        .output()
        .await?;
    if !commit_hash.status.success() {
        let error = from_utf8(commit_hash.stderr, "git log stderr")?;
        if error.contains("not a git repository") {
            return Ok(CommandResult::RepoNotFound);
        };
        let code = commit_hash.status.code().unwrap_or(1);
        return Err(SourceControlError::GitCommand(code, error).into());
    };

    let status = Command::new("git").args(["status", "-sb"]).output().await?;
    if !status.status.success() {
        let error = from_utf8(status.stderr, "git status stderr")?;
        let code = status.status.code().unwrap_or(1);
        return Err(SourceControlError::GitCommand(code, error).into());
    };
    Ok(CommandResult::Ok(format!(
        "Git base commit hash:\n{}\nGit status:\n{}",
        from_utf8(commit_hash.stdout, "git log stdout")?,
        from_utf8(status.stdout, "git status stdout")?
    )))
}

fn from_utf8(result: Vec<u8>, subject: &str) -> anyhow::Result<String> {
    String::from_utf8(result).context(SourceControlError::Utf8(subject.to_owned()))
}