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

use anyhow::Context;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_events::dispatch::span_async;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;

use crate::ctx::BaseServerCommandContext;
use crate::ctx::ServerCommandContext;

pub(crate) async fn materialize_command(
    context: &ServerCommandContext<'_>,
    req: buck2_cli_proto::MaterializeRequest,
) -> anyhow::Result<buck2_cli_proto::MaterializeResponse> {
    let metadata = context.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::MaterializeCommandStart {}.into()),
    };
    span_async(start_event, async move {
        let result = materialize(&context.base_context, req.paths)
            .await
            .map(|()| buck2_cli_proto::MaterializeResponse {})
            .context("Failed to materialize paths");
        let end_event = command_end(metadata, &result, buck2_data::MaterializeCommandEnd {});
        (result, end_event)
    })
    .await
}

async fn materialize(
    server_ctx: &BaseServerCommandContext,
    paths: Vec<String>,
) -> anyhow::Result<()> {
    let mut project_paths = Vec::new();
    for path in paths {
        project_paths.push(ProjectRelativePath::new(&path)?.to_owned())
    }
    server_ctx
        .materializer
        .ensure_materialized(project_paths)
        .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.
 */

#[derive(Clone, Debug)]
pub struct Counters {
    pub bytes_sent: u64,
    pub bytes_recv: u64,
}

#[cfg(any(target_os = "macos", target_os = "linux"))]
mod collector {
    use std::collections::HashMap;
    use std::sync::Arc;
    use std::sync::Mutex;

    use anyhow::Context;
    use dupe::Dupe;
    use psutil::network::NetIoCountersCollector;

    use super::*;

    /// Collects system network I/O stats (tx, rx) on a per-NIC basis. Returns delta
    /// counters in between collection periods. Most useful when done on a periodic
    /// basis (e.g. in the daemon heartbeat process, which happens every 1s).
    #[derive(Clone, Debug, Dupe)]
    pub struct SystemNetworkIoCollector {
        collector: Arc<Mutex<NetIoCountersCollector>>,
    }

    impl SystemNetworkIoCollector {
        pub fn new() -> Self {
            Self {
                collector: Arc::new(Mutex::new(NetIoCountersCollector::default())),
            }
        }

        /// Collect stats for each NIC. Some notes:
        /// * If a new NIC appears between collection periods, we'll start keeping
        ///   track of it.
        /// * If a NIC *disappears*, then we stop reporting on its stats.
        pub fn collect(&self) -> anyhow::Result<Option<HashMap<String, Counters>>> {
            let mut collector = self.collector.lock().expect("poisoned lock");
            let counters: HashMap<_, _> = collector
                .net_io_counters_pernic()
                .context("collecting old counters")?
                .into_iter()
                .map(|(nic, counters)| {
                    (
                        nic,
                        Counters {
                            bytes_sent: counters.bytes_sent(),
                            bytes_recv: counters.bytes_recv(),
                        },
                    )
                })
                .collect();

            Ok(Some(counters))
        }
    }
}

// psutil network stats aren't implemented on windows or other unix-likes.
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
mod collector {
    use std::collections::HashMap;

    use dupe::Dupe;

    use super::*;

    #[derive(Clone, Debug, Dupe)]
    pub struct SystemNetworkIoCollector;

    impl SystemNetworkIoCollector {
        pub fn new() -> Self {
            Self
        }

        pub fn collect(&self) -> anyhow::Result<Option<HashMap<String, Counters>>> {
            Ok(None)
        }
    }
}

pub use collector::SystemNetworkIoCollector;


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

use std::path::Path;
use std::slice;
use std::sync::Arc;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_analysis::analysis::calculation::profile_analysis;
use buck2_analysis::analysis::calculation::profile_analysis_recursively;
use buck2_cli_proto::profile_request::ProfileOpts;
use buck2_cli_proto::target_profile::Action;
use buck2_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::pattern::resolve::resolve_target_patterns;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::PackageSpec;
use buck2_core::target::label::TargetLabel;
use buck2_interpreter::dice::starlark_profiler::StarlarkProfilerConfiguration;
use buck2_interpreter::starlark_profiler::StarlarkProfileDataAndStats;
use buck2_interpreter::starlark_profiler::StarlarkProfiler;
use buck2_interpreter::starlark_profiler::StarlarkProfilerOrInstrumentation;
use buck2_interpreter_for_build::interpreter::dice_calculation_delegate::HasCalculationDelegate;
use buck2_node::target_calculation::ConfiguredTargetCalculation;
use buck2_profile::get_profile_response;
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::pattern::parse_patterns_from_cli_args;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceTransaction;
use dupe::Dupe;

async fn generate_profile_analysis(
    ctx: DiceTransaction,
    package: PackageLabel,
    spec: PackageSpec<TargetPatternExtra>,
    global_target_platform: Option<TargetLabel>,
    profile_mode: &StarlarkProfilerConfiguration,
) -> anyhow::Result<Arc<StarlarkProfileDataAndStats>> {
    let (target, TargetPatternExtra) = match spec {
        PackageSpec::Targets(targets) => one(targets).context("Invalid targets"),
        PackageSpec::All => Err(anyhow::Error::msg("Cannot use a package")),
    }
    .context("Did not find exactly one target")?;

    let label = TargetLabel::new(package.dupe(), target.as_ref());

    let configured_target = ctx
        .get_configured_target(&label, global_target_platform.as_ref())
        .await?;

    match profile_mode {
        StarlarkProfilerConfiguration::ProfileLastAnalysis(profile_mode) => {
            profile_analysis(&ctx, &configured_target, profile_mode)
                .await
                .context("Analysis failed")
        }
        StarlarkProfilerConfiguration::ProfileAnalysisRecursively(_) => {
            profile_analysis_recursively(&ctx, &configured_target)
                .await
                .context("Analysis failed")
                .map(Arc::new)
        }
        _ => Err(anyhow::anyhow!("Incorrect profile mode (internal error)")),
    }
}

async fn generate_profile_loading(
    ctx: DiceTransaction,
    package: PackageLabel,
    spec: PackageSpec<TargetPatternExtra>,
    profile_mode: &StarlarkProfilerConfiguration,
) -> anyhow::Result<Arc<StarlarkProfileDataAndStats>> {
    match spec {
        PackageSpec::Targets(..) => {
            return Err(anyhow::Error::msg("Must use a package"));
        }
        PackageSpec::All => {}
    }

    let calculation = ctx
        .get_interpreter_calculator(package.cell_name(), BuildFileCell::new(package.cell_name()))
        .await?;

    let mut profiler = StarlarkProfiler::new(profile_mode.profile_last_loading()?.dupe(), false);

    calculation
        .eval_build_file(
            package,
            &mut StarlarkProfilerOrInstrumentation::for_profiler(&mut profiler),
        )
        .await?;

    profiler.finish().map(Arc::new)
}

pub async fn profile_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
    req: buck2_cli_proto::ProfileRequest,
) -> anyhow::Result<buck2_cli_proto::ProfileResponse> {
    run_server_command(ProfileServerCommand { req }, ctx, partial_result_dispatcher).await
}

struct ProfileServerCommand {
    req: buck2_cli_proto::ProfileRequest,
}

#[async_trait]
impl ServerCommandTemplate for ProfileServerCommand {
    type StartEvent = buck2_data::ProfileCommandStart;
    type EndEvent = buck2_data::ProfileCommandEnd;
    type Response = buck2_cli_proto::ProfileResponse;
    type PartialResult = NoPartialResult;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        _partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        let output = AbsPath::new(Path::new(&self.req.destination_path))?;

        let profile_mode = starlark_profiler_configuration_from_request(&self.req)?;

        match self
            .req
            .profile_opts
            .as_ref()
            .expect("Target profile not populated")
        {
            ProfileOpts::TargetProfile(opts) => {
                let action = buck2_cli_proto::target_profile::Action::from_i32(opts.action)
                    .context("Invalid action")?;
                let target_pattern = opts
                    .target_pattern
                    .as_ref()
                    .context("Missing target pattern")?;

                let context = self
                    .req
                    .context
                    .as_ref()
                    .context("Missing client context")?;

                let profile_data = generate_profile(
                    server_ctx,
                    ctx,
                    context,
                    target_pattern,
                    action,
                    &profile_mode,
                )
                .await?;

                get_profile_response(profile_data, &self.req, output)
            }
            _ => {
                return Err(anyhow::anyhow!(
                    "Expected target profile opts, not BXL profile opts"
                ));
            }
        }
    }

    fn is_success(&self, _response: &Self::Response) -> bool {
        // No response if we failed.
        true
    }
}

async fn generate_profile(
    server_ctx: &dyn ServerCommandContextTrait,
    mut ctx: DiceTransaction,
    client_ctx: &ClientContext,
    pattern: &buck2_data::TargetPattern,
    action: Action,
    profile_mode: &StarlarkProfilerConfiguration,
) -> anyhow::Result<Arc<StarlarkProfileDataAndStats>> {
    let cells = ctx.get_cell_resolver().await?;

    let global_target_platform =
        target_platform_from_client_context(client_ctx, server_ctx, &ctx).await?;

    let parsed_patterns = parse_patterns_from_cli_args::<TargetPatternExtra>(
        &mut ctx,
        slice::from_ref(pattern),
        server_ctx.working_dir(),
    )
    .await?;

    let resolved_pattern =
        resolve_target_patterns(&cells, &parsed_patterns, &ctx.file_ops()).await?;

    let (package, spec) =
        one(resolved_pattern.specs).context("Did not find exactly one pattern")?;

    match action {
        Action::Analysis => {
            generate_profile_analysis(ctx, package, spec, global_target_platform, profile_mode)
                .await
        }
        Action::Loading => generate_profile_loading(ctx, package, spec, profile_mode).await,
    }
}

fn one<T>(it: impl IntoIterator<Item = T>) -> anyhow::Result<T> {
    let mut it = it.into_iter();
    let val = it.next().context("No value found")?;
    if it.next().is_some() {
        return Err(anyhow::Error::msg("More than one value found"));
    }
    Ok(val)
}


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

use anyhow::Context as _;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_core::io_counters::IoCounterKey;
use buck2_events::EventSink;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::re::manager::ReConnectionManager;
use buck2_util::process_stats::process_stats;
use buck2_util::system_stats::UnixSystemStats;
use dice::Dice;
use dupe::Dupe;

use crate::jemalloc_stats::get_allocator_stats;
use crate::net_io::SystemNetworkIoCollector;

/// Stores state handles necessary to produce snapshots.
#[derive(Clone, Dupe)]
pub struct SnapshotCollector {
    re_client_manager: Arc<ReConnectionManager>,
    blocking_executor: Arc<dyn BlockingExecutor>,
    daemon_start_time: Instant,
    dice: Arc<Dice>,
    materializer: Arc<dyn Materializer>,
    event_sink: Option<Arc<dyn EventSink>>,
    net_io_collector: SystemNetworkIoCollector,
    /// This is only used to obtain statistics from the HTTP client.
    http_client: CountingHttpClient,
}

impl SnapshotCollector {
    pub fn new(
        re_client_manager: Arc<ReConnectionManager>,
        blocking_executor: Arc<dyn BlockingExecutor>,
        daemon_start_time: Instant,
        dice: Arc<Dice>,
        materializer: Arc<dyn Materializer>,
        event_sink: Option<Arc<dyn EventSink>>,
        http_client: CountingHttpClient,
    ) -> SnapshotCollector {
        SnapshotCollector {
            re_client_manager,
            blocking_executor,
            daemon_start_time,
            dice,
            materializer,
            event_sink,
            net_io_collector: SystemNetworkIoCollector::new(),
            http_client,
        }
    }

    /// We emit a Snapshot before a BaseCommandContext is made available.
    /// Initializes snapshot with all information we don't get off a BaseCommandContext.
    /// This lets us send our first Snapshot before fully initializing/syncing.
    pub fn pre_initialization_snapshot(daemon_start_time: Instant) -> buck2_data::Snapshot {
        let mut snapshot = buck2_data::Snapshot::default();
        add_system_metrics(&mut snapshot, daemon_start_time);
        snapshot
    }

    /// Create a new Snapshot.
    pub fn create_snapshot(&self) -> buck2_data::Snapshot {
        let mut snapshot = Self::pre_initialization_snapshot(self.daemon_start_time);
        self.add_daemon_metrics(&mut snapshot);
        self.add_re_metrics(&mut snapshot);
        self.add_http_metrics(&mut snapshot);
        self.add_io_metrics(&mut snapshot);
        self.add_dice_metrics(&mut snapshot);
        self.add_materializer_metrics(&mut snapshot);
        self.add_sink_metrics(&mut snapshot);
        self.add_net_io_metrics(&mut snapshot);
        snapshot
    }

    fn add_daemon_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        snapshot.blocking_executor_io_queue_size = self.blocking_executor.queue_size() as u64;
    }

    fn add_io_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        // Using loop here to make sure no key is forgotten.
        for key in IoCounterKey::ALL {
            let pointer = match key {
                IoCounterKey::Stat => &mut snapshot.io_in_flight_stat,
                IoCounterKey::Copy => &mut snapshot.io_in_flight_copy,
                IoCounterKey::Symlink => &mut snapshot.io_in_flight_symlink,
                IoCounterKey::Hardlink => &mut snapshot.io_in_flight_hardlink,
                IoCounterKey::MkDir => &mut snapshot.io_in_flight_mk_dir,
                IoCounterKey::ReadDir => &mut snapshot.io_in_flight_read_dir,
                IoCounterKey::ReadDirEden => &mut snapshot.io_in_flight_read_dir_eden,
                IoCounterKey::RmDir => &mut snapshot.io_in_flight_rm_dir,
                IoCounterKey::RmDirAll => &mut snapshot.io_in_flight_rm_dir_all,
                IoCounterKey::StatEden => &mut snapshot.io_in_flight_stat_eden,
                IoCounterKey::Chmod => &mut snapshot.io_in_flight_chmod,
                IoCounterKey::ReadLink => &mut snapshot.io_in_flight_read_link,
                IoCounterKey::Remove => &mut snapshot.io_in_flight_remove,
                IoCounterKey::Rename => &mut snapshot.io_in_flight_rename,
                IoCounterKey::Read => &mut snapshot.io_in_flight_read,
                IoCounterKey::Write => &mut snapshot.io_in_flight_write,
                IoCounterKey::Canonicalize => &mut snapshot.io_in_flight_canonicalize,
                IoCounterKey::EdenSettle => &mut snapshot.io_in_flight_eden_settle,
            };
            *pointer = key.get();
        }
    }

    fn add_re_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        fn inner(
            snapshot: &mut buck2_data::Snapshot,
            re: &ReConnectionManager,
        ) -> anyhow::Result<()> {
            let stats = match re
                .get_network_stats()
                .context("Error collecting network stats")?
            {
                Some(stats) => stats,
                None => return Ok(()),
            };

            snapshot.re_download_bytes = stats.downloaded;
            snapshot.re_upload_bytes = stats.uploaded;
            snapshot.re_uploads_started = stats.uploads.started;
            snapshot.re_uploads_finished_successfully = stats.uploads.finished_successfully;
            snapshot.re_uploads_finished_with_error = stats.uploads.finished_with_error;
            snapshot.re_downloads_started = stats.downloads.started;
            snapshot.re_downloads_finished_successfully = stats.downloads.finished_successfully;
            snapshot.re_downloads_finished_with_error = stats.downloads.finished_with_error;
            snapshot.re_action_cache_started = stats.action_cache.started;
            snapshot.re_action_cache_finished_successfully =
                stats.action_cache.finished_successfully;
            snapshot.re_action_cache_finished_with_error = stats.action_cache.finished_with_error;
            snapshot.re_executes_started = stats.executes.started;
            snapshot.re_executes_finished_successfully = stats.executes.finished_successfully;
            snapshot.re_executes_finished_with_error = stats.executes.finished_with_error;
            snapshot.re_materializes_started = stats.materializes.started;
            snapshot.re_materializes_finished_successfully =
                stats.materializes.finished_successfully;
            snapshot.re_materializes_finished_with_error = stats.materializes.finished_with_error;
            snapshot.re_write_action_results_started = stats.write_action_results.started;
            snapshot.re_write_action_results_finished_successfully =
                stats.write_action_results.finished_successfully;
            snapshot.re_write_action_results_finished_with_error =
                stats.write_action_results.finished_with_error;
            snapshot.re_get_digest_expirations_started = stats.get_digest_expirations.started;
            snapshot.re_get_digest_expirations_finished_successfully =
                stats.get_digest_expirations.finished_successfully;
            snapshot.re_get_digest_expirations_finished_with_error =
                stats.get_digest_expirations.finished_with_error;

            Ok(())
        }

        // Nothing we can do if we get an error, unfortunately.
        if let Err(e) = inner(snapshot, &self.re_client_manager) {
            tracing::debug!("Error collecting network stats: {:#}", e);
        }
    }

    fn add_http_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        let stats = self.http_client.get_updated_http_network_stats();
        snapshot.http_download_bytes = stats.downloaded_bytes;
    }

    fn add_dice_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        let metrics = self.dice.metrics();
        snapshot.dice_key_count = metrics.key_count as u64;
        snapshot.dice_currently_active_key_count = metrics.currently_active_key_count as u64;
        snapshot.dice_active_transaction_count = metrics.active_transaction_count;
    }

    fn add_materializer_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        self.materializer.add_snapshot_stats(snapshot);
    }

    fn add_sink_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        if let Some(metrics) = self.event_sink.as_ref().and_then(|sink| sink.stats()) {
            snapshot.sink_successes = Some(metrics.successes);
            snapshot.sink_failures = Some(metrics.failures);
            snapshot.sink_buffer_depth = Some(metrics.buffered);
            snapshot.sink_dropped = Some(metrics.dropped);
        }
    }

    fn add_net_io_metrics(&self, snapshot: &mut buck2_data::Snapshot) {
        if let Ok(Some(mut net_io_counters_per_nic)) = self.net_io_collector.collect() {
            net_io_counters_per_nic
                .retain(|k, _| ["en", "eth"].iter().any(|prefix| k.starts_with(prefix)));
            snapshot.network_interface_stats = net_io_counters_per_nic
                .into_iter()
                .map(|(nic, counters)| {
                    (
                        nic,
                        buck2_data::NetworkInterfaceStats {
                            tx_bytes: counters.bytes_sent,
                            rx_bytes: counters.bytes_recv,
                        },
                    )
                })
                .collect();
        } else {
            snapshot.network_interface_stats = HashMap::new();
        }
    }
}

fn add_system_metrics(snapshot: &mut buck2_data::Snapshot, daemon_start_time: Instant) {
    let process_stats = process_stats();
    if let Some(max_rss_bytes) = process_stats.max_rss_bytes {
        snapshot.buck2_max_rss = max_rss_bytes;
    }
    if let Some(user_cpu_us) = process_stats.user_cpu_us {
        snapshot.buck2_user_cpu_us = user_cpu_us;
    }
    if let Some(system_cpu_us) = process_stats.system_cpu_us {
        snapshot.buck2_system_cpu_us = system_cpu_us;
    }
    snapshot.daemon_uptime_s = daemon_start_time.elapsed().as_secs();
    snapshot.buck2_rss = process_stats.rss_bytes;
    let allocator_stats = get_allocator_stats().ok();
    if let Some(alloc_stats) = allocator_stats {
        snapshot.malloc_bytes_active = alloc_stats.bytes_active;
        snapshot.malloc_bytes_allocated = alloc_stats.bytes_allocated;
    }

    if let Some(UnixSystemStats {
        load1,
        load5,
        load15,
    }) = UnixSystemStats::get()
    {
        snapshot.unix_system_stats = Some(buck2_data::UnixSystemStats {
            load1,
            load5,
            load15,
        });
    }
}


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

use std::time::Duration;

use anyhow::Context as _;
use buck2_events::dispatch::span_async;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::streaming_request_handler::StreamingRequestHandler;
use futures::future::FutureExt;
use gazebo::prelude::*;
use tokio::time::MissedTickBehavior;

use crate::active_commands;

pub(crate) async fn run_subscription_server_command(
    ctx: &dyn ServerCommandContextTrait,
    mut partial_result_dispatcher: PartialResultDispatcher<
        buck2_cli_proto::SubscriptionResponseWrapper,
    >,
    mut req: StreamingRequestHandler<buck2_cli_proto::SubscriptionRequestWrapper>,
) -> anyhow::Result<buck2_cli_proto::SubscriptionCommandResponse> {
    let metadata = ctx.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::SubscriptionCommandStart {}.into()),
    };
    span_async(start_event, async move {
        let result: anyhow::Result<buck2_cli_proto::SubscriptionCommandResponse> = try {
            // NOTE: Long term if we expose more things here then we should probably move this error to
            // only occur when we try to actually interact with materializer subscriptioons
            let materializer = ctx
                .materializer();

            let materializer = materializer
                .as_deferred_materializer_extension()
                .context("Subscriptions only work with the deferred materializer")?;

            let mut materializer_subscription = materializer
                .create_subscription()
                .await
                .context("Error creating a materializer subscription")?;

            let mut wants_active_commands = false;

            let mut ticker = tokio::time::interval(Duration::from_millis(100));
            ticker.set_missed_tick_behavior(MissedTickBehavior::Skip);

            let disconnect = loop {
                futures::select! {
                    message = req.message().fuse() => {
                        use buck2_subscription_proto::subscription_request::Request;

                        match message?.request.context("Empty message")?.request.context("Empty request")? {
                            Request::Disconnect(disconnect) => {
                                break disconnect;
                            }
                            Request::SubscribeToPaths(buck2_subscription_proto::SubscribeToPaths { paths }) => {
                                let paths = paths.into_try_map(|path| path.try_into())?;
                                materializer_subscription.subscribe_to_paths(paths);
                            }
                            Request::UnsubscribeFromPaths(buck2_subscription_proto::UnsubscribeFromPaths { paths }) => {
                                let paths = paths.into_try_map(|path| path.try_into())?;
                                materializer_subscription.unsubscribe_from_paths(paths);
                            }
                            Request::SubscribeToActiveCommands(buck2_subscription_proto::SubscribeToActiveCommands {}) => {
                                wants_active_commands = true;
                            }
                        }
                    }
                    path = materializer_subscription.next_materialization().fuse() => {
                        let path = path.context("Materializer hung up")?;
                        partial_result_dispatcher.emit(buck2_cli_proto::SubscriptionResponseWrapper {
                            response: Some(buck2_subscription_proto::SubscriptionResponse {
                                response: Some(buck2_subscription_proto::Materialized { path: path.to_string() }.into())
                            })
                        });
                    }
                    _ = ticker.tick().fuse() => {
                        if wants_active_commands {
                            let snapshot = active_commands_snapshot();
                            partial_result_dispatcher.emit(buck2_cli_proto::SubscriptionResponseWrapper {
                                response: Some(buck2_subscription_proto::SubscriptionResponse {
                                    response: Some(snapshot.into())
                                })
                            });
                        }
                    }
                }
            };

            partial_result_dispatcher.emit(buck2_cli_proto::SubscriptionResponseWrapper {
                response: Some(buck2_subscription_proto::SubscriptionResponse {
                    response: Some(buck2_subscription_proto::Goodbye {
                        reason: disconnect.reason,
                        ok: disconnect.ok,
                    }.into())
                })
            });

            buck2_cli_proto::SubscriptionCommandResponse {}
        };

        let end_event = command_end(metadata, &result, buck2_data::SubscriptionCommandEnd {});
        (result, end_event)
    })
    .await
}

fn active_commands_snapshot() -> buck2_subscription_proto::ActiveCommandsSnapshot {
    let active_commands = active_commands::active_commands()
        .iter()
        .map(|(trace_id, handle)| {
            let state = handle.state();
            let spans = state.spans();

            buck2_subscription_proto::ActiveCommand {
                trace_id: trace_id.to_string(),
                argv: state.argv.clone(),
                stats: Some(buck2_subscription_proto::ActiveCommandStats {
                    open_spans: spans.open,
                    closed_spans: spans.closed,
                    pending_spans: spans.pending,
                }),
            }
        })
        .collect();

    buck2_subscription_proto::ActiveCommandsSnapshot { active_commands }
}


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

use anyhow::Context;
use buck2_cli_proto::trace_io_request;
use buck2_cli_proto::trace_io_response;
use buck2_common::file_ops::RawSymlink;
use buck2_common::io::trace::TracingIoProvider;
use buck2_events::dispatch::span_async;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;

use crate::ctx::ServerCommandContext;

pub(crate) async fn trace_io_command(
    context: &ServerCommandContext<'_>,
    req: buck2_cli_proto::TraceIoRequest,
) -> anyhow::Result<buck2_cli_proto::TraceIoResponse> {
    let metadata = context.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::TraceIoCommandStart {}.into()),
    };
    span_async(start_event, async move {
        let tracing_provider = &context
            .base_context
            .io
            .as_any()
            .downcast_ref::<TracingIoProvider>();
        let respond_with_trace = matches!(
            req.read_state,
            Some(trace_io_request::ReadIoTracingState { with_trace: true })
        );

        let result = match (tracing_provider, respond_with_trace) {
            (Some(provider), true) => build_response_with_trace(&context, provider).await,
            (Some(_), false) => Ok(buck2_cli_proto::TraceIoResponse {
                enabled: true,
                trace: Vec::new(),
                external_entries: Vec::new(),
                relative_symlinks: Vec::new(),
                external_symlinks: Vec::new(),
            }),
            (None, _) => Ok(buck2_cli_proto::TraceIoResponse {
                enabled: false,
                trace: Vec::new(),
                external_entries: Vec::new(),
                relative_symlinks: Vec::new(),
                external_symlinks: Vec::new(),
            }),
        };

        let end_event = command_end(metadata, &result, buck2_data::TraceIoCommandEnd {});
        (result, end_event)
    })
    .await
}

async fn build_response_with_trace(
    context: &ServerCommandContext<'_>,
    provider: &TracingIoProvider,
) -> anyhow::Result<buck2_cli_proto::TraceIoResponse> {
    // Materialize buck-out paths so they can be archived.
    let buck_out_entries: Vec<_> = provider.trace().buck_out_entries();
    context
        .materializer()
        .ensure_materialized(buck_out_entries.clone())
        .await
        .context("Error materializing buck-out paths for trace")?;

    let mut entries = provider.trace().project_entries();
    entries.extend(buck_out_entries);

    let mut relative_symlinks = Vec::new();
    let mut external_symlinks = Vec::new();
    for link in provider.trace().symlinks.iter() {
        match &link.to {
            RawSymlink::Relative(to) => {
                relative_symlinks.push(trace_io_response::RelativeSymlink {
                    link: link.at.to_string(),
                    target: to.to_string(),
                });
            }
            RawSymlink::External(external) => {
                external_symlinks.push(trace_io_response::ExternalSymlink {
                    link: link.at.to_string(),
                    target: external.target_str().to_owned(),
                    remaining_path: external.remaining_path().map(|path| path.to_string()),
                });
            }
        }
    }

    Ok(buck2_cli_proto::TraceIoResponse {
        enabled: true,
        trace: entries.into_iter().map(|path| path.to_string()).collect(),
        external_entries: provider
            .trace()
            .external_entries()
            .into_iter()
            .map(|path| path.to_string())
            .collect(),
        relative_symlinks,
        external_symlinks,
    })
}


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

use buck2_interpreter_for_build::interpreter::global_interpreter_state::GlobalInterpreterState;
use maplit::hashset;
use once_cell::sync::Lazy;
use starlark::docs::Doc;
use starlark::docs::DocItem;
use starlark::docs::DocMember;
use starlark::docs::DocObject;
use starlark_map::smallmap;

use crate::builtin_docs::docs::builtin_doc;

// List of BXL global function names.
// (@wendyy) I don't think this is a nice implementation, but this offers us finer
// grain control over documentation generated by globals. Let's at least get the
// documentation out there before coming up with a cleaner solution
static BXL_GLOBAL_METHOD_NAME_SET: Lazy<HashSet<&str>> = Lazy::new(|| {
    hashset![
        "sub_target",
        "configured_sub_target",
        "target_set",
        "get_path_without_materialization",
        "get_paths_without_materialization"
    ]
});

// Globals that are in BXL. For BXL, we will probably only have global functions (as opposed
// to objects and properties).
pub(crate) fn get_builtin_bxl_docs(
    interpreter_state: Arc<GlobalInterpreterState>,
) -> anyhow::Result<Vec<Doc>> {
    let b_o = interpreter_state.bxl_file_global_env.documentation();
    let mut docs = vec![];
    for member in b_o.members {
        if BXL_GLOBAL_METHOD_NAME_SET.contains(member.0.as_str()) {
            match member.1 {
                DocMember::Function(function) => {
                    // Convert the function to its own object so it gets its own markdown file name.
                    // Otherwise, it will be named `native.md`, which is confusing in static docs.
                    let doc_item = DocItem::Object(DocObject {
                        docs: None,
                        members: smallmap! {member.0.clone() => DocMember::Function(function)},
                    });

                    docs.push(builtin_doc(member.0.as_str(), "bxl", doc_item));
                }
                _ => (),
            }
        }
    }

    Ok(docs)
}


/*
 * 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::collections::HashSet;
use std::path::Path;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::unstable_docs_request;
use buck2_cli_proto::UnstableDocsRequest;
use buck2_cli_proto::UnstableDocsResponse;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::CellAliasResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_interpreter::parse_import::parse_import_with_config;
use buck2_interpreter::parse_import::ParseImportOptions;
use buck2_interpreter::prelude_path::prelude_path;
use buck2_interpreter_for_build::interpreter::global_interpreter_state::GlobalInterpreterState;
use buck2_interpreter_for_build::interpreter::global_interpreter_state::HasGlobalInterpreterState;
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::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;
use starlark::collections::SmallMap;
use starlark::docs::get_registered_starlark_docs;
use starlark::docs::Doc;
use starlark::docs::DocItem;
use starlark::docs::DocModule;
use starlark::docs::Identifier;
use starlark::docs::Location;
use starlark::environment::Globals;

use super::bxl_docs::get_builtin_bxl_docs;
use crate::builtin_docs::markdown::generate_markdown_files;

#[derive(Debug, thiserror::Error)]
enum DocsError {
    #[error("Unknown format requested (internal error)")]
    UnknownFormat,
}

fn parse_import_paths(
    cell_resolver: &CellAliasResolver,
    current_dir: &CellPath,
    current_cell: BuildFileCell,
    symbol_patterns: &[String],
) -> anyhow::Result<HashSet<ImportPath>> {
    const PARSE_OPTIONS: ParseImportOptions = ParseImportOptions {
        allow_missing_at_symbol: true,
        allow_relative_imports: true,
    };

    symbol_patterns
        .iter()
        .map(|symbol_pattern| {
            let path = parse_import_with_config(
                cell_resolver,
                current_dir,
                symbol_pattern,
                &PARSE_OPTIONS,
            )?;
            ImportPath::new_with_build_file_cells(path, current_cell)
        })
        .collect()
}

pub(crate) fn builtin_doc<S: ToString>(name: S, directory: &str, item: DocItem) -> Doc {
    let mut custom_attrs = HashMap::new();
    if !directory.is_empty() {
        custom_attrs.insert("directory".to_owned(), directory.to_owned());
    }

    Doc {
        id: Identifier {
            name: name.to_string(),
            location: None,
        },
        item,
        custom_attrs,
    }
}

fn get_builtin_global_starlark_docs() -> Doc {
    let globals = Globals::extended();
    builtin_doc(
        "globals",
        "standard",
        DocItem::Module(globals.documentation()),
    )
}

/// Globals that are in the interpreter, but none of the starlark global symbols.
fn get_builtin_build_docs(interpreter_state: Arc<GlobalInterpreterState>) -> anyhow::Result<Doc> {
    let mut b_o = interpreter_state.extension_file_global_env.documentation();
    let globals = Globals::extended();
    let global_symbols: HashSet<_> = globals.names().map(|s| s.as_str()).collect();
    b_o.members = b_o
        .members
        .into_iter()
        .filter(|(name, _)| !global_symbols.contains(&name.as_str()))
        .collect();
    Ok(builtin_doc("globals", "", DocItem::Module(b_o)))
}

pub fn get_builtin_docs(
    interpreter_state: Arc<GlobalInterpreterState>,
) -> anyhow::Result<Vec<Doc>> {
    let mut all_builtins = vec![
        get_builtin_global_starlark_docs(),
        get_builtin_build_docs(interpreter_state.dupe())?,
    ];

    all_builtins.extend(get_builtin_bxl_docs(interpreter_state)?);
    all_builtins.extend(get_registered_starlark_docs());

    Ok(all_builtins)
}

/// Get the documentation for exported symbols in the prelude
///
/// Creates top level docs for member functions of "native" too,
/// presuming that those symbols don't already exist in `existing_globals`
/// (to avoid re-exporting and overriding the real builtins if there is conflict)
pub async fn get_prelude_docs(
    ctx: &DiceTransaction,
    existing_globals: &HashSet<&str>,
) -> anyhow::Result<Vec<Doc>> {
    let cell_resolver = ctx.get_cell_resolver().await?;
    let prelude_path = prelude_path(&cell_resolver)?;
    get_docs_from_module(ctx, prelude_path.import_path(), Some(existing_globals)).await
}

async fn get_docs_from_module(
    ctx: &DiceComputations,
    import_path: &ImportPath,
    // If we want to promote `native`, what should we exclude
    promote_native: Option<&HashSet<&str>>,
) -> anyhow::Result<Vec<Doc>> {
    // Do this so that we don't get the '@' in the display if we're printing targets from a
    // different cell root. i.e. `//foo:bar.bzl`, rather than `//foo:bar.bzl @ cell`
    let import_path_string = format!(
        "{}:{}",
        import_path.path().parent().unwrap(),
        import_path.path().path().file_name().unwrap()
    );
    let module = ctx.get_loaded_module_from_import_path(import_path).await?;
    let frozen_module = module.env();
    let mut module_docs = frozen_module.documentation();

    // For the prelude, we want to promote `native` symbol up one level
    if let Some(existing_globals) = promote_native {
        if let Some(native) = frozen_module.get_option("native")? {
            if let Some(DocItem::Object(native)) = native.value().documentation() {
                for (k, v) in native.members {
                    if !existing_globals.contains(k.as_str())
                        && !module_docs.members.contains_key(&k)
                    {
                        module_docs.members.insert(k, v);
                    }
                }
            }
        }
    }

    let mut docs = vec![];

    if let Some(module_doc) = module_docs.docs {
        docs.push(Doc {
            id: Identifier {
                name: import_path_string.clone(),
                location: Some(starlark::docs::Location {
                    path: import_path_string.clone(),
                    position: None,
                }),
            },
            item: DocItem::Module(DocModule {
                docs: Some(module_doc),
                members: SmallMap::new(),
            }),
            custom_attrs: Default::default(),
        });
    }
    docs.extend(module_docs.members.into_iter().map(|(symbol, d)| {
        Doc {
            // TODO(nmj): Map this back into the codemap to get a line/column
            id: Identifier {
                name: symbol,
                location: Some(Location {
                    path: import_path_string.clone(),
                    position: None,
                }),
            },
            item: d.to_doc_item(),
            custom_attrs: Default::default(),
        }
    }));

    Ok(docs)
}

pub async fn docs_command(
    context: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<NoPartialResult>,
    req: UnstableDocsRequest,
) -> anyhow::Result<UnstableDocsResponse> {
    run_server_command(
        DocsServerCommand { req },
        context,
        partial_result_dispatcher,
    )
    .await
}

struct DocsServerCommand {
    req: UnstableDocsRequest,
}

#[async_trait]
impl ServerCommandTemplate for DocsServerCommand {
    type StartEvent = buck2_data::DocsCommandStart;
    type EndEvent = buck2_data::DocsCommandEnd;
    type Response = UnstableDocsResponse;
    type PartialResult = NoPartialResult;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        _partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        docs(server_ctx, ctx, &self.req).await
    }

    fn is_success(&self, _response: &Self::Response) -> bool {
        // No response if we failed.
        true
    }
}

enum Format {
    Json,
    Markdown,
}

impl Format {
    fn from_proto(request: &UnstableDocsRequest) -> anyhow::Result<Format> {
        let format = unstable_docs_request::Format::from_i32(request.format)
            .context("incorrect enum value")?;
        match format {
            unstable_docs_request::Format::Json => Ok(Format::Json),
            unstable_docs_request::Format::Markdown => Ok(Format::Markdown),
            unstable_docs_request::Format::Unknown => Err(DocsError::UnknownFormat.into()),
        }
    }
}

async fn docs(
    server_ctx: &dyn ServerCommandContextTrait,
    dice_ctx: DiceTransaction,
    request: &UnstableDocsRequest,
) -> anyhow::Result<UnstableDocsResponse> {
    let format = Format::from_proto(request)?;

    let cell_resolver = dice_ctx.get_cell_resolver().await?;
    let current_cell_path = cell_resolver.get_cell_path(server_ctx.working_dir())?;
    let current_cell = BuildFileCell::new(current_cell_path.cell());

    let cell_alias_resolver = cell_resolver
        .get(current_cell_path.cell())?
        .cell_alias_resolver();

    let lookups = parse_import_paths(
        cell_alias_resolver,
        &current_cell_path,
        current_cell,
        &request.symbol_patterns,
    )?;

    let mut docs = if request.retrieve_builtins {
        get_builtin_docs(dice_ctx.get_global_interpreter_state().await?.dupe())?
    } else {
        vec![]
    };

    if request.retrieve_prelude {
        let builtin_names = docs.iter().map(|d| d.id.name.as_str()).collect();
        let prelude_docs = get_prelude_docs(&dice_ctx, &builtin_names).await?;
        docs.extend(prelude_docs);
    }

    let module_calcs: Vec<_> = lookups
        .iter()
        .map(|import_path| get_docs_from_module(&dice_ctx, import_path, None))
        .collect();

    let modules_docs = futures::future::try_join_all(module_calcs).await?;
    docs.extend(modules_docs.into_iter().flatten());

    let json_output = match format {
        Format::Json => Some(serde_json::to_string(&docs)?),
        Format::Markdown => {
            let path = AbsPath::new(Path::new(request.markdown_output_path.as_ref().context(
                "`markdown_output_path` must be set when requesting markdown (internal error)",
            )?))?;
            let starlark_subdir = Path::new(&request.markdown_starlark_subdir);
            let native_subdir = Path::new(&request.markdown_native_subdir);
            generate_markdown_files(path, starlark_subdir, native_subdir, docs)?;
            None
        }
    };

    Ok(UnstableDocsResponse { json_output })
}


/*
 * 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::cmp::Ordering;
use std::collections::HashMap;
use std::path::Path;
use std::path::PathBuf;

use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_events::dispatch::console_message;
use itertools::Itertools;
use starlark::docs::Doc;
use starlark::docs::DocItem;
use starlark::docs::MarkdownFlavor;
use starlark::docs::RenderMarkdown;

use crate::lsp::output_subdir_for_doc;

#[derive(Clone, Debug, Default)]
struct MarkdownOutput {
    sections: Vec<String>,
}

impl MarkdownOutput {
    fn write_to_file(&self, path: &AbsPath) -> anyhow::Result<String> {
        let mut contents = self.sections.join("\n\n---\n");
        if !contents.ends_with('\n') {
            contents.push('\n');
        }
        if let Some(p) = path.parent() {
            fs_util::create_dir_all(p)?;
        }
        fs_util::write(path, &contents)?;

        Ok(contents)
    }

    /// Convert a buck style path (foo//bar:baz.bzl) to path, failing if someone attempts to traverse upward
    fn path_from_location(location: &str) -> anyhow::Result<PathBuf> {
        Ok(Path::new(&location.replace("//", "/").replace(':', "/")).to_path_buf())
    }

    /// Get the output path for the markdown for a given [`Doc`], whether it's in a starlark file, or a native symbol.
    fn markdown_path_for_doc(
        starlark_subdir: &Path,
        native_subdir: &Path,
        doc: &Doc,
    ) -> anyhow::Result<PathBuf> {
        let subdir = output_subdir_for_doc(doc)?;
        let path = match &doc.id.location {
            Some(loc) => starlark_subdir
                .join(subdir.as_path())
                .join(Self::path_from_location(&loc.path)?),
            None => match &doc.item {
                // Functions all go in one file.
                // Objects get their on file (e.g. each provider, Artifact, etc)
                DocItem::Function(_) | DocItem::Property(_) => {
                    native_subdir.join(subdir.as_path()).join("native")
                }
                DocItem::Module(_) | DocItem::Object(_) => {
                    native_subdir.join(subdir.as_path()).join(&doc.id.name)
                }
            },
        };
        let path = path.with_extension(match path.extension() {
            None => "md".to_owned(),
            Some(e) => format!("{}.md", e.to_str().expect("path if not UTF-8")),
        });
        Ok(path)
    }
}

/// Does the heavy work of processing the docs and writing them to markdown files.
pub(crate) fn generate_markdown_files(
    destination_dir: &AbsPath,
    starlark_subdir: &Path,
    native_subdir: &Path,
    docs: Vec<Doc>,
) -> anyhow::Result<()> {
    let mut outputs = HashMap::new();

    fn item_ordering(l: &Doc, r: &Doc) -> Ordering {
        match (&l.item, &r.item) {
            (DocItem::Module(_), DocItem::Module(_)) => l.id.name.cmp(&r.id.name),
            (DocItem::Module(_), _) => Ordering::Less,
            (_, DocItem::Module(_)) => Ordering::Greater,
            _ => l.id.name.cmp(&r.id.name),
        }
    }

    for doc in docs.into_iter().sorted_by(item_ordering) {
        let markdown_path =
            MarkdownOutput::markdown_path_for_doc(starlark_subdir, native_subdir, &doc)?;
        let markdown_file = outputs
            .entry(markdown_path)
            .or_insert_with(MarkdownOutput::default);
        if let Some(docs) = doc.render_markdown_opt(MarkdownFlavor::DocFile) {
            markdown_file.sections.push(docs);
        }
    }

    for (relative_path, markdown_file) in outputs.iter() {
        let path = destination_dir.join(relative_path);
        console_message(format!("Writing to {}", path.display()));

        markdown_file.write_to_file(&path)?;
    }

    Ok(())
}


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

mod bxl_docs;
pub mod docs;
pub(crate) mod markdown;


/*
 * 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 rand::distributions::Alphanumeric;
use rand::distributions::DistString;

/// Verify that our working directory is still here. We often run on Eden, and if Eden restarts
/// ungracefully, our working dir will become unreadable and we are just about done.
pub fn check_working_dir() -> anyhow::Result<()> {
    use std::fs;
    use std::io;

    // Looks like we need to get a name that the OS isn't likely to have seen before for this to
    // work reliably.
    let name = Alphanumeric.sample_string(&mut rand::thread_rng(), 16);

    let err = match fs::metadata(name) {
        Ok(..) => return Ok(()),
        Err(e) => e,
    };

    if err.kind() == io::ErrorKind::NotConnected {
        let err = "Buck2 is running in an Eden mount but Eden restarted uncleanly. \
            This error is unrecoverable and you should restart Buck using `buck2 killall`.";
        return Err(anyhow::anyhow!(err));
    }

    if err.kind() != io::ErrorKind::NotFound {
        tracing::warn!(
            "Buck2 is unable to read its current working directory: {}. Consider restarting",
            err
        );
    }

    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 anyhow::Context as _;
use buck2_cli_proto::client_context::HostPlatformOverride;
use buck2_cli_proto::common_build_options::ExecutionStrategy;
use buck2_common::executor_config::CacheUploadBehavior;
use buck2_common::executor_config::CommandExecutorConfig;
use buck2_common::executor_config::CommandGenerationOptions;
use buck2_common::executor_config::Executor;
use buck2_common::executor_config::HybridExecutionLevel;
use buck2_common::executor_config::LocalExecutorOptions;
use buck2_common::executor_config::PathSeparatorKind;
use buck2_common::executor_config::RemoteEnabledExecutor;
use buck2_common::executor_config::RemoteExecutorOptions;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::project::ProjectRoot;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::execute::cache_uploader::NoOpCacheUploader;
use buck2_execute::execute::dice_data::CommandExecutorResponse;
use buck2_execute::execute::dice_data::HasCommandExecutor;
use buck2_execute::execute::prepared::NoOpCommandExecutor;
use buck2_execute::execute::prepared::PreparedCommandExecutor;
use buck2_execute::execute::request::ExecutorPreference;
use buck2_execute::knobs::ExecutorGlobalKnobs;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::re::manager::ReConnectionHandle;
use buck2_execute_impl::executors::action_cache::ActionCacheChecker;
use buck2_execute_impl::executors::caching::CacheUploader;
use buck2_execute_impl::executors::hybrid::HybridExecutor;
use buck2_execute_impl::executors::local::LocalExecutor;
use buck2_execute_impl::executors::re::ReExecutor;
use buck2_execute_impl::executors::worker::WorkerPool;
use buck2_execute_impl::low_pass_filter::LowPassFilter;
use buck2_forkserver::client::ForkserverClient;
use buck2_util::collections::sorted_map::SortedMap;
use dupe::Dupe;
use host_sharing::HostSharingBroker;
use once_cell::sync::OnceCell;
use remote_execution as RE;

pub fn parse_concurrency(requested: u32) -> anyhow::Result<usize> {
    let mut ret = requested.try_into().context("Invalid concurrency")?;

    if ret == 0 {
        ret = num_cpus::get();
    }

    Ok(ret)
}

/// For each buck invocations, we'll have a single CommandExecutorFactory. This contains shared
/// state used by all command executor strategies.
pub struct CommandExecutorFactory {
    pub re_connection: Arc<ReConnectionHandle>,
    // TODO(cjhopman): This should probably be a global limit, otherwise simultaneous commands may
    // use more resources than intended (this might no longer be accurate since only instances
    // sharing the same DICE context should be allowed to proceed concurrently, and we only have
    // one CommandExecutorFactory per DICE context).
    pub host_sharing_broker: Arc<HostSharingBroker>,
    pub low_pass_filter: Arc<LowPassFilter>,
    pub materializer: Arc<dyn Materializer>,
    pub blocking_executor: Arc<dyn BlockingExecutor>,
    pub strategy: ExecutionStrategy,
    pub executor_global_knobs: ExecutorGlobalKnobs,
    pub upload_all_actions: bool,
    pub forkserver: Option<ForkserverClient>,
    pub skip_cache_read: bool,
    pub skip_cache_write: bool,
    project_root: ProjectRoot,
    worker_pool: Arc<WorkerPool>,
}

impl CommandExecutorFactory {
    pub fn new(
        re_connection: Arc<ReConnectionHandle>,
        host_sharing_broker: HostSharingBroker,
        low_pass_filter: LowPassFilter,
        materializer: Arc<dyn Materializer>,
        blocking_executor: Arc<dyn BlockingExecutor>,
        strategy: ExecutionStrategy,
        executor_global_knobs: ExecutorGlobalKnobs,
        upload_all_actions: bool,
        forkserver: Option<ForkserverClient>,
        skip_cache_read: bool,
        skip_cache_write: bool,
        project_root: ProjectRoot,
        worker_pool: Arc<WorkerPool>,
    ) -> Self {
        Self {
            re_connection,
            host_sharing_broker: Arc::new(host_sharing_broker),
            low_pass_filter: Arc::new(low_pass_filter),
            materializer,
            blocking_executor,
            strategy,
            executor_global_knobs,
            upload_all_actions,
            forkserver,
            skip_cache_read,
            skip_cache_write,
            project_root,
            worker_pool,
        }
    }
}

impl HasCommandExecutor for CommandExecutorFactory {
    fn get_command_executor(
        &self,
        artifact_fs: &ArtifactFs,
        executor_config: &CommandExecutorConfig,
    ) -> anyhow::Result<CommandExecutorResponse> {
        let local_executor_new = |options: &LocalExecutorOptions| {
            let worker_pool = if options.use_persistent_workers {
                Some(self.worker_pool.dupe())
            } else {
                None
            };
            LocalExecutor::new(
                artifact_fs.clone(),
                self.materializer.dupe(),
                self.blocking_executor.dupe(),
                self.host_sharing_broker.dupe(),
                self.project_root.root().to_owned(),
                self.forkserver.dupe(),
                self.executor_global_knobs.dupe(),
                worker_pool,
            )
        };

        if !buck2_core::is_open_source() && !cfg!(fbcode_build) {
            static WARN: OnceCell<()> = OnceCell::new();
            WARN.get_or_init(|| {
                tracing::warn!("Cargo build detected: disabling remote execution and caching!")
            });

            if self.strategy.ban_local() {
                return Err(anyhow::anyhow!(
                    "The desired execution strategy (`{:?}`) is incompatible with the local executor",
                    self.strategy,
                ));
            }

            return Ok(CommandExecutorResponse {
                executor: Arc::new(local_executor_new(&LocalExecutorOptions::default())),
                platform: Default::default(),
                cache_checker: Arc::new(NoOpCommandExecutor {}),
                cache_uploader: Arc::new(NoOpCacheUploader {}),
            });
        }

        let remote_executor_new = |options: &RemoteExecutorOptions,
                                   re_use_case: &RemoteExecutorUseCase,
                                   re_action_key: &Option<String>,
                                   remote_cache_enabled: bool| {
            // 30GB is the max RE can currently support.
            const DEFAULT_RE_MAX_INPUT_FILE_BYTES: u64 = 30 * 1024 * 1024 * 1024;

            ReExecutor {
                artifact_fs: artifact_fs.clone(),
                project_fs: self.project_root.clone(),
                materializer: self.materializer.dupe(),
                re_client: self.re_connection.get_client(),
                re_use_case: *re_use_case,
                re_action_key: re_action_key.clone(),
                re_max_input_files_bytes: options
                    .re_max_input_files_bytes
                    .unwrap_or(DEFAULT_RE_MAX_INPUT_FILE_BYTES),
                re_max_queue_time_ms: options.re_max_queue_time_ms,
                knobs: self.executor_global_knobs.dupe(),
                skip_cache_read: self.skip_cache_read || !remote_cache_enabled,
                skip_cache_write: self.skip_cache_write || !remote_cache_enabled,
            }
        };

        let response = match &executor_config.executor {
            Executor::Local(local) => {
                if self.strategy.ban_local() {
                    None
                } else {
                    Some(CommandExecutorResponse {
                        executor: Arc::new(local_executor_new(local)),
                        platform: Default::default(),
                        cache_checker: Arc::new(NoOpCommandExecutor {}),
                        cache_uploader: Arc::new(NoOpCacheUploader {}),
                    })
                }
            }
            Executor::RemoteEnabled {
                executor,
                re_properties,
                re_use_case,
                re_action_key,
                cache_upload_behavior,
                remote_cache_enabled,
            } => {
                let executor: Option<Arc<dyn PreparedCommandExecutor>> = match &executor {
                    RemoteEnabledExecutor::Local(local) if !self.strategy.ban_local() => {
                        Some(Arc::new(local_executor_new(local)))
                    }
                    RemoteEnabledExecutor::Remote(remote) if !self.strategy.ban_remote() => {
                        Some(Arc::new(remote_executor_new(
                            remote,
                            re_use_case,
                            re_action_key,
                            *remote_cache_enabled,
                        )))
                    }
                    RemoteEnabledExecutor::Hybrid {
                        local,
                        remote,
                        level,
                    } if !self.strategy.ban_hybrid() => Some(Arc::new(HybridExecutor {
                        local: local_executor_new(local),
                        remote: remote_executor_new(
                            remote,
                            re_use_case,
                            re_action_key,
                            *remote_cache_enabled,
                        ),
                        level: *level,
                        executor_preference: self.strategy.hybrid_preference(),
                        low_pass_filter: self.low_pass_filter.dupe(),
                    })),
                    _ => None,
                };

                // NOTE: While we now have a legit flag for this, we keep the env var. This has been used
                // in remediating prod incidents in the past, and this is the kind of thing that can easily
                // become tribal knowledge. Keeping this does not hurt us.
                static DISABLE_CACHING: EnvHelper<bool> =
                    EnvHelper::new("BUCK2_TEST_DISABLE_CACHING");

                let disable_caching = DISABLE_CACHING
                    .get_copied()?
                    .unwrap_or(self.skip_cache_read);

                let (cache_checker, cache_uploader) = if disable_caching || !remote_cache_enabled {
                    (
                        Arc::new(NoOpCommandExecutor {}) as _,
                        Arc::new(NoOpCacheUploader {}) as _,
                    )
                } else {
                    let cache_checker = Arc::new(ActionCacheChecker {
                        artifact_fs: artifact_fs.clone(),
                        materializer: self.materializer.dupe(),
                        re_client: self.re_connection.get_client(),
                        re_use_case: *re_use_case,
                        re_action_key: re_action_key.clone(),
                        upload_all_actions: self.upload_all_actions,
                        knobs: self.executor_global_knobs.dupe(),
                    }) as _;
                    let cache_uploader =
                        if let CacheUploadBehavior::Enabled { max_bytes } = cache_upload_behavior {
                            Arc::new(CacheUploader {
                                artifact_fs: artifact_fs.clone(),
                                materializer: self.materializer.dupe(),
                                re_client: self.re_connection.get_client(),
                                re_use_case: *re_use_case,
                                knobs: self.executor_global_knobs.dupe(),
                                max_bytes: *max_bytes,
                            }) as _
                        } else {
                            Arc::new(NoOpCacheUploader {}) as _
                        };
                    (cache_checker, cache_uploader)
                };

                let platform = RE::Platform {
                    properties: re_properties
                        .iter()
                        .map(|(k, v)| RE::Property {
                            name: k.clone(),
                            value: v.clone(),
                        })
                        .collect(),
                };

                executor.map(|executor| CommandExecutorResponse {
                    executor,
                    platform,
                    cache_checker,
                    cache_uploader,
                })
            }
        };

        let response = response
            .with_context(|| format!(
"The desired execution strategy (`{:?}`) is incompatible with the executor config that was selected: {:?}",
self.strategy, executor_config))?;

        Ok(response)
    }
}

trait ExecutionStrategyExt {
    fn ban_local(&self) -> bool;
    fn ban_remote(&self) -> bool;
    fn ban_hybrid(&self) -> bool;
    fn hybrid_preference(&self) -> ExecutorPreference;
}

impl ExecutionStrategyExt for ExecutionStrategy {
    fn ban_local(&self) -> bool {
        match self {
            Self::RemoteOnly | Self::NoExecution => true,
            _ => false,
        }
    }

    fn ban_remote(&self) -> bool {
        match self {
            Self::LocalOnly | Self::NoExecution => true,
            _ => false,
        }
    }

    fn ban_hybrid(&self) -> bool {
        match self {
            Self::NoExecution => true,
            _ => false,
        }
    }

    fn hybrid_preference(&self) -> ExecutorPreference {
        match self {
            Self::HybridPreferLocal => ExecutorPreference::LocalPreferred,
            Self::HybridPreferRemote => ExecutorPreference::RemotePreferred,
            Self::LocalOnly => ExecutorPreference::LocalRequired,
            Self::RemoteOnly => ExecutorPreference::RemoteRequired,
            _ => ExecutorPreference::Default,
        }
    }
}

/// This is used when execution platforms are not configured.
pub fn get_default_executor_config(host_platform: HostPlatformOverride) -> CommandExecutorConfig {
    let executor = if buck2_core::is_open_source() {
        Executor::Local(LocalExecutorOptions::default())
    } else {
        Executor::RemoteEnabled {
            executor: RemoteEnabledExecutor::Hybrid {
                local: LocalExecutorOptions::default(),
                remote: RemoteExecutorOptions::default(),
                level: HybridExecutionLevel::Limited,
            },
            re_properties: get_default_re_properties(host_platform),
            re_use_case: RemoteExecutorUseCase::buck2_default(),
            re_action_key: None,
            cache_upload_behavior: CacheUploadBehavior::Disabled,
            remote_cache_enabled: true,
        }
    };

    CommandExecutorConfig {
        executor,
        options: CommandGenerationOptions {
            path_separator: PathSeparatorKind::system_default(),
            output_paths_behavior: Default::default(),
        },
    }
}

fn get_default_re_properties(host_platform: HostPlatformOverride) -> SortedMap<String, String> {
    let linux = &[("platform", "linux-remote-execution")];
    let macos = &[("platform", "mac"), ("subplatform", "any")];
    let windows = &[("platform", "windows")];

    let props = match host_platform {
        HostPlatformOverride::Linux => linux.as_slice(),
        HostPlatformOverride::MacOs => macos.as_slice(),
        HostPlatformOverride::Windows => windows.as_slice(),
        HostPlatformOverride::DefaultPlatform => match std::env::consts::OS {
            "linux" => linux.as_slice(),
            "macos" => macos.as_slice(),
            "windows" => windows.as_slice(),
            v => unimplemented!("no support yet for operating system `{}`", v),
        },
    };

    props
        .iter()
        .map(|(k, v)| ((*k).to_owned(), (*v).to_owned()))
        .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.
 */

pub fn create_listener() -> anyhow::Result<(
    buck2_common::buckd_connection::ConnectionType,
    std::net::TcpListener,
)> {
    use std::net::Ipv4Addr;
    use std::net::SocketAddr;

    use buck2_common::buckd_connection::ConnectionType;

    let addr = SocketAddr::new(Ipv4Addr::LOCALHOST.into(), 0);
    let tcp_listener = std::net::TcpListener::bind(addr)?;
    tcp_listener.set_nonblocking(true)?;
    let local_addr = tcp_listener.local_addr()?;
    Ok((
        ConnectionType::Tcp {
            port: local_addr.port(),
        },
        tcp_listener,
    ))
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use buck2_common::buckd_connection::ConnectionType;

    use crate::daemon::daemon_tcp::create_listener;

    #[test]
    fn test_create_listener() {
        let (connection_type, _tcp_listener) = create_listener().unwrap();
        assert_matches!(connection_type, ConnectionType::Tcp { .. });
    }
}


/*
 * 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::File;
use std::io::BufWriter;
use std::path::Path;
use std::sync::Arc;

use anyhow::Context as _;
use bincode::Options;
use buck2_cli_proto::unstable_dice_dump_request::DiceDumpFormat;
use dice::Dice;
use dupe::Dupe;
use flate2::write::GzEncoder;
use flate2::Compression;

pub(crate) async fn dice_dump_spawn(
    dice: &Arc<Dice>,
    path: &Path,
    format: DiceDumpFormat,
) -> anyhow::Result<()> {
    let dice = dice.dupe();
    let path = path.to_path_buf();
    tokio::task::spawn_blocking(move || dice_dump(&dice, &path, format))
        .await
        .context("Failed to spawn")?
        .context("Failed to dump")?;
    Ok(())
}

pub(crate) fn dice_dump(
    dice: &Arc<Dice>,
    path: &Path,
    format: DiceDumpFormat,
) -> anyhow::Result<()> {
    match format {
        DiceDumpFormat::Tsv => dice_dump_tsv(dice, path),
        DiceDumpFormat::Bincode => dice_dump_bincode(dice, path),
        DiceDumpFormat::JsonPretty => dice_dump_json_pretty(dice, path),
    }
}

pub(crate) fn tar_dice_dump(dice_dump_folder: &Path) -> anyhow::Result<()> {
    let tar_gz = File::create(format!("{}.tar.gz", dice_dump_folder.display()))?;
    let enc = GzEncoder::new(tar_gz, Compression::default());
    let mut tar = tar::Builder::new(enc);
    let files = vec!["nodes.gz", "edges.gz", "nodes_currently_running.gz"];
    for file_name in files {
        let mut file = File::open(dice_dump_folder.join(file_name)).context(format!(
            "Failed to open file `{}` for compressing",
            file_name
        ))?;
        tar.append_file(file_name, &mut file)
            .context(format!("Failed to write file `{}` to archive", file_name))?;
    }

    tar.finish()
        .context("Failed to generate DICE dump archive")?;

    Ok(())
}

fn dice_dump_tsv(dice: &Arc<Dice>, path: &Path) -> anyhow::Result<()> {
    let path = path.to_path_buf();
    let nodes_path = path.join("nodes.gz");
    let edges_path = path.join("edges.gz");
    let nodes_currently_running_path = path.join("nodes_currently_running.gz");

    std::fs::create_dir_all(path).context("Failed to create directory")?;

    let nodes = File::create(&nodes_path).context(format!(
        "Failed to open DICE node dumpfile {:?}",
        &nodes_path
    ))?;
    let mut nodes = GzEncoder::new(BufWriter::new(nodes), Compression::default());

    let edges = File::create(&edges_path).context(format!(
        "Failed to open DICE edge dumpfile {:?}",
        &edges_path
    ))?;
    let mut edges = GzEncoder::new(BufWriter::new(edges), Compression::default());

    let nodes_currently_running = File::create(&nodes_currently_running_path).context(format!(
        "Failed to open DICE node currently running dumpfile {:?}",
        &nodes_currently_running_path
    ))?;
    let mut nodes_currently_running = GzEncoder::new(
        BufWriter::new(nodes_currently_running),
        Compression::default(),
    );

    dice.serialize_tsv(&mut nodes, &mut edges, &mut nodes_currently_running)
        .context("Failed to serialize")?;

    nodes
        .try_finish()
        .context(format!("Failed to flush DICE nodes to {:?}", &nodes_path))?;
    edges
        .try_finish()
        .context(format!("Failed to flush DICE edges to {:?}", &edges_path))?;
    nodes_currently_running.try_finish().context(format!(
        "Failed to flush DICE nodes currently running to {:?}",
        &nodes_currently_running_path
    ))?;

    Ok(())
}

fn dice_dump_bincode(dice: &Arc<Dice>, path: &Path) -> anyhow::Result<()> {
    let path = path.to_path_buf();
    std::fs::create_dir_all(path.parent().unwrap()).context("Failed to create directory")?;
    let out =
        File::create(&path).context(format!("Failed to open serde DICE dumpfile {:?}", &path))?;
    let out = GzEncoder::new(BufWriter::new(out), Compression::default());

    let mut writer = bincode::Serializer::new(
        out,
        bincode::config::DefaultOptions::new()
            .with_fixint_encoding()
            .allow_trailing_bytes(),
    );
    dice.serialize_serde(&mut writer)?;
    Ok(())
}

fn dice_dump_json_pretty(dice: &Arc<Dice>, path: &Path) -> anyhow::Result<()> {
    let path = path.to_path_buf();
    std::fs::create_dir_all(path.parent().unwrap()).context("Failed to create directory")?;
    let out =
        File::create(&path).context(format!("Failed to open serde DICE dumpfile {:?}", &path))?;
    let out = GzEncoder::new(BufWriter::new(out), Compression::default());

    let mut writer = serde_json::Serializer::pretty(out);
    dice.serialize_serde(&mut writer)?;
    Ok(())
}


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

use std::collections::HashMap;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_common::invocation_paths::InvocationPaths;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::rollout_percentage::RolloutPercentage;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::materialize::materializer::MaterializationMethod;
use buck2_execute_impl::materializers::deferred::DeferredMaterializerConfigs;
use buck2_execute_impl::materializers::sqlite::MaterializerState;
use buck2_execute_impl::materializers::sqlite::MaterializerStateSqliteDb;
use buck2_execute_impl::materializers::sqlite::DB_SCHEMA_VERSION;

use crate::daemon::server::BuckdServerInitPreferences;

#[derive(Allocative)]
pub struct DiskStateOptions {
    pub sqlite_materializer_state: bool,
    // In future, this will include the config for dep files on disk
}

impl DiskStateOptions {
    pub fn new(
        root_config: &LegacyBuckConfig,
        materialization_method: MaterializationMethod,
    ) -> anyhow::Result<Self> {
        let sqlite_materializer_state = matches!(
            // We can only enable materializer state on sqlite if you use deferred materializer
            materialization_method,
            MaterializationMethod::Deferred | MaterializationMethod::DeferredSkipFinalArtifacts
        ) && root_config
            .parse::<RolloutPercentage>("buck2", "sqlite_materializer_state")?
            .unwrap_or_else(RolloutPercentage::never)
            .roll();
        Ok(Self {
            sqlite_materializer_state,
        })
    }
}

pub(crate) async fn maybe_initialize_materializer_sqlite_db(
    options: &DiskStateOptions,
    paths: &InvocationPaths,
    io_executor: Arc<dyn BlockingExecutor>,
    root_config: &LegacyBuckConfig,
    deferred_materializer_configs: &DeferredMaterializerConfigs,
    fs: ProjectRoot,
    digest_config: DigestConfig,
    init_ctx: &BuckdServerInitPreferences,
) -> anyhow::Result<(Option<MaterializerStateSqliteDb>, Option<MaterializerState>)> {
    if !options.sqlite_materializer_state {
        // When sqlite materializer state is disabled, we should always delete the materializer state db.
        // Otherwise, artifacts in buck-out will diverge from the state stored in db.
        io_executor
            .execute_io_inline(|| fs.remove_path_recursive(&paths.materializer_state_path()))
            .await?;
        return Ok((None, None));
    }

    let metadata = buck2_events::metadata::collect();

    let mut versions = HashMap::from([
        ("schema_version".to_owned(), DB_SCHEMA_VERSION.to_string()),
        (
            "defer_write_actions".to_owned(),
            deferred_materializer_configs
                .defer_write_actions
                .to_string(),
        ),
    ]);
    if let Some(buckconfig_version) =
        root_config.parse("buck2", "sqlite_materializer_state_version")?
    {
        versions.insert("buckconfig_version".to_owned(), buckconfig_version);
    }
    if let Some(hostname) = metadata.get("hostname") {
        versions.insert("hostname".to_owned(), hostname.to_owned());
    }

    // Most things in the rest of `metadata` should go in the metadata sqlite table.
    // TODO(scottcao): Narrow down what metadata we need and and insert them into the
    // metadata table before a feature rollout.
    let (db, load_result) = MaterializerStateSqliteDb::initialize(
        paths.materializer_state_path(),
        versions,
        metadata,
        io_executor,
        digest_config,
        init_ctx.reject_materializer_state.as_ref(),
    )
    .await?;

    let materializer_state = match load_result {
        Ok(s) => Some(s),
        // We know path not found or version mismatch is normal, but some sqlite failures
        // are worth logging here. TODO(scottcao): Refine our error types and figure out what
        // errors to log
        Err(_e) => None,
    };
    Ok((Some(db), materializer_state))
}

// Once we start storing disk state in the cache directory, we need to make sure
// buck2 always deletes the cache directory if the cache is disabled.
// Otherwise, buck-out state can diverge from the state of on-disk cache when
// cache is disabled, causing buck2 to use stale cache when reading from the
// cache is re-enabled. One way this can happen is that someone can build on
// an older revision with a buck2 that doesn't understand the cache directory
// in between 2 builds on newer revisions with buck2 that reads from the cache
// (for ex., as a part of a bisect), then the state can become stale.
// There are 2 (not foolproof) mitigations planned:
// 1) Read from the logs what the last buck2 invocation was and check that the
// last buck2 supported on-disk state. If not, delete the disk state.
// 2) Start always deleting the cache directory now until we add support for disk
// state in buck2.
// The following implements mitigation #2 by always deleting disk state.

/// Recursively deletes all elements under `cache_dir_path`, except for known dirs
/// listed in `known_dir_names`.
pub(crate) fn delete_unknown_disk_state(
    cache_dir_path: &AbsNormPath,
    known_dir_names: &[&FileName],
    fs: ProjectRoot,
) -> anyhow::Result<()> {
    let res: anyhow::Result<()> = try {
        if cache_dir_path.exists() {
            for entry in fs_util::read_dir(cache_dir_path)? {
                let entry = entry?;
                let filename = entry.file_name();
                let filename = filename
                    .to_str()
                    .context("Filename is not UTF-8")
                    .and_then(FileName::new)?;

                // known_dir_names is always small, so this contains isn't expensive
                if !known_dir_names.contains(&filename) || !entry.path().is_dir() {
                    fs.remove_path_recursive(&cache_dir_path.join(filename))?;
                }
            }
        }
    };

    res.with_context(|| {
        format!(
            "deleting unrecognized caches in {} to prevent them from going stale",
            &cache_dir_path
        )
    })
}

#[cfg(test)]
mod tests {
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
    use buck2_core::fs::project::ProjectRootTemp;
    use buck2_core::fs::project_rel_path::ProjectRelativePath;
    use dupe::Dupe;

    use super::*;

    #[test]
    fn test_delete_all_from_cache_dir() {
        let fs_temp = ProjectRootTemp::new().unwrap();
        let fs = fs_temp.path();
        let cache_dir_path = fs.resolve(ProjectRelativePath::unchecked_new("buck-out/v2/cache"));
        let materializer_state_db = cache_dir_path.join(ForwardRelativePath::unchecked_new(
            "materializer_state/db.sqlite",
        ));
        let command_hashes_db = cache_dir_path.join(ForwardRelativePath::unchecked_new(
            "command_hashes/db.sqlite",
        ));
        fs.create_file(&materializer_state_db, false).unwrap();
        fs.create_file(&command_hashes_db, false).unwrap();
        assert!(materializer_state_db.exists());
        assert!(command_hashes_db.exists());

        delete_unknown_disk_state(&cache_dir_path, &[], fs.dupe()).unwrap();

        assert!(!materializer_state_db.exists());
        assert!(!command_hashes_db.exists());
    }

    #[test]
    fn test_delete_from_cache_dir_with_known_dirs() {
        let fs_temp = ProjectRootTemp::new().unwrap();
        let fs = fs_temp.path();
        let cache_dir_path = fs.resolve(ProjectRelativePath::unchecked_new("buck-out/v2/cache"));
        let materializer_state_db = cache_dir_path.join(ForwardRelativePath::unchecked_new(
            "materializer_state/db.sqlite",
        ));
        let command_hashes_db = cache_dir_path.join(ForwardRelativePath::unchecked_new(
            "command_hashes/db.sqlite",
        ));
        fs.create_file(&materializer_state_db, false).unwrap();
        fs.create_file(&command_hashes_db, false).unwrap();
        assert!(materializer_state_db.exists());
        assert!(command_hashes_db.exists());

        delete_unknown_disk_state(
            &cache_dir_path,
            &[FileName::unchecked_new("materializer_state")],
            fs.dupe(),
        )
        .unwrap();

        assert!(materializer_state_db.exists());
        assert!(!command_hashes_db.exists());
    }
}


/*
 * 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_common::legacy_configs::LegacyBuckConfig;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_forkserver::client::ForkserverClient;

#[cfg(unix)]
pub async fn maybe_launch_forkserver(
    root_config: &LegacyBuckConfig,
    forkserver_state_dir: &AbsNormPath,
) -> anyhow::Result<Option<ForkserverClient>> {
    use anyhow::Context;
    use buck2_core::rollout_percentage::RolloutPercentage;

    let config = root_config
        .parse::<RolloutPercentage>("buck2", "forkserver")?
        .unwrap_or_else(RolloutPercentage::always);

    if !config.roll() {
        return Ok(None);
    }

    let exe = std::env::current_exe().context("Cannot access current_exe")?;
    Some(
        buck2_forkserver::unix::launch_forkserver(exe, &["forkserver"], forkserver_state_dir).await,
    )
    .transpose()
}

#[cfg(not(unix))]
pub async fn maybe_launch_forkserver(
    _root_config: &LegacyBuckConfig,
    _forkserver_state_dir: &AbsNormPath,
) -> anyhow::Result<Option<ForkserverClient>> {
    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 std::sync::Arc;

use buck2_common::cas_digest::CasDigestConfig;
use buck2_common::io::fs::FsIoProvider;
use buck2_common::io::trace::TracingIoProvider;
use buck2_common::io::IoProvider;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_core::fs::project::ProjectRoot;

pub async fn create_io_provider(
    fb: fbinit::FacebookInit,
    project_fs: ProjectRoot,
    root_config: Option<&LegacyBuckConfig>,
    cas_digest_config: CasDigestConfig,
    trace_io: bool,
    eden_io_v2: bool,
) -> anyhow::Result<Arc<dyn IoProvider>> {
    #[cfg(any(fbcode_build, cargo_internal_build))]
    {
        use buck2_core::rollout_percentage::RolloutPercentage;

        let allow_eden_io_default =
            RolloutPercentage::from_bool(cfg!(any(target_os = "macos", target_os = "windows")));

        let allow_eden_io = root_config
            .and_then(|c| c.parse("buck2", "allow_eden_io").transpose())
            .transpose()?
            .unwrap_or(allow_eden_io_default)
            .roll();

        if allow_eden_io {
            if let Some(eden) = buck2_eden::io_provider::EdenIoProvider::new(
                fb,
                &project_fs,
                cas_digest_config,
                eden_io_v2,
            )
            .await?
            {
                return if trace_io {
                    Ok(Arc::new(TracingIoProvider::new(Box::new(eden))))
                } else {
                    Ok(Arc::new(eden))
                };
            }
        }
    }

    let _allow_unused = (fb, root_config, eden_io_v2);

    if trace_io {
        Ok(Arc::new(TracingIoProvider::new(Box::new(
            FsIoProvider::new(project_fs, cas_digest_config),
        ))))
    } else {
        Ok(Arc::new(FsIoProvider::new(project_fs, cas_digest_config)))
    }
}


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

pub mod check_working_dir;
pub mod common;
pub mod daemon_tcp;
pub mod dice_dump;
pub mod disk_state;
pub mod forkserver;
pub(crate) mod io_provider;
mod multi_event_stream;
pub mod panic;
pub mod server;
pub(crate) mod server_allocative;
pub mod state;


/*
 * 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::pin::Pin;
use std::task::Context;
use std::task::Poll;

use buck2_cli_proto::CommandProgress;
use buck2_cli_proto::MultiCommandProgress;
use futures::stream::Stream;
use pin_project::pin_project;
use prost::Message;

const PREFERRED_MESSAGE_SIZE_BYTES: usize = 32768; // This is an approximation.

/// Buffer CommandProgress into MultiCommandProgress batches.
#[pin_project]
pub struct MultiEventStream<S, E> {
    #[pin]
    inner: S,
    done: bool,
    buffered_err: Option<E>,
}

impl<S, E> MultiEventStream<S, E> {
    pub fn new(inner: S) -> Self {
        Self {
            inner,
            done: false,
            buffered_err: None,
        }
    }
}

impl<S, E> Stream for MultiEventStream<S, E>
where
    S: Stream<Item = Result<CommandProgress, E>>,
{
    type Item = Result<MultiCommandProgress, E>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        let mut current_size = 0;
        let mut messages = Vec::new();

        loop {
            if *this.done {
                break;
            };

            if current_size > PREFERRED_MESSAGE_SIZE_BYTES {
                break;
            }

            if this.buffered_err.is_some() {
                break;
            }

            match this.inner.as_mut().poll_next(cx) {
                Poll::Ready(Some(Ok(item))) => {
                    // The current_size is an approximation.
                    let len = item.encoded_len();
                    current_size += len + prost::length_delimiter_len(len);
                    messages.push(item);
                }
                Poll::Ready(Some(Err(e))) => {
                    *this.buffered_err = Some(e);
                }
                Poll::Ready(None) => {
                    *this.done = true;
                }
                Poll::Pending => {
                    break;
                }
            }
        }

        if !messages.is_empty() {
            tracing::trace!(n = messages.len(), "multi");
            return Poll::Ready(Some(Ok(MultiCommandProgress { messages })));
        }

        if let Some(err) = this.buffered_err.take() {
            return Poll::Ready(Some(Err(err)));
        }

        if *this.done {
            Poll::Ready(None)
        } else {
            Poll::Pending
        }
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;
    use buck2_data::BuckEvent;
    use futures::stream::poll_fn;
    use futures::stream::Stream;
    use futures::stream::StreamExt;

    use super::*;

    fn event(span_id: u64) -> CommandProgress {
        CommandProgress {
            progress: Some(buck2_cli_proto::command_progress::Progress::Event(
                Box::new(BuckEvent {
                    span_id,
                    ..Default::default()
                }),
            )),
        }
    }

    fn ready_event(span_id: u64) -> Poll<Option<Result<CommandProgress, ()>>> {
        Poll::Ready(Some(Ok(event(span_id))))
    }

    fn test_stream(
        mut events: Vec<Poll<Option<Result<CommandProgress, ()>>>>,
    ) -> impl Stream<Item = Result<CommandProgress, ()>> {
        events.push(Poll::Ready(None));
        events.reverse();
        poll_fn(move |_| events.pop().expect("Polled after completion!"))
    }

    #[tokio::test]
    async fn test_batch_and_done() {
        let s = test_stream(vec![ready_event(1), ready_event(2)]);

        let s = MultiEventStream::new(s);

        assert_eq!(
            s.collect::<Vec<_>>().await,
            vec![Ok(MultiCommandProgress {
                messages: vec![event(1), event(2)]
            })]
        );
    }

    #[tokio::test]
    async fn test_err() {
        let s = test_stream(vec![Poll::Ready(Some(Err(())))]);

        let s = MultiEventStream::new(s);

        assert_eq!(s.collect::<Vec<_>>().await, vec![Err(())]);
    }

    #[tokio::test]
    async fn test_err_buffered() {
        let s = test_stream(vec![ready_event(1), Poll::Ready(Some(Err(())))]);

        let s = MultiEventStream::new(s);

        assert_eq!(
            s.collect::<Vec<_>>().await,
            vec![
                Ok(MultiCommandProgress {
                    messages: vec![event(1)]
                }),
                Err(())
            ]
        );
    }

    #[tokio::test]
    async fn test_batch_on_pending() {
        let s = test_stream(vec![ready_event(1), Poll::Pending, ready_event(2)]);
        let s = MultiEventStream::new(s);
        futures::pin_mut!(s);

        assert_eq!(
            s.next().await,
            Some(Ok(MultiCommandProgress {
                messages: vec![event(1)]
            })),
        );
    }

    #[tokio::test]
    async fn test_batch_limit() {
        let s = poll_fn(|_| ready_event(0));
        let s = MultiEventStream::new(s);
        futures::pin_mut!(s);

        // Make sure this actually finishes and lands on a reasonable message size. Note that
        // given the test messages are *very* small this will be off by quite a bit.
        assert_matches!(s.next().await, Some(Ok(msg)) => {
            assert!(msg.encoded_len() > PREFERRED_MESSAGE_SIZE_BYTES);
            assert!(msg.encoded_len() < 2 * PREFERRED_MESSAGE_SIZE_BYTES);
        })
    }
}


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

//! Daemon-only panic hooks.
//!
//! This module sets up a panic hook to send the panic message to open CLIs.

use std::env::temp_dir;
use std::panic;
use std::panic::PanicInfo;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;

use buck2_cli_proto::unstable_dice_dump_request::DiceDumpFormat;
use buck2_core::env_helper::EnvHelper;
use buck2_wrapper_common::invocation_id::TraceId;
use once_cell::sync::OnceCell;

use crate::daemon::dice_dump::tar_dice_dump;

pub trait DaemonStatePanicDiceDump: Send + Sync + 'static {
    fn dice_dump(&self, path: &Path, format: DiceDumpFormat) -> anyhow::Result<()>;
}

fn get_panic_dump_dir() -> PathBuf {
    temp_dir().join("buck2-dumps")
}

async fn remove_old_panic_dumps() -> anyhow::Result<()> {
    const MAX_PANIC_AGE: Duration = Duration::from_secs(60 * 60 * 24); // 1 day
    let dump_dir = get_panic_dump_dir();
    let now = SystemTime::now();
    if let Ok(dir_result) = std::fs::read_dir(dump_dir) {
        let dumps = dir_result.filter_map(Result::ok).collect::<Vec<_>>();
        for record in dumps {
            let metadata = record.metadata()?;
            if now.duration_since(metadata.modified().unwrap_or(SystemTime::UNIX_EPOCH))?
                > MAX_PANIC_AGE
            {
                match metadata.is_dir() {
                    true => std::fs::remove_dir_all(record.path()),
                    false => std::fs::remove_file(record.path()),
                }
                .ok();
            }
        }
    }
    Ok(())
}

/// Initializes the panic hook.
pub fn initialize(daemon_state: Arc<dyn DaemonStatePanicDiceDump>) {
    let hook = panic::take_hook();
    panic::set_hook(Box::new(move |info| {
        daemon_panic_hook(&daemon_state, info);
        hook(info);
    }));
    tokio::spawn(remove_old_panic_dumps());
}

/// If DICE panics while holding a write lock, that lock will be poisoned and
/// the DICE dump will panic when it tries to acquire that lock.
/// This cell prevents a circular set of panics if this happens.
static ALREADY_DUMPED_DICE: OnceCell<()> = OnceCell::new();

static DICE_DUMP_ON_PANIC: EnvHelper<bool> = EnvHelper::new("BUCK2_DICE_DUMP_ON_PANIC");

fn daemon_panic_hook(daemon_state: &Arc<dyn DaemonStatePanicDiceDump>, info: &PanicInfo) {
    if !buck2_core::is_open_source()
        && DICE_DUMP_ON_PANIC
            .get_copied()
            .unwrap_or_default()
            .unwrap_or_default()
        && ALREADY_DUMPED_DICE.set(()).is_ok()
    {
        let panic_id = TraceId::new();
        maybe_dice_dump(daemon_state, info, &panic_id);
    }
}

fn maybe_dice_dump(
    daemon_state: &Arc<dyn DaemonStatePanicDiceDump>,
    info: &PanicInfo,
    panic_id: &TraceId,
) {
    let is_dice_panic = info.location().map_or(false, |loc| {
        loc.file().split(&['/', '\\']).any(|x| x == "dice")
    });
    if is_dice_panic {
        let dice_dump_folder = get_panic_dump_dir().join(format!("dice-dump-{}", panic_id));
        eprintln!(
            "Buck2 panicked and DICE may be responsible. Please be patient as we try to dump DICE graph to `{:?}` and create an archive file",
            dice_dump_folder
        );
        if let Err(e) = daemon_state.dice_dump(&dice_dump_folder, DiceDumpFormat::Tsv) {
            eprintln!("Failed to dump DICE graph: {:#}", e);
        } else {
            if let Err(e) = tar_dice_dump(&dice_dump_folder) {
                eprintln!(
                    "Failed to create an archive file of DICE dump. You can try manually creating it using `tar -zcf {}.tar.gz {}`. Error: {:#}",
                    dice_dump_folder.display(),
                    dice_dump_folder.display(),
                    e
                );
            }

            let maybe_report_msg = if cfg!(fbcode_build) {
                format!(
                    "Please upload the report via `jf upload {}.tar.gz` and then report to https://fb.workplace.com/groups/buck2users. ",
                    dice_dump_folder.display()
                )
            } else {
                "".to_owned()
            };
            eprintln!(
                "DICE graph dumped to `{:?}`. {}DICE dumps can take up a lot of disk space, you should delete the dump after reporting.",
                dice_dump_folder, maybe_report_msg
            );
        }
    }
}