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

use std::time::Duration;
use std::time::Instant;

use anyhow::Context;
use buck2_cli_proto::daemon_api_client::*;
use buck2_cli_proto::*;
use buck2_wrapper_common::kill;
use sysinfo::Pid;
use sysinfo::PidExt;
use sysinfo::ProcessExt;
use sysinfo::ProcessRefreshKind;
use sysinfo::System;
use sysinfo::SystemExt;
use tonic::codegen::InterceptedService;
use tonic::transport::Channel;
use tonic::Request;

use crate::daemon::client::connect::BuckAddAuthTokenInterceptor;

#[derive(Debug, thiserror::Error)]
enum KillError {
    #[error("Daemon pid {} did not die after kill within {:.1}s (status: {})", _0, _1.as_secs_f32(), _2)]
    DidNotDie(u32, Duration, String),
}

const GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(4);
/// Kill request does not wait for the process to exit.
const KILL_REQUEST_TIMEOUT: Duration = Duration::from_secs(3);
const FORCE_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);

pub struct KillResponse {
    pid: u32,
}

impl KillResponse {
    pub fn log(&self) -> anyhow::Result<()> {
        crate::eprintln!("Buck2 daemon pid {} has exited", self.pid)?;
        Ok(())
    }
}

pub async fn kill(
    client: &mut DaemonApiClient<InterceptedService<Channel, BuckAddAuthTokenInterceptor>>,
    info: &DaemonProcessInfo,
    reason: &str,
) -> anyhow::Result<KillResponse> {
    let pid = info.pid;
    let pid: u32 = pid
        .try_into()
        .with_context(|| format!("Integer overflow converting pid {}", pid))?;
    let callers = get_callers_for_kill();

    let request_fut = client.kill(Request::new(KillRequest {
        reason: reason.to_owned(),
        timeout: Some(GRACEFUL_SHUTDOWN_TIMEOUT.try_into()?),
        callers,
    }));
    let time_to_kill = GRACEFUL_SHUTDOWN_TIMEOUT + FORCE_SHUTDOWN_TIMEOUT;
    let time_req_sent = Instant::now();
    // First we send a Kill request
    match tokio::time::timeout(KILL_REQUEST_TIMEOUT, request_fut).await {
        Ok(inner_result) => {
            match inner_result {
                Ok(_) => loop {
                    if !kill::process_exists(pid)? {
                        return Ok(KillResponse { pid });
                    }
                    if time_req_sent.elapsed() > GRACEFUL_SHUTDOWN_TIMEOUT {
                        crate::eprintln!(
                            "Timed out waiting for graceful shutdown of buck2 daemon pid {}",
                            pid
                        )?;
                        break;
                    }
                    tokio::time::sleep(Duration::from_millis(100)).await;
                },
                Err(e) => {
                    // The kill request can fail if the server is in a bad state and we cannot
                    // authenticate to it.
                    crate::eprintln!(
                        "Error requesting graceful shutdown of buck2 daemon pid {}: {}",
                        pid,
                        e
                    )?;
                }
            }
        }
        Err(e) => {
            let _assert_type: tokio::time::error::Elapsed = e;
            crate::eprintln!(
                "Timed out requesting graceful shutdown of buck2 daemon pid {}",
                pid
            )?;
        }
    };

    hard_kill_impl(pid, time_req_sent, time_to_kill).await
}

pub async fn hard_kill(info: &DaemonProcessInfo) -> anyhow::Result<KillResponse> {
    let pid = info
        .pid
        .try_into()
        .with_context(|| format!("Integer overflow converting pid {}", info.pid))?;

    hard_kill_impl(pid, Instant::now(), FORCE_SHUTDOWN_TIMEOUT).await
}

async fn hard_kill_impl(
    pid: u32,
    start_at: Instant,
    deadline: Duration,
) -> anyhow::Result<KillResponse> {
    tracing::info!(
        "Killing PID {} with status {}",
        pid,
        kill::get_sysinfo_status(pid)
            .as_deref()
            .unwrap_or("<unknown>")
    );

    let handle = kill::kill(pid)?;
    let timestamp_after_kill = Instant::now();
    while start_at.elapsed() < deadline {
        if handle.has_exited()? {
            return Ok(KillResponse { pid });
        }
        tokio::time::sleep(Duration::from_millis(100)).await;
    }

    // Last chance: we do logging this time.
    let status = kill::get_sysinfo_status(pid);
    let status = status.unwrap_or_else(|| "<unknown>".to_owned());
    if handle.has_exited()? {
        return Ok(KillResponse { pid });
    }

    Err(KillError::DidNotDie(pid, timestamp_after_kill.elapsed(), status).into())
}

#[cfg(unix)]
mod os_specific {

    use std::time::Duration;

    use sysinfo::Process;
    use sysinfo::ProcessExt;

    pub(super) fn process_creation_time(process: &Process) -> Option<Duration> {
        // Returns process creation time with 1 second precision.
        Some(Duration::from_secs(process.start_time()))
    }
}

#[cfg(windows)]
mod os_specific {
    use std::time::Duration;

    use sysinfo::PidExt;
    use sysinfo::Process;
    use sysinfo::ProcessExt;

    pub(super) fn process_creation_time(process: &Process) -> Option<Duration> {
        buck2_wrapper_common::kill::os_specific::process_creation_time(process.pid().as_u32())
    }
}

fn get_callers_for_kill() -> Vec<String> {
    /// Add a process to our parts and return its parent PID.
    fn push_process(
        pid: Pid,
        creation_time: Duration,
        system: &mut System,
        process_tree: &mut Vec<String>,
    ) -> Option<(Pid, Duration)> {
        // Specifics about this process need to be refreshed by this time.
        let proc = system.process(pid)?;
        let title = shlex::join(proc.cmd().iter().map(|s| s.as_str()));
        process_tree.push(title);
        let parent_pid = proc.parent()?;
        system.refresh_process_specifics(parent_pid, ProcessRefreshKind::new());
        let parent_proc = system.process(parent_pid)?;
        let parent_creation_time = os_specific::process_creation_time(parent_proc)?;
        if parent_creation_time <= creation_time {
            Some((parent_pid, parent_creation_time))
        } else {
            None
        }
    }

    let mut system = System::new();
    let mut process_tree = Vec::new();

    let pid = Pid::from_u32(std::process::id());
    system.refresh_process_specifics(pid, ProcessRefreshKind::new());
    let mut curr = system
        .process(pid)
        .and_then(|proc| Some((pid, os_specific::process_creation_time(proc)?)));
    while let Some((pid, creation_time)) = curr {
        curr = push_process(pid, creation_time, &mut system, &mut process_tree);
    }

    process_tree.into_iter().rev().collect()
}


/*
 * 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::fs::create_dir_all;
use std::fs::File;
use std::time::Duration;

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::daemon_api_client::*;
use buck2_cli_proto::*;
use buck2_common::daemon_dir::DaemonDir;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use fs4::FileExt;
use futures::future::BoxFuture;
use futures::pin_mut;
use futures::stream;
use futures::Stream;
use futures::StreamExt;
use futures::TryStreamExt;
use tonic::codegen::InterceptedService;
use tonic::transport::Channel;
use tonic::Request;
use tonic::Status;

use crate::command_outcome::CommandOutcome;
use crate::console_interaction_stream::ConsoleInteractionStream;
use crate::daemon::client::connect::BuckAddAuthTokenInterceptor;
use crate::events_ctx::EventsCtx;
use crate::events_ctx::FileTailers;
use crate::events_ctx::PartialResultCtx;
use crate::events_ctx::PartialResultHandler;
use crate::stream_value::StreamValue;
use crate::subscribers::observer::ErrorCause;
use crate::subscribers::observer::ErrorObserver;

pub mod connect;
pub mod kill;

use crate::startup_deadline::StartupDeadline;

#[derive(Debug, thiserror::Error)]
enum LifecycleError {
    #[error("Missing `{}` file in `{}` directory", BuckdLifecycleLock::BUCKD_LIFECYCLE, _0.display())]
    MissingLifecycle(AbsNormPathBuf),
}

/// We need to make sure that all calls to the daemon in buckd flush the tailers after completion.
/// The connector wraps all buckd calls with flushing.
pub struct BuckdClientConnector<'a> {
    client: BuckdClient<'a>,
}

impl<'a> BuckdClientConnector<'a> {
    pub fn with_flushing(&mut self) -> FlushingBuckdClient<'_, 'a> {
        FlushingBuckdClient {
            inner: &mut self.client,
        }
    }

    pub fn daemon_constraints(&self) -> &buck2_cli_proto::DaemonConstraints {
        &self.client.constraints
    }

    pub async fn kill(&mut self, reason: &str) -> anyhow::Result<kill::KillResponse> {
        kill::kill(&mut self.client.client, &self.client.info, reason).await
    }

    pub fn error_observers(&self) -> impl Iterator<Item = &dyn ErrorObserver> {
        self.client
            .events_ctx
            .subscribers
            .iter()
            .filter_map(|s| s.as_error_observer())
    }

    pub fn collect_error_cause(&self) -> ErrorCause {
        if let Some(obs) = self.error_observers().next() {
            return obs.error_cause();
        }

        ErrorCause::Unknown
    }
}

pub struct BuckdLifecycleLock {
    daemon_dir: DaemonDir,
    lock_file: File,
}

impl BuckdLifecycleLock {
    const BUCKD_LIFECYCLE: &'static str = "buckd.lifecycle";

    pub async fn lock_with_timeout(
        daemon_dir: DaemonDir,
        deadline: StartupDeadline,
    ) -> anyhow::Result<BuckdLifecycleLock> {
        create_dir_all(&daemon_dir.path)?;
        let lifecycle_path = daemon_dir.path.as_path().join(Self::BUCKD_LIFECYCLE);
        let file = File::create(lifecycle_path)?;
        deadline
            .retrying(
                "locking buckd lifecycle",
                Duration::from_millis(5),
                Duration::from_millis(100),
                async || Ok(file.try_lock_exclusive()?),
            )
            .await?;

        Ok(BuckdLifecycleLock {
            lock_file: file,
            daemon_dir,
        })
    }

    /// Remove everything except `buckd.lifecycle` file which is the lock file.
    pub fn clean_daemon_dir(&self) -> anyhow::Result<()> {
        let mut seen_lifecycle = false;
        for p in fs_util::read_dir(&self.daemon_dir.path)? {
            let p = p?;
            if p.file_name() == Self::BUCKD_LIFECYCLE {
                seen_lifecycle = true;
                continue;
            }
            fs_util::remove_all(p.path())?;
        }
        if !seen_lifecycle {
            // Self-check.
            return Err(LifecycleError::MissingLifecycle(self.daemon_dir.path.clone()).into());
        }
        Ok(())
    }

    pub fn daemon_dir(&self) -> &DaemonDir {
        &self.daemon_dir
    }
}

impl Drop for BuckdLifecycleLock {
    fn drop(&mut self) {
        self.lock_file
            .unlock()
            .expect("Unexpected failure to unlock buckd.lifecycle file.")
    }
}

/// This provides a thin wrapper around the proto-generated DaemonApiClient and hides
/// some of the complexity/verbosity of making calls with that. For example, the user
/// doesn't need to deal with tonic::Response/Request and this may provide functions
/// that take more primitive types than the protobuf structure itself.
pub struct BuckdClient<'a> {
    client: DaemonApiClient<InterceptedService<Channel, BuckAddAuthTokenInterceptor>>,
    constraints: buck2_cli_proto::DaemonConstraints,
    info: DaemonProcessInfo,
    daemon_dir: DaemonDir,
    // TODO(brasselsprouts): events_ctx should own tailers
    tailers: Option<FileTailers>,
    pub(crate) events_ctx: EventsCtx<'a>,
}

#[derive(Debug, thiserror::Error)]
enum GrpcToStreamError {
    #[error("buck daemon returned an empty CommandProgress")]
    EmptyCommandProgress,
}

/// Translates a tonic streaming response into a stream of StreamValues, the set of things that can flow across the gRPC
/// event stream.
fn grpc_to_stream(
    response: anyhow::Result<tonic::Response<tonic::Streaming<MultiCommandProgress>>>,
) -> impl Stream<Item = anyhow::Result<StreamValue>> {
    let stream = match response {
        Ok(response) => response.into_inner(),
        Err(e) => return futures::stream::once(futures::future::ready(Err(e))).left_stream(),
    };

    let stream = stream
        .map_ok(|e| stream::iter(e.messages.into_iter().map(anyhow::Ok)))
        .try_flatten();

    stream::unfold(stream, |mut stream| async {
        let msg = match stream.try_next().await {
            Ok(Some(msg)) => msg,
            Ok(None) => return None,
            Err(e) => return Some((Err(e), stream)),
        };
        let value = match msg.progress {
            Some(command_progress::Progress::Event(e)) => Some(Ok(StreamValue::Event(e))),
            Some(command_progress::Progress::Result(res)) => Some(Ok(StreamValue::Result(res))),
            Some(command_progress::Progress::PartialResult(res)) => {
                Some(Ok(StreamValue::PartialResult(res)))
            }
            None => Some(Err(GrpcToStreamError::EmptyCommandProgress.into())),
        };

        value.map(|v| (v, stream))
    })
    .right_stream()
}

impl<'a> BuckdClient<'a> {
    fn open_tailers(&mut self) -> anyhow::Result<()> {
        let tailers = FileTailers::new(&self.daemon_dir)?;
        self.tailers = Some(tailers);

        Ok(())
    }

    /// Some commands stream events back from the server.
    /// For these commands, we want to be able to manipulate CLI state.
    async fn stream<'i, T, Res, Handler, Command>(
        &mut self,
        command: Command,
        request: T,
        partial_result_handler: &mut Handler,
        console_interaction: Option<ConsoleInteractionStream<'i>>,
    ) -> anyhow::Result<CommandOutcome<Res>>
    where
        Command: for<'b> FnOnce(
            &'b mut DaemonApiClient<InterceptedService<Channel, BuckAddAuthTokenInterceptor>>,
            Request<T>,
        ) -> BoxFuture<
            'b,
            Result<tonic::Response<tonic::Streaming<MultiCommandProgress>>, Status>,
        >,
        Res: TryFrom<command_result::Result, Error = command_result::Result>,
        Handler: PartialResultHandler,
    {
        let Self {
            client, events_ctx, ..
        } = self;

        let response = command(client, Request::new(request))
            .await
            .context("Error dispatching request");
        let stream = grpc_to_stream(response);
        pin_mut!(stream);
        events_ctx
            .unpack_stream(
                partial_result_handler,
                stream,
                self.tailers.take(),
                console_interaction,
            )
            .await
    }

    pub async fn status(&mut self, snapshot: bool) -> anyhow::Result<StatusResponse> {
        let outcome = self
            .events_ctx
            // Safe to unwrap tailers here because they are instantiated prior to a command being called.
            .unpack_oneshot(&mut self.tailers, || {
                self.client.status(Request::new(StatusRequest { snapshot }))
            })
            .await;
        // TODO(nmj): We have a number of things that wish to use status() and return an anyhow::Result,
        // for now we'll just turn a "CommandMessage" into a error, but that's really not what we
        // want long term.
        match outcome? {
            CommandOutcome::Success(r) => Ok(r),
            CommandOutcome::Failure(_) => {
                Err(anyhow::anyhow!("Unexpected failure message in status()"))
            }
        }
    }

    pub async fn set_log_filter(&mut self, req: SetLogFilterRequest) -> anyhow::Result<()> {
        self.client.set_log_filter(Request::new(req)).await?;

        Ok(())
    }
}

pub struct FlushingBuckdClient<'a, 'b> {
    inner: &'a mut BuckdClient<'b>,
}

impl<'a, 'b> FlushingBuckdClient<'a, 'b> {
    fn enter(&mut self) -> anyhow::Result<()> {
        self.inner.open_tailers()?;
        Ok(())
    }

    async fn exit(&mut self) -> anyhow::Result<()> {
        self.inner.events_ctx.flush(&mut self.inner.tailers).await?;

        Ok(())
    }
}

pub enum NoPartialResult {}

impl TryFrom<buck2_cli_proto::partial_result::PartialResult> for NoPartialResult {
    type Error = buck2_cli_proto::partial_result::PartialResult;

    fn try_from(v: buck2_cli_proto::partial_result::PartialResult) -> Result<Self, Self::Error> {
        Err(v)
    }
}

pub struct NoPartialResultHandler;

#[async_trait]
impl PartialResultHandler for NoPartialResultHandler {
    type PartialResult = NoPartialResult;

    async fn handle_partial_result(
        &mut self,
        _ctx: PartialResultCtx<'_, '_>,
        partial_res: Self::PartialResult,
    ) -> anyhow::Result<()> {
        match partial_res {}
    }
}

/// Receives StdoutBytes, writes them to stdout.
pub struct StdoutPartialResultHandler;

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

    async fn handle_partial_result(
        &mut self,
        mut ctx: PartialResultCtx<'_, '_>,
        partial_res: Self::PartialResult,
    ) -> anyhow::Result<()> {
        ctx.stdout(&partial_res.data).await
    }
}

/// Implement a streaming method with full event reporting.
macro_rules! stream_method {
    ($method: ident, $req: ty, $res: ty, $message: ty) => {
        stream_method!($method, $method, $req, $res, $message);
    };

    ($method: ident, $grpc_method: ident, $req: ty, $res: ty, $message: ty) => {
        pub async fn $method(
            &mut self,
            req: $req,
            console_interaction: Option<ConsoleInteractionStream<'_>>,
            handler: &mut impl PartialResultHandler<PartialResult = $message>,
        ) -> anyhow::Result<CommandOutcome<$res>> {
            self.enter()?;
            let res = self
                .inner
                .stream(
                    |d, r| Box::pin(DaemonApiClient::$grpc_method(d, r)),
                    req,
                    // For now we only support handlers that can be constructed like so, and we
                    // don't let anything go out. Eventually if we wanted to stream structured
                    // data, that could change.
                    handler,
                    console_interaction,
                )
                .await;
            self.exit().await?;
            res
        }
    };
}

/// Implement a bi-directional streaming method with full event reporting.
macro_rules! bidirectional_stream_method {
    ($method: ident, $req: ty, $res: ty, $message: ty) => {
        bidirectional_stream_method!($method, $method, $req, $res, $message);
    };

    ($method: ident, $grpc_method: ident, $req: ty, $res: ty, $message: ty) => {
        pub async fn $method(
            &mut self,
            context: ClientContext,
            requests: impl Stream<Item = $req> + Send + Sync + 'static,
            handler: &mut impl PartialResultHandler<PartialResult = $message>,
        ) -> anyhow::Result<CommandOutcome<$res>> {
            self.enter()?;
            let req = create_client_stream(context, requests);
            let res = self
                .inner
                .stream(
                    |d, r| Box::pin(DaemonApiClient::$method(d, r)),
                    req,
                    handler,
                    None,
                )
                .await;
            self.exit().await?;
            res
        }
    };
}

/// Implement a oneshot method with full event reporting.
macro_rules! oneshot_method {
    ($method: ident, $req: ty, $res: ty) => {
        oneshot_method!($method, $method, $req, $res);
    };

    ($method: ident, $grpc_method: ident, $req: ty, $res: ty) => {
        pub async fn $method(&mut self, req: $req) -> anyhow::Result<CommandOutcome<$res>> {
            self.enter()?;
            let res = self
                .inner
                .events_ctx
                .unpack_oneshot(&mut self.inner.tailers, || {
                    self.inner.client.$method(Request::new(req))
                })
                .await;
            self.exit().await?;
            res
        }
    };
}

/// Implement a method that does not produce a CommandResult and does not produce any events.
macro_rules! debug_method {
    ($method: ident, $req: ty, $res: ty) => {
        debug_method!($method, $method, $req, $res);
    };

    ($method: ident, $grpc_method: ident, $req: ty, $res: ty) => {
        pub async fn $method(&mut self, req: $req) -> anyhow::Result<$res> {
            self.enter()?;
            let out = self.inner.client.$method(Request::new(req)).await;
            self.exit().await?;
            Ok(out?.into_inner())
        }
    };
}

/// Wrap a method that exists on the BuckdClient, with flushing.
macro_rules! wrap_method {
     ($method: ident ($($param: ident : $param_type: ty),*), $res: ty) => {
         pub async fn $method(&mut self, $($param: $param_type)*) -> anyhow::Result<$res> {
             self.enter()?;
             let out = self
                 .inner
                 .$method($($param)*)
                 .await;
             self.exit().await?;
             out
         }
     };
 }

impl<'a, 'b> FlushingBuckdClient<'a, 'b> {
    stream_method!(
        aquery,
        AqueryRequest,
        AqueryResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        cquery,
        CqueryRequest,
        CqueryResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        uquery,
        UqueryRequest,
        UqueryResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        targets,
        TargetsRequest,
        TargetsResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        targets_show_outputs,
        TargetsRequest,
        TargetsShowOutputsResponse,
        NoPartialResult
    );
    stream_method!(
        ctargets,
        ConfiguredTargetsRequest,
        ConfiguredTargetsResponse,
        NoPartialResult
    );
    stream_method!(build, BuildRequest, BuildResponse, NoPartialResult);
    stream_method!(bxl, BxlRequest, BxlResponse, buck2_cli_proto::StdoutBytes);
    stream_method!(test, TestRequest, TestResponse, NoPartialResult);
    stream_method!(install, InstallRequest, InstallResponse, NoPartialResult);
    stream_method!(
        audit,
        GenericRequest,
        GenericResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        starlark,
        GenericRequest,
        GenericResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        materialize,
        MaterializeRequest,
        MaterializeResponse,
        NoPartialResult
    );
    stream_method!(
        clean_stale,
        CleanStaleRequest,
        CleanStaleResponse,
        NoPartialResult
    );
    stream_method!(
        file_status,
        FileStatusRequest,
        GenericResponse,
        buck2_cli_proto::StdoutBytes
    );
    stream_method!(
        unstable_docs,
        UnstableDocsRequest,
        UnstableDocsResponse,
        NoPartialResult
    );
    stream_method!(
        profile,
        profile2,
        ProfileRequest,
        ProfileResponse,
        NoPartialResult
    );
    stream_method!(
        allocative,
        AllocativeRequest,
        AllocativeResponse,
        NoPartialResult
    );

    bidirectional_stream_method!(lsp, LspRequest, LspResponse, LspMessage);
    bidirectional_stream_method!(dap, DapRequest, DapResponse, DapMessage);
    bidirectional_stream_method!(
        subscription,
        SubscriptionRequestWrapper,
        SubscriptionCommandResponse,
        SubscriptionResponseWrapper
    );

    oneshot_method!(flush_dep_files, FlushDepFilesRequest, GenericResponse);

    debug_method!(unstable_crash, UnstableCrashRequest, UnstableCrashResponse);
    debug_method!(segfault, SegfaultRequest, SegfaultResponse);
    debug_method!(
        unstable_heap_dump,
        UnstableHeapDumpRequest,
        UnstableHeapDumpResponse
    );
    debug_method!(
        unstable_allocator_stats,
        UnstableAllocatorStatsRequest,
        UnstableAllocatorStatsResponse
    );
    debug_method!(
        unstable_dice_dump,
        UnstableDiceDumpRequest,
        UnstableDiceDumpResponse
    );

    wrap_method!(status(snapshot: bool), StatusResponse);
    wrap_method!(set_log_filter(log_filter: SetLogFilterRequest), ());
    stream_method!(trace_io, TraceIoRequest, TraceIoResponse, NoPartialResult);
}

/// Create a stream that is sent over as a parameter via GRPC to the daemon.
///
/// Ensures that we send a proper ClientContext message, and that the inner type is wrapped
/// properly into a [`StreamingRequest`]
fn create_client_stream<
    T: Into<StreamingRequest>,
    InStream: Stream<Item = T> + Send + Sync + 'static,
>(
    context: ClientContext,
    requests: InStream,
) -> impl Stream<Item = StreamingRequest> + Send + Sync + 'static {
    let init_req = StreamingRequest {
        request: Some(streaming_request::Request::Context(context)),
    };
    stream::once(async move { init_req }).chain(requests.map(|request| request.into()))
}

#[cfg(test)]
mod tests {
    use std::io::Write;

    use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;

    use super::*;
    use crate::events_ctx::FileTailerEvent;
    use crate::file_tailer::FileTailer;
    use crate::file_tailer::StdoutOrStderr;

    #[tokio::test]
    async fn test_tailer() -> anyhow::Result<()> {
        let mut file = tempfile::NamedTempFile::new()?;
        file.write_all(b"before\n")?;

        // If we could control the interval for tailer polling, we could reliably
        // test more of the behavior. For now, just test a simple case.
        let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel();
        let tailer = FileTailer::tail_file(
            AbsNormPathBuf::new(file.path().to_owned())?,
            sender,
            StdoutOrStderr::Stdout,
        )?;

        let ok_line = b"after\n";
        let invalid_utf8_line = b"\xc3\x28\n";

        file.write_all(ok_line.as_slice())?;
        file.write_all(invalid_utf8_line.as_slice())?;

        // have to sleep long enough for a read or else this test is racy.
        tokio::time::sleep(Duration::from_millis(250)).await;
        std::mem::drop(tailer);
        assert_eq!(
            FileTailerEvent::Stdout((*ok_line).into()),
            receiver.recv().await.unwrap()
        );
        assert_eq!(
            FileTailerEvent::Stdout((*invalid_utf8_line).into()),
            receiver.recv().await.unwrap()
        );

        Ok(())
    }
}


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

use std::sync::Arc;

use async_trait::async_trait;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_events::BuckEvent;

use crate::subscribers::subscriber::EventSubscriber;

pub struct BuildIdWriter {
    path: AbsPathBuf,
}

impl BuildIdWriter {
    pub fn new(path: AbsPathBuf) -> Self {
        Self { path }
    }
}

#[async_trait]
impl EventSubscriber for BuildIdWriter {
    async fn handle_events(&mut self, events: &[Arc<BuckEvent>]) -> anyhow::Result<()> {
        for event in events {
            if event.command_start()?.is_some() {
                tokio::fs::write(&self.path, event.trace_id()?.to_string()).await?;
            }
        }
        Ok(())
    }
}


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

use std::sync::atomic::AtomicU64;
use std::sync::Arc;

use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_event_observer::event_observer::NoopEventObserverExtra;
use buck2_event_observer::verbosity::Verbosity;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;

use crate::argv::SanitizedArgv;
use crate::client_ctx::ClientCommandContext;
use crate::common::CommonDaemonCommandOptions;
use crate::common::ConsoleType;
use crate::path_arg::PathArg;
use crate::subscribers::build_id_writer::BuildIdWriter;
use crate::subscribers::event_log::subscriber::EventLog;
use crate::subscribers::re_log::ReLog;
use crate::subscribers::simpleconsole::SimpleConsole;
use crate::subscribers::subscriber::EventSubscriber;
use crate::subscribers::subscriber_unpack::UnpackingEventSubscriberAsEventSubscriber;
use crate::subscribers::superconsole::StatefulSuperConsole;
use crate::subscribers::superconsole::SuperConsoleConfig;

/// Given a command name and the command arguments, create a default console / superconsole.
pub fn get_console_with_root(
    trace_id: TraceId,
    console_type: ConsoleType,
    verbosity: Verbosity,
    show_waiting_message: bool,
    replay_speed: Option<f64>,
    command_name: &str,
    config: SuperConsoleConfig,
    isolation_dir: FileNameBuf,
) -> anyhow::Result<Option<Box<dyn EventSubscriber>>> {
    match console_type {
        ConsoleType::Simple => Ok(Some(Box::new(UnpackingEventSubscriberAsEventSubscriber(
            SimpleConsole::<NoopEventObserverExtra>::autodetect(
                trace_id,
                isolation_dir,
                verbosity,
                show_waiting_message,
            ),
        )))),
        ConsoleType::SimpleNoTty => Ok(Some(Box::new(UnpackingEventSubscriberAsEventSubscriber(
            SimpleConsole::<NoopEventObserverExtra>::without_tty(
                trace_id,
                isolation_dir,
                verbosity,
                show_waiting_message,
            ),
        )))),
        ConsoleType::SimpleTty => Ok(Some(Box::new(UnpackingEventSubscriberAsEventSubscriber(
            SimpleConsole::<NoopEventObserverExtra>::with_tty(
                trace_id,
                isolation_dir,
                verbosity,
                show_waiting_message,
            ),
        )))),
        ConsoleType::Super => Ok(Some(Box::new(UnpackingEventSubscriberAsEventSubscriber(
            StatefulSuperConsole::new_with_root_forced(
                trace_id,
                command_name,
                verbosity,
                show_waiting_message,
                replay_speed,
                None,
                config,
                isolation_dir,
            )?,
        )))),
        ConsoleType::Auto => {
            match StatefulSuperConsole::new_with_root(
                trace_id.dupe(),
                command_name,
                verbosity,
                show_waiting_message,
                replay_speed,
                config,
                isolation_dir.clone(),
            )? {
                Some(super_console) => Ok(Some(Box::new(
                    UnpackingEventSubscriberAsEventSubscriber(super_console),
                ))),
                None => Ok(Some(Box::new(UnpackingEventSubscriberAsEventSubscriber(
                    SimpleConsole::<NoopEventObserverExtra>::autodetect(
                        trace_id,
                        isolation_dir,
                        verbosity,
                        show_waiting_message,
                    ),
                )))),
            }
        }
        ConsoleType::None => Ok(None),
    }
}

/// Given the command arguments, conditionally create an event log.
pub(crate) fn try_get_event_log_subscriber<'a>(
    event_log_opts: &CommonDaemonCommandOptions,
    sanitized_argv: SanitizedArgv,
    ctx: &ClientCommandContext<'a>,
    log_size_counter_bytes: Option<Arc<AtomicU64>>,
    user_event_log: &Option<PathArg>,
) -> anyhow::Result<Option<Box<dyn EventSubscriber + 'a>>> {
    if event_log_opts.no_event_log {
        return Ok(None);
    }
    let logdir = ctx.paths()?.log_dir();
    let log = EventLog::new(
        logdir,
        ctx.working_dir.clone(),
        event_log_opts
            .event_log
            .as_ref()
            .map(|p| p.resolve(&ctx.working_dir)),
        user_event_log.as_ref().map(|p| p.resolve(&ctx.working_dir)),
        sanitized_argv,
        ctx.async_cleanup_context().dupe(),
        ctx.command_name.clone(),
        log_size_counter_bytes,
    )?;
    Ok(Some(Box::new(log)))
}

pub(crate) fn try_get_re_log_subscriber<'a>(
    ctx: &ClientCommandContext<'a>,
) -> anyhow::Result<Option<Box<dyn EventSubscriber + 'a>>> {
    let log = ReLog::new(
        ctx.paths()?.isolation.clone(),
        ctx.async_cleanup_context().dupe(),
    );
    Ok(Some(Box::new(log)))
}

pub(crate) fn try_get_build_id_writer<'a>(
    opts: &CommonDaemonCommandOptions,
    ctx: &ClientCommandContext<'a>,
) -> anyhow::Result<Option<Box<dyn EventSubscriber + 'a>>> {
    if let Some(file_loc) = opts.write_build_id.as_ref() {
        Ok(Some(Box::new(BuildIdWriter::new(
            file_loc.resolve(&ctx.working_dir),
        ))))
    } else {
        Ok(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 buck2_core::env_helper::EnvHelper;

pub(crate) mod build_id_writer;
pub mod event_log;
pub mod get;
pub(crate) mod observer;
pub mod re_log;
pub mod recorder;
pub(crate) mod simpleconsole;
pub mod stdout_stderr_forwarder;
pub mod subscriber;
pub mod subscriber_unpack;
pub mod superconsole;

pub fn should_upload_log() -> anyhow::Result<bool> {
    if buck2_core::is_open_source() {
        return Ok(false);
    }
    static DISABLE_LOG_UPLOAD: EnvHelper<bool> = EnvHelper::new("BUCK2_TEST_DISABLE_LOG_UPLOAD");
    Ok(!DISABLE_LOG_UPLOAD.get()?.copied().unwrap_or_default())
}


/*
 * 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 trait for such event subscribers that are watching a specific set of
/// errors and keeping the record of them for later use.
pub trait ErrorObserver {
    fn error_cause(&self) -> ErrorCause {
        ErrorCause::Unknown
    }

    /// Whether this observer thinks that the daemon needs killing to work again.
    fn daemon_in_memory_state_is_corrupted(&self) -> bool {
        false
    }

    /// Whether this observer thinks that the daemon needs to dump its materializer state to work
    /// again.
    fn daemon_materializer_state_is_corrupted(&self) -> bool {
        false
    }

    fn restarter_is_enabled(&self) -> bool {
        false
    }
}

pub enum ErrorCause {
    Unknown,
    Infra,
    User,
    DaemonIsBusy,
}


/*
 * 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::atomic::AtomicU64;
use std::sync::Arc;

use dupe::Dupe;

use crate::build_count::BuildCountManager;
use crate::client_ctx::ClientCommandContext;
use crate::common::CommonDaemonCommandOptions;

mod imp {
    use std::cmp;
    use std::collections::HashMap;
    use std::collections::HashSet;
    use std::future::Future;
    use std::io::Write;
    use std::sync::atomic::AtomicU64;
    use std::sync::atomic::Ordering;
    use std::sync::Arc;
    use std::time::Duration;
    use std::time::Instant;
    use std::time::SystemTime;

    use anyhow::Context;
    use async_trait::async_trait;
    use buck2_common::convert::ProstDurationExt;
    use buck2_core::fs::fs_util;
    use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
    use buck2_core::fs::paths::abs_path::AbsPathBuf;
    use buck2_event_observer::action_stats;
    use buck2_event_observer::last_command_execution_kind;
    use buck2_event_observer::last_command_execution_kind::LastCommandExecutionKind;
    use buck2_events::sink::scribe::new_thrift_scribe_sink_if_enabled;
    use buck2_events::BuckEvent;
    use buck2_wrapper_common::invocation_id::TraceId;
    use dupe::Dupe;
    use fbinit::FacebookInit;
    use futures::FutureExt;
    use termwiz::istty::IsTty;

    use crate::build_count::BuildCountManager;
    use crate::cleanup_ctx::AsyncCleanupContext;
    use crate::subscribers::observer::ErrorCause;
    use crate::subscribers::observer::ErrorObserver;
    use crate::subscribers::recorder::system_memory_stats;
    use crate::subscribers::subscriber::EventSubscriber;

    pub struct InvocationRecorder<'a> {
        fb: FacebookInit,
        write_to_path: Option<AbsPathBuf>,
        command_name: &'static str,
        cli_args: Vec<String>,
        isolation_dir: String,
        start_time: Instant,
        async_cleanup_context: AsyncCleanupContext<'a>,
        build_count_manager: BuildCountManager,
        trace_id: TraceId,
        command_start: Option<buck2_data::CommandStart>,
        command_end: Option<buck2_data::CommandEnd>,
        command_critical_start: Option<buck2_data::CommandCriticalStart>,
        command_critical_end: Option<buck2_data::CommandCriticalEnd>,
        command_duration: Option<prost_types::Duration>,
        re_session_id: Option<String>,
        re_experiment_name: Option<String>,
        critical_path_duration: Option<Duration>,
        tags: Vec<String>,
        run_local_count: u64,
        run_remote_count: u64,
        run_action_cache_count: u64,
        run_skipped_count: u64,
        run_fallback_count: u64,
        local_actions_executed_via_worker: u64,
        first_snapshot: Option<buck2_data::Snapshot>,
        last_snapshot: Option<buck2_data::Snapshot>,
        min_build_count_since_rebase: u64,
        cache_upload_count: u64,
        cache_upload_attempt_count: u64,
        parsed_target_patterns: Option<buck2_data::ParsedTargetPatterns>,
        #[allow(dead_code)]
        invocation_root_path: AbsNormPathBuf,
        filesystem: Option<String>,
        watchman_version: Option<String>,
        eden_version: Option<String>,
        test_info: Option<String>,
        eligible_for_full_hybrid: bool,
        max_event_client_delay: Option<Duration>,
        max_malloc_bytes_active: Option<u64>,
        max_malloc_bytes_allocated: Option<u64>,
        run_command_failure_count: u64,
        event_count: u64,
        time_to_first_action_execution: Option<Duration>,
        materialization_output_size: u64,
        initial_materializer_entries_from_sqlite: Option<u64>,
        time_to_command_start: Option<Duration>,
        time_to_command_critical_section: Option<Duration>,
        time_to_first_analysis: Option<Duration>,
        time_to_load_first_build_file: Option<Duration>,
        time_to_first_command_execution_start: Option<Duration>,
        system_total_memory_bytes: Option<u64>,
        file_watcher_stats: Option<buck2_data::FileWatcherStats>,
        time_to_last_action_execution_end: Option<Duration>,
        initial_sink_success_count: Option<u64>,
        initial_sink_failure_count: Option<u64>,
        initial_sink_dropped_count: Option<u64>,
        sink_max_buffer_depth: u64,
        soft_error_categories: HashSet<String>,
        concurrent_command_blocking_duration: Option<prost_types::Duration>,
        metadata: HashMap<String, String>,
        analysis_count: u64,
        total_concurrent_commands: Option<u32>,
        exit_when_different_state: bool,
        daemon_in_memory_state_is_corrupted: bool,
        daemon_materializer_state_is_corrupted: bool,
        enable_restarter: bool,
        restarted_trace_id: Option<TraceId>,
        has_command_result: bool,
        has_end_of_stream: bool,
        compressed_event_log_size_bytes: Option<Arc<AtomicU64>>,
        critical_path_backend: Option<String>,
        instant_command_is_success: Option<bool>,
        bxl_ensure_artifacts_duration: Option<prost_types::Duration>,
        initial_re_upload_bytes: Option<u64>,
        initial_re_download_bytes: Option<u64>,
    }

    impl<'a> InvocationRecorder<'a> {
        pub fn new(
            fb: FacebookInit,
            async_cleanup_context: AsyncCleanupContext<'a>,
            write_to_path: Option<AbsPathBuf>,
            mut command_name: &'static str,
            sanitized_argv: Vec<String>,
            trace_id: TraceId,
            isolation_dir: String,
            build_count_manager: BuildCountManager,
            invocation_root_path: AbsNormPathBuf,
            restarted_trace_id: Option<TraceId>,
            log_size_counter_bytes: Option<Arc<AtomicU64>>,
        ) -> Self {
            // FIXME: Figure out if we can replace this. We used to log this this way in Ingress :/
            if command_name == "uquery" {
                command_name = "query";
            }

            Self {
                fb,
                write_to_path,
                command_name,
                cli_args: sanitized_argv,
                isolation_dir,
                start_time: Instant::now(),
                async_cleanup_context,
                build_count_manager,
                trace_id,
                command_start: None,
                command_end: None,
                command_critical_start: None,
                command_critical_end: None,
                command_duration: None,
                re_session_id: None,
                re_experiment_name: None,
                critical_path_duration: None,
                tags: vec![],
                run_local_count: 0,
                run_remote_count: 0,
                run_action_cache_count: 0,
                run_skipped_count: 0,
                run_fallback_count: 0,
                local_actions_executed_via_worker: 0,
                first_snapshot: None,
                last_snapshot: None,
                min_build_count_since_rebase: 0,
                cache_upload_count: 0,
                cache_upload_attempt_count: 0,
                parsed_target_patterns: None,
                invocation_root_path,
                filesystem: None,
                watchman_version: None,
                eden_version: None,
                test_info: None,
                eligible_for_full_hybrid: false,
                max_event_client_delay: None,
                max_malloc_bytes_active: None,
                max_malloc_bytes_allocated: None,
                run_command_failure_count: 0,
                event_count: 0,
                time_to_first_action_execution: None,
                materialization_output_size: 0,
                initial_materializer_entries_from_sqlite: None,
                time_to_command_start: None,
                time_to_command_critical_section: None,
                time_to_first_analysis: None,
                time_to_load_first_build_file: None,
                time_to_first_command_execution_start: None,
                system_total_memory_bytes: Some(system_memory_stats()),
                file_watcher_stats: None,
                time_to_last_action_execution_end: None,
                initial_sink_success_count: None,
                initial_sink_failure_count: None,
                initial_sink_dropped_count: None,
                sink_max_buffer_depth: 0,
                soft_error_categories: HashSet::new(),
                concurrent_command_blocking_duration: None,
                metadata: buck2_events::metadata::collect(),
                analysis_count: 0,
                total_concurrent_commands: None,
                exit_when_different_state: false,
                daemon_in_memory_state_is_corrupted: false,
                daemon_materializer_state_is_corrupted: false,
                enable_restarter: false,
                restarted_trace_id,
                has_command_result: false,
                has_end_of_stream: false,
                compressed_event_log_size_bytes: log_size_counter_bytes,
                critical_path_backend: None,
                instant_command_is_success: None,
                bxl_ensure_artifacts_duration: None,
                initial_re_upload_bytes: None,
                initial_re_download_bytes: None,
            }
        }

        pub fn instant_command_outcome(&mut self, is_success: bool) {
            self.instant_command_is_success = Some(is_success);
        }

        async fn build_count(
            &mut self,
            target_patterns: &[buck2_data::TargetPattern],
        ) -> anyhow::Result<u64> {
            if let Some(stats) = &self.file_watcher_stats {
                if let Some(merge_base) = &stats.branched_from_revision {
                    return self
                        .build_count_manager
                        .min_build_count(
                            merge_base,
                            self.parsed_target_patterns
                                .as_ref()
                                .map_or(target_patterns, |d| &d.target_patterns[..]),
                        )
                        .await
                        .context("Error recording build count");
                }
            }

            Ok(0)
        }

        fn send_it(&mut self) -> Option<impl Future<Output = ()> + 'static + Send> {
            let mut sink_success_count = None;
            let mut sink_failure_count = None;
            let mut sink_dropped_count = None;
            let mut re_upload_bytes = None;
            let mut re_download_bytes = None;
            if let Some(snapshot) = &self.last_snapshot {
                sink_success_count = calculate_diff_if_some(
                    &snapshot.sink_successes,
                    &self.initial_sink_success_count,
                );
                sink_failure_count = calculate_diff_if_some(
                    &snapshot.sink_failures,
                    &self.initial_sink_failure_count,
                );
                sink_dropped_count = calculate_diff_if_some(
                    &snapshot.sink_dropped,
                    &self.initial_sink_dropped_count,
                );
                re_upload_bytes = calculate_diff_if_some(
                    &Some(snapshot.re_upload_bytes),
                    &self.initial_re_upload_bytes,
                );
                re_download_bytes = calculate_diff_if_some(
                    &Some(snapshot.re_download_bytes),
                    &self.initial_re_download_bytes,
                );
            }

            let mut metadata = Self::default_metadata();
            metadata.strings.extend(std::mem::take(&mut self.metadata));

            let record = buck2_data::InvocationRecord {
                command_name: Some(self.command_name.to_owned()),
                command_start: self.command_start.take(),
                command_end: self.command_end.take(),
                command_critical_start: self.command_critical_start.take(),
                command_critical_end: self.command_critical_end.take(),
                command_duration: self.command_duration.take(),
                client_walltime: self.start_time.elapsed().try_into().ok(),
                re_session_id: self.re_session_id.take().unwrap_or_default(),
                re_experiment_name: self.re_experiment_name.take().unwrap_or_default(),
                cli_args: self.cli_args.clone(),
                critical_path_duration: self.critical_path_duration.and_then(|x| x.try_into().ok()),
                metadata: Some(metadata),
                tags: self.tags.drain(..).collect(),
                run_local_count: self.run_local_count,
                run_remote_count: self.run_remote_count,
                run_action_cache_count: self.run_action_cache_count,
                run_skipped_count: self.run_skipped_count,
                run_fallback_count: Some(self.run_fallback_count),
                local_actions_executed_via_worker: Some(self.local_actions_executed_via_worker),
                first_snapshot: self.first_snapshot.take(),
                last_snapshot: self.last_snapshot.take(),
                min_build_count_since_rebase: self.min_build_count_since_rebase,
                cache_upload_count: self.cache_upload_count,
                cache_upload_attempt_count: self.cache_upload_attempt_count,
                parsed_target_patterns: self.parsed_target_patterns.take(),
                filesystem: self.filesystem.take().unwrap_or_default(),
                watchman_version: self.watchman_version.take(),
                eden_version: self.eden_version.take(),
                test_info: self.test_info.take(),
                eligible_for_full_hybrid: Some(self.eligible_for_full_hybrid),
                max_event_client_delay_ms: self
                    .max_event_client_delay
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                max_malloc_bytes_active: self.max_malloc_bytes_active.take(),
                max_malloc_bytes_allocated: self.max_malloc_bytes_allocated.take(),
                run_command_failure_count: Some(self.run_command_failure_count),
                event_count: Some(self.event_count),
                time_to_first_action_execution_ms: self
                    .time_to_first_action_execution
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                materialization_output_size: Some(self.materialization_output_size),
                initial_materializer_entries_from_sqlite: self
                    .initial_materializer_entries_from_sqlite,
                time_to_command_start_ms: self
                    .time_to_command_start
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                time_to_command_critical_section_ms: self
                    .time_to_command_critical_section
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                time_to_first_analysis_ms: self
                    .time_to_first_analysis
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                time_to_load_first_build_file_ms: self
                    .time_to_load_first_build_file
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                time_to_first_command_execution_start_ms: self
                    .time_to_first_command_execution_start
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                system_total_memory_bytes: self.system_total_memory_bytes,
                file_watcher_stats: self.file_watcher_stats.take(),
                time_to_last_action_execution_end_ms: self
                    .time_to_last_action_execution_end
                    .and_then(|d| u64::try_from(d.as_millis()).ok()),
                isolation_dir: Some(self.isolation_dir.clone()),
                sink_success_count,
                sink_failure_count,
                sink_dropped_count,
                sink_max_buffer_depth: Some(self.sink_max_buffer_depth),
                soft_error_categories: std::mem::take(&mut self.soft_error_categories)
                    .into_iter()
                    .collect(),
                concurrent_command_blocking_duration: self
                    .concurrent_command_blocking_duration
                    .take(),
                analysis_count: Some(self.analysis_count),
                total_concurrent_commands: self.total_concurrent_commands,
                exit_when_different_state: Some(self.exit_when_different_state),
                restarted_trace_id: self.restarted_trace_id.as_ref().map(|t| t.to_string()),
                has_command_result: Some(self.has_command_result),
                has_end_of_stream: Some(self.has_end_of_stream),
                // At this point we expect the event log writer to have finished
                compressed_event_log_size_bytes: Some(
                    self.compressed_event_log_size_bytes
                        .as_ref()
                        .map(|x| x.load(Ordering::Relaxed))
                        .unwrap_or_default(),
                ),
                critical_path_backend: self.critical_path_backend.take(),
                instant_command_is_success: self.instant_command_is_success.take(),
                bxl_ensure_artifacts_duration: self.bxl_ensure_artifacts_duration.take(),
                re_upload_bytes,
                re_download_bytes,
            };

            let event = BuckEvent::new(
                SystemTime::now(),
                self.trace_id.dupe(),
                None,
                None,
                buck2_data::RecordEvent {
                    data: Some((Box::new(record)).into()),
                }
                .into(),
            );

            if let Some(path) = &self.write_to_path {
                let res = (|| {
                    let out = fs_util::create_file(path).context("Error opening")?;
                    let mut out = std::io::BufWriter::new(out);
                    serde_json::to_writer(&mut out, event.event()).context("Error writing")?;
                    out.flush().context("Error flushing")?;
                    anyhow::Ok(())
                })();

                if let Err(e) = &res {
                    tracing::warn!(
                        "Failed to write InvocationRecord to `{}`: {:#}",
                        path.as_path().display(),
                        e
                    );
                }
            }

            if let Ok(Some(scribe_sink)) =
                new_thrift_scribe_sink_if_enabled(self.fb, 1, Duration::from_millis(500), 5, None)
            {
                tracing::info!("Recording invocation to Scribe: {:?}", &event);
                Some(async move {
                    scribe_sink.send_now(event).await;
                })
            } else {
                tracing::info!("Invocation record is not sent to Scribe: {:?}", &event);
                None
            }
        }

        // Collects client-side state and data, suitable for telemetry.
        // NOTE: If data is visible from the daemon, put it in cli::metadata::collect()
        fn default_metadata() -> buck2_data::TypedMetadata {
            let mut ints = HashMap::new();
            ints.insert("is_tty".to_owned(), std::io::stderr().is_tty() as i64);
            buck2_data::TypedMetadata {
                ints,
                strings: HashMap::new(),
            }
        }

        fn handle_command_start(
            &mut self,
            command: &buck2_data::CommandStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let mut command = command.clone();
            self.metadata.extend(std::mem::take(&mut command.metadata));
            self.command_start = Some(command);
            self.time_to_command_start = Some(self.start_time.elapsed());
            Ok(())
        }

        async fn handle_command_end(
            &mut self,
            command: &buck2_data::CommandEnd,
            event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let mut command = command.clone();
            self.metadata.extend(std::mem::take(&mut command.metadata));

            // Awkwardly unpacks the SpanEnd event so we can read its duration.
            let command_end = match event.data() {
                buck2_data::buck_event::Data::SpanEnd(ref end) => end.clone(),
                _ => {
                    return Err(anyhow::anyhow!(
                        "handle_command_end was passed a CommandEnd not contained in a SpanEndEvent"
                    ));
                }
            };
            self.command_duration = command_end.duration;
            self.min_build_count_since_rebase =
                match command.data.as_ref().context("Missing command data")? {
                    buck2_data::command_end::Data::Build(cmd) => {
                        self.build_count(&cmd.unresolved_target_patterns).await?
                    }
                    buck2_data::command_end::Data::Test(cmd) => {
                        self.build_count(&cmd.unresolved_target_patterns).await?
                    }
                    buck2_data::command_end::Data::Install(cmd) => {
                        self.build_count(&cmd.unresolved_target_patterns).await?
                    }
                    buck2_data::command_end::Data::Targets(cmd) => {
                        self.build_count(&cmd.unresolved_target_patterns).await?
                    }
                    // other events don't have target patterns
                    _ => 0,
                };
            self.command_end = Some(command);
            #[cfg(any(fbcode_build, cargo_internal_build))]
            {
                let root = std::path::Path::to_owned(self.invocation_root_path.as_ref());
                if detect_eden::is_eden(root).unwrap_or(false) {
                    self.filesystem = Some("eden".to_owned());
                } else {
                    self.filesystem = Some("default".to_owned());
                }
            }
            #[cfg(not(any(fbcode_build, cargo_internal_build)))]
            {
                self.filesystem = Some("default".to_owned());
            }
            Ok(())
        }
        fn handle_command_critical_start(
            &mut self,
            command: &buck2_data::CommandCriticalStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let mut command = command.clone();
            self.metadata.extend(std::mem::take(&mut command.metadata));
            self.command_critical_start = Some(command);
            self.time_to_command_critical_section = Some(self.start_time.elapsed());
            Ok(())
        }
        fn handle_command_critical_end(
            &mut self,
            command: &buck2_data::CommandCriticalEnd,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let mut command = command.clone();
            self.metadata.extend(std::mem::take(&mut command.metadata));
            self.command_critical_end = Some(command);
            Ok(())
        }

        fn handle_action_execution_start(
            &mut self,
            _action: &buck2_data::ActionExecutionStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            if self.time_to_first_action_execution.is_none() {
                self.time_to_first_action_execution = Some(self.start_time.elapsed());
            }
            Ok(())
        }
        fn handle_action_execution_end(
            &mut self,
            action: &buck2_data::ActionExecutionEnd,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            if action.kind == buck2_data::ActionKind::Run as i32 {
                if action_stats::was_fallback_action(action) {
                    self.run_fallback_count += 1;
                }

                match last_command_execution_kind::get_last_command_execution_kind(action) {
                    LastCommandExecutionKind::Local => {
                        self.run_local_count += 1;
                    }
                    LastCommandExecutionKind::LocalWorker => {
                        self.run_local_count += 1;
                        self.local_actions_executed_via_worker += 1;
                    }
                    LastCommandExecutionKind::Cached => {
                        self.run_action_cache_count += 1;
                    }
                    LastCommandExecutionKind::Remote => {
                        self.run_remote_count += 1;
                    }
                    LastCommandExecutionKind::NoCommand => {
                        self.run_skipped_count += 1;
                    }
                }
            }

            if action.eligible_for_full_hybrid.unwrap_or_default() {
                self.eligible_for_full_hybrid = true;
            }

            if action.commands.iter().any(|c| {
                matches!(
                    c.status,
                    Some(buck2_data::command_execution::Status::Failure(..))
                )
            }) {
                self.run_command_failure_count += 1;
            }

            self.time_to_last_action_execution_end = Some(self.start_time.elapsed());

            Ok(())
        }

        fn handle_analysis_start(
            &mut self,
            _analysis: &buck2_data::AnalysisStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            self.time_to_first_analysis
                .get_or_insert_with(|| self.start_time.elapsed());
            Ok(())
        }

        fn handle_load_start(
            &mut self,
            _eval: &buck2_data::LoadBuildFileStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            self.time_to_load_first_build_file
                .get_or_insert_with(|| self.start_time.elapsed());
            Ok(())
        }

        fn handle_executor_stage_start(
            &mut self,
            executor_stage: &buck2_data::ExecutorStageStart,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            match &executor_stage.stage {
                Some(buck2_data::executor_stage_start::Stage::Re(re_stage)) => {
                    match &re_stage.stage {
                        Some(buck2_data::re_stage::Stage::Execute(_)) => {
                            self.time_to_first_command_execution_start
                                .get_or_insert_with(|| self.start_time.elapsed());
                        }
                        _ => {}
                    }
                }
                Some(buck2_data::executor_stage_start::Stage::Local(local_stage)) => {
                    match &local_stage.stage {
                        Some(buck2_data::local_stage::Stage::Execute(_)) => {
                            self.time_to_first_command_execution_start
                                .get_or_insert_with(|| self.start_time.elapsed());
                        }
                        _ => {}
                    }
                }
                _ => {}
            }
            Ok(())
        }

        fn handle_cache_upload_end(
            &mut self,
            cache_upload: &buck2_data::CacheUploadEnd,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            if cache_upload.success {
                self.cache_upload_count += 1;
            }
            self.cache_upload_attempt_count += 1;
            Ok(())
        }

        fn handle_re_session_created(
            &mut self,
            session: &buck2_data::RemoteExecutionSessionCreated,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            self.re_session_id = Some(session.session_id.clone());
            self.re_experiment_name = Some(session.experiment_name.clone());
            Ok(())
        }

        fn handle_materialization_end(
            &mut self,
            materialization: &buck2_data::MaterializationEnd,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            self.materialization_output_size += materialization.total_bytes;
            Ok(())
        }

        fn handle_materializer_state_info(
            &mut self,
            materializer_state_info: &buck2_data::MaterializerStateInfo,
        ) -> anyhow::Result<()> {
            self.initial_materializer_entries_from_sqlite =
                Some(materializer_state_info.num_entries_from_sqlite);
            Ok(())
        }

        fn handle_bxl_ensure_artifacts_end(
            &mut self,
            _bxl_ensure_artifacts_end: &buck2_data::BxlEnsureArtifactsEnd,
            event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let bxl_ensure_artifacts_end = match event.data() {
                buck2_data::buck_event::Data::SpanEnd(ref end) => end.clone(),
                _ => {
                    return Err(anyhow::anyhow!(
                        "handle_bxl_ensure_artifacts_end was passed a BxlEnsureArtifacts not contained in a SpanEndEvent"
                    ));
                }
            };

            self.bxl_ensure_artifacts_duration = bxl_ensure_artifacts_end.duration;
            Ok(())
        }

        fn handle_test_discovery(
            &mut self,
            test_info: &buck2_data::TestDiscovery,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            match &test_info.data {
                Some(buck2_data::test_discovery::Data::Session(session_info)) => {
                    self.test_info = Some(session_info.info.clone());
                }
                Some(buck2_data::test_discovery::Data::Tests(..)) | None => {}
            }

            Ok(())
        }

        fn handle_build_graph_info(
            &mut self,
            info: &buck2_data::BuildGraphExecutionInfo,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let mut duration = Duration::default();

            for node in &info.critical_path {
                if let Some(d) = &node.duration {
                    duration += d.try_into_duration()?;
                }
            }

            for node in &info.critical_path2 {
                if let Some(d) = &node.duration {
                    duration += d.try_into_duration()?;
                }
            }

            self.critical_path_duration = Some(duration);
            self.critical_path_backend = info.backend_name.clone();
            Ok(())
        }

        fn handle_io_provider_info(
            &mut self,
            io_provider_info: &buck2_data::IoProviderInfo,
        ) -> anyhow::Result<()> {
            self.eden_version = io_provider_info.eden_version.to_owned();
            Ok(())
        }

        fn handle_dice_concurrent_commands(
            &mut self,
            dice_concurrent_commands: &buck2_data::DiceConcurrentCommands,
        ) -> anyhow::Result<()> {
            self.total_concurrent_commands =
                Some(dice_concurrent_commands.total_concurrent_commands);
            Ok(())
        }

        fn handle_exit_when_different_state(
            &mut self,
            _exit_when_different_state: &buck2_data::ExitWhenDifferentState,
        ) -> anyhow::Result<()> {
            self.exit_when_different_state = true;
            Ok(())
        }

        fn handle_tag(&mut self, tag: &buck2_data::TagEvent) -> anyhow::Result<()> {
            self.tags.extend(tag.tags.iter().cloned());
            Ok(())
        }

        fn handle_snapshot(
            &mut self,
            update: &buck2_data::Snapshot,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            self.max_malloc_bytes_active =
                cmp::max(self.max_malloc_bytes_active, update.malloc_bytes_active);
            self.max_malloc_bytes_allocated = cmp::max(
                self.max_malloc_bytes_allocated,
                update.malloc_bytes_allocated,
            );
            if self.first_snapshot.is_none() {
                self.first_snapshot = Some(update.clone());
            } else {
                self.last_snapshot = Some(update.clone());
            }
            if self.initial_sink_success_count.is_none() {
                self.initial_sink_success_count = update.sink_successes;
            }
            if self.initial_sink_failure_count.is_none() {
                self.initial_sink_failure_count = update.sink_failures;
            }
            if self.initial_sink_dropped_count.is_none() {
                self.initial_sink_dropped_count = update.sink_dropped;
            }
            self.sink_max_buffer_depth =
                cmp::max(self.sink_max_buffer_depth, update.sink_buffer_depth());

            if self.initial_re_upload_bytes.is_none() {
                self.initial_re_upload_bytes = Some(update.re_upload_bytes);
            }
            if self.initial_re_download_bytes.is_none() {
                self.initial_re_download_bytes = Some(update.re_download_bytes);
            }

            Ok(())
        }

        fn handle_file_watcher_end(
            &mut self,
            file_watcher: &buck2_data::FileWatcherEnd,
            _event: &BuckEvent,
        ) -> anyhow::Result<()> {
            // We might receive this event twice, so ... deal with it by merging the two.
            // See: https://fb.workplace.com/groups/buck2dev/permalink/3396726613948720/
            self.file_watcher_stats = merge_file_watcher_stats(
                self.file_watcher_stats.take(),
                file_watcher.stats.clone(),
            );

            if let Some(stats) = &file_watcher.stats {
                self.watchman_version = stats.watchman_version.to_owned();
            }
            Ok(())
        }

        fn handle_parsed_target_patterns(
            &mut self,
            patterns: &buck2_data::ParsedTargetPatterns,
        ) -> anyhow::Result<()> {
            self.parsed_target_patterns = Some(patterns.clone());
            Ok(())
        }

        fn handle_structured_error(
            &mut self,
            err: &buck2_data::StructuredError,
        ) -> anyhow::Result<()> {
            if let Some(soft_error_category) = err.soft_error_category.as_ref() {
                self.soft_error_categories
                    .insert(soft_error_category.to_owned());

                if err.daemon_in_memory_state_is_corrupted {
                    self.daemon_in_memory_state_is_corrupted = true;
                }

                if err.daemon_materializer_state_is_corrupted {
                    self.daemon_materializer_state_is_corrupted = true;
                }
            }

            Ok(())
        }

        fn handle_dice_block_concurrent_command_end(
            &mut self,
            _command: &buck2_data::DiceBlockConcurrentCommandEnd,
            event: &BuckEvent,
        ) -> anyhow::Result<()> {
            let block_concurrent_command = match event.data() {
                buck2_data::buck_event::Data::SpanEnd(ref end) => end.clone(),
                _ => {
                    return Err(anyhow::anyhow!(
                        "handle_dice_block_concurrent_command_end was passed a DiceBlockConcurrentCommandEnd not contained in a SpanEndEvent"
                    ));
                }
            };
            self.concurrent_command_blocking_duration = block_concurrent_command.duration;
            Ok(())
        }

        async fn handle_event(&mut self, event: &Arc<BuckEvent>) -> anyhow::Result<()> {
            // TODO(nga): query now once in `EventsCtx`.
            let now = SystemTime::now();
            if let Ok(delay) = now.duration_since(event.timestamp()) {
                self.max_event_client_delay = Some(cmp::max(
                    self.max_event_client_delay.unwrap_or_default(),
                    delay,
                ));
            }
            self.event_count += 1;

            match event.data() {
                buck2_data::buck_event::Data::SpanStart(ref start) => {
                    match start.data.as_ref().context("Missing `start`")? {
                        buck2_data::span_start_event::Data::Command(command) => {
                            self.handle_command_start(command, event)
                        }
                        buck2_data::span_start_event::Data::CommandCritical(command) => {
                            self.handle_command_critical_start(command, event)
                        }
                        buck2_data::span_start_event::Data::ActionExecution(action) => {
                            self.handle_action_execution_start(action, event)
                        }
                        buck2_data::span_start_event::Data::Analysis(analysis) => {
                            self.handle_analysis_start(analysis, event)
                        }
                        buck2_data::span_start_event::Data::Load(eval) => {
                            self.handle_load_start(eval, event)
                        }
                        buck2_data::span_start_event::Data::ExecutorStage(stage) => {
                            self.handle_executor_stage_start(stage, event)
                        }
                        _ => Ok(()),
                    }
                }
                buck2_data::buck_event::Data::SpanEnd(ref end) => {
                    match end.data.as_ref().context("Missing `end`")? {
                        buck2_data::span_end_event::Data::Command(command) => {
                            self.handle_command_end(command, event).await
                        }
                        buck2_data::span_end_event::Data::CommandCritical(command) => {
                            self.handle_command_critical_end(command, event)
                        }
                        buck2_data::span_end_event::Data::ActionExecution(action) => {
                            self.handle_action_execution_end(action, event)
                        }
                        buck2_data::span_end_event::Data::FileWatcher(file_watcher) => {
                            self.handle_file_watcher_end(file_watcher, event)
                        }
                        buck2_data::span_end_event::Data::CacheUpload(cache_upload) => {
                            self.handle_cache_upload_end(cache_upload, event)
                        }
                        buck2_data::span_end_event::Data::Materialization(materialization) => {
                            self.handle_materialization_end(materialization, event)
                        }
                        buck2_data::span_end_event::Data::Analysis(..) => {
                            self.analysis_count += 1;
                            Ok(())
                        }
                        buck2_data::span_end_event::Data::DiceBlockConcurrentCommand(
                            block_concurrent_command,
                        ) => self.handle_dice_block_concurrent_command_end(
                            block_concurrent_command,
                            event,
                        ),
                        buck2_data::span_end_event::Data::BxlEnsureArtifacts(
                            _bxl_ensure_artifacts,
                        ) => self.handle_bxl_ensure_artifacts_end(_bxl_ensure_artifacts, event),
                        _ => Ok(()),
                    }
                }
                buck2_data::buck_event::Data::Instant(ref instant) => {
                    match instant.data.as_ref().context("Missing `data`")? {
                        buck2_data::instant_event::Data::ReSession(session) => {
                            self.handle_re_session_created(session, event)
                        }
                        buck2_data::instant_event::Data::BuildGraphInfo(info) => {
                            self.handle_build_graph_info(info, event)
                        }
                        buck2_data::instant_event::Data::TestDiscovery(discovery) => {
                            self.handle_test_discovery(discovery, event)
                        }
                        buck2_data::instant_event::Data::Snapshot(result) => {
                            self.handle_snapshot(result, event)
                        }
                        buck2_data::instant_event::Data::TagEvent(tag) => self.handle_tag(tag),
                        buck2_data::instant_event::Data::IoProviderInfo(io_provider_info) => {
                            self.handle_io_provider_info(io_provider_info)
                        }
                        buck2_data::instant_event::Data::TargetPatterns(tag) => {
                            self.handle_parsed_target_patterns(tag)
                        }
                        buck2_data::instant_event::Data::MaterializerStateInfo(
                            materializer_state,
                        ) => self.handle_materializer_state_info(materializer_state),
                        buck2_data::instant_event::Data::StructuredError(err) => {
                            self.handle_structured_error(err)
                        }
                        buck2_data::instant_event::Data::DiceConcurrentCommands(
                            dice_concurrent_commands,
                        ) => self.handle_dice_concurrent_commands(dice_concurrent_commands),
                        buck2_data::instant_event::Data::ExitWhenDifferentState(
                            exit_when_different_state,
                        ) => self.handle_exit_when_different_state(exit_when_different_state),
                        buck2_data::instant_event::Data::RestartConfiguration(conf) => {
                            self.enable_restarter = conf.enable_restarter;
                            Ok(())
                        }
                        _ => Ok(()),
                    }
                }
                buck2_data::buck_event::Data::Record(_) => Ok(()),
            }
        }
    }

    impl<'a> Drop for InvocationRecorder<'a> {
        fn drop(&mut self) {
            if let Some(fut) = self.send_it() {
                self.async_cleanup_context
                    .register("sending invocation to Scribe", fut.boxed());
            }
        }
    }

    #[async_trait]
    impl<'a> EventSubscriber for InvocationRecorder<'a> {
        async fn handle_events(&mut self, events: &[Arc<BuckEvent>]) -> anyhow::Result<()> {
            for event in events {
                self.handle_event(event).await?;
            }
            Ok(())
        }

        async fn handle_console_interaction(&mut self, _c: char) -> anyhow::Result<()> {
            self.tags.push("console-interaction".to_owned());
            Ok(())
        }

        async fn handle_command_result(
            &mut self,
            _result: &buck2_cli_proto::CommandResult,
        ) -> anyhow::Result<()> {
            self.has_command_result = true;
            Ok(())
        }

        async fn exit(&mut self) -> anyhow::Result<()> {
            self.has_end_of_stream = true;
            Ok(())
        }

        fn as_error_observer(&self) -> Option<&dyn ErrorObserver> {
            Some(self)
        }
    }

    impl<'a> ErrorObserver for InvocationRecorder<'a> {
        fn error_cause(&self) -> ErrorCause {
            if self.exit_when_different_state {
                // User wants to immediately exit concurrent commands with different states
                return ErrorCause::DaemonIsBusy;
            } else if self.run_command_failure_count > 0 {
                // Action fails likely because of user-defined commands in the action
                return ErrorCause::User;
            }

            ErrorCause::Unknown
        }

        fn daemon_in_memory_state_is_corrupted(&self) -> bool {
            self.daemon_in_memory_state_is_corrupted
        }

        fn daemon_materializer_state_is_corrupted(&self) -> bool {
            self.daemon_materializer_state_is_corrupted
        }

        fn restarter_is_enabled(&self) -> bool {
            self.enable_restarter
        }
    }

    fn calculate_diff_if_some(a: &Option<u64>, b: &Option<u64>) -> Option<u64> {
        match (a, b) {
            (Some(av), Some(bv)) => Some(std::cmp::max(av, bv) - std::cmp::min(av, bv)),
            _ => None,
        }
    }

    fn merge_file_watcher_stats(
        a: Option<buck2_data::FileWatcherStats>,
        b: Option<buck2_data::FileWatcherStats>,
    ) -> Option<buck2_data::FileWatcherStats> {
        let (mut a, b) = match (a, b) {
            (Some(a), Some(b)) => (a, b),
            (a, None) => return a,
            (None, b) => return b,
        };

        a.fresh_instance = a.fresh_instance || b.fresh_instance;
        a.events_total += b.events_total;
        a.events_processed += b.events_processed;
        a.branched_from_revision = a.branched_from_revision.or(b.branched_from_revision);
        a.events.extend(b.events);
        a.incomplete_events_reason = a.incomplete_events_reason.or(b.incomplete_events_reason);
        a.watchman_version = a.watchman_version.or(b.watchman_version);
        Some(a)
    }
}

pub fn try_get_invocation_recorder<'a>(
    ctx: &ClientCommandContext<'a>,
    opts: &CommonDaemonCommandOptions,
    command_name: &'static str,
    sanitized_argv: Vec<String>,
    log_size_counter_bytes: Option<Arc<AtomicU64>>,
) -> anyhow::Result<Box<imp::InvocationRecorder<'a>>> {
    let write_to_path = opts
        .unstable_write_invocation_record
        .as_ref()
        .map(|path| path.resolve(&ctx.working_dir));

    let recorder = imp::InvocationRecorder::new(
        ctx.fbinit(),
        ctx.async_cleanup_context().dupe(),
        write_to_path,
        command_name,
        sanitized_argv,
        ctx.trace_id.dupe(),
        ctx.paths()?.isolation.as_str().to_owned(),
        BuildCountManager::new(ctx.paths()?.build_count_dir()),
        ctx.paths()?.project_root().root().to_buf(),
        ctx.restarted_trace_id.dupe(),
        log_size_counter_bytes,
    );
    Ok(Box::new(recorder))
}

fn system_memory_stats() -> u64 {
    use sysinfo::RefreshKind;
    use sysinfo::System;
    use sysinfo::SystemExt;

    let system = System::new_with_specifics(RefreshKind::new().with_memory());
    system.total_memory()
}

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

    #[test]
    fn get_system_memory_stats() {
        let total_mem = system_memory_stats();
        // sysinfo returns zero when fails to retrieve data
        assert!(total_mem > 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::process::Stdio;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use async_trait::async_trait;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_event_observer::unpack_event::unpack_event;
use buck2_event_observer::unpack_event::UnpackedBuckEvent;
use buck2_events::BuckEvent;
use futures::Future;
use futures::FutureExt;
use tokio::process::Child;

use crate::cleanup_ctx::AsyncCleanupContext;
use crate::subscribers::should_upload_log;
use crate::subscribers::subscriber::EventSubscriber;

pub(crate) struct ReLog<'a> {
    re_session_id: Option<String>,
    isolation_dir: FileNameBuf,
    async_cleanup_context: AsyncCleanupContext<'a>,
}

impl<'a> ReLog<'a> {
    pub(crate) fn new(
        isolation_dir: FileNameBuf,
        async_cleanup_context: AsyncCleanupContext<'a>,
    ) -> Self {
        Self {
            re_session_id: None,
            isolation_dir,
            async_cleanup_context,
        }
    }

    fn log_upload(&mut self) -> impl Future<Output = anyhow::Result<()>> + 'static + Send + Sync {
        // We put `None` in place of re_session_id which means we will only attempt to upload
        // the logs once no matter how many times this function is called
        let session_id = self.re_session_id.take();
        let isolation_dir = self.isolation_dir.clone();
        async move {
            if let Some(s_id) = session_id {
                log_upload_impl(s_id, isolation_dir).await?;
            }
            Ok(())
        }
    }
}

#[async_trait]
impl<'a> EventSubscriber for ReLog<'a> {
    async fn exit(&mut self) -> anyhow::Result<()> {
        self.log_upload().await
    }

    async fn handle_events(&mut self, events: &[Arc<BuckEvent>]) -> anyhow::Result<()> {
        for event in events {
            match unpack_event(event)? {
                UnpackedBuckEvent::Instant(
                    _,
                    _,
                    buck2_data::instant_event::Data::ReSession(session),
                ) => {
                    self.re_session_id = Some(session.session_id.clone());
                }
                _ => {}
            }
        }
        Ok(())
    }
}

impl<'a> Drop for ReLog<'a> {
    fn drop(&mut self) {
        let upload = self.log_upload();
        self.async_cleanup_context.register(
            "RE log upload",
            async move {
                if let Err(e) = upload.await {
                    tracing::warn!("Failed to cleanup ReLog: {:#}", e);
                }
            }
            .boxed(),
        );
    }
}

async fn log_upload_impl(session_id: String, isolation_dir: FileNameBuf) -> anyhow::Result<()> {
    if !should_upload_log()? {
        return Ok(());
    }

    let mut buck = buck2_util::process::async_background_command(std::env::current_exe()?);
    let command = buck
        .arg("--isolation-dir")
        .arg(isolation_dir.as_str())
        .arg("debug")
        .arg("upload-re-logs")
        .arg("--session-id")
        .arg(session_id)
        .stdin(Stdio::null())
        .stdout(Stdio::null());

    let block_on_upload = std::env::var_os("SANDCASTLE").is_some();
    if block_on_upload {
        let child = command.stderr(Stdio::piped()).spawn()?;
        match blocking_upload(child).await {
            Ok(_) => {}
            Err(e) => {
                tracing::warn!("Error uploading RE logs: {:#}", e);
            }
        }
    } else {
        command.stderr(Stdio::null()).spawn()?;
    };

    Ok(())
}

async fn blocking_upload(child: Child) -> anyhow::Result<()> {
    let res = child.wait_with_output();
    let timeout = tokio::time::timeout(Duration::from_secs(20), res)
        .await
        .context("Timed out waiting for RE log upload to manifold")??;
    if !timeout.status.success() {
        let stderr = String::from_utf8_lossy(&timeout.stderr);
        return Err(anyhow::anyhow!(
            "RE log upload exited with {}. Stderr: `{}`",
            timeout.status,
            stderr.trim(),
        ));
    };
    Ok(())
}