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

//! A Sink for forwarding events directly to Scribe.
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;

use buck2_core::env_helper::EnvHelper;
use fbinit::FacebookInit;

#[cfg(fbcode_build)]
mod fbcode {

    use std::time::Duration;
    use std::time::SystemTime;

    use buck2_data::InstantEvent;
    use buck2_data::Location;
    use buck2_data::StructuredError;
    use buck2_util::truncate::truncate;
    use fbinit::FacebookInit;
    use prost::Message;

    use crate::metadata;
    use crate::BuckEvent;
    use crate::ControlEvent;
    use crate::EventSink;
    use crate::EventSinkStats;
    use crate::TraceId;

    // 1 MiB limit
    static SCRIBE_MESSAGE_SIZE_LIMIT: usize = 1024 * 1024;
    // 50k characters
    static TRUNCATED_SCRIBE_MESSAGE_SIZE: usize = 50000;

    /// ThriftScribeSink is a ScribeSink backed by the Thrift-based client in the `buck2_scribe_client` crate.
    pub struct ThriftScribeSink {
        category: String,
        client: scribe_client::ScribeClient,
    }

    impl ThriftScribeSink {
        /// Creates a new ThriftScribeSink that forwards messages onto the Thrift-backed Scribe client.
        pub fn new(
            fb: FacebookInit,
            category: String,
            buffer_size: usize,
            retry_backoff: Duration,
            retry_attempts: usize,
            message_batch_size: Option<usize>,
        ) -> anyhow::Result<ThriftScribeSink> {
            let client = scribe_client::ScribeClient::new(
                fb,
                buffer_size,
                retry_backoff,
                retry_attempts,
                message_batch_size,
            )?;
            Ok(ThriftScribeSink { category, client })
        }

        // Send this event now, bypassing internal message queue.
        pub async fn send_now(&self, event: BuckEvent) {
            let message_key = event.trace_id().unwrap().hash();
            if let Some(bytes) = self.encode_message(event, false) {
                self.client
                    .send_now(scribe_client::Message {
                        category: self.category.clone(),
                        message: bytes,
                        message_key: Some(message_key),
                    })
                    .await;
            }
        }

        // Send this event by placing it on the internal message queue.
        pub fn offer(&self, event: BuckEvent) {
            let message_key = event.trace_id().unwrap().hash();
            if let Some(bytes) = self.encode_message(event, false) {
                self.client.offer(scribe_client::Message {
                    category: self.category.clone(),
                    message: bytes,
                    message_key: Some(message_key),
                });
            }
        }

        // Encodes message into something scribe understands.
        fn encode_message(&self, mut event: BuckEvent, is_truncated: bool) -> Option<Vec<u8>> {
            Self::smart_truncate_event(event.data_mut());
            let proto: Box<buck2_data::BuckEvent> = event.into();

            // Add a header byte to indicate this is _not_ base64 encoding.
            let mut buf = Vec::with_capacity(proto.encoded_len() + 1);
            buf.push(b'!');
            let mut proto_bytes = proto.encode_to_vec();
            buf.append(&mut proto_bytes);

            if buf.len() > SCRIBE_MESSAGE_SIZE_LIMIT {
                // if this BuckEvent is already a truncated one but the buffer byte size exceeds the limit,
                // do not send Scribe another truncated version
                if is_truncated {
                    return None;
                }
                let json = serde_json::to_string(&proto).unwrap();

                self.encode_message(
                    BuckEvent::new(
                        SystemTime::now(),
                        TraceId::new(),
                        None,
                        None,
                        buck2_data::buck_event::Data::Instant(InstantEvent {
                            data: Some(
                                StructuredError {
                                    location: Some(Location {
                                        file: file!().to_string(),
                                        line: line!(),
                                        column: column!(),
                                    }),
                                    payload: format!("Soft Error: oversized_scribe: Message is oversized. Event data: {}. Original message size: {}", truncate(&json, TRUNCATED_SCRIBE_MESSAGE_SIZE),
                                    buf.len()),
                                    metadata: metadata::collect(),
                                    backtrace: Vec::new(),
                                    quiet: false,
                                    task: Some(true),
                                    soft_error_category: Some("oversized_scribe".to_owned()),
                                    daemon_in_memory_state_is_corrupted: false,
                                    daemon_materializer_state_is_corrupted: false,
                                    action_cache_is_corrupted: false,
                                }
                                .into(),
                            ),
                        }),
                    ),
                    true,
                )
            } else {
                Some(buf)
            }
        }

        fn smart_truncate_event(d: &mut buck2_data::buck_event::Data) {
            use buck2_data::buck_event::Data;

            match d {
                Data::SpanEnd(ref mut s) => {
                    use buck2_data::span_end_event::Data;

                    match &mut s.data {
                        Some(Data::ActionExecution(ref mut action_execution)) => {
                            Self::truncate_action_execution_end(action_execution);
                        }
                        Some(Data::Command(ref mut command_end)) => {
                            Self::truncate_command_end(command_end, false);
                        }
                        _ => {}
                    };
                }
                Data::Instant(ref mut inst) => {
                    use buck2_data::instant_event::Data;
                    match &mut inst.data {
                        Some(Data::TestResult(ref mut test_result)) => {
                            Self::truncate_test_result(test_result);
                        }
                        Some(Data::TargetPatterns(ref mut target_patterns)) => {
                            Self::truncate_target_patterns(&mut target_patterns.target_patterns);
                        }
                        _ => {}
                    }
                }
                Data::Record(ref mut rec) => {
                    if let Some(buck2_data::record_event::Data::InvocationRecord(
                        ref mut invocation_record,
                    )) = rec.data
                    {
                        if let Some(ref mut file_watcher_stats) =
                            invocation_record.file_watcher_stats
                        {
                            Self::truncate_file_watcher_stats(file_watcher_stats);
                        }
                        if let Some(ref mut resolved_target_patterns) =
                            invocation_record.parsed_target_patterns
                        {
                            Self::truncate_target_patterns(
                                &mut resolved_target_patterns.target_patterns,
                            );
                            // Clear `unresolved_traget_patterns` to save bandwidth. It has less information
                            // than `resolved` one does, and will never be used if `resolved` one is available.
                            if let Some(ref mut command_end) = invocation_record.command_end {
                                Self::truncate_command_end(command_end, true);
                            }
                        } else if let Some(ref mut command_end) = invocation_record.command_end {
                            Self::truncate_command_end(command_end, false);
                        }
                    }
                }
                _ => {}
            };
        }

        fn truncate_action_execution_end(
            action_execution_end: &mut buck2_data::ActionExecutionEnd,
        ) {
            // truncate(...) can panic if asked to truncate too short.
            const MIN_CMD_TRUNCATION: usize = 20;
            let per_command_size_budget =
                ((500 * 1024) / action_execution_end.commands.len().max(1)).max(MIN_CMD_TRUNCATION);

            let truncate_cmd = |cmd: &mut buck2_data::CommandExecution, truncate_all: bool| {
                if let Some(details) = &mut cmd.details {
                    details.stderr = if truncate_all {
                        "<<omitted>>".to_owned()
                    } else {
                        truncate(&details.stderr, per_command_size_budget)
                    };
                }
            };

            if let Some((last_command, retries)) = action_execution_end.commands.split_last_mut() {
                for retried in retries {
                    truncate_cmd(retried, false);
                }
                // Current Scribe tailers don't read stderr of successful actions.
                // Save some bytes.
                truncate_cmd(last_command, !action_execution_end.failed);
            }
        }

        fn truncate_command_end(
            command_end: &mut buck2_data::CommandEnd,
            clear_target_patterns: bool,
        ) {
            use buck2_data::command_end::Data;

            if let Some(ref mut target_patterns) = match &mut command_end.data {
                Some(Data::Build(build_command_end)) => {
                    Some(&mut build_command_end.unresolved_target_patterns)
                }
                Some(Data::Test(test_command_end)) => {
                    Some(&mut test_command_end.unresolved_target_patterns)
                }
                Some(Data::Install(install_command_end)) => {
                    Some(&mut install_command_end.unresolved_target_patterns)
                }
                Some(Data::Targets(targets_command_end)) => {
                    Some(&mut targets_command_end.unresolved_target_patterns)
                }
                _ => None,
            } {
                if clear_target_patterns {
                    target_patterns.clear();
                } else {
                    Self::truncate_target_patterns(target_patterns);
                }
            }

            if let Some(buck2_data::command_end::Data::Audit(ref mut audit_command_end)) =
                command_end.data
            {
                const TRUNCATED_ARGS_LENGTH: usize = 512 * 1024; // 512Kb
                audit_command_end.args = truncate(&audit_command_end.args, TRUNCATED_ARGS_LENGTH);
            }
        }

        fn truncate_file_watcher_stats(file_watcher_stats: &mut buck2_data::FileWatcherStats) {
            const MAX_FILE_CHANGE_BYTES: usize = 100 * 1024;
            let mut bytes: usize = 0;
            for (index, ev) in file_watcher_stats.events.iter().enumerate() {
                bytes += ev.path.len();
                if bytes > MAX_FILE_CHANGE_BYTES {
                    file_watcher_stats.events.truncate(index);
                    file_watcher_stats.incomplete_events_reason = Some(format!(
                        "Too long file change records ({} bytes, max {} bytes)",
                        bytes, MAX_FILE_CHANGE_BYTES
                    ));
                    break;
                }
            }
        }

        fn truncate_test_result(test_result: &mut buck2_data::TestResult) {
            const TRUNCATED_DETAILS_LENGTH: usize = 512 * 1024; // 512Kb
            test_result.details = truncate(&test_result.details, TRUNCATED_DETAILS_LENGTH);
        }

        fn truncate_target_patterns(target_patterns: &mut Vec<buck2_data::TargetPattern>) {
            const MAX_TARGET_PATTERNS_BYTES: usize = 512 * 1024;
            let orig_len = target_patterns.len();
            let mut bytes: usize = 0;
            for (index, target) in target_patterns.iter().enumerate() {
                bytes += target.value.len();
                if bytes > MAX_TARGET_PATTERNS_BYTES {
                    target_patterns.truncate(index);
                    let warn = format!("<<Truncated (reported {} / {})>>", index, orig_len);
                    target_patterns.push(buck2_data::TargetPattern { value: warn });
                    break;
                }
            }
        }
    }

    impl EventSink for ThriftScribeSink {
        fn send(&self, event: BuckEvent) {
            if !should_send_event(event.data()) {
                return;
            }
            self.offer(event);
        }

        fn send_control(&self, _control_event: ControlEvent) {}

        fn stats(&self) -> Option<EventSinkStats> {
            self.client
                .export_counters()
                .map(|counters| EventSinkStats {
                    successes: counters.successes,
                    failures: counters.failures,
                    buffered: counters.queue_depth,
                    dropped: counters.dropped,
                })
        }
    }

    fn should_send_event(d: &buck2_data::buck_event::Data) -> bool {
        use buck2_data::buck_event::Data;
        use buck2_data::ActionKind;

        match d {
            Data::SpanStart(s) => {
                use buck2_data::span_start_event::Data;

                match &s.data {
                    Some(Data::Command(..)) => true,
                    Some(Data::ActionExecution(a)) => match ActionKind::from_i32(a.kind) {
                        // Simple actions are not useful for most log analysis cases
                        Some(ActionKind::Copy)
                        | Some(ActionKind::SymlinkedDir)
                        | Some(ActionKind::Write)
                        | Some(ActionKind::WriteMacrosToFile) => false,
                        _ => true,
                    },
                    None => false,
                    _ => false,
                }
            }
            Data::SpanEnd(s) => {
                use buck2_data::span_end_event::Data;
                use buck2_data::ActionExecutionKind;

                match &s.data {
                    Some(Data::Command(..)) => true,
                    Some(Data::ActionExecution(a)) => {
                        match ActionExecutionKind::from_i32(a.execution_kind) {
                            // Not useful for most log analysis cases
                            Some(ActionExecutionKind::Simple) => false,
                            _ => true,
                        }
                    }
                    Some(Data::Analysis(..)) => true,
                    Some(Data::Load(..)) => true,
                    Some(Data::CacheUpload(..)) => true,
                    Some(Data::Materialization(..)) => true,
                    None => false,
                    _ => false,
                }
            }
            Data::Instant(i) => {
                use buck2_data::instant_event::Data;

                match i.data {
                    Some(Data::BuildGraphInfo(..)) => true,
                    Some(Data::RageResult(..)) => true,
                    Some(Data::ReSession(..)) => true,
                    Some(Data::StructuredError(..)) => true,
                    Some(Data::TestResult(..)) => true,
                    None => false,
                    _ => false,
                }
            }
            Data::Record(r) => {
                use buck2_data::record_event::Data;

                match r.data {
                    Some(Data::InvocationRecord(..)) => true,
                    None => false,
                }
            }
        }
    }

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

        fn make_invocation_record(
            data: buck2_data::InvocationRecord,
        ) -> buck2_data::buck_event::Data {
            buck2_data::buck_event::Data::Record(buck2_data::RecordEvent {
                data: Some(buck2_data::record_event::Data::InvocationRecord(Box::new(
                    data,
                ))),
            })
        }

        fn make_action_execution_end(
            data: buck2_data::ActionExecutionEnd,
        ) -> buck2_data::buck_event::Data {
            buck2_data::buck_event::Data::SpanEnd(buck2_data::SpanEndEvent {
                data: Some(buck2_data::span_end_event::Data::ActionExecution(Box::new(
                    data,
                ))),
                ..Default::default()
            })
        }

        fn make_command_end(data: buck2_data::CommandEnd) -> buck2_data::buck_event::Data {
            buck2_data::buck_event::Data::SpanEnd(buck2_data::SpanEndEvent {
                data: Some(buck2_data::span_end_event::Data::Command(data)),
                ..Default::default()
            })
        }

        fn make_build_command_end(
            unresolved_target_patterns: Vec<buck2_data::TargetPattern>,
        ) -> buck2_data::CommandEnd {
            buck2_data::CommandEnd {
                data: Some(buck2_data::command_end::Data::Build(
                    buck2_data::BuildCommandEnd {
                        unresolved_target_patterns,
                    },
                )),
                ..Default::default()
            }
        }

        fn make_audit_command_end(args: String) -> buck2_data::CommandEnd {
            buck2_data::CommandEnd {
                data: Some(buck2_data::command_end::Data::Audit(
                    buck2_data::AuditCommandEnd {
                        args,
                        ..Default::default()
                    },
                )),
                ..Default::default()
            }
        }

        fn make_command_execution_with_stderr(stderr: String) -> buck2_data::CommandExecution {
            buck2_data::CommandExecution {
                details: Some(buck2_data::CommandExecutionDetails {
                    stderr,
                    ..Default::default()
                }),
                ..Default::default()
            }
        }

        #[test]
        fn smart_truncate_resolved_target_patterns_clears_unresolved_one() {
            let mut record = buck2_data::InvocationRecord::default();
            let mut record_expected = record.clone();

            let resolved_target_patterns = vec![buck2_data::TargetPattern {
                value: "some_resolved_target".to_owned(),
            }];
            record.parsed_target_patterns = Some(buck2_data::ParsedTargetPatterns {
                target_patterns: resolved_target_patterns.clone(),
            });
            // resolved_target_patterns is expected to be unchanged.
            record_expected.parsed_target_patterns = Some(buck2_data::ParsedTargetPatterns {
                target_patterns: resolved_target_patterns,
            });

            let unresolved_target_patterns = vec![buck2_data::TargetPattern {
                value: "some_unresolved_target".to_owned(),
            }];
            record.command_end = Some(make_build_command_end(unresolved_target_patterns));

            // unresolved_target_patterns is expected to be empty.
            record_expected.command_end = Some(make_build_command_end(vec![]));

            let mut event_data = make_invocation_record(record);
            let event_data_expected = make_invocation_record(record_expected);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_unresolved_target_used_when_resolved_one_unavailable() {
            let mut record = buck2_data::InvocationRecord::default();
            let mut record_expected = record.clone();

            record.parsed_target_patterns = None;
            record_expected.parsed_target_patterns = None;

            let unresolved_target_patterns = vec![buck2_data::TargetPattern {
                value: "some_unresolved_target".to_owned(),
            }];
            let command_end = make_build_command_end(unresolved_target_patterns);

            record.command_end = Some(command_end.clone());
            // unresolved_target_patterns is expected to be unchanged.
            record_expected.command_end = Some(command_end);

            let mut event_data = make_invocation_record(record);
            let event_data_expected = make_invocation_record(record_expected);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_action_execution_end_one_last_command_truncated() {
            let command_execution_with_stderr =
                make_command_execution_with_stderr("this is a test".to_owned());
            let command_execution_stderr_omitted =
                make_command_execution_with_stderr("<<omitted>>".to_owned());

            let action_execution_end_with_stderrs = buck2_data::ActionExecutionEnd {
                commands: vec![command_execution_with_stderr],
                ..Default::default()
            };
            let action_execution_end_last_stderr_omitted = buck2_data::ActionExecutionEnd {
                commands: vec![command_execution_stderr_omitted],
                ..Default::default()
            };
            let mut event_data = make_action_execution_end(action_execution_end_with_stderrs);
            let event_data_expected =
                make_action_execution_end(action_execution_end_last_stderr_omitted);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_action_execution_end_long_stderr_command_truncated() {
            let command_execution_with_stderr =
                make_command_execution_with_stderr("this is a test".to_owned());
            let mut over_sized_str = "0123456789".repeat(10 * 1024);
            over_sized_str.push_str("0123456789"); // 100k + 10; 10-byte over
            let command_execution_with_long_stderr =
                make_command_execution_with_stderr(over_sized_str);
            let mut omitted_str = "0123456789".repeat(10 * 1024);
            omitted_str.replace_range((50 * 1024 - 6)..(50 * 1024 + 6), "<<omitted>>");
            let command_execution_stderr_partially_omitted =
                make_command_execution_with_stderr(omitted_str);
            let command_execution_stderr_all_omitted =
                make_command_execution_with_stderr("<<omitted>>".to_owned());

            let action_execution_end_with_stderrs = buck2_data::ActionExecutionEnd {
                commands: vec![
                    command_execution_with_stderr.clone(),
                    command_execution_with_long_stderr.clone(),
                    command_execution_with_stderr.clone(),
                    command_execution_with_long_stderr,
                    command_execution_with_stderr.clone(),
                ],
                ..Default::default()
            };
            let action_execution_end_last_stderr_omitted = buck2_data::ActionExecutionEnd {
                commands: vec![
                    command_execution_with_stderr.clone(),
                    command_execution_stderr_partially_omitted.clone(),
                    command_execution_with_stderr,
                    command_execution_stderr_partially_omitted,
                    command_execution_stderr_all_omitted,
                ],
                ..Default::default()
            };
            let mut event_data = make_action_execution_end(action_execution_end_with_stderrs);
            let event_data_expected =
                make_action_execution_end(action_execution_end_last_stderr_omitted);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_build_command_end_short_target_patterns_not_truncated() {
            let unresolved_target_patterns = vec![
                buck2_data::TargetPattern {
                    value: "hello".to_owned(),
                },
                buck2_data::TargetPattern {
                    value: "world".to_owned(),
                },
                buck2_data::TargetPattern {
                    value: "!\n".to_owned(),
                },
            ];
            let command_end = make_build_command_end(unresolved_target_patterns);

            let mut event_data = make_command_end(command_end);
            let event_data_expected = event_data.clone();

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_build_command_end_long_target_patterns_truncated() {
            let unresolved_target_patterns = vec![
                buck2_data::TargetPattern {
                    value: "0123456789".repeat(20 * 1024),
                },
                buck2_data::TargetPattern {
                    value: "0123456789".repeat(20 * 1024),
                },
                buck2_data::TargetPattern {
                    value: "0123456789".repeat(20 * 1024), // 600k in total; 88k-byte over
                },
            ];
            let command_end = make_build_command_end(unresolved_target_patterns);

            let unresolved_target_patterns_truncated = vec![
                buck2_data::TargetPattern {
                    value: "0123456789".repeat(20 * 1024),
                },
                buck2_data::TargetPattern {
                    value: "0123456789".repeat(20 * 1024),
                },
                buck2_data::TargetPattern {
                    value: "<<Truncated (reported 2 / 3)>>".to_owned(),
                },
            ];
            let command_end_truncated =
                make_build_command_end(unresolved_target_patterns_truncated);

            let mut event_data = make_command_end(command_end);
            let event_data_expected = make_command_end(command_end_truncated);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_long_file_watcher_stats_truncated() {
            let file_watcher_event = buck2_data::FileWatcherEvent {
                path: "0123456789".repeat(3 * 1024),
                ..Default::default()
            };
            let file_watcher_stats = buck2_data::FileWatcherStats {
                events: vec![
                    file_watcher_event.clone(),
                    file_watcher_event.clone(),
                    file_watcher_event.clone(),
                    file_watcher_event.clone(), // 120k in total; 20k-byte over
                ],
                ..Default::default()
            };
            let file_watcher_stats_truncated = buck2_data::FileWatcherStats {
                events: vec![
                    file_watcher_event.clone(),
                    file_watcher_event.clone(),
                    file_watcher_event,
                ],
                incomplete_events_reason: Some(format!(
                    "Too long file change records ({} bytes, max {} bytes)",
                    120 * 1024,
                    100 * 1024
                )),
                ..Default::default()
            };
            let record = buck2_data::InvocationRecord {
                file_watcher_stats: Some(file_watcher_stats),
                ..Default::default()
            };
            let record_truncated = buck2_data::InvocationRecord {
                file_watcher_stats: Some(file_watcher_stats_truncated),
                ..Default::default()
            };
            let mut event_data = make_invocation_record(record);
            let event_data_expected = make_invocation_record(record_truncated);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_short_file_watcher_stats_not_truncated() {
            let file_watcher_event = buck2_data::FileWatcherEvent {
                path: "this is a test".to_owned(),
                ..Default::default()
            };
            let file_watcher_stats = buck2_data::FileWatcherStats {
                events: vec![
                    file_watcher_event.clone(),
                    file_watcher_event.clone(),
                    file_watcher_event,
                ],
                ..Default::default()
            };
            let record = buck2_data::InvocationRecord {
                file_watcher_stats: Some(file_watcher_stats),
                ..Default::default()
            };
            let mut event_data = make_invocation_record(record);
            let event_data_expected = event_data.clone();

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_long_audit_args_truncated() {
            let long_args = "0123456789".repeat(100 * 1024); // 1000KB
            let mut long_args_truncated = long_args.clone();
            // See buck2_util::truncate::truncate to know the formula here
            long_args_truncated.replace_range(
                ((512 * 1024 - 11) / 2)..(1000 * 1024 - ((512 * 1024 - 11) / 2)),
                "<<omitted>>",
            );
            let audit_command_end = make_audit_command_end(long_args);
            let audit_command_end_truncated = make_audit_command_end(long_args_truncated);
            let record = buck2_data::InvocationRecord {
                command_end: Some(audit_command_end),
                ..Default::default()
            };
            let record_truncated = buck2_data::InvocationRecord {
                command_end: Some(audit_command_end_truncated),
                ..Default::default()
            };
            let mut event_data = make_invocation_record(record);
            let event_data_expected = make_invocation_record(record_truncated);

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }

        #[test]
        fn smart_truncate_short_audit_args_not_truncated() {
            let audit_command_end = make_audit_command_end("some args".to_owned());
            let record = buck2_data::InvocationRecord {
                command_end: Some(audit_command_end),
                ..Default::default()
            };
            let mut event_data = make_invocation_record(record);
            let event_data_expected = event_data.clone();

            ThriftScribeSink::smart_truncate_event(&mut event_data);

            assert_eq!(event_data, event_data_expected);
        }
    }
}

#[cfg(not(fbcode_build))]
mod fbcode {
    use crate::BuckEvent;
    use crate::ControlEvent;
    use crate::EventSink;
    use crate::EventSinkStats;

    pub struct ThriftScribeSink;

    impl ThriftScribeSink {
        pub async fn send_now(&self, _event: BuckEvent) {}
    }

    impl EventSink for ThriftScribeSink {
        fn send(&self, _event: BuckEvent) {}

        fn send_control(&self, _control_event: ControlEvent) {}

        fn stats(&self) -> Option<EventSinkStats> {
            None
        }
    }
}

pub use fbcode::*;

fn new_thrift_scribe_sink_if_fbcode(
    fb: FacebookInit,
    buffer_size: usize,
    retry_backoff: Duration,
    retry_attempts: usize,
    message_batch_size: Option<usize>,
) -> anyhow::Result<Option<ThriftScribeSink>> {
    #[cfg(fbcode_build)]
    {
        Ok(Some(ThriftScribeSink::new(
            fb,
            scribe_category()?,
            buffer_size,
            retry_backoff,
            retry_attempts,
            message_batch_size,
        )?))
    }
    #[cfg(not(fbcode_build))]
    {
        let _ = (
            fb,
            buffer_size,
            retry_backoff,
            retry_attempts,
            message_batch_size,
        );
        Ok(None)
    }
}

pub fn new_thrift_scribe_sink_if_enabled(
    fb: FacebookInit,
    buffer_size: usize,
    retry_backoff: Duration,
    retry_attempts: usize,
    message_batch_size: Option<usize>,
) -> anyhow::Result<Option<ThriftScribeSink>> {
    if is_enabled() {
        new_thrift_scribe_sink_if_fbcode(
            fb,
            buffer_size,
            retry_backoff,
            retry_attempts,
            message_batch_size,
        )
    } else {
        Ok(None)
    }
}

/// Whether or not Scribe logging is enabled for this process. It must be explicitly disabled via `disable()`.
static SCRIBE_ENABLED: AtomicBool = AtomicBool::new(true);

/// Returns whether this process should actually write to Scribe, even if it is fully supported by the platform and
/// binary.
pub fn is_enabled() -> bool {
    SCRIBE_ENABLED.load(Ordering::Relaxed)
}

/// Disables Scribe logging for this process. Scribe logging must be disabled explicitly on startup, otherwise it is
/// on by default.
pub fn disable() {
    SCRIBE_ENABLED.store(false, Ordering::Relaxed);
}

pub fn scribe_category() -> anyhow::Result<String> {
    const DEFAULT_SCRIBE_CATEGORY: &str = "buck2_events";
    // Note that both daemon and client are emitting events, and that changing this variable has
    // no effect on the daemon until buckd is restarted but has effect on the client.
    static SCRIBE_CATEGORY: EnvHelper<String> = EnvHelper::new("BUCK2_SCRIBE_CATEGORY");
    Ok(SCRIBE_CATEGORY
        .get()?
        .map_or_else(|| DEFAULT_SCRIBE_CATEGORY.to_owned(), |c| c.clone()))
}


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

use crate::BuckEvent;
use crate::ControlEvent;
use crate::EventSink;
use crate::EventSinkStats;

/// A Sink implementation that wraps two EventSinks and sends events to both of them.
pub struct TeeSink<A, B>(A, B);

impl<A, B> TeeSink<A, B> {
    pub fn new(a: A, b: B) -> TeeSink<A, B> {
        TeeSink(a, b)
    }
}

impl<A: EventSink, B: EventSink> EventSink for TeeSink<A, B> {
    fn send(&self, event: BuckEvent) {
        self.0.send(event.clone());
        self.1.send(event);
    }

    fn send_control(&self, control_event: ControlEvent) {
        self.0.send_control(control_event.clone());
        self.1.send_control(control_event);
    }

    fn stats(&self) -> Option<EventSinkStats> {
        match (self.0.stats(), self.1.stats()) {
            (Some(stats0), Some(stats1)) => Some(stats0.aggregate(&stats1)),
            (Some(stats0), None) => Some(stats0),
            (None, Some(stats1)) => Some(stats1),
            _ => None,
        }
    }
}


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

use crate::last_command_execution_kind::get_last_command_execution_kind;
use crate::last_command_execution_kind::LastCommandExecutionKind;

/// Records the number of actions depending on how they executed.
/// There's no overlap between the actions - summing them all up
/// gives the total number of actions. `local_actions` + `remote_actions`
/// provides the total number of actually executed actions while
/// `cached_actions` provides number of actions which we found
/// in the action cache.  `fallback_actions` provides the number of actions
/// that had its command run more than once (hence, using fallback to run).
///
/// These stats only track executions/commands.
#[derive(Default, Clone, Dupe)]
pub struct ActionStats {
    pub local_actions: u64,
    pub remote_actions: u64,
    pub cached_actions: u64,
    pub fallback_actions: u64,
}

impl ActionStats {
    pub fn action_cache_hit_percentage(&self) -> u8 {
        // We want special semantics for the return value: the terminal values (0% and 100%)
        // should _only_ be used when there are exactly no cache hits and full cache hits.
        // So, even if we have 99.6% cache hits, we want to display 99% and conversely,
        // if the value is 0.1% cache hits, we want to display 1%.
        //
        // This allows us to have special meaning for 0% and 100%: complete cache-divergence
        // and fully cacheable builds.
        let total_actions = self.total_executed_and_cached_actions();
        if total_actions == 0 || self.cached_actions == total_actions {
            100
        } else if self.cached_actions == 0 {
            0
        } else {
            let hit_percent = ((self.cached_actions as f64) / (total_actions as f64)) * 100f64;
            let integral_percent = (hit_percent.round()) as u8;
            integral_percent.clamp(1, 99)
        }
    }

    pub fn total_executed_actions(&self) -> u64 {
        self.local_actions + self.remote_actions
    }

    pub fn total_executed_and_cached_actions(&self) -> u64 {
        self.local_actions + self.remote_actions + self.cached_actions
    }

    pub fn update(&mut self, action: &buck2_data::ActionExecutionEnd) {
        if was_fallback_action(action) {
            self.fallback_actions += 1;
        }
        match get_last_command_execution_kind(action) {
            LastCommandExecutionKind::Local | LastCommandExecutionKind::LocalWorker => {
                self.local_actions += 1;
            }
            LastCommandExecutionKind::Cached => {
                self.cached_actions += 1;
            }
            LastCommandExecutionKind::Remote => {
                self.remote_actions += 1;
            }
            LastCommandExecutionKind::NoCommand => {}
        }
    }

    pub fn log_stats(&self) -> bool {
        self.total_executed_and_cached_actions() > 0
    }
}

impl fmt::Display for ActionStats {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut action_stats_message = format!(
            "Cache hits: {}%. Commands: {} (cached: {}, remote: {}, local: {})",
            self.action_cache_hit_percentage(),
            self.total_executed_and_cached_actions(),
            self.cached_actions,
            self.remote_actions,
            self.local_actions
        );
        if self.fallback_actions > 0 {
            action_stats_message += format!(
                ". Fallback: {}/{}",
                self.fallback_actions,
                self.total_executed_actions()
            )
            .as_str();
        }
        write!(f, "{}", action_stats_message)
    }
}

/// Identify whether an action was a fallback action. A fallback action is an action that executed
/// two commands, unless one of those was a Cancelled, which just means hybrid execution
/// cancelled the local run (and which is not a fallback).
pub fn was_fallback_action(action: &buck2_data::ActionExecutionEnd) -> bool {
    use buck2_data::command_execution::Status;

    action
        .commands
        .iter()
        .filter(|c| !matches!(c.status, Some(Status::Cancelled(..))))
        .count()
        > 1
}


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

use std::collections::BTreeMap;
use std::collections::VecDeque;
use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;

use buck2_events::BuckEvent;
use gazebo::variants::VariantName;

use crate::unpack_event::unpack_event;
use crate::unpack_event::UnpackedBuckEvent;

static NUM_DELAYS_FOR_AVERAGE: usize = 10;

#[derive(Default)]
pub struct SpanData {
    pub started: u64,
    pub finished: u64,
    pub total_duration: Duration,
    pub total_poll_time: Duration,
    // We display this as average max poll time, but easiest to track it as total.
    pub total_max_poll_time: Duration,
}

#[derive(Default)]
pub struct InstantData {
    pub count: u64,
}

pub struct DebugEventsState {
    pub event_count: u64,
    pub spans: BTreeMap<String, SpanData>,
    pub instants: BTreeMap<String, InstantData>,
    /// events we receive have SystemTime timestamps and may be based on a different time than
    /// we are running in. This is used to adjust them to our local Instant-based time.
    pub events_start_time: Option<SystemTime>,
    pub recent_delays: VecDeque<Duration>,
    pub max_delay: Duration,
}

impl DebugEventsState {
    pub fn new() -> Self {
        Self {
            event_count: 0,
            spans: BTreeMap::new(),
            instants: BTreeMap::new(),
            events_start_time: None,
            recent_delays: VecDeque::new(),
            max_delay: Duration::ZERO,
        }
    }

    pub fn handle_event(&mut self, start_time: Instant, event: &BuckEvent) -> anyhow::Result<()> {
        let event_time: SystemTime = event.timestamp();
        let elapsed = start_time.elapsed();
        let events_start_time = match self.events_start_time {
            Some(v) => v,
            None => {
                let start_time = event_time - elapsed;
                self.events_start_time = Some(start_time);
                start_time
            }
        };

        self.event_count += 1;
        // event_elapsed represents the time from start_time to when this event was timestamped in the server
        let event_elapsed = event_time.duration_since(events_start_time)?;
        // if events are being delayed, the real elapsed time will be greater than the "event_elapsed" time
        let delay = elapsed.checked_sub(event_elapsed).unwrap_or(Duration::ZERO);

        self.max_delay = std::cmp::max(delay, self.max_delay);
        self.recent_delays.push_back(delay);
        while self.recent_delays.len() > NUM_DELAYS_FOR_AVERAGE {
            self.recent_delays.pop_front();
        }

        match unpack_event(event)? {
            UnpackedBuckEvent::SpanStart(_, _, data) => {
                let name = data.variant_name();

                let entry = {
                    match self.spans.get_mut(name) {
                        Some(v) => v,
                        None => {
                            self.spans.insert(name.to_owned(), SpanData::default());
                            self.spans.get_mut(name).unwrap()
                        }
                    }
                };
                entry.started += 1;
            }
            UnpackedBuckEvent::SpanEnd(_, span_end, data) => {
                // Right now, matching these end events to the start events depends on the field names in the protobufs
                // matching. That works but is fragile, if it breaks at some point we can do the match and explicitly
                // match them ourselves.
                let name = data.variant_name();

                let entry = {
                    match self.spans.get_mut(name) {
                        Some(v) => v,
                        None => {
                            self.spans.insert(name.to_owned(), SpanData::default());
                            self.spans.get_mut(name).unwrap()
                        }
                    }
                };
                entry.finished += 1;
                if let Some(v) = &span_end.stats {
                    entry.total_poll_time += Duration::from_micros(v.max_poll_time_us);
                    entry.total_max_poll_time += Duration::from_micros(v.max_poll_time_us);
                }
            }
            UnpackedBuckEvent::Instant(_, _, data) => {
                let name = data.variant_name();

                let entry = {
                    match self.instants.get_mut(name) {
                        Some(v) => v,
                        None => {
                            self.instants
                                .insert(name.to_owned(), InstantData::default());
                            self.instants.get_mut(name).unwrap()
                        }
                    }
                };

                entry.count += 1;
            }
        }

        Ok(())
    }
}


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

use std::collections::BTreeMap;

use buck2_data::DiceKeyState;
use buck2_data::DiceStateSnapshot;

pub struct DiceState {
    key_states: BTreeMap<String, DiceKeyState>,
}

impl DiceState {
    pub fn new() -> Self {
        Self {
            key_states: BTreeMap::new(),
        }
    }

    pub fn update(&mut self, update: &DiceStateSnapshot) {
        for (k, v) in &update.key_states {
            self.key_states.insert(k.clone(), v.clone());
        }
    }

    pub fn key_states(&self) -> &BTreeMap<String, DiceKeyState> {
        &self.key_states
    }
}


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

// TODO(brasselsprouts): move this onto the original core types and convert in events

use std::borrow::Cow;
use std::fmt;
use std::fmt::Write;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context as _;
use buck2_common::convert::ProstDurationExt;
use buck2_data::action_key;
use buck2_data::span_start_event::Data;
use buck2_data::ActionKey;
use buck2_data::ActionName;
use buck2_data::AnonTarget;
use buck2_data::BxlFunctionKey;
use buck2_data::BxlFunctionLabel;
use buck2_data::ConfiguredTargetLabel;
use buck2_data::TargetLabel;
use buck2_events::BuckEvent;
use buck2_test_api::data::TestStatus;
use buck2_util::collections::ordered_set::OrderedSet;
use buck2_util::commas::commas;
use buck2_util::truncate::truncate;
use dupe::Dupe;
use superconsole::style::Stylize;
use superconsole::Line;
use superconsole::Lines;
use superconsole::Span;
use thiserror::Error;

use crate::fmt_duration;
use crate::verbosity::Verbosity;

#[derive(Copy, Clone, Dupe)]
pub struct TargetDisplayOptions {
    with_configuration: bool,
}

impl TargetDisplayOptions {
    pub fn for_log() -> Self {
        Self {
            with_configuration: true,
        }
    }

    pub fn for_console(with_configuration: bool) -> Self {
        Self { with_configuration }
    }

    pub fn for_chrome_trace() -> Self {
        Self {
            with_configuration: false,
        }
    }
}

pub fn display_configured_target_label(
    ctl: &ConfiguredTargetLabel,
    opts: TargetDisplayOptions,
) -> anyhow::Result<String> {
    if let ConfiguredTargetLabel {
        label: Some(TargetLabel { package, name }),
        configuration: Some(configuration),
        // We never display execution configurations at the moment
        execution_configuration: _,
    } = ctl
    {
        Ok(if opts.with_configuration {
            format!("{}:{} ({})", package, name, configuration.full_name)
        } else {
            format!("{}:{}", package, name)
        })
    } else {
        Err(ParseEventError::InvalidConfiguredTargetLabel.into())
    }
}

pub fn display_anon_target(ctl: &AnonTarget) -> anyhow::Result<String> {
    if let AnonTarget {
        name: Some(TargetLabel { package, name }),
        // We currently never display execution configurations, only normal configurations
        execution_configuration: _,
        hash,
    } = ctl
    {
        Ok(format!("{}:{}@{}", package, name, hash))
    } else {
        Err(ParseEventError::InvalidAnonTarget.into())
    }
}

pub fn display_analysis_target(
    target: &buck2_data::analysis_start::Target,
    opts: TargetDisplayOptions,
) -> anyhow::Result<String> {
    use buck2_data::analysis_start::Target;
    match target {
        Target::StandardTarget(ctl) => display_configured_target_label(ctl, opts),
        Target::AnonTarget(anon) => display_anon_target(anon),
    }
}

pub fn display_bxl_key(ctl: &BxlFunctionKey) -> anyhow::Result<String> {
    if let BxlFunctionKey {
        label: Some(BxlFunctionLabel { bxl_path, name }),
    } = ctl
    {
        Ok(format!("{}:{}", bxl_path, name))
    } else {
        Err(ParseEventError::MissingBxlFunctionLabel.into())
    }
}

pub fn display_action_owner(
    owner: &action_key::Owner,
    opts: TargetDisplayOptions,
) -> anyhow::Result<String> {
    match owner {
        action_key::Owner::TargetLabel(target_label)
        | action_key::Owner::TestTargetLabel(target_label)
        | action_key::Owner::LocalResourceSetup(target_label) => {
            display_configured_target_label(target_label, opts)
        }
        action_key::Owner::BxlKey(bxl_key) => display_bxl_key(bxl_key),
        action_key::Owner::AnonTarget(anon_target) => display_anon_target(anon_target),
    }
}

pub fn display_action_key(
    action_key: &ActionKey,
    opts: TargetDisplayOptions,
) -> anyhow::Result<String> {
    if let ActionKey {
        owner: Some(owner), ..
    } = action_key
    {
        display_action_owner(owner, opts)
    } else {
        Err(ParseEventError::MissingActionOwner.into())
    }
}

fn display_action_name_opt(name: Option<&ActionName>) -> String {
    match name {
        Some(name) if name.identifier.is_empty() => name.category.clone(),
        Some(name) => format!("{} {}", name.category, name.identifier),
        _ => "unknown".to_owned(),
    }
}

pub fn display_action_identity(
    action_key: Option<&ActionKey>,
    name: Option<&ActionName>,
    opts: TargetDisplayOptions,
) -> anyhow::Result<String> {
    let key_string = match action_key {
        Some(key) => display_action_key(key, opts),
        None => Err(ParseEventError::MissingActionKey.into()),
    }?;
    let action_string = match name {
        Some(ActionName {
            category,
            identifier,
        }) if !identifier.is_empty() => format!(" ({} {})", category, identifier),
        Some(ActionName { category, .. }) => format!(" ({})", category),
        None => String::new(),
    };

    Ok(format!("{}{}", key_string, action_string))
}

/// Formats event payloads for display.
pub fn display_event(event: &BuckEvent, opts: TargetDisplayOptions) -> anyhow::Result<String> {
    let res: anyhow::Result<_> = try {
        let data = match event.data() {
            buck2_data::buck_event::Data::SpanStart(ref start) => start.data.as_ref().unwrap(),
            _ => Err(anyhow::Error::from(ParseEventError::UnexpectedEvent))?,
        };

        let res: anyhow::Result<_> = match data {
            Data::ActionExecution(action) => match &action.key {
                Some(key) => {
                    let string = display_action_key(key, opts)?;
                    let action_descriptor = display_action_name_opt(action.name.as_ref());
                    Ok(format!("{} -- action ({})", string, action_descriptor))
                }
                None => Err(ParseEventError::MissingActionKey.into()),
            },
            Data::FinalMaterialization(materialization) => {
                let build = &materialization
                    .artifact
                    .as_ref()
                    .ok_or(ParseEventError::MissingArtifact)?;

                let key = display_action_key(
                    build
                        .key
                        .as_ref()
                        .ok_or(ParseEventError::MissingActionKey)?,
                    opts,
                )?;
                let path = {
                    if build.path.is_empty() {
                        Err(ParseEventError::MissingMaterializationPath)
                    } else {
                        Ok(&build.path)
                    }
                }?;
                Ok(format!("{} -- materializing `{}`", key, path))
            }
            Data::Analysis(analysis) => match &analysis.target {
                Some(target) => {
                    let target = display_analysis_target(target, opts)?;
                    Ok(format!("{} -- running analysis", target))
                }
                None => Err(ParseEventError::MissingConfiguredTargetLabel.into()),
            },
            Data::AnalysisStage(info) => {
                let stage = info.stage.as_ref().context("analysis stage is missing")?;
                let stage = display_analysis_stage(stage);
                Ok(stage.into())
            }
            Data::LoadPackage(load) => Ok(format!("{} -- loading package file tree", load.path)),
            Data::Load(load) => Ok(format!("{} -- evaluating build file", load.module_id)),
            Data::ExecutorStage(info) => {
                let stage = info.stage.as_ref().context("executor stage is missing")?;
                let stage = display_executor_stage(stage).context("unknown executor stage")?;
                Ok(stage.into())
            }
            Data::TestDiscovery(discovery) => Ok(format!(
                "Test {} -- discovering tests",
                discovery.suite_name
            )),
            Data::TestStart(start) => match &start.suite {
                Some(suite) => {
                    let tests = suite.test_names.join(" ");
                    Ok(format!("Test {} -- {}", suite.suite_name, tests))
                }
                None => Err(ParseEventError::MissingSuiteName.into()),
            },
            Data::CommandCritical(..) => Err(ParseEventError::UnexpectedEvent.into()),
            Data::Command(..) => Err(ParseEventError::UnexpectedEvent.into()),
            Data::FileWatcher(x) => Ok(format!(
                "Syncing file changes via {}",
                display_file_watcher(x.provider)
            )),
            Data::MatchDepFiles(buck2_data::MatchDepFilesStart {}) => Ok("dep_files".to_owned()),
            Data::SharedTask(..) => Ok("Waiting on task from another command".to_owned()),
            Data::CacheUpload(..) => Ok("upload".to_owned()),
            Data::CreateOutputSymlinks(..) => Ok("Creating output symlinks".to_owned()),
            Data::InstallEventInfo(info) => Ok(format!(
                "Sending {} at path {}",
                info.artifact_name, info.file_path
            )),
            Data::DiceStateUpdate(..) => Ok("Syncing changes to graph".to_owned()),
            Data::Materialization(..) => Ok("materializing".to_owned()),
            Data::DiceCriticalSection(..) => Err(ParseEventError::UnexpectedEvent.into()),
            Data::DiceBlockConcurrentCommand(cmd) => Ok(format!(
                "Waiting for command [{}] to finish",
                truncate(&cmd.cmd_args, 200),
            )),
            Data::DiceSynchronizeSection(..) => Ok("Synchronizing buck2 internal state".to_owned()),
            Data::DiceCleanup(..) => Ok("Cleaning up graph state".to_owned()),
            Data::ExclusiveCommandWait(buck2_data::ExclusiveCommandWaitStart { command_name }) => {
                if let Some(name) = command_name {
                    Ok(format!("Waiting for command [{}] to finish", name))
                } else {
                    Ok("Waiting for dice".to_owned())
                }
            }
            Data::DeferredPreparationStage(prep) => {
                use buck2_data::deferred_preparation_stage_start::Stage;
                match prep.stage.as_ref().context("Missing `stage`")? {
                    Stage::MaterializedArtifacts(_) => Ok("local_materialize_inputs".to_owned()),
                }
            }
            Data::DynamicLambda(lambda) => {
                use buck2_data::dynamic_lambda_start::Owner;

                let label = match lambda.owner.as_ref().context("Missing `owner`")? {
                    Owner::TargetLabel(target_label) => {
                        display_configured_target_label(target_label, opts)
                    }
                    Owner::BxlKey(bxl_key) => display_bxl_key(bxl_key),
                    Owner::AnonTarget(anon_target) => display_anon_target(anon_target),
                }?;

                Ok(format!("{} -- dynamic analysis", label))
            }
            Data::BxlExecution(execution) => {
                Ok(format!("Executing BXL script `{}`", execution.name))
            }
            Data::BxlDiceInvocation(..) => Ok("Waiting for graph computations".to_owned()),
            Data::ReUpload(..) => Ok("re_upload".to_owned()),
            Data::ConnectToInstaller(buck2_data::ConnectToInstallerStart { tcp_port }) => {
                Ok(format!("Connecting to installer on port {}", tcp_port))
            }
            Data::Fake(fake) => Ok(format!("{} -- speak of the devil", fake.caramba)),
            Data::LocalResources(..) => Ok("Local resources setup".to_owned()),
            Data::ReleaseLocalResources(..) => Ok("Releasing local resources".to_owned()),
            Data::BxlEnsureArtifacts(..) => Err(ParseEventError::UnexpectedEvent.into()),
        };

        // This shouldn't really be necessary, but that's how try blocks work :(
        res?
    };

    res.with_context(|| InvalidBuckEvent(Arc::new(event.clone())))
}

fn display_file_watcher(provider: i32) -> &'static str {
    match buck2_data::FileWatcherProvider::from_i32(provider) {
        Some(buck2_data::FileWatcherProvider::Watchman) => "Watchman",
        Some(buck2_data::FileWatcherProvider::RustNotify) => "notify",
        None => "unknown mechanism",
    }
}

pub fn display_analysis_stage(stage: &buck2_data::analysis_stage_start::Stage) -> &'static str {
    use buck2_data::analysis_stage_start::Stage;

    match stage {
        Stage::ResolveQueries(()) => "resolve_queries",
        Stage::EvaluateRule(()) => "evaluate_rule",
    }
}

pub fn display_file_watcher_end(file_watcher_end: &buck2_data::FileWatcherEnd) -> Vec<String> {
    const MAX_PRINT_MESSAGES: usize = 3;
    let mut res = Vec::new();

    if let Some(stats) = &file_watcher_end.stats {
        // The `FileWatcherEvent` contain no duplicates. However, there can be two distinct events
        // for the same file, e.g. `foo` was deleted as a file and then created as a directory.
        //
        // It looks really odd to print "File changed: foo" twice, so we dedupe user messages.
        //
        // If there are more file change records than we passed over, and some of the omitted ones are
        // duplicates on the same file, then our "additional file change events" count is slightly high.
        // Shouldn't be a big deal in practice, since it is rare, and fairly big numbers already.

        let mut to_print = OrderedSet::new();
        for x in &stats.events {
            to_print.insert(&x.path);
        }
        for path in to_print.iter().take(MAX_PRINT_MESSAGES) {
            res.push(format!("File changed: {}", path));
        }
        let unprinted_paths =
            // those we have the names of but didn't print
            to_print.len().saturating_sub(MAX_PRINT_MESSAGES) +
            // plus those we didn't get the names for
            (stats.events_processed as usize).saturating_sub(stats.events.len());
        if unprinted_paths > 0 {
            res.push(format!("{} additional file change events", unprinted_paths));
        }

        if let Some(fresh_instance) = &stats.fresh_instance_data {
            let file_watcher = if stats.watchman_version.is_some() {
                "Watchman"
            } else {
                "File Watcher"
            };

            let mut msg = format!("{} fresh instance: ", file_watcher);
            let mut comma = commas();
            if fresh_instance.new_mergebase {
                comma(&mut msg).unwrap();
                write!(&mut msg, "new mergebase").unwrap();
            }
            if fresh_instance.cleared_dice {
                comma(&mut msg).unwrap();
                write!(&mut msg, "cleared graph state").unwrap();
            }
            if fresh_instance.cleared_dep_files {
                comma(&mut msg).unwrap();
                write!(&mut msg, "cleared dep files").unwrap();
            }
            res.push(msg);
        }
    }

    res
}

pub fn display_executor_stage(
    stage: &buck2_data::executor_stage_start::Stage,
) -> Option<&'static str> {
    use buck2_data::executor_stage_start::Stage;

    let label = match stage {
        Stage::Prepare(..) => "prepare",
        Stage::CacheQuery(..) => "re_action_cache",
        Stage::CacheHit(..) => "re_download",
        Stage::Re(re) => {
            use buck2_data::re_stage::Stage;

            match re.stage.as_ref()? {
                Stage::Execute(..) => "re_execute",
                Stage::Download(..) => "re_download",
                Stage::Queue(..) => "re_queued",
                Stage::WorkerDownload(..) => "re_worker_download",
                Stage::WorkerUpload(..) => "re_worker_upload",
                Stage::Unknown(..) => "re_unknown",
            }
        }
        Stage::Local(local) => {
            use buck2_data::local_stage::Stage;

            match local.stage.as_ref()? {
                Stage::Queued(..) => "local_queued",
                Stage::Execute(..) => "local_execute",
                Stage::MaterializeInputs(..) => "local_materialize_inputs",
                Stage::PrepareOutputs(_) => "local_prepare_outputs",
                Stage::AcquireLocalResource(_) => "acquire_local_resource",
                Stage::InitializeWorker(_) => "initialize_worker",
                Stage::WorkerExecute(_) => "worker_execute",
            }
        }
    };

    Some(label)
}

#[derive(Error, Debug)]
enum ParseEventError {
    #[error("Missing configured target label")]
    MissingConfiguredTargetLabel,
    #[error("Invalid configured target label")]
    InvalidConfiguredTargetLabel,
    #[error("Invalid anon target")]
    InvalidAnonTarget,
    #[error("Missing action key")]
    MissingActionKey,
    #[error("Missing suite name")]
    MissingSuiteName,
    #[error("Missing artifact")]
    MissingArtifact,
    #[error("Missing materialization path")]
    MissingMaterializationPath,
    #[error("Missing action owner")]
    MissingActionOwner,
    #[error("Missing bxl function label")]
    MissingBxlFunctionLabel,
    #[error("Unexpected event")]
    UnexpectedEvent,
}

#[derive(Error, Debug)]
#[error("Invalid buck event: `{0:?}`")]
pub struct InvalidBuckEvent(pub Arc<BuckEvent>);

pub fn format_test_result(test_result: &buck2_data::TestResult) -> anyhow::Result<Option<Lines>> {
    let buck2_data::TestResult {
        name,
        status,
        duration,
        details,
        ..
    } = test_result;
    let status = TestStatus::try_from(*status)?;

    // Pass results normally have no details, unless the --print-passing-details is set.
    // Do not display anything for passing tests unless details are present to avoid
    // cluttering the UI with unimportant test results.
    if matches!(&status, TestStatus::PASS | TestStatus::LISTING_SUCCESS) && details.is_empty() {
        return Ok(None);
    }

    let prefix = match status {
        TestStatus::FAIL => Span::new_styled("âœ— Fail".to_owned().red()),
        TestStatus::SKIP => Span::new_styled("â†· Skip".to_owned().cyan()),
        TestStatus::OMITTED => Span::new_styled("\u{20E0} Omitted".to_owned().cyan()),
        TestStatus::FATAL => Span::new_styled("âš  Fatal".to_owned().red()),
        TestStatus::TIMEOUT => Span::new_styled("âœ‰ Timeout".to_owned().cyan()),
        TestStatus::PASS => Span::new_styled("âœ“ Pass".to_owned().green()),
        TestStatus::LISTING_SUCCESS => Span::new_styled("âœ“ Listing success".to_owned().green()),
        TestStatus::UNKNOWN => Span::new_styled("? Unknown".to_owned().cyan()),
        TestStatus::RERUN => Span::new_styled("â†» Rerun".to_owned().cyan()),
        TestStatus::LISTING_FAILED => Span::new_styled("âš  Listing failed".to_owned().red()),
    }?;
    let mut base = Line::from_iter([prefix, Span::new_unstyled(format!(": {}", name,))?]);
    if let Some(duration) = duration {
        if let Ok(duration) = Duration::try_from(duration.clone()) {
            base.push(Span::new_unstyled(format!(
                " ({})",
                // Set time_speed parameter as 1.0 because this is taking the duration of something that was measured somewhere else,
                // so it doesn't make sense to apply the speed adjustment.
                fmt_duration::fmt_duration(duration, 1.0)
            ))?);
        }
    }
    // If a test has details, we always show them. It's the test runner's
    // responsibility to withhold details when these are not relevant.
    // For instance, tpx will always withhold details of passing tests
    // unless the --print-passing-details is set.
    let mut lines = vec![base];
    if !details.is_empty() {
        lines.append(&mut Lines::from_multiline_string(details, Default::default()).0);
    }
    Ok(Some(Lines(lines)))
}

pub struct ActionErrorDisplay<'a> {
    pub action_id: String,
    pub reason: String,
    pub command: Option<Cow<'a, buck2_data::CommandExecutionDetails>>,
}

impl<'a> ActionErrorDisplay<'a> {
    pub fn to_static(self) -> ActionErrorDisplay<'static> {
        ActionErrorDisplay {
            action_id: self.action_id,
            reason: self.reason,
            command: self.command.map(|c| Cow::Owned(c.into_owned())),
        }
    }
}

pub fn display_action_error<'a>(
    action: &'a buck2_data::ActionExecutionEnd,
    error: &'a buck2_data::action_execution_end::Error,
    opts: TargetDisplayOptions,
) -> anyhow::Result<ActionErrorDisplay<'a>> {
    use buck2_data::action_execution_end::Error;

    let command = action.commands.last().and_then(|c| c.details.as_ref());

    let reason = match error {
        Error::MissingOutputs(missing_outputs) => {
            format!("Required outputs are missing: {}", missing_outputs.message)
        }
        Error::Unknown(error_string) => {
            format!("Internal error: {}", error_string)
        }
        Error::CommandExecutionError(buck2_data::CommandExecutionError {}) => {
            match action.commands.last() {
                Some(c) => failure_reason_for_command_execution(c)?,
                None => "Unexpected command status".to_owned(),
            }
        }
    };

    Ok(ActionErrorDisplay {
        action_id: display_action_identity(action.key.as_ref(), action.name.as_ref(), opts)?,
        reason,
        command: command.map(Cow::Borrowed),
    })
}

fn failure_reason_for_command_execution(
    command_execution: &buck2_data::CommandExecution,
) -> anyhow::Result<String> {
    use buck2_data::command_execution::Cancelled;
    use buck2_data::command_execution::Error;
    use buck2_data::command_execution::Failure;
    use buck2_data::command_execution::Status;
    use buck2_data::command_execution::Success;
    use buck2_data::command_execution::Timeout;
    use buck2_data::command_execution_details::Command;

    let command = command_execution
        .details
        .as_ref()
        .context("CommandExecution did not include a `command`")?;

    let status = command_execution
        .status
        .as_ref()
        .context("CommandExecution did not include a `status`")?;

    let locality = match command.command {
        Some(Command::RemoteCommand(..)) => "Remote ",
        Some(Command::LocalCommand(..)) | Some(Command::OmittedLocalCommand(..)) => "Local ",
        Some(Command::WorkerInitCommand(..)) => "Local Worker Initialization ",
        Some(Command::WorkerCommand(..)) => "Local Worker ",
        None => "",
    };

    Ok(match status {
        Status::Success(Success {}) => "Unexpected command status".to_owned(),
        Status::Failure(Failure {}) => {
            struct OptionalExitCode {
                code: Option<i32>,
            }

            impl fmt::Display for OptionalExitCode {
                fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
                    match self.code {
                        Some(code) => write!(f, "{}", code),
                        None => write!(f, "<no exit code>"),
                    }
                }
            }

            format!(
                "{}command returned non-zero exit code {}",
                locality,
                OptionalExitCode {
                    code: command.signed_exit_code
                }
            )
        }
        Status::Timeout(Timeout { duration }) => {
            let duration = duration
                .as_ref()
                .context("Timeout did not include a `duration`")?
                .try_into_duration()
                .context("Timeout `duration` was invalid")?;

            format!("Command timed out after {:.3}s", duration.as_secs_f64(),)
        }
        Status::Error(Error { stage, error }) => {
            format!("Internal error (stage: {}): {}", stage, error)
        }
        Status::Cancelled(Cancelled {}) => "Command was cancelled".to_owned(),
    })
}

pub fn success_stderr<'a>(
    action: &'a buck2_data::ActionExecutionEnd,
    verbosity: Verbosity,
) -> anyhow::Result<Option<&'a str>> {
    if !(verbosity.print_success_stderr() || action.always_print_stderr) {
        return Ok(None);
    }

    let stderr = match action.commands.last() {
        Some(command) => {
            &command
                .details
                .as_ref()
                .context("CommandExecution did not include a `command`")?
                .stderr
        }
        None => return Ok(None),
    };

    if stderr.is_empty() {
        return Ok(None);
    }

    Ok(Some(stderr))
}


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

use std::sync::Arc;
use std::time::Instant;

use anyhow::Context;
use buck2_events::BuckEvent;
use buck2_wrapper_common::invocation_id::TraceId;

use crate::action_stats::ActionStats;
use crate::debug_events::DebugEventsState;
use crate::dice_state::DiceState;
use crate::io_state::IoState;
use crate::re_state::ReState;
use crate::session_info::SessionInfo;
use crate::span_tracker::BuckEventSpanTracker;
use crate::starlark_debug::StarlarkDebuggerState;
use crate::test_state::TestState;
use crate::two_snapshots::TwoSnapshots;

pub struct EventObserver<E> {
    pub span_tracker: BuckEventSpanTracker,
    pub action_stats: ActionStats,
    re_state: ReState,
    two_snapshots: TwoSnapshots, // NOTE: We got many more copies of this than we should.
    session_info: SessionInfo,
    io_state: IoState,
    test_state: TestState,
    starlark_debugger_state: StarlarkDebuggerState,
    /// When running without the Superconsole, we skip some state that we don't need. This might be
    /// premature optimization.
    extra: E,
}

impl<E> EventObserver<E>
where
    E: EventObserverExtra,
{
    pub fn new(trace_id: TraceId) -> Self {
        Self {
            span_tracker: BuckEventSpanTracker::new(),
            action_stats: ActionStats::default(),
            re_state: ReState::new(),
            two_snapshots: TwoSnapshots::default(),
            session_info: SessionInfo {
                trace_id,
                test_session: None,
            },
            io_state: IoState::default(),
            test_state: TestState::default(),
            starlark_debugger_state: StarlarkDebuggerState::new(),
            extra: E::new(),
        }
    }

    pub fn observe(&mut self, receive_time: Instant, event: &Arc<BuckEvent>) -> anyhow::Result<()> {
        self.span_tracker.handle_event(receive_time, event)?;

        {
            use buck2_data::buck_event::Data::*;

            match event.data() {
                SpanEnd(end) => {
                    use buck2_data::span_end_event::Data::*;

                    match end.data.as_ref().context("Missing `data` in SpanEnd")? {
                        ActionExecution(action_execution_end) => {
                            self.action_stats.update(action_execution_end);
                        }
                        _ => {}
                    }
                }
                Instant(instant) => {
                    use buck2_data::instant_event::Data::*;

                    match instant
                        .data
                        .as_ref()
                        .context("Missing `data` in `Instant`")?
                    {
                        ReSession(re_session) => {
                            self.re_state.add_re_session(re_session);
                        }
                        Snapshot(snapshot) => {
                            self.re_state.update(event.timestamp(), snapshot);
                            self.two_snapshots.update(event.timestamp(), snapshot);
                            self.io_state.update(event.timestamp(), snapshot);
                        }
                        TestDiscovery(discovery) => {
                            use buck2_data::test_discovery::Data::*;

                            match discovery
                                .data
                                .as_ref()
                                .context("Missing `data` in `TestDiscovery`")?
                            {
                                Session(session) => {
                                    self.session_info.test_session = Some(session.clone());
                                }
                                Tests(tests) => {
                                    self.test_state.discovered += tests.test_names.len() as u64
                                }
                            }
                        }
                        TestResult(result) => {
                            self.test_state.update(result)?;
                        }
                        DebugAdapterSnapshot(snapshot) => {
                            self.starlark_debugger_state
                                .update(event.timestamp(), snapshot)?;
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }

        self.extra.observe(receive_time, event)?;

        Ok(())
    }

    pub fn spans(&self) -> &BuckEventSpanTracker {
        &self.span_tracker
    }

    pub fn action_stats(&self) -> &ActionStats {
        &self.action_stats
    }

    pub fn re_state(&self) -> &ReState {
        &self.re_state
    }

    pub fn two_snapshots(&self) -> &TwoSnapshots {
        &self.two_snapshots
    }

    pub fn session_info(&self) -> &SessionInfo {
        &self.session_info
    }

    pub fn io_state(&self) -> &IoState {
        &self.io_state
    }

    pub fn starlark_debugger_state(&self) -> &StarlarkDebuggerState {
        &self.starlark_debugger_state
    }

    pub fn test_state(&self) -> &TestState {
        &self.test_state
    }

    pub fn extra(&self) -> &E {
        &self.extra
    }
}

pub trait EventObserverExtra: Send {
    fn new() -> Self;

    fn observe(&mut self, receive_time: Instant, event: &Arc<BuckEvent>) -> anyhow::Result<()>;
}

/// This has more fields for debug info. We don't always capture those.
pub struct DebugEventObserverExtra {
    dice_state: DiceState,
    debug_events: DebugEventsState,
}

impl EventObserverExtra for DebugEventObserverExtra {
    fn new() -> Self {
        Self {
            dice_state: DiceState::new(),
            debug_events: DebugEventsState::new(),
        }
    }

    fn observe(&mut self, receive_time: Instant, event: &Arc<BuckEvent>) -> anyhow::Result<()> {
        self.debug_events.handle_event(receive_time, event)?;

        {
            use buck2_data::buck_event::Data::*;

            match event.data() {
                Instant(instant) => {
                    use buck2_data::instant_event::Data::*;

                    match instant
                        .data
                        .as_ref()
                        .context("Missing `data` in `Instant`")?
                    {
                        DiceStateSnapshot(dice) => {
                            self.dice_state.update(dice);
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
        }

        Ok(())
    }
}

impl DebugEventObserverExtra {
    pub fn dice_state(&self) -> &DiceState {
        &self.dice_state
    }

    pub fn debug_events(&self) -> &DebugEventsState {
        &self.debug_events
    }
}

pub struct NoopEventObserverExtra;

impl EventObserverExtra for NoopEventObserverExtra {
    fn new() -> Self {
        Self
    }

    fn observe(&mut self, _receive_time: Instant, _event: &Arc<BuckEvent>) -> anyhow::Result<()> {
        // Noop
        Ok(())
    }
}


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

use std::time::Duration;

pub fn fmt_duration(elapsed: Duration, time_speed: f64) -> String {
    let elapsed = elapsed.mul_f64(time_speed);
    let nanos = elapsed.as_nanos().try_into().unwrap_or(u64::MAX);
    let millis = nanos.saturating_add(50_000_000); // Round up.
    let subsec = millis % 1_000_000_000;
    let secs = millis / 1_000_000_000;
    let mins = secs / 60;
    let secs_of_min = secs % 60;
    let hours = mins / 60;
    let mins_of_hour = mins % 60;
    if hours != 0 {
        format!(
            "{}:{:02}:{:02}.{}s",
            hours,
            mins_of_hour,
            secs_of_min,
            subsec / 100_000_000
        )
    } else if mins != 0 {
        format!("{}:{:02}.{}s", mins, secs_of_min, subsec / 100_000_000)
    } else {
        format!("{}.{}s", secs, subsec / 100_000_000)
    }
}

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

    use crate::fmt_duration::fmt_duration;

    #[test]
    fn test_fmt_duration() {
        fn hmss(h: u64, m: u64, s: u64, ms: u64) -> Duration {
            Duration::from_millis(h * 3_600_000 + m * 60_000 + s * 1000 + ms)
        }

        assert_eq!("0.0s", fmt_duration(hmss(0, 0, 0, 0), 1.0));
        assert_eq!("0.0s", fmt_duration(hmss(0, 0, 0, 49), 1.0));
        assert_eq!("0.1s", fmt_duration(hmss(0, 0, 0, 50), 1.0));
        assert_eq!("0.1s", fmt_duration(hmss(0, 0, 0, 99), 1.0));
        assert_eq!("0.1s", fmt_duration(hmss(0, 0, 0, 100), 1.0));
        assert_eq!("0.9s", fmt_duration(hmss(0, 0, 0, 949), 1.0));
        assert_eq!("1.0s", fmt_duration(hmss(0, 0, 0, 999), 1.0));
        assert_eq!("1.0s", fmt_duration(hmss(0, 0, 1, 0), 1.0));
        assert_eq!("59.9s", fmt_duration(hmss(0, 0, 59, 949), 1.0));
        assert_eq!("1:00.0s", fmt_duration(hmss(0, 0, 59, 999), 1.0));
        assert_eq!("1:00.0s", fmt_duration(hmss(0, 1, 0, 0), 1.0));
        assert_eq!("59:59.9s", fmt_duration(hmss(0, 59, 59, 949), 1.0));
        assert_eq!("1:00:00.0s", fmt_duration(hmss(0, 59, 59, 950), 1.0));
        assert_eq!("1:00:00.0s", fmt_duration(hmss(1, 0, 0, 0), 1.0));
        assert_eq!("9876:54:32.1s", fmt_duration(hmss(9876, 54, 32, 100), 1.0));
    }
}


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

/// Write out a u64 representing bytes as something more readable
pub struct HumanizedBytes {
    bytes: u64,
    fixed_width: bool,
}

impl HumanizedBytes {
    pub fn new(bytes: u64) -> Self {
        HumanizedBytes {
            bytes,
            fixed_width: false,
        }
    }

    pub const FIXED_WIDTH_WIDTH: usize = 7;

    pub fn fixed_width(bytes: u64) -> Self {
        HumanizedBytes {
            bytes,
            fixed_width: true,
        }
    }
}

pub struct HumanizedBytesPerSecond {
    bytes_per_second: u64,
    fixed_width: bool,
}

impl HumanizedBytesPerSecond {
    pub fn _new(bytes_per_second: u64) -> Self {
        HumanizedBytesPerSecond {
            bytes_per_second,
            fixed_width: false,
        }
    }

    pub const FIXED_WIDTH_WIDTH: usize = HumanizedBytes::FIXED_WIDTH_WIDTH + "/s".len();

    pub fn fixed_width(bytes_per_second: u64) -> Self {
        HumanizedBytesPerSecond {
            bytes_per_second,
            fixed_width: true,
        }
    }
}

pub struct HumanizedCount {
    count: u64,
    fixed_width: bool,
}

impl HumanizedCount {
    pub fn new(count: u64) -> Self {
        HumanizedCount {
            count,
            fixed_width: false,
        }
    }

    pub fn fixed_width(count: u64) -> Self {
        HumanizedCount {
            count,
            fixed_width: true,
        }
    }
}

struct Preformat {
    val: f64,
    label: &'static str,
    point: bool,
}

fn preformat(
    value: u64,
    factor: f64,
    one_label: &'static str,
    labels: &[&'static str],
) -> Preformat {
    let mut val = value as f64;
    let mut label = one_label;

    let mut labels = labels.iter();

    loop {
        if val < factor {
            break;
        }

        let next_label = match labels.next() {
            Some(l) => l,
            None => break,
        };

        val /= factor;
        label = next_label;
    }

    if f64::round(val) >= 1000.0 {
        if let Some(next_label) = labels.next() {
            val = 1.0;
            label = next_label;
        }
    }

    Preformat {
        val,
        label,
        point: value >= 1000 && f64::round(val) < 10.0,
    }
}

impl fmt::Display for HumanizedBytes {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let Preformat { val, label, point } =
            preformat(self.bytes, 1024.0, "B", &["KiB", "MiB", "GiB"]);

        match (point, self.fixed_width) {
            (false, false) => write!(f, "{:.0} {}", val, label),
            (true, false) => write!(f, "{:.1} {}", val, label),
            (false, true) => write!(f, "{:>3.0} {:<3}", val, label),
            (true, true) => write!(f, "{:>3.1} {:<3}", val, label),
        }
    }
}

impl fmt::Display for HumanizedBytesPerSecond {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let Preformat { val, label, point } = preformat(
            self.bytes_per_second,
            1024.0,
            "B/s",
            &["KiB/s", "MiB/s", "GiB/s"],
        );

        match (point, self.fixed_width) {
            (false, false) => write!(f, "{:.0} {}", val, label),
            (true, false) => write!(f, "{:.1} {}", val, label),
            (false, true) => write!(f, "{:>3.0} {:<5}", val, label),
            (true, true) => write!(f, "{:>3.1} {:<5}", val, label),
        }
    }
}

impl fmt::Display for HumanizedCount {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let Preformat { val, label, point } = preformat(self.count, 1000.0, "", &["K", "M", "B"]);

        struct SpacePrefixedUnlessEmpty<'a> {
            label: &'a str,
        }

        impl fmt::Display for SpacePrefixedUnlessEmpty<'_> {
            fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
                if self.label.is_empty() {
                    return Ok(());
                }
                write!(f, " {}", self.label)
            }
        }

        match (point, self.fixed_width) {
            (false, false) => write!(f, "{:.0}{}", val, SpacePrefixedUnlessEmpty { label }),
            (true, false) => write!(f, "{:.1}{}", val, SpacePrefixedUnlessEmpty { label }),
            (false, true) => write!(f, "{:>3.0} {:<1}", val, label),
            (true, true) => write!(f, "{:>3.1} {:<1}", val, label),
        }
    }
}

#[cfg(test)]
mod test {
    use super::HumanizedBytes;
    use super::HumanizedBytesPerSecond;
    use super::HumanizedCount;

    #[allow(clippy::identity_op)]
    #[test]
    fn test_humanized() {
        fn t(value: u64, expected: &str, expected_fixed_width: &str) {
            assert_eq!(
                HumanizedBytes {
                    bytes: value,
                    fixed_width: false
                }
                .to_string(),
                expected
            );
            if value <= (1 << 9) {
                assert_eq!(
                    HumanizedBytes::FIXED_WIDTH_WIDTH,
                    expected_fixed_width.len()
                );
            }
            assert_eq!(
                HumanizedBytes {
                    bytes: value,
                    fixed_width: true,
                }
                .to_string(),
                expected_fixed_width
            );
        }

        t(0, "0 B", "  0 B  ");
        t(1, "1 B", "  1 B  ");
        t(10, "10 B", " 10 B  ");
        t(345, "345 B", "345 B  ");
        t(1023, "1.0 KiB", "1.0 KiB");
        t(1024, "1.0 KiB", "1.0 KiB");
        t(1536, "1.5 KiB", "1.5 KiB");
        t(10 * 1024 - 1, "10 KiB", " 10 KiB");
        t(10 * 1024 + 0, "10 KiB", " 10 KiB");
        t(654 * 1024 + 0, "654 KiB", "654 KiB");
        t(1024 * 1024 - 1, "1.0 MiB", "1.0 MiB");
        t(1024 * 1024 + 0, "1.0 MiB", "1.0 MiB");
        t(1024 * 1024 * 1024 - 1, "1.0 GiB", "1.0 GiB");
        t(2034 * 1024 * 1024 + 0, "2.0 GiB", "2.0 GiB");
        t(100500 * 1024 * 1024 * 1024, "100500 GiB", "100500 GiB");
    }

    #[test]
    fn test_humanized_bytes_per_second() {
        fn t(value: u64, expected: &str, expected_fixed_width: &str) {
            assert_eq!(
                HumanizedBytesPerSecond {
                    bytes_per_second: value,
                    fixed_width: false
                }
                .to_string(),
                expected
            );
            if value <= (1 << 9) {
                assert_eq!(
                    HumanizedBytesPerSecond::FIXED_WIDTH_WIDTH,
                    expected_fixed_width.len()
                );
            }
            assert_eq!(
                HumanizedBytesPerSecond {
                    bytes_per_second: value,
                    fixed_width: true,
                }
                .to_string(),
                expected_fixed_width
            );
        }

        t(0, "0 B/s", "  0 B/s  ");
        t(22, "22 B/s", " 22 B/s  ");
        t(1024 * 1024, "1.0 MiB/s", "1.0 MiB/s");
    }

    #[test]
    fn test_humanized_count() {
        fn t(value: u64, expected: &str, expected_fixed_width: &str) {
            assert_eq!(
                HumanizedCount {
                    count: value,
                    fixed_width: false
                }
                .to_string(),
                expected
            );
            assert_eq!(
                HumanizedCount {
                    count: value,
                    fixed_width: true,
                }
                .to_string(),
                expected_fixed_width
            );
        }

        t(0, "0", "  0  ");
        t(22, "22", " 22  ");
        t(1000000, "1.0 M", "1.0 M");
    }
}


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

use buck2_core::io_counters::IoCounterKey;
use gazebo::prelude::VecExt;
use superconsole::DrawMode;
use superconsole::Line;
use superconsole::Lines;

use crate::humanized::HumanizedBytes;
use crate::two_snapshots::TwoSnapshots;

#[derive(Default)]
pub struct IoState {
    two_snapshots: TwoSnapshots,
}

/// Place space-separated words on lines.
fn words_to_lines(words: Vec<String>, width: usize) -> Vec<String> {
    let mut lines = Vec::new();
    let mut current_line = String::new();
    for word in words {
        if current_line.is_empty() {
            current_line = word;
            continue;
        }
        // This works correctly only for ASCII strings.
        if current_line.len() + 1 + word.len() > width {
            lines.push(current_line);
            current_line = word;
        } else {
            current_line.push(' ');
            current_line.push_str(&word);
        }
    }
    if !current_line.is_empty() {
        lines.push(current_line);
    }
    lines
}

pub fn io_in_flight_non_zero_counters(
    snapshot: &buck2_data::Snapshot,
) -> impl Iterator<Item = (IoCounterKey, u32)> + '_ {
    IoCounterKey::ALL
        .iter()
        .map(|key| {
            let value = match key {
                IoCounterKey::Stat => snapshot.io_in_flight_stat,
                IoCounterKey::Copy => snapshot.io_in_flight_copy,
                IoCounterKey::Symlink => snapshot.io_in_flight_symlink,
                IoCounterKey::Hardlink => snapshot.io_in_flight_hardlink,
                IoCounterKey::MkDir => snapshot.io_in_flight_mk_dir,
                IoCounterKey::ReadDir => snapshot.io_in_flight_read_dir,
                IoCounterKey::ReadDirEden => snapshot.io_in_flight_read_dir_eden,
                IoCounterKey::RmDir => snapshot.io_in_flight_rm_dir,
                IoCounterKey::RmDirAll => snapshot.io_in_flight_rm_dir_all,
                IoCounterKey::StatEden => snapshot.io_in_flight_stat_eden,
                IoCounterKey::Chmod => snapshot.io_in_flight_chmod,
                IoCounterKey::ReadLink => snapshot.io_in_flight_read_link,
                IoCounterKey::Remove => snapshot.io_in_flight_remove,
                IoCounterKey::Rename => snapshot.io_in_flight_rename,
                IoCounterKey::Read => snapshot.io_in_flight_read,
                IoCounterKey::Write => snapshot.io_in_flight_write,
                IoCounterKey::Canonicalize => snapshot.io_in_flight_canonicalize,
                IoCounterKey::EdenSettle => snapshot.io_in_flight_eden_settle,
            };
            (*key, value)
        })
        .filter(|(_, value)| *value > 0)
}

impl IoState {
    pub fn update(&mut self, timestamp: SystemTime, snapshot: &buck2_data::Snapshot) {
        self.two_snapshots.update(timestamp, snapshot);
    }

    fn do_render(&self, snapshot: &buck2_data::Snapshot, width: usize) -> anyhow::Result<Lines> {
        let mut lines = Vec::new();
        let mut parts = Vec::new();
        if let Some(buck2_rss) = snapshot.buck2_rss {
            parts.push(format!("RSS = {}", HumanizedBytes::new(buck2_rss)));
        }
        if let Some(cpu) = self.two_snapshots.cpu_percents() {
            parts.push(format!("CPU = {}%", cpu));
        }
        if snapshot.deferred_materializer_queue_size > 0 {
            parts.push(format!(
                "DM Queue = {}",
                snapshot.deferred_materializer_queue_size
            ));
        }
        if snapshot.blocking_executor_io_queue_size > 0 {
            parts.push(format!(
                "IO Queue = {}",
                snapshot.blocking_executor_io_queue_size
            ));
        }
        if !parts.is_empty() {
            lines.push(Line::from_iter([superconsole::Span::new_unstyled(
                parts.join("  "),
            )?]));
        }

        let mut counters = Vec::new();
        for (key, value) in io_in_flight_non_zero_counters(snapshot) {
            counters.push(format!("{:?} = {}", key, value));
        }
        lines.extend(words_to_lines(counters, width).into_try_map(|s| Line::unstyled(&s))?);

        Ok(Lines(lines))
    }

    pub fn render(
        &self,
        draw_mode: DrawMode,
        width: usize,
        enabled: bool,
    ) -> anyhow::Result<Lines> {
        if !enabled {
            return Ok(Lines::new());
        }
        if let DrawMode::Final = draw_mode {
            return Ok(Lines::new());
        }
        if let Some((_, snapshot)) = &self.two_snapshots.last {
            self.do_render(snapshot, width)
        } else {
            Ok(Lines::new())
        }
    }
}

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

    #[test]
    fn test_words_to_lines() {
        assert_eq!(vec![String::new(); 0], words_to_lines(vec![], 5));
        assert_eq!(
            vec!["ab".to_owned()],
            words_to_lines(vec!["ab".to_owned()], 5)
        );
        assert_eq!(
            vec!["ab cd".to_owned()],
            words_to_lines(vec!["ab".to_owned(), "cd".to_owned()], 5)
        );
        assert_eq!(
            vec!["ab".to_owned(), "cd".to_owned()],
            words_to_lines(vec!["ab".to_owned(), "cd".to_owned()], 4)
        );
        assert_eq!(
            vec!["abcd".to_owned()],
            words_to_lines(vec!["abcd".to_owned()], 3)
        );
    }
}


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

pub enum LastCommandExecutionKind {
    Local,
    LocalWorker,
    Remote,
    Cached,
    NoCommand,
}

/// Returns what the execution kind of the last command was in the given action.
/// It tells the execution kind of the commands that actually produced a result
/// or an error, but not the commands that fell back and were retried.
pub fn get_last_command_execution_kind(
    action: &buck2_data::ActionExecutionEnd,
) -> LastCommandExecutionKind {
    use buck2_data::command_execution_details::Command;

    let last_command = action
        .commands
        .last()
        .and_then(|c| c.details.as_ref())
        .and_then(|c| c.command.as_ref());

    match last_command {
        Some(Command::LocalCommand(..)) | Some(Command::OmittedLocalCommand(..)) => {
            LastCommandExecutionKind::Local
        }
        Some(Command::WorkerCommand(_)) | Some(Command::WorkerInitCommand(_)) => {
            LastCommandExecutionKind::LocalWorker
        }
        Some(Command::RemoteCommand(buck2_data::RemoteCommand {
            cache_hit: true, ..
        })) => LastCommandExecutionKind::Cached,
        Some(Command::RemoteCommand(buck2_data::RemoteCommand {
            cache_hit: false, ..
        })) => LastCommandExecutionKind::Remote,
        None => LastCommandExecutionKind::NoCommand,
    }
}


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

#![feature(try_blocks)]

pub mod action_stats;
pub mod debug_events;
pub mod dice_state;
pub mod display;
pub mod event_observer;
pub mod fmt_duration;
pub mod humanized;
pub mod io_state;
pub mod last_command_execution_kind;
pub mod pending_estimate;
pub mod re_state;
pub mod session_info;
pub mod span_tracker;
pub mod starlark_debug;
pub mod test_state;
pub mod two_snapshots;
pub mod unpack_event;
pub mod verbosity;
pub mod what_ran;


/*
 * 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 crate::dice_state::DiceState;
use crate::span_tracker::Roots;
use crate::span_tracker::SpanTrackable;

/// Estimate how many things are still left to do in a build. This is an approximation since our
/// roots and DICE state are not necessarily entirely in sync.
pub fn pending_estimate<T: SpanTrackable>(roots: &Roots<T>, dice: &DiceState) -> u64 {
    let mut total = 0;
    for k in &["BuildKey", "AnalysisKey"] {
        let from_dice = dice.key_states().get(*k).map_or(0, |v| {
            v.started + v.check_deps_started - v.finished - v.check_deps_finished
        });

        let from_roots = roots.dice_counts().get(k).copied().unwrap_or(0);

        total += u64::from(from_dice).saturating_sub(from_roots);
    }

    total
}