/*
 * 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::future;
use std::io;
use std::path::Path;
use std::pin::Pin;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::task::Context;
use std::task::Poll;
use std::thread;
use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_build_api::configure_dice::configure_dice_for_buck;
use buck2_build_api::spawner::BuckSpawner;
use buck2_cli_proto::daemon_api_server::*;
use buck2_cli_proto::*;
use buck2_common::buckd_connection::BUCK_AUTH_TOKEN_HEADER;
use buck2_common::events::HasEvents;
use buck2_common::invocation_paths::InvocationPaths;
use buck2_common::io::trace::TracingIoProvider;
use buck2_common::io::IoProvider;
use buck2_common::legacy_configs::cells::DaemonStartupConfig;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_common::memory;
use buck2_core::env_helper::EnvHelper;
use buck2_core::error::reload_hard_error_config;
use buck2_core::error::reset_soft_error_counters;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_core::logging::LogConfigurationReloadHandle;
use buck2_events::dispatch::EventDispatcher;
use buck2_events::ControlEvent;
use buck2_events::Event;
use buck2_events::EventSource;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute_impl::materializers::sqlite::MaterializerStateIdentity;
use buck2_interpreter::dice::starlark_profiler::StarlarkProfilerConfiguration;
use buck2_profile::starlark_profiler_configuration_from_request;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::NoPartialResult;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::streaming_request_handler::StreamingRequestHandler;
use buck2_server_starlark_debug::run::run_dap_server_command;
use dice::DetectCycles;
use dice::Dice;
use dice::WhichDice;
use dupe::Dupe;
use futures::channel::mpsc;
use futures::channel::mpsc::UnboundedReceiver;
use futures::channel::mpsc::UnboundedSender;
use futures::future::BoxFuture;
use futures::stream;
use futures::Future;
use futures::FutureExt;
use futures::Stream;
use futures::StreamExt;
use futures::TryFutureExt;
use more_futures::cancellation::ExplicitCancellationContext;
use more_futures::drop::DropTogether;
use more_futures::spawn::spawn_cancellable;
use rand::RngCore;
use rand::SeedableRng;
use tokio::sync::oneshot;
use tonic::service::interceptor;
use tonic::service::Interceptor;
use tonic::transport::Server;
use tonic::Code;
use tonic::Request;
use tonic::Response;
use tonic::Status;

use crate::active_commands::ActiveCommand;
use crate::active_commands::ActiveCommandStateWriter;
use crate::clean_stale::clean_stale_command;
use crate::ctx::ServerCommandContext;
use crate::daemon::multi_event_stream::MultiEventStream;
use crate::daemon::server_allocative::spawn_allocative;
use crate::daemon::state::DaemonState;
use crate::file_status::file_status_command;
use crate::lsp::run_lsp_server_command;
use crate::materialize::materialize_command;
use crate::snapshot;
use crate::subscription::run_subscription_server_command;
use crate::trace_io::trace_io_command;

// TODO(cjhopman): Figure out a reasonable value for this.
static DEFAULT_KILL_TIMEOUT: Duration = Duration::from_millis(500);

static DEFAULT_INACTIVITY_TIMEOUT: Duration = Duration::from_secs(4 * 86400);

pub trait BuckdServerDelegate: Allocative + Send + Sync {
    fn force_shutdown_with_timeout(&self, reason: String, timeout: Duration);
}

#[derive(Allocative)]
struct DaemonShutdown {
    delegate: Box<dyn BuckdServerDelegate>,

    /// This channel is used to trigger a graceful shutdown of the grpc server. After
    /// an item is sent on this channel, the server will start rejecting new requests
    /// and once current requests are finished the server will shutdown.
    #[allocative(skip)]
    shutdown_channel: UnboundedSender<()>,
}

impl DaemonShutdown {
    /// Trigger a graceful server shutdown with a timeout. After the timeout expires, a hard shutdown
    /// will be triggered.
    ///
    /// As we might be processing a `kill()` (or other) request, we cannot wait for the server to actually
    /// shutdown (as it will wait for current requests to finish), so this returns immediately.
    fn start_shutdown(&self, reason: buck2_data::DaemonShutdown, timeout: Option<Duration>) {
        crate::active_commands::broadcast_shutdown(&reason);

        let timeout = timeout.unwrap_or(DEFAULT_KILL_TIMEOUT);

        // Ignore errors on shutdown_channel as that would mean we've already started shutdown;
        let _ = self.shutdown_channel.unbounded_send(());
        self.delegate
            .force_shutdown_with_timeout(reason.to_string(), timeout);
    }
}

#[derive(Allocative)]
pub struct BuckdServerInitPreferences {
    pub detect_cycles: Option<DetectCycles>,
    pub which_dice: Option<WhichDice>,
    pub enable_trace_io: bool,
    pub reject_materializer_state: Option<MaterializerStateIdentity>,
    pub daemon_startup_config: DaemonStartupConfig,
}

impl BuckdServerInitPreferences {
    pub async fn construct_dice(
        &self,
        io: Arc<dyn IoProvider>,
        digest_config: DigestConfig,
        root_config: &LegacyBuckConfig,
    ) -> anyhow::Result<Arc<Dice>> {
        configure_dice_for_buck(
            io,
            digest_config,
            Some(root_config),
            self.detect_cycles,
            self.which_dice,
        )
        .await
    }
}

/// Access to functions which live outside of `buck2_server` crate.
#[async_trait]
pub trait BuckdServerDependencies: Send + Sync + 'static {
    async fn test(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: TestRequest,
    ) -> anyhow::Result<TestResponse>;
    async fn build(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: BuildRequest,
    ) -> anyhow::Result<BuildResponse>;
    async fn install(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: InstallRequest,
    ) -> anyhow::Result<InstallResponse>;
    async fn bxl(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::BxlRequest,
    ) -> anyhow::Result<BxlResponse>;
    async fn audit(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::GenericRequest,
    ) -> anyhow::Result<buck2_cli_proto::GenericResponse>;
    async fn starlark(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::GenericRequest,
    ) -> anyhow::Result<buck2_cli_proto::GenericResponse>;
    async fn profile(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: buck2_cli_proto::ProfileRequest,
    ) -> anyhow::Result<buck2_cli_proto::ProfileResponse>;
    async fn uquery(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::UqueryRequest,
    ) -> anyhow::Result<buck2_cli_proto::UqueryResponse>;
    async fn cquery(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::CqueryRequest,
    ) -> anyhow::Result<CqueryResponse>;
    async fn aquery(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: buck2_cli_proto::AqueryRequest,
    ) -> anyhow::Result<buck2_cli_proto::AqueryResponse>;
    async fn targets(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
        req: TargetsRequest,
    ) -> anyhow::Result<TargetsResponse>;
    async fn ctargets(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: ConfiguredTargetsRequest,
    ) -> anyhow::Result<ConfiguredTargetsResponse>;
    async fn targets_show_outputs(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: TargetsRequest,
    ) -> anyhow::Result<TargetsShowOutputsResponse>;
    async fn docs(
        &self,
        ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
        req: buck2_cli_proto::UnstableDocsRequest,
    ) -> anyhow::Result<buck2_cli_proto::UnstableDocsResponse>;
}

#[derive(Clone)]
struct BuckCheckAuthTokenInterceptor {
    auth_token: String,
}

impl Interceptor for BuckCheckAuthTokenInterceptor {
    fn call(&mut self, request: Request<()>) -> Result<Request<()>, Status> {
        let token = match request.metadata().get(BUCK_AUTH_TOKEN_HEADER) {
            Some(token) => token,
            None => return Err(Status::unauthenticated("missing auth token")),
        };
        if !constant_time_eq::constant_time_eq(
            token.as_bytes(),
            self.auth_token.as_str().as_bytes(),
        ) {
            return Err(Status::unauthenticated("invalid auth token"));
        }

        static FAIL_AUTH: EnvHelper<bool> = EnvHelper::new("BUCK2_TEST_FAIL_BUCKD_AUTH");
        if FAIL_AUTH
            .get_copied()
            .unwrap_or_default()
            .unwrap_or_default()
        {
            return Err(Status::unauthenticated("injected auth error"));
        }

        Ok(request)
    }
}

#[derive(Allocative)]
pub(crate) struct BuckdServerData {
    /// The flag that is set to true when server is shutting down.
    stop_accepting_requests: AtomicBool,
    #[allocative(skip)]
    process_info: DaemonProcessInfo,
    base_daemon_constraints: buck2_cli_proto::DaemonConstraints,
    start_time: prost_types::Timestamp,
    start_instant: Instant,
    daemon_shutdown: DaemonShutdown,
    daemon_state: Arc<DaemonState>,
    #[allocative(skip)]
    command_channel: UnboundedSender<()>,
    #[allocative(skip)]
    callbacks: &'static dyn BuckdServerDependencies,
    #[allocative(skip)]
    log_reload_handle: Arc<dyn LogConfigurationReloadHandle>,
}

/// The BuckdServer implements the DaemonApi.
///
/// Simple endpoints are implemented here and complex things will be implemented in a sibling
/// module taking just a ServerCommandContext.
#[derive(Allocative)]
pub struct BuckdServer(Arc<BuckdServerData>);

impl BuckdServer {
    pub async fn run<I>(
        fb: fbinit::FacebookInit,
        log_reload_handle: Arc<dyn LogConfigurationReloadHandle>,
        paths: InvocationPaths,
        delegate: Box<dyn BuckdServerDelegate>,
        init_ctx: BuckdServerInitPreferences,
        process_info: DaemonProcessInfo,
        base_daemon_constraints: buck2_cli_proto::DaemonConstraints,
        listener: I,
        callbacks: &'static dyn BuckdServerDependencies,
    ) -> anyhow::Result<()>
    where
        I: Stream<Item = Result<tokio::net::TcpStream, io::Error>>,
    {
        let now = SystemTime::now();
        let now = now.duration_since(SystemTime::UNIX_EPOCH)?;

        let (shutdown_channel, shutdown_receiver): (UnboundedSender<()>, _) = mpsc::unbounded();
        let (command_channel, command_receiver): (UnboundedSender<()>, _) = mpsc::unbounded();

        let daemon_state = Arc::new(DaemonState::new(fb, paths, init_ctx).await);

        let auth_token = process_info.auth_token.clone();
        let api_server = BuckdServer(Arc::new(BuckdServerData {
            stop_accepting_requests: AtomicBool::new(false),
            process_info,
            base_daemon_constraints,
            start_time: prost_types::Timestamp {
                seconds: now.as_secs() as i64,
                nanos: now.subsec_nanos() as i32,
            },
            start_instant: Instant::now(),
            daemon_shutdown: DaemonShutdown {
                delegate,
                shutdown_channel,
            },
            daemon_state,
            command_channel,
            callbacks,
            log_reload_handle,
        }));

        let shutdown = server_shutdown_signal(command_receiver, shutdown_receiver)?;
        let server = Server::builder()
            .layer(interceptor(BuckCheckAuthTokenInterceptor { auth_token }))
            .add_service(
                DaemonApiServer::new(api_server)
                    .max_encoding_message_size(usize::MAX)
                    .max_decoding_message_size(usize::MAX),
            )
            .serve_with_incoming_shutdown(listener, shutdown);

        server.await?;

        Ok(())
    }

    /// Run a request that does bidirectional streaming.
    ///
    /// This mostly just ensures that a client context has been sent first, and passes a client
    /// stream to `func` that converts to the correct type (or returns an error and shuts the
    /// stream down)
    async fn run_bidirectional<Req, Res, PartialRes, F>(
        &self,
        req: Request<tonic::Streaming<StreamingRequest>>,
        opts: impl StreamingCommandOptions<StreamingRequest>,
        func: F,
    ) -> Result<Response<ResponseStream>, Status>
    where
        F: for<'a> FnOnce(
                &'a ServerCommandContext,
                PartialResultDispatcher<PartialRes>,
                &ClientContext,
                StreamingRequestHandler<Req>,
            ) -> BoxFuture<'a, anyhow::Result<Res>>
            + Send
            + 'static,
        Req: TryFrom<StreamingRequest, Error = anyhow::Error> + Send + Sync + 'static,
        Res: Into<command_result::Result> + Send + 'static,
        PartialRes: Into<partial_result::PartialResult> + Send + 'static,
    {
        let mut req = req.into_inner();
        let init_request = match req.message().await? {
            Some(
                m @ StreamingRequest {
                    request: Some(buck2_cli_proto::streaming_request::Request::Context(_)),
                },
            ) => Ok(m),
            _ => Err(Status::failed_precondition(
                "no client context message was received",
            )),
        }?;

        let init_request = Request::new(init_request);
        self.run_streaming(
            init_request,
            opts,
            |ctx, partial_result_dispatcher, init_req| {
                // TODO: Use the PartialResultDispatcher instead of writing events.
                func(
                    ctx,
                    partial_result_dispatcher,
                    init_req
                        .client_context()
                        .expect("already checked for a valid context"),
                    StreamingRequestHandler::new(req),
                )
            },
        )
        .await
    }

    async fn run_streaming_anyhow<Req, Res, PartialRes, F>(
        &self,
        req: Request<Req>,
        opts: impl StreamingCommandOptions<Req>,
        func: F,
    ) -> anyhow::Result<Response<ResponseStream>>
    where
        F: for<'a> FnOnce(
                &'a ServerCommandContext,
                PartialResultDispatcher<PartialRes>,
                Req,
            ) -> BoxFuture<'a, anyhow::Result<Res>>
            + Send
            + 'static,
        Req: HasClientContext + HasBuildOptions + Send + Sync + 'static,
        Res: Into<command_result::Result> + Send + 'static,
        PartialRes: Into<partial_result::PartialResult> + Send + 'static,
    {
        let client_ctx = req.get_ref().client_context()?;

        // This will reset counters incorrectly if commands are running concurrently.
        // This is fine.
        reset_soft_error_counters();

        reload_hard_error_config(&client_ctx.buck2_hard_error)?;

        OneshotCommandOptions::pre_run(&opts, self)?;

        let daemon_state = self.0.daemon_state.dupe();
        let trace_id = client_ctx.trace_id.parse()?;
        let (events, dispatch) = daemon_state.prepare_events(trace_id).await?;
        let ActiveCommand {
            guard,
            daemon_shutdown_channel,
            state,
        } = ActiveCommand::new(&dispatch, client_ctx);
        let data = daemon_state.data()?;

        dispatch.instant_event(Box::new(
            snapshot::SnapshotCollector::pre_initialization_snapshot(data.start_time),
        ));

        let resp = streaming(
            req,
            events,
            state,
            dispatch.dupe(),
            daemon_shutdown_channel,
            move |req, cancellations| {
                async move {
                    let result: anyhow::Result<Res> = try {
                        let base_context =
                            daemon_state.prepare_command(dispatch.dupe(), guard).await?;

                        let context = ServerCommandContext::new(
                            base_context,
                            req.client_context()?,
                            opts.starlark_profiler_instrumentation_override(&req)?,
                            req.build_options(),
                            daemon_state.paths.buck_out_dir(),
                            cancellations,
                        )?;

                        func(&context, PartialResultDispatcher::new(dispatch.dupe()), req).await?
                    };

                    let result: CommandResult = result_to_command_result(result);
                    dispatch.control_event(ControlEvent::CommandResult(Box::new(result)));
                }
                .boxed()
            },
        )
        .await;
        Ok(resp)
    }

    /// Runs a single command (given by the func F). Prior to running the command, calls the
    /// `opts`'s `pre_run` hook.  then bootstraps an event source and command context so that the
    /// invoked function has the ability to stream events to the caller.
    async fn run_streaming<Req, Res, PartialRes, F>(
        &self,
        req: Request<Req>,
        opts: impl StreamingCommandOptions<Req>,
        func: F,
    ) -> Result<Response<ResponseStream>, Status>
    where
        F: for<'a> FnOnce(
                &'a ServerCommandContext,
                PartialResultDispatcher<PartialRes>,
                Req,
            ) -> BoxFuture<'a, anyhow::Result<Res>>
            + Send
            + 'static,
        Req: HasClientContext + HasBuildOptions + Send + Sync + 'static,
        Res: Into<command_result::Result> + Send + 'static,
        PartialRes: Into<partial_result::PartialResult> + Send + 'static,
    {
        // send signal to register new command time
        _ = self.0.command_channel.unbounded_send(());

        Ok(self
            .run_streaming_anyhow(req, opts, func)
            .await
            .unwrap_or_else(error_to_response_stream))
    }

    async fn oneshot<
        Req,
        Res: Into<command_result::Result>,
        Fut: Future<Output = anyhow::Result<Res>> + Send,
        F: FnOnce(Req) -> Fut,
    >(
        &self,
        req: Request<Req>,
        opts: impl OneshotCommandOptions,
        func: F,
    ) -> Result<Response<CommandResult>, Status> {
        opts.pre_run(self)?;

        let req = req.into_inner();
        let result = func(req).await;
        Ok(Response::new(result_to_command_result(result)))
    }

    /// Checks if the server is accepting requests.
    fn check_if_accepting_requests(&self) -> Result<(), Status> {
        if self.0.stop_accepting_requests.load(Ordering::Relaxed) {
            Err(Status::failed_precondition(
                "Failed to run command, `buckd` is shutting down soon!",
            ))
        } else {
            Ok(())
        }
    }
}

fn convert_positive_duration(proto_duration: &prost_types::Duration) -> Result<Duration, Status> {
    if proto_duration.seconds < 0 || proto_duration.nanos < 0 {
        return Err(Status::new(
            Code::Unknown,
            format!("received invalid timeout: `{:?}`", proto_duration),
        ));
    }
    Ok(Duration::from_secs(proto_duration.seconds as u64)
        + Duration::from_nanos(proto_duration.nanos as u64))
}

fn error_to_command_result(e: anyhow::Error) -> CommandResult {
    let messages = vec![format!("{:?}", e)];

    CommandResult {
        result: Some(command_result::Result::Error(CommandError { messages })),
    }
}

fn result_to_command_result<R: Into<command_result::Result>>(
    result: anyhow::Result<R>,
) -> CommandResult {
    match result {
        Ok(result) => CommandResult {
            result: Some(result.into()),
        },
        Err(e) => error_to_command_result(e),
    }
}

fn error_to_command_progress(e: anyhow::Error) -> CommandProgress {
    CommandProgress {
        progress: Some(command_progress::Progress::Result(Box::new(
            error_to_command_result(e),
        ))),
    }
}

fn error_to_response_stream(e: anyhow::Error) -> Response<ResponseStream> {
    tonic::Response::new(Box::pin(stream::once(future::ready(Ok(
        buck2_cli_proto::MultiCommandProgress {
            messages: vec![error_to_command_progress(e)],
        },
    )))))
}

/// tonic requires the response for a streaming api to be a Sync Stream. With async/await, that requirement is really difficult
/// to meet. This simple wrapper allows us to wrap a non-Sync stream into a Sync one (the inner stream is never accessed in a
/// non-exclusive manner).
struct SyncStream<T: Stream + Send> {
    // SyncWrapper provides a Sync type that only allows (statically checked) exclusive access to
    // the underlying object, this allows using a non-Sync object where a Sync one is required
    // but is never accessed from multiple threads.
    // See https://internals.rust-lang.org/t/what-shall-sync-mean-across-an-await/12020/31
    // and https://github.com/hyperium/tonic/issues/117
    wrapped: sync_wrapper::SyncWrapper<T>,
}

impl<T: Stream + Send> Stream for SyncStream<T> {
    type Item = <T as Stream>::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        // This is a safe pin projection. See https://doc.rust-lang.org/std/pin/index.html#projections-and-structural-pinning
        // Specifically see the requirements when pinning is structural for a field here: https://doc.rust-lang.org/std/pin/index.html#pinning-is-structural-for-field
        unsafe { self.map_unchecked_mut(|a| a.wrapped.get_mut()) }.poll_next(cx)
    }
}

fn pump_events<E: EventSource>(
    mut events: E,
    mut state: ActiveCommandStateWriter,
    output_send: tokio::sync::mpsc::UnboundedSender<
        Result<buck2_cli_proto::CommandProgress, tonic::Status>,
    >,
) {
    while let Some(next_event) = events.receive() {
        // Note that writes to `output_send` have their errors explicitly ignored here. There is only one reason
        // for a write to a `mpsc::channel` to fail: the receiving end of the channel has already been closed.
        //
        // This function returns the receiving channel back to `tonic` as part of a streaming response. Tonic can
        // drop the stream before it is fully resolved if, for example, the gRPC client disconnects during the
        // command. In this case, we explicitly ignore write errors and let them float off into the void, since no
        // client is listening.
        //
        // TODO(swgillespie) - We should handle client disconnects better.
        match next_event {
            Event::Control(control_event) => {
                // A control event. This event isn't going to be sent to gRPC, but we do need to react to it. In
                // this case, the CommandResult event indicates that the spawned computation has produced a result
                // and will not be producing any more events.
                match control_event {
                    ControlEvent::CommandResult(result) => {
                        let _ignore = output_send.send(Ok(CommandProgress {
                            progress: Some(command_progress::Progress::Result(result)),
                        }));
                        return;
                    }
                    ControlEvent::PartialResult(result) => {
                        let _ignore = output_send.send(Ok(CommandProgress {
                            progress: Some(command_progress::Progress::PartialResult(Box::new(
                                result,
                            ))),
                        }));
                    }
                }
            }
            Event::Buck(buck_event) => {
                state.peek_event(&buck_event);

                // A buck event. These events should be forwarded directly to gRPC.
                let _ignore = output_send.send(Ok(CommandProgress {
                    progress: Some(command_progress::Progress::Event(buck_event.into())),
                }));
            }
        }
    }
}

/// Dispatches a request to the given function and returns a stream of responses, suitable for streaming to a client.
#[allow(clippy::mut_mut)] // select! does this internally
async fn streaming<
    Req: Send + Sync + 'static,
    F: for<'a> FnOnce(Req, &'a ExplicitCancellationContext) -> BoxFuture<'a, ()>,
    E: EventSource + 'static,
>(
    req: Request<Req>,
    events: E,
    state: ActiveCommandStateWriter,
    dispatcher: EventDispatcher,
    daemon_shutdown_channel: oneshot::Receiver<buck2_data::DaemonShutdown>,
    func: F,
) -> Response<ResponseStream>
where
    F: Send + 'static,
    E: Sync,
{
    // This function is responsible for receiving all events coming into an EventSource and reacting accordingly. There
    // are two categories events that can be seen:
    // 1. Control events, which are not to be sent across the gRPC boundary but instruct this function to do something.
    // 2. Buck events, which are to be sent across the gRPC boundary.
    //
    // The function `func` is the computation that we are going to run. It communicates its success or failure using
    // control events; our first step is to spawn it.

    struct EventsCtx {
        dispatcher: EventDispatcher,
    }
    impl HasEvents for EventsCtx {
        fn get_dispatcher(&self) -> &EventDispatcher {
            &self.dispatcher
        }
    }

    let trace_id = dispatcher.trace_id().dupe();

    let req = req.into_inner();
    let events_ctx = EventsCtx { dispatcher };
    let spawned = spawn_cancellable(
        |cancellations| func(req, cancellations),
        &BuckSpawner,
        &events_ctx,
    );
    let (output_send, output_recv) = tokio::sync::mpsc::unbounded_channel();

    // We run the event consumer on a totally separate tokio runtime to avoid the consumer task from getting stuck behind
    // another tokio task in its lifo task slot. See T96012305 and https://github.com/tokio-rs/tokio/issues/4323 for more
    // information.
    let merge_task = thread::Builder::new()
        .name("pump-events".to_owned())
        .spawn(move || {
            pump_events(events, state, output_send);
        });
    let _merge_task = match merge_task {
        Ok(merge_task) => merge_task,
        Err(e) => {
            return error_to_response_stream(
                anyhow::Error::new(e).context("failed to spawn pump-events"),
            );
        }
    };

    //
    // Note that while this is an event, we don't send it through our normal event
    // processing. The reason for that is that we dont want this event to queue behind any other
    // events in the (2) unbounded channels that form our event pipeline. So, we inject this one
    // directly where Tonic is polling for responses (which, unlike the rest of the pipeline, is
    // not unbounded, and has backpressure).

    let daemon_shutdown_stream = daemon_shutdown_channel
        .map_ok(move |shutdown| CommandProgress {
            progress: Some(command_progress::Progress::Event(Box::new(
                buck2_data::BuckEvent {
                    timestamp: Some(SystemTime::now().into()),
                    trace_id: trace_id.to_string(),
                    span_id: 0,
                    parent_id: 0,
                    data: Some(
                        buck2_data::InstantEvent {
                            data: Some(shutdown.into()),
                        }
                        .into(),
                    ),
                },
            ))),
        })
        .into_stream()
        .filter_map(|e| {
            // If the channel yields an Err, that means we didnt shut down, so for us that is
            // simply something we want to drop from the stream.
            futures::future::ready(e.ok().map(Ok))
        });

    // The stream we ultimately return is the receiving end of the channel that the above task is
    // writing to, plus the shutdown channel.
    let events = futures::stream::select(
        tokio_stream::wrappers::UnboundedReceiverStream::new(output_recv),
        daemon_shutdown_stream,
    );

    let events = MultiEventStream::new(events);

    Response::new(Box::pin(SyncStream {
        wrapped: sync_wrapper::SyncWrapper::new(DropTogether::new(
            events,
            spawned.into_drop_cancel(),
        )),
    }))
}

type ResponseStream =
    Pin<Box<dyn Stream<Item = Result<MultiCommandProgress, Status>> + Send + Sync>>;
#[async_trait]
impl DaemonApi for BuckdServer {
    async fn kill(&self, req: Request<KillRequest>) -> Result<Response<CommandResult>, Status> {
        struct KillRunCommandOptions;

        impl OneshotCommandOptions for KillRunCommandOptions {
            /// kill should be always available
            fn pre_run(&self, _server: &BuckdServer) -> Result<(), Status> {
                Ok(())
            }
        }

        self.oneshot(req, KillRunCommandOptions, move |req| async move {
            self.0
                .stop_accepting_requests
                .store(true, Ordering::Relaxed);

            let timeout = req
                .timeout
                .as_ref()
                .map(convert_positive_duration)
                .transpose()?;

            let reason = buck2_data::DaemonShutdown {
                reason: req.reason,
                callers: req.callers,
            };

            self.0.daemon_shutdown.start_shutdown(reason, timeout);
            Ok(KillResponse {})
        })
        .await
    }

    async fn ping(&self, req: Request<PingRequest>) -> Result<Response<CommandResult>, Status> {
        self.oneshot(req, DefaultCommandOptions, move |req| async move {
            match &req.delay {
                Some(delay) => {
                    let delay = convert_positive_duration(delay)?;
                    tokio::time::sleep(delay).await;
                }
                _ => {}
            }

            let mut payload = vec![
                0;
                req.response_payload_size
                    .try_into()
                    .context("requested payload too large")?
            ];
            rand::rngs::SmallRng::seed_from_u64(10).fill_bytes(&mut payload);

            Ok(PingResponse { payload })
        })
        .await
    }

    async fn status(&self, req: Request<StatusRequest>) -> Result<Response<CommandResult>, Status> {
        let daemon_state = self.0.daemon_state.dupe();

        self.oneshot(req, DefaultCommandOptions, move |req| async move {
            let snapshot = if req.snapshot {
                let data = daemon_state.data()?;
                Some(
                    snapshot::SnapshotCollector::new(
                        data.re_client_manager.dupe(),
                        data.blocking_executor.dupe(),
                        data.start_time,
                        data.dice_manager.unsafe_dice().dupe(),
                        data.materializer.dupe(),
                        data.scribe_sink.dupe() as _,
                        data.http_client.dupe(),
                    )
                    .create_snapshot(),
                )
            } else {
                None
            };

            let extra_constraints = daemon_state.data().as_ref().ok().map(|state| {
                buck2_cli_proto::ExtraDaemonConstraints {
                    trace_io_enabled: state.io.as_any().is::<TracingIoProvider>(),
                    materializer_state_identity: state
                        .materializer_state_identity
                        .as_ref()
                        .map(|i| i.to_string()),
                }
            });

            let mut daemon_constraints = self.0.base_daemon_constraints.clone();
            daemon_constraints.extra = extra_constraints;

            let uptime = self.0.start_instant.elapsed();
            let base = StatusResponse {
                process_info: Some(self.0.process_info.clone()),
                start_time: Some(self.0.start_time.clone()),
                uptime: Some(uptime.try_into()?),
                snapshot,
                daemon_constraints: Some(daemon_constraints),
                project_root: daemon_state.paths.project_root().to_string(),
                isolation_dir: daemon_state.paths.isolation.to_string(),
                forkserver_pid: daemon_state
                    .data
                    .as_ref()
                    .ok()
                    .and_then(|state| state.forkserver.as_ref().map(|f| f.pid())),
                supports_vpnless: daemon_state
                    .data()
                    .as_ref()
                    .ok()
                    .map(|state| state.http_client.supports_vpnless()),
                ..Default::default()
            };
            Ok(base)
        })
        .await
    }

    async fn flush_dep_files(
        &self,
        req: Request<FlushDepFilesRequest>,
    ) -> Result<Response<CommandResult>, Status> {
        self.oneshot(req, DefaultCommandOptions, move |req| async move {
            let FlushDepFilesRequest {} = req;
            buck2_file_watcher::dep_files::flush_dep_files();
            Ok(GenericResponse {})
        })
        .await
    }

    type FileStatusStream = ResponseStream;
    async fn file_status(
        &self,
        req: Request<FileStatusRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |context, partial_result_dispatcher, req| {
                file_status_command(context, partial_result_dispatcher, req).boxed()
            },
        )
        .await
    }

    type BuildStream = ResponseStream;
    async fn build(&self, req: Request<BuildRequest>) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.build(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type BxlStream = ResponseStream;
    async fn bxl(&self, req: Request<BxlRequest>) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.bxl(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type TestStream = ResponseStream;
    async fn test(&self, req: Request<TestRequest>) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.test(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type AqueryStream = ResponseStream;
    async fn aquery(
        &self,
        req: Request<AqueryRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.aquery(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type UqueryStream = ResponseStream;
    async fn uquery(
        &self,
        req: Request<UqueryRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.uquery(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type CqueryStream = ResponseStream;
    async fn cquery(
        &self,
        req: Request<CqueryRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.cquery(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type TargetsStream = ResponseStream;
    async fn targets(
        &self,
        req: Request<TargetsRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.targets(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type CtargetsStream = ResponseStream;
    async fn ctargets(
        &self,
        req: Request<ConfiguredTargetsRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.ctargets(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type TargetsShowOutputsStream = ResponseStream;
    async fn targets_show_outputs(
        &self,
        req: Request<TargetsRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.targets_show_outputs(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type AuditStream = ResponseStream;
    async fn audit(
        &self,
        req: Request<GenericRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.audit(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type StarlarkStream = ResponseStream;
    async fn starlark(
        &self,
        req: Request<GenericRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.starlark(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type InstallStream = ResponseStream;
    async fn install(
        &self,
        req: Request<InstallRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.install(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    async fn unstable_crash(
        &self,
        _req: Request<UnstableCrashRequest>,
    ) -> Result<Response<UnstableCrashResponse>, Status> {
        panic!("explicitly requested panic (via unstable_crash)");
    }

    async fn segfault(
        &self,
        _req: Request<SegfaultRequest>,
    ) -> Result<Response<SegfaultResponse>, Status> {
        unsafe {
            std::ptr::null_mut::<&'static str>()
                .write("Explicitly requested segfault (via `segfault`)")
        };
        unreachable!()
    }

    async fn unstable_heap_dump(
        &self,
        req: Request<UnstableHeapDumpRequest>,
    ) -> Result<Response<UnstableHeapDumpResponse>, Status> {
        self.check_if_accepting_requests()?;

        let heap_dump = memory::write_heap_to_file(&req.into_inner().destination_path);
        match heap_dump {
            Ok(_) => Ok(Response::new(UnstableHeapDumpResponse {})),
            Err(e) => Err(Status::invalid_argument(format!(
                "failed to perform heap dump: {}",
                e
            ))),
        }
    }

    async fn unstable_allocator_stats(
        &self,
        req: Request<UnstableAllocatorStatsRequest>,
    ) -> Result<Response<UnstableAllocatorStatsResponse>, Status> {
        self.check_if_accepting_requests()?;

        let response = memory::allocator_stats(&req.into_inner().options)
            .context("Failed to retrieve allocator stats");

        match response {
            Ok(response) => Ok(Response::new(UnstableAllocatorStatsResponse { response })),
            Err(e) => Err(Status::invalid_argument(format!("{:#}", e))),
        }
    }

    async fn unstable_dice_dump(
        &self,
        req: Request<UnstableDiceDumpRequest>,
    ) -> Result<Response<UnstableDiceDumpResponse>, Status> {
        self.check_if_accepting_requests()?;

        let inner = req.into_inner();
        let path = inner.destination_path;
        let res: anyhow::Result<_> = try {
            let path = Path::new(&path);
            let format_proto =
                buck2_cli_proto::unstable_dice_dump_request::DiceDumpFormat::from_i32(inner.format)
                    .context("Invalid DICE dump format")?;

            self.0
                .daemon_state
                .data()?
                .spawn_dice_dump(path, format_proto)
                .await
                .with_context(|| format!("Failed to perform dice dump to {}", path.display()))?;

            UnstableDiceDumpResponse {}
        };

        res.map(Response::new)
            .map_err(|e| Status::internal(format!("{:#}", e)))
    }

    type AllocativeStream = ResponseStream;
    async fn allocative(
        &self,
        req: Request<AllocativeRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        self.check_if_accepting_requests()?;

        let res: anyhow::Result<_> = try {
            let client_ctx = req.get_ref().client_context()?;
            let trace_id = client_ctx.trace_id.parse()?;
            let (event_source, dispatcher) = self.0.daemon_state.prepare_events(trace_id).await?;
            let active_command = ActiveCommand::new(&dispatcher, client_ctx);
            (event_source, dispatcher, active_command)
        };

        let (event_source, dispatcher, active_command) = match res {
            Ok(v) => v,
            Err(e) => return Ok(error_to_response_stream(e)),
        };

        let ActiveCommand {
            guard,
            daemon_shutdown_channel,
            state,
        } = active_command;

        let this = self.0.dupe();
        Ok(streaming(
            req,
            event_source,
            state,
            dispatcher.dupe(),
            daemon_shutdown_channel,
            move |req, _| {
                async move {
                    let result = try {
                        spawn_allocative(
                            this,
                            AbsPathBuf::try_from(req.output_path)?,
                            dispatcher.dupe(),
                        )
                        .await?;
                        AllocativeResponse {}
                    };

                    let result: CommandResult = result_to_command_result(result);
                    dispatcher.control_event(ControlEvent::CommandResult(Box::new(result)));

                    drop(guard);
                }
                .boxed()
            },
        )
        .await)
    }

    type UnstableDocsStream = ResponseStream;
    async fn unstable_docs(
        &self,
        req: Request<UnstableDocsRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.docs(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type Profile2Stream = ResponseStream;
    async fn profile2(
        &self,
        req: Request<ProfileRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        struct ProfileCommandOptions;

        impl OneshotCommandOptions for ProfileCommandOptions {}

        impl StreamingCommandOptions<ProfileRequest> for ProfileCommandOptions {
            fn starlark_profiler_instrumentation_override(
                &self,
                req: &ProfileRequest,
            ) -> anyhow::Result<StarlarkProfilerConfiguration> {
                starlark_profiler_configuration_from_request(req)
            }
        }

        let callbacks = self.0.callbacks;
        self.run_streaming(
            req,
            ProfileCommandOptions,
            |ctx, partial_result_dispatcher, req| {
                callbacks.profile(ctx, partial_result_dispatcher, req)
            },
        )
        .await
    }

    type MaterializeStream = ResponseStream;
    async fn materialize(
        &self,
        req: Request<MaterializeRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |context, _: PartialResultDispatcher<NoPartialResult>, req| {
                materialize_command(context, req).boxed()
            },
        )
        .await
    }

    type CleanStaleStream = ResponseStream;
    async fn clean_stale(
        &self,
        req: Request<CleanStaleRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |context, partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>, req| {
                clean_stale_command(context, partial_result_dispatcher, req).boxed()
            },
        )
        .await
    }

    type LspStream = ResponseStream;
    async fn lsp(
        &self,
        req: Request<tonic::Streaming<StreamingRequest>>,
    ) -> Result<Response<Self::LspStream>, Status> {
        self.run_bidirectional(
            req,
            DefaultCommandOptions,
            |ctx,
             partial_result_dispatcher,
             _client_ctx,
             req: StreamingRequestHandler<LspRequest>| {
                run_lsp_server_command(ctx, partial_result_dispatcher, req).boxed()
            },
        )
        .await
    }

    type SubscriptionStream = ResponseStream;
    async fn subscription(
        &self,
        req: Request<tonic::Streaming<StreamingRequest>>,
    ) -> Result<Response<Self::SubscriptionStream>, Status> {
        self.run_bidirectional(
            req,
            DefaultCommandOptions,
            |ctx,
             partial_result_dispatcher,
             _client_ctx,
             req: StreamingRequestHandler<SubscriptionRequestWrapper>| {
                run_subscription_server_command(ctx, partial_result_dispatcher, req).boxed()
            },
        )
        .await
    }

    type DapStream = ResponseStream;
    async fn dap(
        &self,
        req: Request<tonic::Streaming<StreamingRequest>>,
    ) -> Result<Response<Self::DapStream>, Status> {
        self.run_bidirectional(
            req,
            DefaultCommandOptions,
            |ctx,
             partial_result_dispatcher,
             _client_ctx,
             req: StreamingRequestHandler<DapRequest>| {
                run_dap_server_command(ctx, partial_result_dispatcher, req).boxed()
            },
        )
        .await
    }

    async fn set_log_filter(
        &self,
        req: Request<SetLogFilterRequest>,
    ) -> Result<Response<SetLogFilterResponse>, Status> {
        let req = req.into_inner();

        if req.daemon {
            self.0
                .log_reload_handle
                .update_log_filter(&req.log_filter)
                .context("Error updating daemon log filter")
                .map_err(|e| Status::invalid_argument(format!("{:#}", e)))?;
        }

        if req.forkserver {
            if let Ok(data) = self.0.daemon_state.data() {
                if let Some(forkserver) = data.forkserver.as_ref() {
                    forkserver
                        .set_log_filter(req.log_filter)
                        .await
                        .context("Error forwarding daemon log filter to forkserver")
                        .map_err(|e| Status::invalid_argument(format!("{:#}", e)))?;
                }
            }
        }

        Ok(Response::new(SetLogFilterResponse {}))
    }

    type TraceIoStream = ResponseStream;
    async fn trace_io(
        &self,
        req: Request<TraceIoRequest>,
    ) -> Result<Response<ResponseStream>, Status> {
        self.run_streaming(
            req,
            DefaultCommandOptions,
            |context, _: PartialResultDispatcher<NoPartialResult>, req| {
                trace_io_command(context, req).boxed()
            },
        )
        .await
    }
}

/// Options to configure the execution of a oneshot command (i.e. what happens in `oneshot()`).
trait OneshotCommandOptions: Send + Sync + 'static {
    fn pre_run(&self, server: &BuckdServer) -> Result<(), Status> {
        server.check_if_accepting_requests()
    }
}

/// Options to configure the execution of a streaming command (i.e. what happens in `run_streaming()`).
trait StreamingCommandOptions<Req>: OneshotCommandOptions {
    fn starlark_profiler_instrumentation_override(
        &self,
        _req: &Req,
    ) -> anyhow::Result<StarlarkProfilerConfiguration> {
        Ok(StarlarkProfilerConfiguration::None)
    }
}

fn server_shutdown_signal(
    command_receiver: UnboundedReceiver<()>,
    mut shutdown_receiver: UnboundedReceiver<()>,
) -> anyhow::Result<impl Future<Output = ()>> {
    static TESTING_INACTIVITY_TIMEOUT: EnvHelper<bool> =
        EnvHelper::new("BUCK2_TESTING_INACTIVITY_TIMEOUT");

    let mut duration = DEFAULT_INACTIVITY_TIMEOUT;
    if *TESTING_INACTIVITY_TIMEOUT.get()?.unwrap_or(&false) {
        duration = Duration::from_secs(1);
    }

    Ok(async move {
        let timeout = inactivity_timeout(command_receiver, duration);
        let shutdown = shutdown_receiver.next();

        futures::pin_mut!(shutdown);
        futures::pin_mut!(timeout);

        futures::future::select(timeout, shutdown).await;
    })
}

async fn inactivity_timeout(mut command_receiver: UnboundedReceiver<()>, duration: Duration) {
    // this restarts the timer everytime there is a new command
    loop {
        let command = command_receiver.next();
        let timer = tokio::time::sleep(duration);

        futures::pin_mut!(command);
        futures::pin_mut!(timer);

        match futures::future::select(command, timer).await {
            futures::future::Either::Left(_) => continue,
            futures::future::Either::Right(_) => break,
        };
    }
}

/// No-op set of command options.
struct DefaultCommandOptions;

impl OneshotCommandOptions for DefaultCommandOptions {}
impl<Req> StreamingCommandOptions<Req> for DefaultCommandOptions {}


/*
 * 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 allocative::FlameGraph;
use allocative::FlameGraphBuilder;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_path::AbsPathBuf;
use buck2_events::dispatch::EventDispatcher;
use buck2_util::process_stats::process_stats;

use crate::daemon::server::BuckdServerData;
use crate::jemalloc_stats::get_allocator_stats;

/// In `FlameGraph` nodes do not have names, only child keys.
/// This helper makes it a bit easier to manipulate the tree.
struct NamedFlameGraph {
    name: allocative::Key,
    flamegraph: FlameGraph,
}

impl NamedFlameGraph {
    #[must_use]
    fn layer_on_top(self, name: &'static str, size: usize) -> NamedFlameGraph {
        // Each layer should be larger than the previous,
        // but if it is not, round up the difference to zero.
        let rem_size = size.saturating_sub(self.flamegraph.total_size());
        let mut new_fg = FlameGraph::default();
        new_fg.add_self(rem_size);
        new_fg.add_child(self.name, self.flamegraph);
        NamedFlameGraph {
            name: allocative::Key::new(name),
            flamegraph: new_fg,
        }
    }

    fn into_flamegraph(self) -> FlameGraph {
        let mut new_fg = FlameGraph::default();
        new_fg.add_child(self.name, self.flamegraph);
        new_fg
    }
}

fn wrap_flamegraph_with_system_stats(fg: &FlameGraph) -> FlameGraph {
    let mut fg = NamedFlameGraph {
        name: allocative::Key::new("allocative"),
        flamegraph: fg.clone(),
    };

    if let Ok(stats) = get_allocator_stats() {
        // This is technically incorrect because not everything allocative visits
        // is allocated by malloc, but it's close enough.
        if let Some(allocated) = stats.bytes_allocated {
            fg = fg.layer_on_top("jemalloc.allocated", allocated as usize);
        }
        if let Some(active) = stats.bytes_active {
            fg = fg.layer_on_top("jemalloc.active", active as usize);
        }
    }

    if let Some(rss_bytes) = process_stats().rss_bytes {
        // This is also incorrect because memory can be swapped out.
        fg = fg.layer_on_top("rss", rss_bytes as usize);
    }

    fg.into_flamegraph()
}

pub(crate) async fn spawn_allocative(
    buckd_server_data: Arc<BuckdServerData>,
    path: AbsPathBuf,
    dispatcher: EventDispatcher,
) -> anyhow::Result<()> {
    tokio::task::spawn_blocking(move || {
        let mut graph = FlameGraphBuilder::default();
        dispatcher.console_message(
            "Starting allocative profiling. It may take a while to finish...".to_owned(),
        );
        // TODO(nga): Emit some progress.
        dispatcher.console_message("Visiting global roots...".to_owned());
        graph.visit_global_roots();
        dispatcher.console_message("Visiting buckd...".to_owned());
        graph.visit_root(&buckd_server_data);
        let fg = graph.finish();
        fs_util::create_dir_if_not_exists(&path)?;
        dispatcher.console_message(format!("Writing allocative to `{}`...", path.display()));
        let final_fg = wrap_flamegraph_with_system_stats(fg.flamegraph());
        fs_util::write(path.join("flamegraph.src"), final_fg.write())?;
        let mut fg_svg = Vec::new();
        inferno::flamegraph::from_reader(
            &mut inferno::flamegraph::Options::default(),
            final_fg.write().as_bytes(),
            &mut fg_svg,
        )?;
        fs_util::write(path.join("flamegraph.svg"), &fg_svg)?;

        fs_util::write(path.join("warnings.txt"), fg.warnings())?;

        dispatcher.console_message("Profile written.".to_owned());

        anyhow::Ok(())
    })
    .await?
}


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

use std::collections::HashMap;
use std::path::Path;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use allocative::Allocative;
use anyhow::Context;
use buck2_cli_proto::unstable_dice_dump_request::DiceDumpFormat;
use buck2_common::cas_digest::DigestAlgorithm;
use buck2_common::cas_digest::DigestAlgorithmKind;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_common::http::http_client;
use buck2_common::ignores::ignore_set::IgnoreSet;
use buck2_common::invocation_paths::InvocationPaths;
use buck2_common::io::IoProvider;
use buck2_common::legacy_configs::cells::BuckConfigBasedCells;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_core::cells::name::CellName;
use buck2_core::env_helper::EnvHelper;
use buck2_core::facebook_only;
use buck2_core::fs::fs_util;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::rollout_percentage::RolloutPercentage;
use buck2_core::tag_result;
use buck2_events::dispatch::EventDispatcher;
use buck2_events::sink::scribe;
use buck2_events::sink::tee::TeeSink;
use buck2_events::EventSink;
use buck2_events::EventSource;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::execute::blocking::BuckBlockingExecutor;
use buck2_execute::materialize::materializer::MaterializationMethod;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::re::manager::ReConnectionManager;
use buck2_execute_impl::materializers::deferred::DeferredMaterializer;
use buck2_execute_impl::materializers::deferred::DeferredMaterializerConfigs;
use buck2_execute_impl::materializers::deferred::TtlRefreshConfiguration;
use buck2_execute_impl::materializers::immediate::ImmediateMaterializer;
use buck2_execute_impl::materializers::sqlite::MaterializerState;
use buck2_execute_impl::materializers::sqlite::MaterializerStateIdentity;
use buck2_execute_impl::materializers::sqlite::MaterializerStateSqliteDb;
use buck2_file_watcher::file_watcher::FileWatcher;
use buck2_forkserver::client::ForkserverClient;
use buck2_re_configuration::RemoteExecutionStaticMetadata;
use buck2_re_configuration::RemoteExecutionStaticMetadataImpl;
use buck2_server_ctx::concurrency::ConcurrencyHandler;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use fbinit::FacebookInit;
use gazebo::prelude::*;
use gazebo::variants::VariantName;
use tokio::sync::Mutex;

use crate::active_commands::ActiveCommandDropGuard;
use crate::ctx::BaseServerCommandContext;
use crate::daemon::check_working_dir;
use crate::daemon::disk_state::delete_unknown_disk_state;
use crate::daemon::disk_state::maybe_initialize_materializer_sqlite_db;
use crate::daemon::disk_state::DiskStateOptions;
use crate::daemon::forkserver::maybe_launch_forkserver;
use crate::daemon::io_provider::create_io_provider;
use crate::daemon::panic::DaemonStatePanicDiceDump;
use crate::daemon::server::BuckdServerInitPreferences;
/// For a buckd process there is a single DaemonState created at startup and never destroyed.
#[derive(Allocative)]
pub struct DaemonState {
    #[allocative(skip)]
    fb: fbinit::FacebookInit,

    pub paths: InvocationPaths,

    /// This holds the main data shared across different commands.
    pub(crate) data: SharedResult<Arc<DaemonStateData>>,
}

/// DaemonStateData is the main shared data across all commands. It's lazily initialized on
/// the first command that requires it.
#[derive(Allocative)]
pub struct DaemonStateData {
    /// The Dice computation graph. Generally, we shouldn't add things to the DaemonStateData
    /// (or DaemonState) itself and instead they should be represented on the computation graph.
    ///
    /// The DICE graph is held by the concurrency handler to manage locking for concurrent commands
    pub(crate) dice_manager: ConcurrencyHandler,

    /// Synced every time we run a command.
    file_watcher: Arc<dyn FileWatcher>,

    /// Settled every time we run a command.
    pub io: Arc<dyn IoProvider>,

    /// The RE connection, managed such that all build commands that are concurrently active uses
    /// the same connection. Once there are no active build commands, the connection will be
    /// terminated
    pub re_client_manager: Arc<ReConnectionManager>,

    /// Executor responsible for coordinating and rate limiting I/O.
    pub blocking_executor: Arc<dyn BlockingExecutor>,

    /// Most materializations go through the materializer, providing a single point
    /// where the most expensive network and fs IO operations are performed. It
    /// needs access to the `ReConnectionManager` to download from RE. It must
    /// live for the entire lifetime of the daemon, in order to allow deferred
    /// materializations to work properly between distinct build commands.
    pub(crate) materializer: Arc<dyn Materializer>,

    pub(crate) forkserver: Option<ForkserverClient>,

    #[allocative(skip)]
    pub scribe_sink: Option<Arc<dyn EventSink>>,

    /// Whether or not to hash all commands
    pub hash_all_commands: bool,

    /// Whether to consult the offline-cache buck-out dir for network action
    /// outputs prior to running them. If no cached output exists, the action
    /// (download_file, cas_artifact) will execute normally.
    ///
    /// This supports fully-offline builds, where network actions like
    /// download_file have an execution component that is inherently non-local (
    /// e.g. making a HEAD request against the remote artifact to determine if
    /// it needs to be downloaded again).
    pub use_network_action_output_cache: bool,

    /// What buck2 state to store on disk, ex. materializer state on sqlite
    pub disk_state_options: DiskStateOptions,

    pub start_time: Instant,

    #[allocative(skip)]
    pub create_unhashed_outputs_lock: Arc<Mutex<()>>,

    /// A unique identifier for the materializer state.
    pub materializer_state_identity: Option<MaterializerStateIdentity>,

    /// Whether to enable the restarter. This controls whether the client will attempt to restart
    /// the daemon when we hit an error.
    pub enable_restarter: bool,

    /// Http client used for materializer and RunAction implementations.
    pub http_client: CountingHttpClient,

    /// Did we enable eden I/O v2?
    pub eden_io_v2: bool,
}

impl DaemonStateData {
    pub fn dice_dump(&self, path: &Path, format: DiceDumpFormat) -> anyhow::Result<()> {
        crate::daemon::dice_dump::dice_dump(self.dice_manager.unsafe_dice(), path, format)
    }

    pub async fn spawn_dice_dump(&self, path: &Path, format: DiceDumpFormat) -> anyhow::Result<()> {
        crate::daemon::dice_dump::dice_dump_spawn(self.dice_manager.unsafe_dice(), path, format)
            .await
    }
}

impl DaemonStatePanicDiceDump for DaemonStateData {
    fn dice_dump(&self, path: &Path, format: DiceDumpFormat) -> anyhow::Result<()> {
        self.dice_dump(path, format)
    }
}

impl DaemonState {
    #[tracing::instrument(name = "daemon_listener", skip_all)]
    pub async fn new(
        fb: fbinit::FacebookInit,
        paths: InvocationPaths,
        init_ctx: BuckdServerInitPreferences,
    ) -> Self {
        let data = Self::init_data(fb, &paths, init_ctx)
            .await
            .context("Error initializing DaemonStateData");
        if let Ok(data) = &data {
            crate::daemon::panic::initialize(data.dupe());
        }

        tracing::info!("Daemon state is ready.");

        let data = data.shared_error();

        DaemonState { fb, paths, data }
    }

    // Creates the initial DaemonStateData.
    // Starts up the watchman query.
    async fn init_data(
        fb: fbinit::FacebookInit,
        paths: &InvocationPaths,
        init_ctx: BuckdServerInitPreferences,
    ) -> anyhow::Result<Arc<DaemonStateData>> {
        let fs = paths.project_root().clone();

        tracing::info!("Reading config...");
        let legacy_cells = BuckConfigBasedCells::parse(&fs)?;

        tracing::info!("Starting...");

        let (legacy_configs, cells) = (legacy_cells.configs_by_name, legacy_cells.cell_resolver);

        let root_config = legacy_configs
            .get(cells.root_cell())
            .context("No config for root cell")?;

        // This really should belong in  in DaemonCommand::run, but we can't read configs there. In
        // practice, while changing cwd after starting is not a great idea, it's ... fine.

        fs_util::create_dir_all(paths.buck_out_path()).context("Error creating buck_out_path")?;

        let materialization_method =
            MaterializationMethod::try_new_from_config(legacy_configs.get(cells.root_cell()).ok())?;

        if !matches!(materialization_method, MaterializationMethod::Eden) {
            fs_util::set_current_dir(paths.buck_out_path()).context("Error changing dirs")?;
            buck2_core::fs::cwd::cwd_will_not_change().context("Error initializing static cwd")?;
        }

        static DEFAULT_DIGEST_ALGORITHM: EnvHelper<DigestAlgorithmKind> =
            EnvHelper::new("BUCK_DEFAULT_DIGEST_ALGORITHM");

        let default_digest_algorithm =
            DEFAULT_DIGEST_ALGORITHM.get_copied()?.unwrap_or_else(|| {
                if buck2_core::is_open_source() {
                    DigestAlgorithmKind::Sha256
                } else {
                    DigestAlgorithmKind::Sha1
                }
            });

        let digest_algorithms = init_ctx
            .daemon_startup_config
            .digest_algorithms
            .as_ref()
            .map(|algos| {
                algos
                    .split(',')
                    .map(DigestAlgorithmKind::from_str)
                    .collect::<Result<_, _>>()
            })
            .transpose()
            .context("Invalid digest_algorithms")?
            .unwrap_or_else(|| vec![default_digest_algorithm])
            .into_try_map(convert_algorithm_kind)?;

        let preferred_source_algorithm = init_ctx
            .daemon_startup_config
            .source_digest_algorithm
            .as_deref()
            .map(|a| convert_algorithm_kind(a.parse()?))
            .transpose()
            .context("Invalid source_digest_algorithm")?;

        let digest_config = DigestConfig::leak_new(digest_algorithms, preferred_source_algorithm)
            .context("Error initializing DigestConfig")?;

        // TODO(rafaelc): merge configs from all cells once they are consistent
        let static_metadata = Arc::new(RemoteExecutionStaticMetadata::from_legacy_config(
            root_config,
        )?);

        let ignore_specs: HashMap<CellName, IgnoreSet> = legacy_configs
            .iter()
            .map(|(cell, config)| {
                Ok((
                    cell,
                    IgnoreSet::from_ignore_spec(
                        config.get("project", "ignore").unwrap_or(""),
                        cells.is_root_cell(cell),
                    )?,
                ))
            })
            .collect::<anyhow::Result<_>>()?;

        let disk_state_options = DiskStateOptions::new(root_config, materialization_method.dupe())?;
        let blocking_executor = Arc::new(BuckBlockingExecutor::default_concurrency(fs.dupe())?);
        let cache_dir_path = paths.cache_dir_path();
        let valid_cache_dirs = paths.valid_cache_dirs();
        let fs_duped = fs.dupe();

        let deferred_materializer_configs = {
            let defer_write_actions = root_config
                .parse::<RolloutPercentage>("buck2", "defer_write_actions")?
                .unwrap_or_else(RolloutPercentage::never)
                .roll();

            // RE will refresh any TTL < 1 hour, so we check twice an hour and refresh any TTL
            // < 1 hour.
            let ttl_refresh_frequency = root_config
                .parse("buck2", "ttl_refresh_frequency_seconds")?
                .unwrap_or(1800);

            let ttl_refresh_min_ttl = root_config
                .parse("buck2", "ttl_refresh_min_ttl_seconds")?
                .unwrap_or(3600);

            let ttl_refresh_enabled = root_config
                .parse::<RolloutPercentage>("buck2", "ttl_refresh_enabled")?
                .unwrap_or_else(RolloutPercentage::never)
                .roll();

            DeferredMaterializerConfigs {
                materialize_final_artifacts: matches!(
                    materialization_method,
                    MaterializationMethod::Deferred
                ),
                defer_write_actions,
                ttl_refresh: TtlRefreshConfiguration {
                    frequency: std::time::Duration::from_secs(ttl_refresh_frequency),
                    min_ttl: chrono::Duration::seconds(ttl_refresh_min_ttl),
                    enabled: ttl_refresh_enabled,
                },
            }
        };

        let eden_io_v2 = init_ctx
            .daemon_startup_config
            .eden_io_v2
            .as_deref()
            .map(RolloutPercentage::from_str)
            .transpose()
            .context("Invalid eden_io_v2")?
            .unwrap_or_else(RolloutPercentage::never)
            .roll();

        let (io, _, (materializer_db, materializer_state)) = futures::future::try_join3(
            create_io_provider(
                fb,
                fs.dupe(),
                legacy_configs.get(cells.root_cell()).ok(),
                digest_config.cas_digest_config(),
                init_ctx.enable_trace_io,
                eden_io_v2,
            ),
            (blocking_executor.dupe() as Arc<dyn BlockingExecutor>).execute_io_inline(|| {
                // Using `execute_io_inline` is just out of convenience.
                // It doesn't really matter what's used here since there's no IO-heavy
                // operations on daemon startup
                delete_unknown_disk_state(&cache_dir_path, &valid_cache_dirs, fs_duped)
            }),
            maybe_initialize_materializer_sqlite_db(
                &disk_state_options,
                paths,
                blocking_executor.dupe() as Arc<dyn BlockingExecutor>,
                root_config,
                &deferred_materializer_configs,
                fs,
                digest_config,
                &init_ctx,
            ),
        )
        .await?;

        let allow_vpnless = root_config
            .parse("buck2", "allow_vpnless")?
            .unwrap_or(false);
        let http_client = http_client(allow_vpnless)?;

        let materializer_state_identity = materializer_db.as_ref().map(|d| d.identity().clone());

        let re_client_manager = Arc::new(ReConnectionManager::new(
            fb,
            false,
            10,
            static_metadata,
            Some(paths.re_logs_dir()),
            paths.buck_out_path(),
        ));
        let materializer = Self::create_materializer(
            fb,
            io.project_root().dupe(),
            digest_config,
            paths.buck_out_dir(),
            re_client_manager.dupe(),
            blocking_executor.dupe(),
            materialization_method,
            deferred_materializer_configs,
            materializer_db,
            materializer_state,
            http_client.dupe(),
        )?;

        // Create this after the materializer because it'll want to write to buck-out, and an Eden
        // materializer would create buck-out now.
        let forkserver =
            maybe_launch_forkserver(root_config, &paths.forkserver_state_dir()).await?;

        let dice = init_ctx
            .construct_dice(io.dupe(), digest_config, root_config)
            .await?;

        // TODO(cjhopman): We want to use Expr::True here, but we need to workaround
        // https://github.com/facebook/watchman/issues/911. Adding other filetypes to
        // this list should be safe until we can revert it to Expr::True.

        let file_watcher = <dyn FileWatcher>::new(
            paths.project_root(),
            root_config,
            cells.dupe(),
            ignore_specs,
        )
        .with_context(|| {
            format!(
                "Error creating a FileWatcher for project root `{}`",
                paths.project_root()
            )
        })?;

        let hash_all_commands = root_config
            .parse::<RolloutPercentage>("buck2", "hash_all_commands")?
            .unwrap_or_else(RolloutPercentage::never)
            .roll();

        let use_network_action_output_cache = root_config
            .parse("buck2", "use_network_action_output_cache")?
            .unwrap_or(false);

        let create_unhashed_outputs_lock = Arc::new(Mutex::new(()));

        let buffer_size = root_config
            .parse("buck2", "event_log_buffer_size")?
            .unwrap_or(10000);
        let retry_backoff = Duration::from_millis(
            root_config
                .parse("buck2", "event_log_retry_backoff_duration_ms")?
                .unwrap_or(500),
        );
        let retry_attempts = root_config
            .parse("buck2", "event_log_retry_attempts")?
            .unwrap_or(5);
        let message_batch_size = root_config.parse("buck2", "event_log_message_batch_size")?;
        let scribe_sink = Self::init_scribe_sink(
            fb,
            buffer_size,
            retry_backoff,
            retry_attempts,
            message_batch_size,
        )
        .context("failed to init scribe sink")?;

        let enable_restarter = root_config
            .parse::<RolloutPercentage>("buck2", "restarter")?
            .unwrap_or_else(RolloutPercentage::never)
            .roll();

        // Kick off an initial sync eagerly. This gets Watchamn to start watching the path we care
        // about (potentially kicking off an initial crawl).

        // disable the eager spawn for watchman until we fix dice commit to avoid a panic TODO(bobyf)
        // tokio::task::spawn(watchman_query.sync());
        Ok(Arc::new(DaemonStateData {
            dice_manager: ConcurrencyHandler::new(dice),
            file_watcher,
            io,
            re_client_manager,
            blocking_executor,
            materializer,
            forkserver,
            scribe_sink,
            hash_all_commands,
            use_network_action_output_cache,
            disk_state_options,
            start_time: std::time::Instant::now(),
            create_unhashed_outputs_lock,
            materializer_state_identity,
            enable_restarter,
            http_client,
            eden_io_v2,
        }))
    }

    fn create_materializer(
        fb: FacebookInit,
        fs: ProjectRoot,
        digest_config: DigestConfig,
        buck_out_path: ProjectRelativePathBuf,
        re_client_manager: Arc<ReConnectionManager>,
        blocking_executor: Arc<dyn BlockingExecutor>,
        materialization_method: MaterializationMethod,
        deferred_materializer_configs: DeferredMaterializerConfigs,
        materializer_db: Option<MaterializerStateSqliteDb>,
        materializer_state: Option<MaterializerState>,
        http_client: CountingHttpClient,
    ) -> anyhow::Result<Arc<dyn Materializer>> {
        match materialization_method {
            MaterializationMethod::Immediate => Ok(Arc::new(ImmediateMaterializer::new(
                fs,
                digest_config,
                re_client_manager,
                blocking_executor,
                http_client,
            ))),
            MaterializationMethod::Deferred | MaterializationMethod::DeferredSkipFinalArtifacts => {
                Ok(Arc::new(DeferredMaterializer::new(
                    fs,
                    digest_config,
                    buck_out_path,
                    re_client_manager,
                    blocking_executor,
                    deferred_materializer_configs,
                    materializer_db,
                    materializer_state,
                    http_client,
                )?))
            }
            MaterializationMethod::Eden => {
                #[cfg(any(fbcode_build, cargo_internal_build))]
                {
                    use buck2_execute::materialize::eden_api::EdenBuckOut;
                    use buck2_execute_impl::materializers::eden::EdenMaterializer;

                    let buck_out_mount = fs.root().join(&buck_out_path);

                    if cfg!(unix) {
                        Ok(Arc::new(
                            EdenMaterializer::new(
                                fs,
                                digest_config,
                                re_client_manager.dupe(),
                                blocking_executor,
                                EdenBuckOut::new(
                                    fb,
                                    buck_out_path,
                                    buck_out_mount,
                                    re_client_manager,
                                )
                                .context("Failed to create EdenFS-based buck-out")?,
                                http_client,
                            )
                            .context("Failed to create Eden materializer")?,
                        ))
                    } else {
                        Err(anyhow::anyhow!(
                            "`eden` materialization method is not supported on Windows"
                        ))
                    }
                }
                #[cfg(not(any(fbcode_build, cargo_internal_build)))]
                {
                    let _unused = buck_out_path;
                    let _unused = fs;
                    let _unused = fb;
                    Err(anyhow::anyhow!(
                        "`eden` materialization method is only supported in Meta internal builds"
                    ))
                }
            }
        }
    }

    pub fn init_scribe_sink(
        fb: FacebookInit,
        buffer_size: usize,
        retry_backoff: Duration,
        retry_attempts: usize,
        message_batch_size: Option<usize>,
    ) -> anyhow::Result<Option<Arc<dyn EventSink>>> {
        facebook_only();
        scribe::new_thrift_scribe_sink_if_enabled(
            fb,
            buffer_size,
            retry_backoff,
            retry_attempts,
            message_batch_size,
        )
        .map(|maybe_scribe| maybe_scribe.map(|scribe| Arc::new(scribe) as _))
    }

    /// Prepares an event stream for a request by bootstrapping an event source and EventDispatcher pair. The given
    /// EventDispatcher will log to the returned EventSource and (optionally) to Scribe if enabled via buckconfig.
    pub async fn prepare_events(
        &self,
        trace_id: TraceId,
    ) -> SharedResult<(impl EventSource, EventDispatcher)> {
        // facebook only: logging events to Scribe.
        facebook_only();
        let (events, sink) = buck2_events::create_source_sink_pair();
        let data = self.data()?;
        let dispatcher = if let Some(scribe_sink) = data.scribe_sink.dupe() {
            EventDispatcher::new(trace_id, TeeSink::new(scribe_sink, sink))
        } else {
            EventDispatcher::new(trace_id, sink)
        };
        Ok((events, dispatcher))
    }

    /// Prepares a ServerCommandContext for processing a complex command (that accesses the dice computation graph, for example).
    ///
    /// This initializes (if necessary) the shared daemon state and syncs the watchman query (to flush any recent filesystem events).
    pub async fn prepare_command(
        &self,
        dispatcher: EventDispatcher,
        drop_guard: ActiveCommandDropGuard,
    ) -> SharedResult<BaseServerCommandContext> {
        let data = self.data();

        dispatcher.instant_event(buck2_data::RestartConfiguration {
            enable_restarter: data.as_ref().map_or(false, |d| d.enable_restarter),
        });

        tag_result!(
            "eden_not_connected",
            check_working_dir::check_working_dir(),
            quiet: true,
            daemon_in_memory_state_is_corrupted: true,
            task: false
        )?;

        self.validate_buck_out_mount()
            .context("Error validating buck-out mount")?;

        let data = data?;

        let tags = vec![
            format!(
                "dice-detect-cycles:{}",
                data.dice_manager
                    .unsafe_dice()
                    .detect_cycles()
                    .variant_name()
            ),
            format!(
                "which-dice:{}",
                data.dice_manager.unsafe_dice().which_dice().variant_name()
            ),
            format!("hash-all-commands:{}", data.hash_all_commands),
            format!(
                "sqlite-materializer-state:{}",
                data.disk_state_options.sqlite_materializer_state
            ),
            format!("eden-io-v2:{}", data.eden_io_v2),
        ];

        dispatcher.instant_event(buck2_data::TagEvent { tags });

        // Sync any FS changes and invalidate DICE state if necessary.  Get the Eden
        // version of the underlying system in parallel if available.
        let (_, eden_version) =
            futures::future::try_join(data.io.settle(), data.io.eden_version()).await?;

        dispatcher.instant_event(buck2_data::IoProviderInfo { eden_version });

        Ok(BaseServerCommandContext {
            _fb: self.fb,
            project_root: self.paths.project_root().clone(),
            dice_manager: data.dice_manager.dupe(),
            io: data.io.dupe(),
            re_client_manager: data.re_client_manager.dupe(),
            blocking_executor: data.blocking_executor.dupe(),
            materializer: data.materializer.dupe(),
            file_watcher: data.file_watcher.dupe(),
            events: dispatcher,
            forkserver: data.forkserver.dupe(),
            hash_all_commands: data.hash_all_commands,
            use_network_action_output_cache: data.use_network_action_output_cache,
            _drop_guard: drop_guard,
            daemon_start_time: data.start_time,
            create_unhashed_outputs_lock: data.create_unhashed_outputs_lock.dupe(),
            http_client: data.http_client.dupe(),
        })
    }

    pub fn data(&self) -> anyhow::Result<Arc<DaemonStateData>> {
        Ok(self.data.dupe()?)
    }

    fn validate_buck_out_mount(&self) -> anyhow::Result<()> {
        #[cfg(any(fbcode_build, cargo_internal_build))]
        {
            use buck2_core::soft_error;

            let project_root = self.paths.project_root().root();
            if !detect_eden::is_eden(project_root.to_path_buf())? {
                return Ok(());
            }

            let buck_out_root = project_root.join(InvocationPaths::buck_out_dir_prefix());

            if let Some(buck_out_root_meta) = fs_util::symlink_metadata_if_exists(buck_out_root)? {
                // If buck-out is a symlink, we'll be happy with that.
                if buck_out_root_meta.is_symlink() {
                    return Ok(());
                }

                // If we are on UNIX, then buck-out could also be on a different device from the repo.
                // We don't check which kind of device, we just assume it's not mounted completely
                // wrong.
                #[cfg(unix)]
                {
                    use std::os::unix::fs::MetadataExt;

                    let project_device = fs_util::symlink_metadata(project_root)?.dev();
                    let buck_out_device = buck_out_root_meta.dev();

                    if project_device != buck_out_device {
                        return Ok(());
                    }
                }
            }

            soft_error!(
                "eden_buck_out",
                anyhow::anyhow!(
                    "Buck is running in an Eden repository, but `buck-out` is not redirected. \
                     This will likely lead to failed or slow builds. \
                     To remediate, run `eden redirect fixup`."
                )
            )?;
        }

        Ok(())
    }
}

fn convert_algorithm_kind(kind: DigestAlgorithmKind) -> anyhow::Result<DigestAlgorithm> {
    anyhow::Ok(match kind {
        DigestAlgorithmKind::Sha1 => DigestAlgorithm::Sha1,
        DigestAlgorithmKind::Sha256 => DigestAlgorithm::Sha256,
        DigestAlgorithmKind::Blake3 => DigestAlgorithm::Blake3,
        DigestAlgorithmKind::Blake3Keyed => {
            #[cfg(fbcode_build)]
            {
                let key = blake3_constants::BLAKE3_HASH_KEY
                    .as_bytes()
                    .try_into()
                    .context("BLAKE3_HASH_KEY is the wrong size")?;

                DigestAlgorithm::Blake3Keyed { key }
            }

            #[cfg(not(fbcode_build))]
            {
                // We probably should just add it as a separate buckconfig, there is
                // zero reason not to.
                return Err(anyhow::anyhow!(
                    "{} is not supported in the open source build",
                    kind
                ));
            }
        }
    })
}


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

/// Serialized JSON.
/// It can be output as is to a JSON file.
pub(crate) struct QuotedJson(String);

impl QuotedJson {
    pub(crate) fn as_str(&self) -> &str {
        &self.0
    }

    pub(crate) fn quote_str(s: &str) -> QuotedJson {
        QuotedJson(serde_json::ser::to_string(s).unwrap())
    }

    pub(crate) fn quote_display(d: impl Display) -> QuotedJson {
        Self::quote_str(&d.to_string())
    }

    pub(crate) fn list(items: impl IntoIterator<Item = QuotedJson>) -> QuotedJson {
        let mut s = String::new();
        s.push('[');
        for (i, item) in items.into_iter().enumerate() {
            if i != 0 {
                s.push(',');
            }
            s.push_str(item.as_str());
        }
        s.push(']');
        QuotedJson(s)
    }

    pub(crate) fn from_serde_json_value(v: serde_json::Value) -> QuotedJson {
        QuotedJson(serde_json::ser::to_string(&v).unwrap())
    }
}

#[cfg(test)]
mod tests {
    use crate::json::QuotedJson;

    #[test]
    fn test_quote_string() {
        assert_eq!("\"\"", QuotedJson::quote_str("").as_str());
        assert_eq!("\"\\\"\"", QuotedJson::quote_str("\"").as_str());
        assert_eq!("\"\\\\\"", QuotedJson::quote_str("\\").as_str());
        assert_eq!(
            "\"a\\n\\r\\t\\b\\f\"",
            QuotedJson::quote_str("a\n\r\t\x08\x0c").as_str()
        );
    }

    #[test]
    pub fn test_list() {
        assert_eq!("[]", QuotedJson::list(Vec::new()).as_str());
        assert_eq!(
            "[\"a\"]",
            QuotedJson::list(vec![QuotedJson::quote_str("a")]).as_str()
        );
        assert_eq!(
            "[\"a\",\"b\"]",
            QuotedJson::list(vec![QuotedJson::quote_str("a"), QuotedJson::quote_str("b")]).as_str()
        );
    }
}


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

//! Implementation of several server commands.

#![feature(async_closure)]
#![feature(box_patterns)]
#![feature(try_blocks)]

pub mod commands;
pub mod dot;
pub(crate) mod json;
pub mod target_hash;

pub fn init_late_bindings() {
    commands::query::printer::init_print_action_node();
}