/*
 * 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::BTreeSet;
use std::collections::HashMap;
use std::collections::HashSet;
use std::io::BufWriter;
use std::io::Write;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;

use anyhow::Context;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::stream_value::StreamValue;
use buck2_client_ctx::subscribers::event_log::file_names::retrieve_nth_recent_log;
use buck2_client_ctx::subscribers::event_log::read::EventLogPathBuf;
use buck2_client_ctx::subscribers::event_log::utils::Invocation;
use buck2_common::convert::ProstDurationExt;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_event_observer::display;
use buck2_event_observer::display::TargetDisplayOptions;
use buck2_events::BuckEvent;
use dupe::Dupe;
use futures::TryStreamExt;
use serde::Serialize;
use serde_json::json;

#[derive(Debug, clap::Parser)]
pub struct ChromeTraceCommand {
    #[clap(
        long,
        help = "Where to write the chrome trace JSON. If a directory is passed, the filename of the event log will be used as a base filename."
    )]
    pub trace_path: PathArg,
    /// The path to read the event log from.
    #[clap(
        long,
        help = "A path to an event-log file to read from. Only works for log files with a single command in them. If no event-log is passed, the most recent one will be used.",
        group = "event_log",
        value_name = "PATH"
    )]
    pub path: Option<PathArg>,

    /// Which recent command to read the event log from.
    #[clap(
        long,
        help = "Use the event-log from the Nth most recent command (`--recent 0` is the most recent).",
        group = "event_log",
        value_name = "NUMBER"
    )]
    pub recent: Option<usize>,
}

struct ChromeTraceFirstPass {
    /// Track assignment needs to know, when it sees a SpanStart, whether that
    /// span is going to be included in the final trace.
    /// But some spans need to be filtered based on later events, like:
    /// 1. We shouldn't assign tracks to StartLoad events whose SpanEnd records
    ///    a really short duration.
    /// 2. We shouldn't assign tracks to ActionExecutionStart events who have
    ///    no child LocalStage spans.
    /// 3. (eventually) We should assign tracks to ActionExecutionStart events
    ///    only if they appear in the CriticalPath, but the CriticalPath is one
    ///    of the last events.
    /// So this first pass builds up several lists of "interesting" span IDs.
    pub long_analyses: HashSet<buck2_events::span::SpanId>,
    pub long_loads: HashSet<buck2_events::span::SpanId>,
    pub local_actions: HashSet<buck2_events::span::SpanId>,
    pub critical_path_action_keys: HashSet<buck2_data::ActionKey>,
    pub critical_path_span_ids: HashSet<u64>,
}

impl ChromeTraceFirstPass {
    const LONG_ANALYSIS_CUTOFF: Duration = Duration::from_millis(50);
    const LONG_LOAD_CUTOFF: Duration = Duration::from_millis(50);
    fn new() -> Self {
        Self {
            long_analyses: HashSet::new(),
            long_loads: HashSet::new(),
            local_actions: HashSet::new(),
            critical_path_action_keys: HashSet::new(),
            critical_path_span_ids: HashSet::new(),
        }
    }

    fn handle_event(&mut self, event: &BuckEvent) -> anyhow::Result<()> {
        match event.data() {
            buck2_data::buck_event::Data::SpanStart(ref start) => {
                match start.data.as_ref() {
                    Some(buck2_data::span_start_event::Data::ExecutorStage(exec)) => {
                        // A local stage means that we want to show the entire action execution.
                        if let Some(buck2_data::executor_stage_start::Stage::Local(_)) = exec.stage
                        {
                            self.local_actions.insert(event.parent_id().unwrap());
                        }
                    }
                    _ => {}
                }
            }
            buck2_data::buck_event::Data::SpanEnd(ref end) => {
                match end.data.as_ref() {
                    Some(buck2_data::span_end_event::Data::Analysis(_)) => {
                        if end
                            .duration
                            .as_ref()
                            .expect("Analysis SpanEnd missing duration")
                            .try_into_duration()?
                            > Self::LONG_ANALYSIS_CUTOFF
                        {
                            self.long_analyses.insert(event.span_id().unwrap());
                        }
                    }
                    Some(buck2_data::span_end_event::Data::Load(_)) => {
                        if end
                            .duration
                            .as_ref()
                            .expect("Load SpanEnd missing duration")
                            .try_into_duration()?
                            > Self::LONG_LOAD_CUTOFF
                        {
                            self.long_loads.insert(event.span_id().unwrap());
                        }
                    }
                    _ => {}
                };
            }
            buck2_data::buck_event::Data::Instant(ref instant) => match instant.data.as_ref() {
                Some(buck2_data::instant_event::Data::BuildGraphInfo(info)) => {
                    self.critical_path_action_keys = info
                        .critical_path
                        .iter()
                        .map(|entry| entry.action_key.clone().unwrap())
                        .collect();

                    self.critical_path_span_ids = info
                        .critical_path2
                        .iter()
                        .flat_map(|entry| entry.span_ids.iter().copied())
                        .collect()
                }
                _ => {}
            },
            buck2_data::buck_event::Data::Record(_) => {}
        };
        Ok(())
    }
}

enum SpanTrackAssignment {
    Owned(TrackId),
    Inherited(TrackId),
}

impl SpanTrackAssignment {
    fn get_track_id(&self) -> TrackId {
        match self {
            Self::Owned(tid) => *tid,
            Self::Inherited(tid) => *tid,
        }
    }
}

struct ChromeTraceOpenSpan {
    name: String,
    start: SystemTime,
    process_id: u64,
    track: SpanTrackAssignment,
    categories: Vec<&'static str>,
    // Any misc. per-event unstructured data.
    args: serde_json::Value,
}

struct ChromeTraceClosedSpan {
    open: ChromeTraceOpenSpan,
    duration: Duration,
}

impl ChromeTraceClosedSpan {
    fn to_json(self) -> anyhow::Result<serde_json::Value> {
        Ok(json!(
            {
                "name": self.open.name,
                "ts": self.open.start.duration_since(SystemTime::UNIX_EPOCH)?.as_micros() as u64,
                "dur": self.duration.as_micros() as u64,
                "ph": "X", // Chrome trace "complete event"
                "pid": self.open.process_id,
                "tid": String::from(self.open.track.get_track_id()),
                "cat": self.open.categories.join(","),
                "args": self.open.args,
            }
        ))
    }
}

/// Spans are directed to a category, like "critical-path" or "misc". Spans in a
/// category that would overlap are put on different tracks within that category.
#[derive(Clone, Copy, Dupe)]
struct TrackId(&'static str, u64);

impl From<TrackId> for String {
    fn from(tid: TrackId) -> String {
        // Outputs like "misc-00", "misc-01", ...
        format!("{}-{:02}", tid.0, tid.1)
    }
}

struct TrackIdAllocator {
    unused_track_ids: BTreeSet<u64>,
    // Used to extend |unused_track_ids| when it's empty.
    lowest_never_used: u64,
}

impl TrackIdAllocator {
    pub fn new() -> Self {
        Self {
            unused_track_ids: BTreeSet::new(),
            lowest_never_used: 0,
        }
    }

    fn get_smallest(&mut self) -> u64 {
        let maybe_smallest = self.unused_track_ids.iter().next().copied();
        if let Some(n) = maybe_smallest {
            self.unused_track_ids.remove(&n);
            n
        } else {
            let n = self.lowest_never_used;
            self.lowest_never_used += 1;
            n
        }
    }

    pub fn mark_unused(&mut self, tid: u64) {
        self.unused_track_ids.insert(tid);
    }
}

struct SimpleCounters<T> {
    name: &'static str,
    // timeseries are flushed every BUCKET_DURATION, if any changed.
    next_flush: SystemTime,
    /// Stores the current value of each timeseries.
    /// Set to None when we output a zero, so we can save a bit of filesize
    /// by omitting them from the JSON output.
    counters: HashMap<String, Option<T>>,
    start_value: T,
    trace_events: Vec<serde_json::Value>,
}

impl<T> SimpleCounters<T>
where
    T: std::ops::Sub<Output = T>
        + std::cmp::PartialEq
        + std::ops::Add<Output = T>
        + std::marker::Copy
        + Serialize,
{
    const BUCKET_DURATION: Duration = Duration::from_millis(10);
    pub fn new(name: &'static str, start_value: T) -> Self {
        Self {
            name,
            next_flush: SystemTime::UNIX_EPOCH,
            counters: HashMap::new(),
            trace_events: vec![],
            start_value,
        }
    }

    /// Process the given timestamp and flush if needed and update next_flush accordingly
    fn process_timestamp(&mut self, timestamp: SystemTime) -> anyhow::Result<()> {
        if self.next_flush == SystemTime::UNIX_EPOCH {
            self.next_flush = timestamp + Self::BUCKET_DURATION;
        }
        if timestamp > self.next_flush + Self::BUCKET_DURATION {
            self.flush()?;
            self.next_flush = timestamp - Duration::from_micros(1);
        }
        Ok(())
    }

    /// If the given key is new to the map, initialize it to self.start_value and flush
    /// Return the value stored at the given key
    fn initialize_first_entry_if_needed(
        &mut self,
        timestamp: SystemTime,
        key: &str,
    ) -> anyhow::Result<T> {
        // If counter is being bumped from zero, we need to output its zero count
        // immediately so the line graph won't interpolate from the last time it was zero.
        let entry = *self.counters.entry(key.to_owned()).or_insert(None);
        if entry.is_none() {
            // Add a zero output immediately before the counter changes from zero.
            self.next_flush = timestamp - Duration::from_micros(1);
            self.counters.insert(key.to_owned(), Some(self.start_value));
            self.flush()?;
            self.next_flush = timestamp;
        } else if timestamp > self.next_flush {
            self.flush()?;
            self.next_flush = timestamp + Self::BUCKET_DURATION;
        }
        Ok(entry.unwrap_or(self.start_value))
    }

    fn set(&mut self, timestamp: SystemTime, key: &str, amount: T) -> anyhow::Result<()> {
        self.process_timestamp(timestamp)?;
        self.counters.insert(key.to_owned(), Some(amount));
        Ok(())
    }

    fn bump(&mut self, timestamp: SystemTime, key: &str, amount: T) -> anyhow::Result<()> {
        self.process_timestamp(timestamp)?;
        let entry = self.initialize_first_entry_if_needed(timestamp, key);
        self.counters
            .insert(key.to_owned(), Some(entry.unwrap() + amount));
        Ok(())
    }

    fn subtract(&mut self, timestamp: SystemTime, key: &str, amount: T) -> anyhow::Result<()> {
        self.process_timestamp(timestamp)?;
        let entry = self.initialize_first_entry_if_needed(timestamp, key);
        self.counters
            .insert(key.to_owned(), Some(entry.unwrap() - amount));
        Ok(())
    }

    fn flush(&mut self) -> anyhow::Result<()> {
        // Output size optimization: omit counters that were previously, and still are, zero.
        let mut output_counters = json!({});
        for (key, value) in self.counters.iter_mut() {
            if let Some(v) = value {
                output_counters[key] = json!(v);
                if *v == self.start_value {
                    *value = None;
                }
            }
        }

        self.trace_events.push(json!(
            {
                "name": self.name,
                "pid": 0,
                "tid": "counters",
                "ph": "C",
                "ts": self.next_flush
                    .duration_since(SystemTime::UNIX_EPOCH)?
                    .as_micros() as u64,
                "args": output_counters,
            }
        ));
        self.next_flush += Self::BUCKET_DURATION;
        Ok(())
    }

    pub fn flush_all_to(&mut self, output: &mut Vec<serde_json::Value>) -> anyhow::Result<()> {
        self.flush()?;
        output.append(&mut self.trace_events);
        Ok(())
    }
}

struct TimestampAndAmount {
    timestamp: SystemTime,
    amount: u64,
}

struct AverageRateOfChangeCounters {
    counters: SimpleCounters<f32>,
    previous_timestamp_and_amount_by_key: HashMap<String, TimestampAndAmount>,
}

impl AverageRateOfChangeCounters {
    pub fn new(name: &'static str) -> Self {
        Self {
            previous_timestamp_and_amount_by_key: HashMap::new(),
            counters: SimpleCounters::<f32>::new(name, 0.0),
        }
    }

    fn set_average_rate_of_change_per_s(
        &mut self,
        timestamp: SystemTime,
        key: &str,
        amount: u64,
    ) -> anyhow::Result<()> {
        // We only plot if there exists a previous item to compute the rate of change off of
        if let Some(previous) = self.previous_timestamp_and_amount_by_key.get(key) {
            let secs_since_last_datapoint =
                timestamp.duration_since(previous.timestamp)?.as_secs_f32();
            let value_change_since_last_datapoint = (amount - previous.amount) as f32;
            if secs_since_last_datapoint > 0.0 {
                self.counters.set(
                    timestamp,
                    key,
                    value_change_since_last_datapoint / secs_since_last_datapoint,
                )?;
            }
        }
        self.previous_timestamp_and_amount_by_key
            .insert(key.to_owned(), TimestampAndAmount { timestamp, amount });

        Ok(())
    }
}

struct SpanCounters {
    counter: SimpleCounters<i32>,
    // Stores how current open spans contribute to counter values.
    open_spans: HashMap<buck2_events::span::SpanId, (&'static str, i32)>,
}

impl SpanCounters {
    pub fn new(name: &'static str) -> Self {
        Self {
            counter: SimpleCounters::new(name, 0),
            open_spans: HashMap::new(),
        }
    }

    fn bump_counter_while_span(
        &mut self,
        event: &BuckEvent,
        key: &'static str,
        amount: i32,
    ) -> anyhow::Result<()> {
        self.open_spans
            .insert(event.span_id().unwrap(), (key, amount));
        self.counter.bump(event.timestamp(), key, amount)
    }

    fn handle_event_end(
        &mut self,
        _end: &buck2_data::SpanEndEvent,
        event: &BuckEvent,
    ) -> anyhow::Result<()> {
        if let Some((key, value)) = self.open_spans.remove(&event.span_id().unwrap()) {
            self.counter.subtract(event.timestamp(), key, value)?;
        }
        Ok(())
    }
}

struct ChromeTraceWriter {
    trace_events: Vec<serde_json::Value>,
    open_spans: HashMap<buck2_events::span::SpanId, ChromeTraceOpenSpan>,
    invocation: Invocation,
    first_pass: ChromeTraceFirstPass,
    span_counters: SpanCounters,
    unused_track_ids: HashMap<&'static str, TrackIdAllocator>,
    // Wrappers to contain values from InstantEvent.Data.Snapshot as a timeseries
    snapshot_counters: SimpleCounters<u64>,
    max_rss_gigabytes_counter: SimpleCounters<f64>,
    rate_of_change_counters: AverageRateOfChangeCounters,
}

impl ChromeTraceWriter {
    const UNCATEGORIZED: &'static str = "uncategorized";
    const CRITICAL_PATH: &'static str = "critical-path";
    const BYTES_PER_GIGABYTE: f64 = 1000000000.0;

    pub fn new(invocation: Invocation, first_pass: ChromeTraceFirstPass) -> Self {
        Self {
            trace_events: vec![],
            open_spans: HashMap::new(),
            invocation,
            first_pass,
            unused_track_ids: HashMap::new(),
            span_counters: SpanCounters::new("spans"),
            snapshot_counters: SimpleCounters::<u64>::new("snapshot_counters", 0),
            max_rss_gigabytes_counter: SimpleCounters::<f64>::new("max_rss", 0.0),
            rate_of_change_counters: AverageRateOfChangeCounters::new("rate_of_change_counters"),
        }
    }

    fn assign_track_for_span(
        &mut self,
        track_key: &'static str,
        event: &BuckEvent,
    ) -> anyhow::Result<SpanTrackAssignment> {
        let parent_track_id = event.parent_id().and_then(|parent_id| {
            self.open_spans
                .get(&parent_id)
                .map(|open_span| open_span.track.get_track_id())
        });

        match parent_track_id {
            None => Ok(SpanTrackAssignment::Owned(TrackId(
                track_key,
                self.unused_track_ids
                    .entry(track_key)
                    .or_insert_with(TrackIdAllocator::new)
                    .get_smallest(),
            ))),
            Some(track_id) => Ok(SpanTrackAssignment::Inherited(track_id)),
        }
    }

    pub fn to_writer<W>(mut self, file: W) -> anyhow::Result<()>
    where
        W: Write,
    {
        self.span_counters
            .counter
            .flush_all_to(&mut self.trace_events)?;
        self.snapshot_counters
            .flush_all_to(&mut self.trace_events)?;
        self.max_rss_gigabytes_counter
            .flush_all_to(&mut self.trace_events)?;
        self.rate_of_change_counters
            .counters
            .flush_all_to(&mut self.trace_events)?;

        serde_json::to_writer(
            file,
            &json!({
                "traceEvents": self.trace_events
            }),
        )?;
        Ok(())
    }

    fn open_span(&mut self, event: &BuckEvent, span: ChromeTraceOpenSpan) -> anyhow::Result<()> {
        self.open_spans.insert(event.span_id().unwrap(), span);
        Ok(())
    }

    fn open_named_span(
        &mut self,
        event: &BuckEvent,
        name: String,
        track_key: &'static str,
    ) -> anyhow::Result<()> {
        // Allocate this span to its parent's track or to a new track.
        let track = self.assign_track_for_span(track_key, event)?;
        self.open_span(
            event,
            ChromeTraceOpenSpan {
                name,
                start: event.timestamp(),
                process_id: 0,
                track,
                categories: vec!["buck2"],
                args: json!({
                    "span_id": event.span_id(),
                }),
            },
        )
    }

    fn handle_event(&mut self, event: &Arc<BuckEvent>) -> anyhow::Result<()> {
        match event.data() {
            buck2_data::buck_event::Data::SpanStart(buck2_data::SpanStartEvent {
                data: Some(start_data),
            }) => {
                let on_critical_path = event.span_id().map_or(false, |span_id| {
                    self.first_pass
                        .critical_path_span_ids
                        .contains(&span_id.into())
                });

                enum Categorization<'a> {
                    /// Show this node on a speciifc tack
                    Show {
                        category: &'static str,
                        name: Cow<'a, str>,
                    },
                    /// Show this node if its parent is being shown.
                    ShowIfParent { name: Cow<'a, str> },
                    /// Do not show this node.
                    Omit,
                }

                let categorization = match start_data {
                    buck2_data::span_start_event::Data::Command(_command) => Categorization::Show {
                        category: Self::UNCATEGORIZED,
                        name: self.invocation.command_line_args.join(" ").into(),
                    },
                    buck2_data::span_start_event::Data::Analysis(analysis) => {
                        self.span_counters
                            .bump_counter_while_span(event, "analysis", 1)?;

                        let category = if on_critical_path {
                            Some(Self::CRITICAL_PATH)
                        } else if self
                            .first_pass
                            .long_analyses
                            .contains(&event.span_id().unwrap())
                        {
                            Some(Self::UNCATEGORIZED)
                        } else {
                            None
                        };

                        match category {
                            Some(category) => {
                                let name = format!(
                                    "analysis {}",
                                    display::display_analysis_target(
                                        analysis
                                            .target
                                            .as_ref()
                                            .expect("AnalysisStart event missing 'target' field"),
                                        TargetDisplayOptions::for_chrome_trace()
                                    )?,
                                );

                                Categorization::Show {
                                    category,
                                    name: name.into(),
                                }
                            }
                            None => Categorization::Omit,
                        }
                    }
                    buck2_data::span_start_event::Data::Load(eval) => {
                        self.span_counters
                            .bump_counter_while_span(event, "load", 1)?;

                        let category = if on_critical_path {
                            Some(Self::CRITICAL_PATH)
                        } else if self
                            .first_pass
                            .long_loads
                            .contains(&event.span_id().unwrap())
                        {
                            Some(Self::UNCATEGORIZED)
                        } else {
                            None
                        };

                        match category {
                            Some(category) => Categorization::Show {
                                category,
                                name: format!("load {}", eval.module_id).into(),
                            },
                            None => Categorization::Omit,
                        }
                    }
                    buck2_data::span_start_event::Data::ActionExecution(action) => {
                        #[allow(clippy::if_same_then_else)]
                        let category = if self
                            .first_pass
                            .critical_path_action_keys
                            .contains(action.key.as_ref().unwrap())
                        {
                            Some(Self::CRITICAL_PATH)
                        } else if on_critical_path {
                            Some(Self::CRITICAL_PATH)
                        } else if self
                            .first_pass
                            .local_actions
                            .contains(&event.span_id().unwrap())
                        {
                            Some(Self::UNCATEGORIZED)
                        } else {
                            None
                        };

                        match category {
                            Some(category) => {
                                let name = display::display_action_identity(
                                    action.key.as_ref(),
                                    action.name.as_ref(),
                                    TargetDisplayOptions::for_chrome_trace(),
                                )?;

                                Categorization::Show {
                                    category,
                                    name: name.into(),
                                }
                            }
                            None => Categorization::Omit,
                        }
                    }
                    buck2_data::span_start_event::Data::ExecutorStage(stage) => {
                        let name = stage
                            .stage
                            .as_ref()
                            .and_then(display::display_executor_stage);

                        match name {
                            Some(name) => {
                                self.span_counters.bump_counter_while_span(event, name, 1)?;
                                Categorization::ShowIfParent { name: name.into() }
                            }
                            None => Categorization::Omit,
                        }
                    }
                    buck2_data::span_start_event::Data::ReUpload(_) => {
                        let name = "re_upload";
                        self.span_counters.bump_counter_while_span(event, name, 1)?;
                        Categorization::ShowIfParent { name: name.into() }
                    }
                    buck2_data::span_start_event::Data::FinalMaterialization(..) => {
                        if on_critical_path {
                            Categorization::Show {
                                category: Self::CRITICAL_PATH,
                                name: "materialization".into(),
                            }
                        } else {
                            Categorization::Omit
                        }
                    }
                    buck2_data::span_start_event::Data::FileWatcher(_file_watcher) => {
                        Categorization::Show {
                            category: Self::CRITICAL_PATH,
                            name: "file_watcher_sync".into(),
                        }
                    }
                    _ if on_critical_path => Categorization::Show {
                        category: Self::CRITICAL_PATH,
                        name: "<unknown>".into(),
                    },
                    _ => Categorization::Omit,
                };

                match categorization {
                    Categorization::Show { category, name } => {
                        self.open_named_span(event, name.into_owned(), category)?;
                    }
                    Categorization::ShowIfParent { name } => {
                        let parent_is_open = event
                            .parent_id()
                            .map_or(false, |id| self.open_spans.contains_key(&id));

                        if parent_is_open {
                            // Inherit the parent's track.
                            self.open_named_span(event, name.into_owned(), Self::UNCATEGORIZED)?;
                        }
                    }

                    Categorization::Omit => {}
                }
            }
            // Data field is oneof and `None` means the event is produced with newer version of `.proto` file
            // which added a variant which is not available in version used when compiling this program.
            buck2_data::buck_event::Data::SpanStart(buck2_data::SpanStartEvent { data: None }) => {}
            buck2_data::buck_event::Data::SpanEnd(ref end) => self.handle_event_end(end, event)?,
            buck2_data::buck_event::Data::Instant(buck2_data::InstantEvent {
                data: Some(ref instant_data),
            }) => {
                if let buck2_data::instant_event::Data::Snapshot(_snapshot) = instant_data {
                    self.max_rss_gigabytes_counter.set(
                        event.timestamp(),
                        "max_rss_gigabyte",
                        (_snapshot.buck2_max_rss) as f64 / Self::BYTES_PER_GIGABYTE,
                    )?;
                    self.rate_of_change_counters
                        .set_average_rate_of_change_per_s(
                            event.timestamp(),
                            "average_user_cpu_in_usecs_per_s",
                            _snapshot.buck2_user_cpu_us,
                        )?;
                    self.rate_of_change_counters
                        .set_average_rate_of_change_per_s(
                            event.timestamp(),
                            "average_system_cpu_in_usecs_per_s",
                            _snapshot.buck2_system_cpu_us,
                        )?;
                    self.snapshot_counters.set(
                        event.timestamp(),
                        "blocking_executor_io_queue_size",
                        _snapshot.blocking_executor_io_queue_size,
                    )?;
                    for (nic, stats) in &_snapshot.network_interface_stats {
                        self.rate_of_change_counters
                            .set_average_rate_of_change_per_s(
                                event.timestamp(),
                                &format!("{}_send_bytes", &nic),
                                stats.tx_bytes,
                            )?;
                        self.rate_of_change_counters
                            .set_average_rate_of_change_per_s(
                                event.timestamp(),
                                &format!("{}_receive_bytes", &nic),
                                stats.rx_bytes,
                            )?;
                    }
                    self.rate_of_change_counters
                        .set_average_rate_of_change_per_s(
                            event.timestamp(),
                            "re_upload_bytes",
                            _snapshot.re_upload_bytes,
                        )?;
                    self.rate_of_change_counters
                        .set_average_rate_of_change_per_s(
                            event.timestamp(),
                            "re_download_bytes",
                            _snapshot.re_download_bytes,
                        )?;
                    self.rate_of_change_counters
                        .set_average_rate_of_change_per_s(
                            event.timestamp(),
                            "http_download_bytes",
                            _snapshot.http_download_bytes,
                        )?;
                }
            }
            // Data field is oneof and `None` means the event is produced with newer version of `.proto` file
            // which added a variant which is not available in version used when compiling this program.
            buck2_data::buck_event::Data::Instant(buck2_data::InstantEvent { data: None }) => {}
            buck2_data::buck_event::Data::Record(_) => {}
        };
        Ok(())
    }

    fn handle_event_end(
        &mut self,
        end: &buck2_data::SpanEndEvent,
        event: &BuckEvent,
    ) -> anyhow::Result<()> {
        self.span_counters.handle_event_end(end, event)?;
        if let Some(open) = self.open_spans.remove(&event.span_id().unwrap()) {
            let duration = end
                .duration
                .as_ref()
                .context("Expected SpanEndEvent to have duration")?
                .try_into_duration()?;
            if let SpanTrackAssignment::Owned(track_id) = &open.track {
                self.unused_track_ids
                    .get_mut(track_id.0)
                    .unwrap()
                    .mark_unused(track_id.1);
            }
            self.trace_events
                .push(ChromeTraceClosedSpan { open, duration }.to_json()?);
        }
        Ok(())
    }
}

impl ChromeTraceCommand {
    async fn load_events(path: AbsPathBuf) -> anyhow::Result<(Invocation, Vec<BuckEvent>)> {
        let log_path = EventLogPathBuf::infer(path)?;
        let (invocation, mut stream_values) = log_path.unpack_stream().await?;

        let mut buck_events = Vec::new();

        while let Some(stream_value) = stream_values.try_next().await? {
            match stream_value {
                StreamValue::Event(e) => {
                    let buck_event_result = BuckEvent::try_from(e);
                    match buck_event_result {
                        Ok(buck_event) => buck_events.push(buck_event),
                        Err(e) => {
                            buck2_client_ctx::eprintln!("Error converting event-log: {:#}", e)?
                        }
                    }
                }
                _ => (),
            }
        }

        Ok((invocation, buck_events))
    }

    fn trace_path_from_dir(dir: AbsPathBuf, log: &std::path::Path) -> anyhow::Result<AbsPathBuf> {
        match log.file_name() {
            None => Err(anyhow::anyhow!(
                "Could not determine filename from event log path: `{:#}`",
                log.display()
            )),
            Some(file_name) => {
                let mut trace_path = dir;
                trace_path.push(file_name);
                trace_path.set_extension("trace");
                Ok(trace_path)
            }
        }
    }

    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let log = match self.path {
            Some(path) => path.resolve(&ctx.working_dir),
            None => retrieve_nth_recent_log(&ctx, self.recent.unwrap_or(0))?
                .path()
                .to_owned(),
        };

        let trace_path = self.trace_path.resolve(&ctx.working_dir);
        let dest_path_result = if trace_path.is_dir() {
            Self::trace_path_from_dir(trace_path, &log)
        } else {
            Ok(trace_path)
        };

        let dest_path = match dest_path_result {
            Ok(dest_path) => dest_path,
            Err(e) => {
                buck2_client_ctx::eprintln!("Could not determine trace path, {:#}", e)?;
                return ExitResult::failure();
            }
        };

        let (invocation, events) = ctx
            .runtime
            .block_on(async move { Self::load_events(log).await })?;

        let mut first_pass = ChromeTraceFirstPass::new();
        for event in events.iter() {
            first_pass
                .handle_event(event)
                .with_context(|| display::InvalidBuckEvent(Arc::new(event.clone())))?;
        }
        let mut writer = ChromeTraceWriter::new(invocation, first_pass);
        for event in events {
            let event = Arc::new(event);
            writer
                .handle_event(&event)
                .with_context(|| display::InvalidBuckEvent(event))?;
        }
        let tracefile = std::fs::OpenOptions::new()
            .create(true)
            .write(true)
            .truncate(true)
            .open(dest_path)?;
        writer.to_writer(BufWriter::new(tracefile))?;
        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::UnstableCrashRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct CrashCommand {}

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

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let _err = buckd
            .with_flushing()
            .unstable_crash(UnstableCrashRequest {})
            .await;
        ExitResult::success()
    }

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

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

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


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

use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;

/// Print buck2 daemon directory (`~/.buckd/xxx`).
#[derive(Debug, clap::Parser)]
pub struct DaemonDirCommand {}

impl DaemonDirCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        buck2_client_ctx::println!("{}", ctx.paths?.daemon_dir()?.path.display())?;
        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::unstable_dice_dump_request::DiceDumpFormat;
use buck2_cli_proto::UnstableDiceDumpRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct DiceDumpCommand {
    /// The path to write the heap dump to.
    #[clap(short, long, value_name = "PATH")]
    path: PathArg,
    #[clap(long, group = "dice_dump_format")]
    serde: bool,
    #[clap(long, group = "dice_dump_format")]
    serde_pretty: bool,
}

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

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let format = if self.serde {
            DiceDumpFormat::Bincode
        } else if self.serde_pretty {
            DiceDumpFormat::JsonPretty
        } else {
            DiceDumpFormat::Tsv
        };
        buckd
            .with_flushing()
            .unstable_dice_dump(UnstableDiceDumpRequest {
                destination_path: self.path.resolve(&ctx.working_dir).into_string()?,
                format: format.into(),
            })
            .await?;
        ExitResult::success()
    }

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

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

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


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

use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;

/// Path to current executable.
#[derive(Debug, clap::Parser)]
pub struct ExeCommand {}

impl ExeCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, _ctx: ClientCommandContext<'_>) -> ExitResult {
        buck2_client_ctx::println!("{}", env::current_exe()?.display())?;
        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::FileStatusRequest;
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::path_arg::PathArg;
use buck2_client_ctx::streaming::StreamingCommand;
use gazebo::prelude::*;

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

    /// Paths to validate
    #[clap(value_name = "PATH", required = true)]
    paths: Vec<PathArg>,

    #[clap(long, short)]
    verbose: bool,
}

#[async_trait]
impl StreamingCommand for FileStatusCommand {
    const COMMAND_NAME: &'static str = "file-status";

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;
        buckd
            .with_flushing()
            .file_status(
                FileStatusRequest {
                    context: Some(context),
                    paths: self
                        .paths
                        .try_map(|x| x.resolve(&ctx.working_dir).into_string())?,
                    verbose: self.verbose,
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut StdoutPartialResultHandler,
            )
            .await??;

        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 async_trait::async_trait;
use buck2_cli_proto::FlushDepFilesRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct FlushDepFilesCommand {}

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

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        buckd
            .with_flushing()
            .flush_dep_files(FlushDepFilesRequest {})
            .await??;
        ExitResult::success()
    }

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

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

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


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

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

#[derive(Debug, clap::Parser)]
pub struct HeapDumpCommand {
    /// The path to write the heap dump to.
    #[clap(short, long, value_name = "PATH")]
    path: String,
}

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

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        buckd
            .with_flushing()
            .unstable_heap_dump(UnstableHeapDumpRequest {
                destination_path: self.path,
            })
            .await?;
        ExitResult::success()
    }

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

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

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


/*
 * 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::version::BuckVersion;

#[derive(Debug, clap::Parser)]
pub struct InternalVersionCommand {}

impl InternalVersionCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, _ctx: ClientCommandContext<'_>) -> ExitResult {
        buck2_client_ctx::println!("buck2 internal-version {}", BuckVersion::get_unique_id())?;
        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::subscribers::event_log::read::ReaderStats;
use tokio_stream::StreamExt;

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

/// This command outputs the most recent log in JSON format
#[derive(Debug, clap::Parser)]
pub struct LogPerfCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,

    /// Stats will be emitted every `interval` events.
    #[clap(long, default_value = "10000")]
    interval: u64,
}

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

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

            let mut total_alloc = 0;

            let stats = ReaderStats::new();
            let (_invocation, mut events) = log_path.unpack_stream_with_stats(&stats).await?;

            let mut i = 0;

            while let Some(event) = events.try_next().await? {
                total_alloc += allocative::size_of_unique(&event);

                if i % interval == 0 {
                    buck2_client_ctx::println!(
                        "{}\t{}\t{}\t{}",
                        i,
                        stats.compressed_bytes(),
                        stats.decompressed_bytes(),
                        total_alloc
                    )?;
                }

                i += 1;
            }
            buck2_client_ctx::println!(
                "{}\t{}\t{}\t{}",
                i,
                stats.compressed_bytes(),
                stats.decompressed_bytes(),
                total_alloc
            )?;

            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::MaterializeRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonCommandOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

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

    /// Paths to materialize, relative to project root
    #[clap(value_name = "PATH")]
    paths: Vec<String>,
}

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

    fn existing_only() -> bool {
        true
    }

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context = ctx.client_context(
            &self.common_opts.config_opts,
            matches,
            ctx.sanitized_argv.argv.clone(),
        )?;
        buckd
            .with_flushing()
            .materialize(
                MaterializeRequest {
                    context: Some(context),
                    paths: self.paths,
                },
                ctx.stdin()
                    .console_interaction_stream(&self.common_opts.console_opts),
                &mut NoPartialResultHandler,
            )
            .await??;

        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 allocator_stats::AllocatorStatsCommand;
use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::BuckSubcommand;
use chrome_trace::ChromeTraceCommand;
use crash::CrashCommand;
use dice_dump::DiceDumpCommand;
use file_status::FileStatusCommand;
use flush_dep_files::FlushDepFilesCommand;
use heap_dump::HeapDumpCommand;
use internal_version::InternalVersionCommand;
use materialize::MaterializeCommand;

use crate::commands::debug::allocative::AllocativeCommand;
use crate::commands::debug::daemon_dir::DaemonDirCommand;
use crate::commands::debug::exe::ExeCommand;
use crate::commands::debug::log_perf::LogPerfCommand;
use crate::commands::debug::persist_event_logs::PersistEventLogsCommand;
use crate::commands::debug::segfault::SegfaultCommand;
use crate::commands::debug::set_log_filter::SetLogFilterCommand;
use crate::commands::debug::trace_io::TraceIoCommand;
use crate::commands::debug::upload_re_logs::UploadReLogsCommand;
use crate::commands::log::debug_last_log::DebugLastLogCommand;
use crate::commands::log::debug_replay::DebugReplayCommand;
use crate::commands::log::debug_what_ran::DebugWhatRanCommand;

mod allocative;
mod allocator_stats;
mod chrome_trace;
mod crash;
mod daemon_dir;
mod dice_dump;
mod exe;
mod file_status;
mod flush_dep_files;
mod heap_dump;
mod internal_version;
mod log_perf;
mod materialize;
mod persist_event_logs;
mod segfault;
mod set_log_filter;
mod trace_io;
mod upload_re_logs;

#[derive(Debug, clap::Parser)]
#[clap(about = "Hidden debug commands useful for testing buck2")]
pub enum DebugCommand {
    /// Deliberately crashes the Buck daemon, for testing purposes.
    Crash(CrashCommand),
    /// Causes a segfault in the daemon.
    ///
    /// Useful to make sure that we're reporting it correctly.
    SegFault(SegfaultCommand),
    /// Performs a heap dump of the running buck daemon and saves it to disk.
    HeapDump(HeapDumpCommand),
    /// Dumps allocator stat
    AllocatorStats(AllocatorStatsCommand),
    /// Dump the DICE graph to a file and saves it to disk.
    DiceDump(DiceDumpCommand),
    #[clap(setting(clap::AppSettings::Hidden))]
    Replay(DebugReplayCommand),
    /// Prints the hash of the buck2 binary
    InternalVersion(InternalVersionCommand),
    /// Renders an event-log to a Chrome trace file for inspection with a browser.
    ChromeTrace(ChromeTraceCommand),
    /// Flushes all dep files known to Buck2.
    FlushDepFiles(FlushDepFilesCommand),
    /// Forces materialization of a path, even on the deferred materializer
    Materialize(MaterializeCommand),
    // Upload RE logs given an RE session ID
    UploadReLogs(UploadReLogsCommand),
    /// Validates that Buck2 and disk agree on the state of files.
    FileStatus(FileStatusCommand),
    /// Shows the commands that buck ran
    #[clap(alias = "whatran", setting(clap::AppSettings::Hidden))]
    WhatRan(DebugWhatRanCommand),
    /// Shows the path to the most recent event log
    #[clap(alias = "lastlog", setting(clap::AppSettings::Hidden))]
    LastLog(DebugLastLogCommand),
    /// Prints buck2 daemon directory (`~/.buckd/xxx`).
    DaemonDir(DaemonDirCommand),
    /// Prints buck2 executable (this executable) path.
    Exe(ExeCommand),
    Allocative(AllocativeCommand),
    SetLogFilter(SetLogFilterCommand),
    /// Make sense of log perf
    LogPerf(LogPerfCommand),
    /// Interact with I/O tracing of the daemon.
    TraceIo(TraceIoCommand),
    #[doc(hidden)]
    PersistEventLogs(PersistEventLogsCommand),
}

impl DebugCommand {
    pub fn exec(self, matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        let matches = matches.subcommand().expect("subcommand not found").1;
        match self {
            DebugCommand::DiceDump(cmd) => cmd.exec(matches, ctx),
            DebugCommand::Crash(cmd) => cmd.exec(matches, ctx),
            DebugCommand::HeapDump(cmd) => cmd.exec(matches, ctx),
            DebugCommand::AllocatorStats(cmd) => cmd.exec(matches, ctx),
            DebugCommand::Replay(cmd) => cmd.exec(matches, ctx),
            DebugCommand::InternalVersion(cmd) => cmd.exec(matches, ctx),
            DebugCommand::ChromeTrace(cmd) => cmd.exec(matches, ctx),
            DebugCommand::SegFault(cmd) => cmd.exec(matches, ctx),
            DebugCommand::FlushDepFiles(cmd) => cmd.exec(matches, ctx),
            DebugCommand::WhatRan(cmd) => cmd.exec(matches, ctx),
            DebugCommand::LastLog(cmd) => cmd.exec(matches, ctx),
            DebugCommand::Materialize(cmd) => cmd.exec(matches, ctx),
            DebugCommand::UploadReLogs(cmd) => cmd.exec(matches, ctx),
            DebugCommand::DaemonDir(cmd) => cmd.exec(matches, ctx),
            DebugCommand::Exe(cmd) => cmd.exec(matches, ctx),
            DebugCommand::Allocative(cmd) => cmd.exec(matches, ctx),
            DebugCommand::SetLogFilter(cmd) => cmd.exec(matches, ctx),
            DebugCommand::FileStatus(cmd) => cmd.exec(matches, ctx),
            DebugCommand::LogPerf(cmd) => cmd.exec(matches, ctx),
            DebugCommand::TraceIo(cmd) => cmd.exec(matches, ctx),
            DebugCommand::PersistEventLogs(cmd) => cmd.exec(matches, ctx),
        }
    }

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


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

use std::time::Duration;

use anyhow::Context;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::manifold;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use thiserror::Error;
use tokio::fs::File;
use tokio::fs::OpenOptions;
use tokio::io;
use tokio::io::AsyncReadExt;
use tokio::io::AsyncSeekExt;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;

static UPLOAD_CHUNK_SIZE: EnvHelper<u64> = EnvHelper::new("BUCK2_TEST_MANIFOLD_CHUNK_BYTES");
static MANIFOLD_TTL_S: EnvHelper<u64> = EnvHelper::new("BUCK2_TEST_MANIFOLD_TTL_S");

#[derive(Debug, Error)]
pub(crate) enum PersistLogError {
    #[error("Read more bytes than are available")]
    ReadBytesOverflow,
}

/// Read binary event log from stdin and simultaneously write it to disk and optionally upload to Manifold.
///
/// This command is launched by the buck client to continue log streaming
/// after client command finishes. It is not intended to be used directly.
#[derive(Debug, clap::Parser)]
pub struct PersistEventLogsCommand {
    #[clap(long, help = "Name this log will take in Manifold")]
    manifold_name: String,
    #[clap(long, help = "Where to write this log to on disk")]
    local_path: AbsPathBuf,
    #[clap(long, help = "If present, only write to disk and don't upload")]
    no_upload: bool,
}

impl PersistEventLogsCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        buck2_core::facebook_only();
        ctx.with_runtime(async move |mut ctx| {
            let mut stdin = io::BufReader::new(ctx.stdin());
            self.write_and_upload(&mut stdin).await
        })
        .context("Error writing or uploading event log")?;
        ExitResult::success()
    }

    async fn write_and_upload(self, stdin: impl io::AsyncBufRead + Unpin) -> anyhow::Result<()> {
        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        let file = Mutex::new(create_log_file(self.local_path).await?);

        let write = write_task(&file, tx, stdin);
        let upload = upload_task(&file, rx, self.manifold_name, self.no_upload);

        // Wait for both tasks to finish. If the upload fails we want to keep writing to disk
        let (write_result, upload_result) = tokio::join!(write, upload);
        write_result?;
        upload_result?;
        Ok(())
    }
}

async fn write_task(
    file_mutex: &Mutex<File>,
    tx: tokio::sync::mpsc::UnboundedSender<u64>,
    mut stdin: impl io::AsyncBufRead + Unpin,
) -> anyhow::Result<()> {
    let mut write_position = 0;
    loop {
        let mut buf = vec![0; 64 * 1024]; // maximum pipe size in linux
        let bytes_read = stdin.read(&mut buf).await?;
        if bytes_read == 0 {
            break; // closed stdin
        }
        let mut file = file_mutex.lock().await;
        write_to_file(&mut file, write_position, &buf[..bytes_read]).await?;
        drop(file);
        write_position += bytes_read as u64;
        let _ignored = tx.send(bytes_read as u64); // If this errors, that means the upload task died.
    }

    Ok(())
}

async fn create_log_file(local_path: AbsPathBuf) -> Result<tokio::fs::File, anyhow::Error> {
    let file = OpenOptions::new()
        .create(true)
        .write(true)
        .read(true)
        .open(&local_path)
        .await
        .with_context(|| {
            format!(
                "Failed to open event log for writing at `{}`",
                local_path.display()
            )
        })?;
    Ok(file)
}

async fn upload_task(
    file_mutex: &Mutex<File>,
    mut rx: tokio::sync::mpsc::UnboundedReceiver<u64>,
    manifold_name: String,
    no_upload: bool,
) -> anyhow::Result<()> {
    if no_upload {
        return Ok(());
    }
    let manifold_client = manifold::ManifoldClient::new()?;
    let manifold_path = format!("flat/{}", manifold_name);
    let upload_chunk_size = UPLOAD_CHUNK_SIZE.get_copied()?.unwrap_or(8 * 1024 * 1024);
    let mut read_position = 0;
    let mut total_bytes = 0_u64;

    while let Some(n) = rx.recv().await {
        total_bytes += n;
        while should_upload_chunk(total_bytes, read_position, upload_chunk_size).await? {
            do_the_upload_and_increment_read_position(
                &manifold_client,
                file_mutex,
                &mut read_position,
                &manifold_path,
                upload_chunk_size,
            )
            .await?;
        }
    }

    // When tx gets dropped, rx will return None
    while should_upload_chunk(total_bytes, read_position, upload_chunk_size).await? {
        do_the_upload_and_increment_read_position(
            &manifold_client,
            file_mutex,
            &mut read_position,
            &manifold_path,
            upload_chunk_size,
        )
        .await?;
    }

    // Last chunk to upload is smaller than UPLOAD_CHUNK_SIZE
    do_the_upload_and_increment_read_position(
        &manifold_client,
        file_mutex,
        &mut read_position,
        &manifold_path,
        upload_chunk_size,
    )
    .await?;

    Ok(())
}

async fn should_upload_chunk(
    total_bytes: u64,
    read_position: u64,
    upload_chunk_size: u64,
) -> anyhow::Result<bool> {
    Ok(total_bytes
        .checked_sub(read_position)
        .ok_or(PersistLogError::ReadBytesOverflow)?
        > upload_chunk_size)
}

async fn do_the_upload_and_increment_read_position(
    manifold_client: &manifold::ManifoldClient,
    file_mutex: &Mutex<File>,
    read_position: &mut u64,
    manifold_path: &str,
    upload_chunk_size: u64,
) -> anyhow::Result<u64> {
    let mut file = file_mutex.lock().await;
    // Upload chunk
    file.seek(io::SeekFrom::Start(*read_position))
        .await
        .context("Failed to seek log file")?;
    let (buf, len) = read_chunk(&mut file, upload_chunk_size).await?;
    drop(file);
    if *read_position == 0 {
        // First chunk
        manifold_client
            .write(
                manifold::Bucket::EVENT_LOGS,
                manifold_path,
                buf.into(),
                MANIFOLD_TTL_S.get_copied()?.map(Duration::from_secs),
            )
            .await?
    } else {
        manifold_client
            .append(
                manifold::Bucket::EVENT_LOGS,
                manifold_path,
                buf.into(),
                *read_position,
            )
            .await?
    }
    *read_position += len;
    Ok(len)
}

async fn read_chunk(read_ref: &mut File, chunk_size: u64) -> Result<(Vec<u8>, u64), anyhow::Error> {
    let mut buf = vec![];
    let mut handle = read_ref.take(chunk_size);
    let len = handle
        .read_to_end(&mut buf)
        .await
        .context("Cannot read log file chunk")?;
    Ok((buf, len as u64))
}

async fn write_to_file(
    file: &mut File,
    write_position: u64,
    buf: &[u8],
) -> Result<(), anyhow::Error> {
    file.seek(io::SeekFrom::Start(write_position))
        .await
        .context("Failed to seek log file")?;
    file.write_all(buf).await?;
    file.flush().await?;
    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 async_trait::async_trait;
use buck2_cli_proto::SegfaultRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::streaming::StreamingCommand;

#[derive(Debug, clap::Parser)]
pub struct SegfaultCommand {}

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

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        _matches: &clap::ArgMatches,
        _ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let _err = buckd.with_flushing().segfault(SegfaultRequest {}).await;
        ExitResult::success()
    }

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

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

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


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

use buck2_cli_proto::SetLogFilterRequest;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::daemon::client::connect::BuckdConnectOptions;
use buck2_client_ctx::exit_result::ExitResult;

/// Change the log filter that's currently applied by the Buck2 daemon.
#[derive(Debug, clap::Parser)]
#[clap()]
pub struct SetLogFilterCommand {
    /// The log filter to apply.
    #[clap()]
    log_filter: String,

    /// Whether not to apply it to the daemon.
    #[clap(long)]
    no_daemon: bool,

    /// Whether not to apply it to the forkserver.
    #[clap(long)]
    no_forkserver: bool,
}

impl SetLogFilterCommand {
    pub fn exec(self, _matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        ctx.with_runtime(async move |ctx| {
            let mut buckd = ctx
                .connect_buckd(BuckdConnectOptions::existing_only_no_console())
                .await?;

            buckd
                .with_flushing()
                .set_log_filter(SetLogFilterRequest {
                    log_filter: self.log_filter,
                    daemon: !self.no_daemon,
                    forkserver: !self.no_forkserver,
                })
                .await?;

            ExitResult::success()
        })
    }
}


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

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::trace_io_request;
use buck2_cli_proto::TraceIoRequest;
use buck2_cli_proto::TraceIoResponse;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::command_outcome::CommandOutcome;
use buck2_client_ctx::common::CommonBuildConfigurationOptions;
use buck2_client_ctx::common::CommonConsoleOptions;
use buck2_client_ctx::common::CommonDaemonCommandOptions;
use buck2_client_ctx::daemon::client::connect::DesiredTraceIoState;
use buck2_client_ctx::daemon::client::BuckdClientConnector;
use buck2_client_ctx::daemon::client::NoPartialResultHandler;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::streaming::StreamingCommand;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_offline_archive::ExternalSymlink;
use buck2_offline_archive::OfflineArchiveManifest;
use buck2_offline_archive::RelativeSymlink;
use buck2_offline_archive::RepositoryMetadata;

/// Enable I/O tracing in the buck daemon so we keep track of which files
/// go into a build.
#[derive(Debug, clap::Parser)]
pub struct TraceIoCommand {
    #[clap(subcommand)]
    trace_io_action: Subcommand,
}

/// Sub-settings of I/O tracing
#[derive(Debug, clap::Subcommand)]
enum Subcommand {
    /// Turn on I/O tracing. Has no effect if tracing is already enabled.
    Enable,
    /// Turn off I/O tracing. Has no effect if tracing is already disabled.
    Disable,
    /// Return whether I/O tracing is enabled.
    Status,
    /// Exports the I/O trace taken by the daemon in a structured manifest format.
    ExportManifest {
        #[clap(short, long, help = "Output path to write manifest to")]
        out: Option<PathArg>,
    },
}

impl TraceIoCommand {
    async fn send_request(
        &self,
        req: TraceIoRequest,
        buckd: &mut BuckdClientConnector<'_>,
        ctx: &mut ClientCommandContext<'_>,
    ) -> anyhow::Result<CommandOutcome<TraceIoResponse>> {
        buckd
            .with_flushing()
            .trace_io(
                req,
                ctx.stdin().console_interaction_stream(self.console_opts()),
                &mut NoPartialResultHandler,
            )
            .await
    }
}

#[async_trait]
impl StreamingCommand for TraceIoCommand {
    const COMMAND_NAME: &'static str = "trace-io";

    async fn exec_impl(
        self,
        buckd: &mut BuckdClientConnector,
        matches: &clap::ArgMatches,
        ctx: &mut ClientCommandContext<'_>,
    ) -> ExitResult {
        let context =
            ctx.client_context(self.common_opts(), matches, ctx.sanitized_argv.argv.clone())?;
        match &self.trace_io_action {
            Subcommand::Status => {
                let req = TraceIoRequest {
                    context: Some(context),
                    read_state: Some(trace_io_request::ReadIoTracingState { with_trace: false }),
                };
                let resp = self.send_request(req, buckd, ctx).await??;
                buck2_client_ctx::println!("I/O tracing status: {}", resp.enabled)?;
            }
            Subcommand::ExportManifest { out } => {
                let req = TraceIoRequest {
                    context: Some(context),
                    read_state: Some(trace_io_request::ReadIoTracingState { with_trace: true }),
                };
                let resp = self.send_request(req, buckd, ctx).await??;

                let manifest = OfflineArchiveManifest {
                    paths: resp
                        .trace
                        .into_iter()
                        // Note: Safe because these are all ProjectRelativePath's on the daemon side.
                        .map(ProjectRelativePathBuf::unchecked_new)
                        .collect(),
                    external_paths: resp
                        .external_entries
                        .into_iter()
                        .map(|path| {
                            AbsNormPathBuf::try_from(path)
                                .expect("got unexpected non-absolute path")
                        })
                        .collect(),
                    relative_symlinks: resp
                        .relative_symlinks
                        .into_iter()
                        .map(|symlink| RelativeSymlink {
                            link: ProjectRelativePathBuf::unchecked_new(symlink.link),
                            target: ProjectRelativePathBuf::unchecked_new(symlink.target),
                        })
                        .collect(),
                    external_symlinks: resp
                        .external_symlinks
                        .into_iter()
                        .map(|symlink| ExternalSymlink {
                            link: ProjectRelativePathBuf::unchecked_new(symlink.link),
                            target: AbsPathBuf::try_from(symlink.target)
                                .expect("got unexpected non-absolute symlink target"),
                            remaining_path: symlink
                                .remaining_path
                                .map(ForwardRelativePathBuf::unchecked_new),
                        })
                        .collect(),
                    repository: RepositoryMetadata::from_cwd()
                        .context("creating repository metadata")?,
                };
                let serialized = serde_json::to_string(&manifest)
                    .context("serializing offline archive manifest to json")?;
                if let Some(output_path) = &out {
                    fs_util::write(output_path.resolve(&ctx.working_dir), &serialized)
                        .context("writing offline archive manifest")?;
                } else {
                    buck2_client_ctx::println!("{}", serialized)?;
                }
            }
            // Subcommand::{Enable, Disable} handled by StreamingCommand::trace_io()
            // via implicit daemon restart.
            _ => {}
        }

        ExitResult::success()
    }

    /// Results in a daemon restart if tracing is not already enabled.
    fn trace_io(&self) -> DesiredTraceIoState {
        match self.trace_io_action {
            Subcommand::Enable => DesiredTraceIoState::Enabled,
            Subcommand::Disable => DesiredTraceIoState::Disabled,
            _ => DesiredTraceIoState::Existing,
        }
    }

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

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

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


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

use async_compression::tokio::bufread::ZstdEncoder;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use buck2_client_ctx::manifold;
use buck2_core::fs::async_fs_util;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use tokio::io::AsyncRead;
use tokio::io::BufReader;

#[derive(Debug, clap::Parser)]
#[clap(about = "upload RE logs")]
pub struct UploadReLogsCommand {
    #[clap(long)]
    session_id: String,
}

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

        ctx.with_runtime(async move |ctx| {
            let re_logs_location = ctx.paths()?.re_logs_dir();
            let logs_path = re_logs_location
                .join(ForwardRelativePath::new(&self.session_id)?)
                .join(ForwardRelativePath::new("REClientFolly.log")?);
            let file = async_fs_util::open(&logs_path).await?;
            let mut encoder =
                ZstdEncoder::with_quality(BufReader::new(file), async_compression::Level::Default);

            self.upload_file(&mut encoder).await?
        })
    }

    async fn upload_file<'a, R>(&self, reader: &'a mut R) -> anyhow::Result<ExitResult>
    where
        R: AsyncRead + Unpin,
    {
        let bucket_path = &format!("{}.log.zst", self.session_id);

        manifold::Upload::new(manifold::Bucket::RE_LOGS, bucket_path)
            .with_default_ttl()
            .from_async_read(reader)?
            .spawn()
            .await?;
        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::fmt;
use std::time::Duration;

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;

/// Show the critical path for a selected build.
///
/// This produces tab-delimited output listing every node on the critical path.
///
/// It includes the kind of node, its name, category and identfier, as well as total duration
/// (runtime of this node), user duration (duration the user can improve) and potential improvement
/// before this node stops being on the critical path.
///
/// All durations are in microseconds.
#[derive(Debug, clap::Parser)]
pub struct CriticalPathCommand {
    #[clap(flatten)]
    event_log: EventLogOptions,
}

impl CriticalPathCommand {
    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?;
            buck2_client_ctx::eprintln!(
                "Showing critical path 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::Instant(instant)) => {
                            match instant.data {
                                Some(buck2_data::instant_event::Data::BuildGraphInfo(
                                    build_graph,
                                )) => {
                                    log_critical_path(&build_graph)?;
                                }
                                _ => {}
                            }
                        }
                        _ => {}
                    },
                    _ => {}
                }
            }

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

        ExitResult::success()
    }
}

fn log_critical_path(critical_path: &buck2_data::BuildGraphExecutionInfo) -> anyhow::Result<()> {
    let target_display_options = TargetDisplayOptions::for_log();

    for entry in &critical_path.critical_path2 {
        use buck2_data::critical_path_entry2::Entry;

        let kind;
        let name;
        let mut category = "";
        let mut identifier = "";

        match &entry.entry {
            Some(Entry::Analysis(analysis)) => {
                use buck2_data::critical_path_entry2::analysis::Target;

                kind = "analysis";

                name = match &analysis.target {
                    Some(Target::StandardTarget(t)) => {
                        display::display_configured_target_label(t, target_display_options)?
                    }
                    None => continue,
                };
            }
            Some(Entry::ActionExecution(action_execution)) => {
                use buck2_data::critical_path_entry2::action_execution::Owner;

                kind = "action";

                name = match &action_execution.owner {
                    Some(Owner::TargetLabel(t)) => {
                        display::display_configured_target_label(t, target_display_options)?
                    }
                    Some(Owner::BxlKey(t)) => display::display_bxl_key(t)?,
                    Some(Owner::AnonTarget(t)) => display::display_anon_target(t)?,
                    None => continue,
                };

                match &action_execution.name {
                    Some(name) => {
                        category = &name.category;
                        identifier = &name.identifier;
                    }
                    None => {}
                }
            }
            Some(Entry::Materialization(materialization)) => {
                use buck2_data::critical_path_entry2::materialization::Owner;

                kind = "materialization";

                name = match &materialization.owner {
                    Some(Owner::TargetLabel(t)) => {
                        display::display_configured_target_label(t, target_display_options)?
                    }
                    Some(Owner::BxlKey(t)) => display::display_bxl_key(t)?,
                    Some(Owner::AnonTarget(t)) => display::display_anon_target(t)?,
                    None => continue,
                };

                identifier = &materialization.path;
            }
            Some(Entry::ComputeCriticalPath(..)) => {
                kind = "compute-critical-path";
                name = "".to_owned();
            }
            Some(Entry::Load(load)) => {
                kind = "load";
                name = load.package.clone();
            }
            None => continue,
        }

        struct OptionalDuration {
            inner: Option<Duration>,
        }

        impl OptionalDuration {
            fn new<T, E>(d: Option<T>) -> Result<Self, E>
            where
                T: TryInto<Duration, Error = E>,
            {
                Ok(Self {
                    inner: d.map(|d| d.try_into()).transpose()?,
                })
            }
        }

        impl fmt::Display for OptionalDuration {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                if let Some(inner) = self.inner {
                    write!(f, "{}", inner.as_micros())?;
                }
                Ok(())
            }
        }

        buck2_client_ctx::println!(
            "{}\t{}\t{}\t{}\t{}\t{}\t{}",
            kind,
            name,
            category,
            identifier,
            OptionalDuration::new(entry.total_duration.clone())?,
            OptionalDuration::new(entry.user_duration.clone())?,
            OptionalDuration::new(entry.potential_improvement_duration.clone())?,
        )?;
    }

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

use crate::commands::log::path_log::PathLogCommand;

#[derive(Debug, thiserror::Error)]
enum DebugLastLogCommandError {
    #[error("`buck2 debug last-log` is deprecated. Use `buck2 log path` instead.")]
    Deprecated,
}

#[derive(clap::Parser, Debug)]
pub struct DebugLastLogCommand {
    #[clap(flatten)]
    path_log: PathLogCommand,
}

impl DebugLastLogCommand {
    pub(crate) fn exec(
        self,
        matches: &clap::ArgMatches,
        ctx: ClientCommandContext<'_>,
    ) -> ExitResult {
        soft_error!(
            "debug_last_log",
            DebugLastLogCommandError::Deprecated.into()
        )?;
        self.path_log.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_core::soft_error;

use crate::commands::log::replay::ReplayCommand;

#[derive(Debug, thiserror::Error)]
enum DebugReplayCommandError {
    #[error("`buck2 debug replay` is deprecated. Use `buck2 log replay` instead.")]
    Deprecated,
}

#[derive(clap::Parser, Debug)]
pub struct DebugReplayCommand {
    #[clap(flatten)]
    replay: ReplayCommand,
}

impl DebugReplayCommand {
    pub(crate) fn exec(
        self,
        matches: &clap::ArgMatches,
        ctx: ClientCommandContext<'_>,
    ) -> ExitResult {
        soft_error!("debug_replay", DebugReplayCommandError::Deprecated.into())?;
        self.replay.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_core::soft_error;

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

#[derive(Debug, thiserror::Error)]
enum DebugWhatRanCommandError {
    #[error("`buck2 debug what-ran` is deprecated. Use `buck2 log what-ran` instead.")]
    Deprecated,
}

#[derive(clap::Parser, Debug)]
pub struct DebugWhatRanCommand {
    #[clap(flatten)]
    what_ran: WhatRanCommand,
}

impl DebugWhatRanCommand {
    pub(crate) fn exec(
        self,
        matches: &clap::ArgMatches,
        ctx: ClientCommandContext<'_>,
    ) -> ExitResult {
        soft_error!(
            "debug_what_ran",
            DebugWhatRanCommandError::Deprecated.into()
        )?;
        self.what_ran.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.
 */

mod critical_path;
pub(crate) mod debug_last_log;
pub(crate) mod debug_replay;
pub(crate) mod debug_what_ran;
pub(crate) mod options;
pub(crate) mod path_log;
mod replay;
mod show_log;
mod show_user_log;
mod what_cmd;
mod what_failed;
mod what_materialized;
pub(crate) mod what_ran;
mod what_up;
mod what_uploaded;

use buck2_client_ctx::argv::Argv;
use buck2_client_ctx::argv::SanitizedArgv;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::exit_result::ExitResult;
use dupe::Dupe;

#[derive(
    Debug,
    serde::Serialize,
    serde::Deserialize,
    Clone,
    Dupe,
    clap::ArgEnum
)]
#[clap(rename_all = "snake_case")]
pub enum LogCommandOutputFormat {
    Tabulated,
    Json,
    Csv,
}

#[derive(Debug, clap::Subcommand)]
#[clap(about = "Commands for interacting with buck2 logs")]
pub enum LogCommand {
    #[clap(alias = "whatran")]
    WhatRan(what_ran::WhatRanCommand),
    #[clap(alias = "whatfailed")]
    WhatFailed(what_failed::WhatFailedCommand),
    #[clap(alias = "last")]
    Path(path_log::PathLogCommand),
    Show(show_log::ShowLogCommand),
    #[clap(alias = "whatcmd", alias = "what-cmd")]
    Cmd(what_cmd::WhatCmdCommand),
    #[clap(alias = "whatup")]
    WhatUp(what_up::WhatUpCommand),
    WhatMaterialized(what_materialized::WhatMaterializedCommand),
    WhatUploaded(what_uploaded::WhatUploadedCommand),
    CriticalPath(critical_path::CriticalPathCommand),
    Replay(replay::ReplayCommand),
    ShowUser(show_user_log::ShowUserLogCommand),
}

impl LogCommand {
    pub fn exec(self, matches: &clap::ArgMatches, ctx: ClientCommandContext<'_>) -> ExitResult {
        match self {
            Self::WhatRan(cmd) => cmd.exec(matches, ctx),
            Self::WhatFailed(cmd) => cmd.exec(matches, ctx),
            Self::Path(cmd) => cmd.exec(matches, ctx),
            Self::Show(cmd) => cmd.exec(matches, ctx),
            Self::Cmd(cmd) => cmd.exec(matches, ctx),
            Self::WhatUp(cmd) => cmd.exec(matches, ctx),
            Self::WhatMaterialized(cmd) => cmd.exec(matches, ctx),
            Self::WhatUploaded(cmd) => cmd.exec(matches, ctx),
            Self::CriticalPath(cmd) => cmd.exec(matches, ctx),
            Self::Replay(cmd) => cmd.exec(matches, ctx),
            Self::ShowUser(cmd) => cmd.exec(matches, ctx),
        }
    }

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


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

use std::process::Stdio;

use anyhow::Context;
use buck2_client_ctx::client_ctx::ClientCommandContext;
use buck2_client_ctx::path_arg::PathArg;
use buck2_client_ctx::subscribers::event_log::file_names::find_log_by_trace_id;
use buck2_client_ctx::subscribers::event_log::file_names::retrieve_nth_recent_log;
use buck2_client_ctx::subscribers::event_log::read::EventLogPathBuf;
use buck2_client_ctx::subscribers::event_log::utils::Encoding;
use buck2_common::temp_path::TempPath;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_util::indent::indent;
use buck2_util::process::async_background_command;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use rand::Rng;

#[derive(Debug, thiserror::Error)]
enum EventLogOptionsError {
    #[error("Manifold failed; stderr:\n{}", indent("  ", _0))]
    ManifoldFailed(String),
    #[error(
        "Log not found locally by trace id `{0}`; try `--allow-remote` to download from manifold"
    )]
    LogNotFoundLocally(TraceId),
}

#[derive(Debug, clap::Parser)]
#[clap(group = clap::ArgGroup::with_name("event_log"))]
pub(crate) struct EventLogOptions {
    /// Open the event-log file from a recent command.
    #[clap(long, group = "event_log", value_name = "NUMBER")]
    recent: Option<usize>,

    /// Show log by trace id.
    #[clap(long, group = "event_log", value_name = "ID")]
    trace_id: Option<TraceId>,

    /// Allow downloading the log from manifold if it's not found locally.
    #[clap(long, requires = "trace-id")]
    allow_remote: bool,

    /// A path to an event-log file to read from.
    #[clap(group = "event_log", value_name = "PATH")]
    path: Option<PathArg>,
}

impl EventLogOptions {
    pub(crate) async fn get(
        &self,
        ctx: &ClientCommandContext<'_>,
    ) -> anyhow::Result<EventLogPathBuf> {
        if let Some(path) = &self.path {
            EventLogPathBuf::infer(path.resolve(&ctx.working_dir))
        } else if let Some(id) = &self.trace_id {
            if let Some(log_path) = find_log_by_trace_id(&ctx.paths()?.log_dir(), id)? {
                Ok(log_path)
            } else if self.allow_remote {
                EventLogPathBuf::infer(self.download_remote_id(id, ctx).await?)
            } else {
                return Err(EventLogOptionsError::LogNotFoundLocally(id.dupe()).into());
            }
        } else {
            retrieve_nth_recent_log(ctx, self.recent.unwrap_or(0))
        }
    }

    fn random_string() -> String {
        let mut s = String::with_capacity(10);
        for _ in 0..10 {
            s.push(rand::thread_rng().gen_range('a'..='z'));
        }
        s
    }

    async fn download_remote_id(
        &self,
        trace_id: &TraceId,
        ctx: &ClientCommandContext<'_>,
    ) -> anyhow::Result<AbsPathBuf> {
        let manifold_file_name = FileNameBuf::try_from(format!(
            "{}{}",
            trace_id,
            // TODO(nga): hardcoded default, should at least use the same default buck2 uses,
            //   or better enumerate all the possible suffixes.
            Encoding::PROTO_ZSTD.extensions[0]
        ))?;

        let log_path = ctx
            .paths()?
            .log_dir()
            .join(FileName::new(&format!("dl-{}", manifold_file_name))?);

        if fs_util::try_exists(&log_path)? {
            return Ok(log_path.into_abs_path_buf());
        }

        let tmp_dir = ctx.paths()?.tmp_dir();
        fs_util::create_dir_if_not_exists(&tmp_dir)?;
        let temp_path = tmp_dir.join(FileName::new(&format!(
            "dl.{}.{}.tmp",
            manifold_file_name,
            Self::random_string()
        ))?);

        // Delete the file on failure.
        let temp_path = TempPath::new_path(temp_path);

        let args = [
            "get",
            &format!("buck2_logs/flat/{}", manifold_file_name),
            temp_path
                .path()
                .as_os_str()
                .to_str()
                .context("temp_path is not valid UTF-8")?,
        ];
        buck2_client_ctx::eprintln!("Spawning: manifold {}", args.join(" "))?;
        let command = async_background_command("manifold")
            .args(args)
            .stdin(Stdio::null())
            .stdout(Stdio::null())
            .stderr(Stdio::piped())
            .spawn()?;

        // No timeout here, just press Ctrl-C if you want it to cancel.
        let result = command.wait_with_output().await?;
        if !result.status.success() {
            return Err(EventLogOptionsError::ManifoldFailed(
                String::from_utf8_lossy(&result.stderr).into_owned(),
            )
            .into());
        }

        fs_util::rename(temp_path.path(), &log_path)?;
        buck2_client_ctx::eprintln!("Downloaded event-log to `{}`", log_path.display())?;

        temp_path.close()?;

        Ok(log_path.into_abs_path_buf())
    }
}


/*
 * 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::subscribers::event_log::file_names::retrieve_all_logs;

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

/// Output the path to the selected log.
#[derive(Debug, clap::Parser)]
pub struct PathLogCommand {
    /// Find the log from the Nth most recent command (`--recent 0` is the most recent).
    #[clap(flatten)]
    event_log_options: EventLogOptions,

    /// List all the logs.
    #[clap(long, group = "event_log")]
    all: bool,
}

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

        ctx.with_runtime(async move |ctx| {
            let paths = if all {
                retrieve_all_logs(&ctx)?
            } else {
                vec![event_log_options.get(&ctx).await?]
            };
            for path in paths {
                buck2_client_ctx::println!("{}", path.path().display())?;
            }
            anyhow::Ok(())
        })?;

        ExitResult::success()
    }
}