/*
 * 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::io::BufWriter;
use std::marker::PhantomData;
use std::ops::Deref;
use std::sync::Arc;
use std::time::Instant;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::actions::impls::run_action_knobs::HasRunActionKnobs;
use buck2_build_api::actions::impls::run_action_knobs::RunActionKnobs;
use buck2_build_api::build::HasCreateUnhashedSymlinkLock;
use buck2_build_api::build_signals::create_build_signals;
use buck2_build_api::build_signals::BuildSignalsInstaller;
use buck2_build_api::build_signals::SetBuildSignals;
use buck2_build_api::context::SetBuildContextData;
use buck2_build_api::keep_going::HasKeepGoing;
use buck2_build_api::spawner::BuckSpawner;
use buck2_build_signals::CriticalPathBackendName;
use buck2_build_signals::HasCriticalPathBackend;
use buck2_cli_proto::client_context::HostArchOverride;
use buck2_cli_proto::client_context::HostPlatformOverride;
use buck2_cli_proto::common_build_options::ExecutionStrategy;
use buck2_cli_proto::ClientContext;
use buck2_cli_proto::CommonBuildOptions;
use buck2_cli_proto::ConfigOverride;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::cycles::CycleDetectorAdapter;
use buck2_common::dice::cycles::PairDiceCycleDetector;
use buck2_common::dice::data::HasIoProvider;
use buck2_common::executor_config::CommandExecutorConfig;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_common::http::SetHttpClient;
use buck2_common::io::trace::TracingIoProvider;
use buck2_common::io::IoProvider;
use buck2_common::legacy_configs::dice::HasLegacyConfigs;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_common::legacy_configs::LegacyBuckConfigs;
use buck2_common::result::SharedError;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_configured::calculation::ConfiguredGraphCycleDescriptor;
use buck2_core::async_once_cell::AsyncOnceCell;
use buck2_core::cells::CellResolver;
use buck2_core::facebook_only;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::fs::working_dir::WorkingDir;
use buck2_core::pattern::pattern_type::ConfiguredProvidersPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::rollout_percentage::RolloutPercentage;
use buck2_events::daemon_id;
use buck2_events::dispatch::EventDispatcher;
use buck2_events::metadata;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::execute::blocking::SetBlockingExecutor;
use buck2_execute::execute::dice_data::set_fallback_executor_config;
use buck2_execute::execute::dice_data::SetCommandExecutor;
use buck2_execute::execute::dice_data::SetReClient;
use buck2_execute::knobs::ExecutorGlobalKnobs;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::materialize::materializer::SetMaterializer;
use buck2_execute::re::client::RemoteExecutionClient;
use buck2_execute::re::manager::ReConnectionHandle;
use buck2_execute::re::manager::ReConnectionManager;
use buck2_execute::re::manager::ReConnectionObserver;
use buck2_execute_impl::executors::worker::WorkerPool;
use buck2_execute_impl::low_pass_filter::LowPassFilter;
use buck2_file_watcher::file_watcher::FileWatcher;
use buck2_file_watcher::mergebase::SetMergebase;
use buck2_forkserver::client::ForkserverClient;
use buck2_interpreter::dice::starlark_debug::SetStarlarkDebugger;
use buck2_interpreter::dice::starlark_profiler::StarlarkProfilerConfiguration;
use buck2_interpreter::extra::xcode::XcodeVersionInfo;
use buck2_interpreter::extra::InterpreterHostArchitecture;
use buck2_interpreter::extra::InterpreterHostPlatform;
use buck2_interpreter::prelude_path::prelude_path;
use buck2_interpreter_for_build::interpreter::configuror::BuildInterpreterConfiguror;
use buck2_interpreter_for_build::interpreter::cycles::LoadCycleDescriptor;
use buck2_interpreter_for_build::interpreter::globals::configure_build_file_globals;
use buck2_interpreter_for_build::interpreter::globals::configure_bxl_file_globals;
use buck2_interpreter_for_build::interpreter::globals::configure_extension_file_globals;
use buck2_interpreter_for_build::interpreter::globals::configure_package_file_globals;
use buck2_interpreter_for_build::interpreter::interpreter_setup::setup_interpreter;
use buck2_server_ctx::concurrency::ConcurrencyHandler;
use buck2_server_ctx::concurrency::DiceDataProvider;
use buck2_server_ctx::concurrency::DiceUpdater;
use buck2_server_ctx::ctx::DiceAccessor;
use buck2_server_ctx::ctx::PrivateStruct;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::stderr_output_guard::StderrOutputGuard;
use buck2_server_ctx::stderr_output_guard::StderrOutputWriter;
use buck2_server_starlark_debug::create_debugger_handle;
use buck2_server_starlark_debug::BuckStarlarkDebuggerHandle;
use buck2_util::truncate::truncate_container;
use dice::DiceComputations;
use dice::DiceData;
use dice::DiceTransactionUpdater;
use dice::UserComputationData;
use dice::UserCycleDetector;
use dupe::Dupe;
use gazebo::prelude::SliceExt;
use host_sharing::HostSharingBroker;
use host_sharing::HostSharingStrategy;
use more_futures::cancellation::ExplicitCancellationContext;
use tokio::sync::Mutex;
use tracing::warn;

use crate::active_commands::ActiveCommandDropGuard;
use crate::configs::parse_legacy_cells;
use crate::daemon::common::get_default_executor_config;
use crate::daemon::common::parse_concurrency;
use crate::daemon::common::CommandExecutorFactory;
use crate::dice_tracker::BuckDiceTracker;
use crate::heartbeat_guard::HeartbeatGuard;
use crate::host_info;

#[derive(Debug, thiserror::Error)]
enum DaemonCommunicationError {
    #[error("Got invalid working directory `{0}`")]
    InvalidWorkingDirectory(String),
}

/// BaseCommandContext provides access to the global daemon state and information specific to a command (like the
/// EventDispatcher). Most commands use a ServerCommandContext which has more command/client-specific information.
pub struct BaseServerCommandContext {
    /// An fbinit token for using things that require fbinit. fbinit is initialized on daemon startup.
    pub _fb: fbinit::FacebookInit,
    /// Absolute path to the project root.
    pub project_root: ProjectRoot,
    /// A reference to the dice graph. Most interesting things are accessible from this (and new interesting things should be
    /// added there rather than as fields here). This has some per-request setup done already (like attaching a per-request
    /// event dispatcher).
    pub dice_manager: ConcurrencyHandler,
    /// A reference to the I/O provider.
    pub io: Arc<dyn IoProvider>,
    /// The RE connection, managed such that all build commands that are concurrently active uses
    /// the same connection.
    pub re_client_manager: Arc<ReConnectionManager>,
    /// Executor responsible for coordinating and rate limiting I/O.
    pub blocking_executor: Arc<dyn BlockingExecutor>,
    /// Object responsible for handling most materializations.
    pub materializer: Arc<dyn Materializer>,
    /// Forkserver connection, if any was started
    pub forkserver: Option<ForkserverClient>,
    /// The event dispatcher for this command context.
    pub events: EventDispatcher,
    /// Removes this command from the set of active commands when dropped.
    pub _drop_guard: ActiveCommandDropGuard,
    /// The file watcher that keeps buck2 up to date with disk changes.
    pub file_watcher: Arc<dyn FileWatcher>,
    /// Whether or not to hash all commands
    pub hash_all_commands: bool,
    /// Whether to try to read from the network action output cache when running
    /// network actions like download_file; only useful for offline builds.
    pub use_network_action_output_cache: bool,
    /// Start time to track daemon uptime
    pub daemon_start_time: Instant,
    /// Mutex for creating symlinks
    pub create_unhashed_outputs_lock: Arc<Mutex<()>>,
    /// Http client used during run actions; shared with materializer.
    pub http_client: CountingHttpClient,
}

/// ServerCommandContext provides access to the global daemon state and information about the calling client for
/// the implementation of DaemonApi endpoints (ex. targets, query, build).
pub struct ServerCommandContext<'a> {
    pub base_context: BaseServerCommandContext,

    /// The working directory of the client. This is used for resolving things in the request in a
    /// working-dir relative way. For example, it's common to resolve target patterns relative to
    /// the working directory and resolving cell aliases there. This should generally only be used
    /// to interpret values that are in the request. We should convert to client-agnostic things early.
    pub working_dir: ProjectRelativePathBuf,

    working_dir_abs: WorkingDir,

    /// The oncall specified by the client, if any. This gets injected into request metadata.
    pub oncall: Option<String>,

    host_platform_override: HostPlatformOverride,
    host_arch_override: HostArchOverride,
    host_xcode_version_override: Option<String>,

    // This ensures that there's only one RE connection during the lifetime of this context. It's possible
    // that we give out other handles, but we don't depend on the lifetimes of those for this guarantee. We
    // also use this to send a RemoteExecutionSessionCreated if the connection is made.
    _re_connection_handle: ReConnectionHandle,

    /// Starlark profiler instrumentation requested throughout the duration of this command. Usually associated with
    /// the `buck2 profile` command.
    pub starlark_profiler_instrumentation_override: StarlarkProfilerConfiguration,

    debugger_handle: Option<BuckStarlarkDebuggerHandle>,

    record_target_call_stacks: bool,
    skip_targets_with_duplicate_names: bool,
    disable_starlark_types: bool,

    pub buck_out_dir: ProjectRelativePathBuf,

    /// Common build options associated with this command.
    build_options: Option<CommonBuildOptions>,

    /// The CellResolver and Configs loader for this command
    cell_configs_loader: Arc<CellConfigLoader>,

    /// Keep emitting heartbeat events while the ServerCommandContext is alive  We put this in an
    /// Option so that we can ensure heartbeat events are cancelled before everything else is
    /// dropped.
    heartbeat_guard_handle: Option<HeartbeatGuard>,

    /// Daemon uuid passed in from the client side to detect nested invocation.
    pub(crate) daemon_uuid_from_client: Option<String>,

    /// Sanitized argument vector from the CLI from the client side.
    pub(crate) sanitized_argv: Vec<String>,

    cancellations: &'a ExplicitCancellationContext,

    exit_when_different_state: bool,
}

impl<'a> ServerCommandContext<'a> {
    pub fn new(
        base_context: BaseServerCommandContext,
        client_context: &ClientContext,
        starlark_profiler_instrumentation_override: StarlarkProfilerConfiguration,
        build_options: Option<&CommonBuildOptions>,
        buck_out_dir: ProjectRelativePathBuf,
        cancellations: &'a ExplicitCancellationContext,
    ) -> anyhow::Result<Self> {
        let working_dir = AbsNormPath::new(&client_context.working_dir)?;

        let working_dir_project_relative = working_dir
            .strip_prefix(base_context.project_root.root())
            .map_err(|_| {
                Into::<anyhow::Error>::into(DaemonCommunicationError::InvalidWorkingDirectory(
                    client_context.working_dir.clone(),
                ))
            })?;

        #[derive(Allocative)]
        struct Observer {
            events: EventDispatcher,
        }

        impl ReConnectionObserver for Observer {
            fn session_created(&self, client: &RemoteExecutionClient) {
                let session_id = client.get_session_id();
                let experiment_name = match client.get_experiment_name() {
                    Ok(Some(exp)) => exp,
                    Ok(None) => "".to_owned(),
                    Err(e) => {
                        tracing::debug!("Failed to access RE experiment name: {:#}", e);
                        "<ffi error>".to_owned()
                    }
                };

                self.events
                    .instant_event(buck2_data::RemoteExecutionSessionCreated {
                        session_id: session_id.to_owned(),
                        experiment_name,
                    })
            }
        }

        let mut re_connection_handle = base_context.re_client_manager.get_re_connection();

        re_connection_handle.set_observer(Arc::new(Observer {
            events: base_context.events.dupe(),
        }));

        // Add argfiles read by client into IO tracing state.
        if let Some(tracing_provider) = base_context.io.as_any().downcast_ref::<TracingIoProvider>()
        {
            let argfiles: anyhow::Result<Vec<_>> = client_context
                .argfiles
                .iter()
                .map(|s| AbsNormPathBuf::new(s.into()))
                .collect();
            tracing_provider.add_config_paths(&base_context.project_root, argfiles?);
        }

        let oncall = if client_context.oncall.is_empty() {
            None
        } else {
            Some(client_context.oncall.clone())
        };

        let heartbeat_guard_handle = HeartbeatGuard::new(&base_context);

        let cell_configs_loader = Arc::new(CellConfigLoader {
            project_root: base_context.project_root.clone(),
            working_dir: working_dir_project_relative.to_buf().into(),
            reuse_current_config: client_context.reuse_current_config,
            config_overrides: client_context.config_overrides.clone(),
            loaded_cell_configs: AsyncOnceCell::new(),
        });

        let debugger_handle = create_debugger_handle(base_context.events.dupe());

        Ok(ServerCommandContext {
            base_context,
            working_dir: working_dir_project_relative.to_buf().into(),
            working_dir_abs: WorkingDir::unchecked_new(working_dir.to_buf()),
            host_platform_override: client_context.host_platform(),
            host_arch_override: client_context.host_arch(),
            host_xcode_version_override: client_context.host_xcode_version.clone(),
            oncall,
            _re_connection_handle: re_connection_handle,
            starlark_profiler_instrumentation_override,
            buck_out_dir,
            build_options: build_options.cloned(),
            cell_configs_loader,
            record_target_call_stacks: client_context.target_call_stacks,
            skip_targets_with_duplicate_names: client_context.skip_targets_with_duplicate_names,
            disable_starlark_types: client_context.disable_starlark_types,
            heartbeat_guard_handle: Some(heartbeat_guard_handle),
            daemon_uuid_from_client: client_context.daemon_uuid.clone(),
            sanitized_argv: client_context.sanitized_argv.clone(),
            debugger_handle,
            cancellations,
            exit_when_different_state: client_context.exit_when_different_state,
        })
    }

    async fn dice_data_constructor(
        &self,
        build_signals: BuildSignalsInstaller,
    ) -> DiceCommandDataProvider {
        let execution_strategy = self
            .build_options
            .as_ref()
            .map(|opts| opts.execution_strategy)
            .map_or(ExecutionStrategy::LocalOnly, |strategy| {
                ExecutionStrategy::from_i32(strategy).expect("execution strategy should be valid")
            });

        let skip_cache_read = self
            .build_options
            .as_ref()
            .map(|opts| opts.skip_cache_read)
            .unwrap_or_default();

        let skip_cache_write = self
            .build_options
            .as_ref()
            .map(|opts| opts.skip_cache_write)
            .unwrap_or_default();

        let mut run_action_knobs = RunActionKnobs {
            hash_all_commands: self.base_context.hash_all_commands,
            use_network_action_output_cache: self.base_context.use_network_action_output_cache,
            ..Default::default()
        };

        if let Some(build_options) = self.build_options.as_ref() {
            run_action_knobs.eager_dep_files = build_options.eager_dep_files;
        }

        let concurrency = self
            .build_options
            .as_ref()
            .and_then(|opts| opts.concurrency.as_ref())
            .map(|obj| parse_concurrency(obj.concurrency))
            .map(|v| v.map_err(SharedError::from));

        let executor_config = get_default_executor_config(self.host_platform_override);
        let blocking_executor: Arc<_> = self.base_context.blocking_executor.dupe();
        let materializer = self.base_context.materializer.dupe();
        let re_connection = Arc::new(self.get_re_connection());

        let forkserver = self.base_context.forkserver.dupe();

        let upload_all_actions = self
            .build_options
            .as_ref()
            .map_or(false, |opts| opts.upload_all_actions);

        let create_unhashed_symlink_lock = self.base_context.create_unhashed_outputs_lock.dupe();

        DiceCommandDataProvider {
            cell_configs_loader: self.cell_configs_loader.dupe(),
            events: self.events().dupe(),
            execution_strategy,
            run_action_knobs,
            concurrency,
            executor_config: Arc::new(executor_config),
            blocking_executor,
            materializer,
            re_connection,
            build_signals,
            forkserver,
            upload_all_actions,
            skip_cache_read,
            skip_cache_write,
            create_unhashed_symlink_lock,
            starlark_debugger: self.debugger_handle.dupe(),
            keep_going: self
                .build_options
                .as_ref()
                .map_or(false, |opts| opts.keep_going),
            http_client: self.base_context.http_client.dupe(),
        }
    }

    async fn dice_updater(&self) -> anyhow::Result<DiceCommandUpdater> {
        let (interpreter_platform, interpreter_architecture, interpreter_xcode_version) =
            host_info::get_host_info(
                self.host_platform_override,
                self.host_arch_override,
                &self.host_xcode_version_override,
            )?;

        Ok(DiceCommandUpdater {
            file_watcher: self.base_context.file_watcher.dupe(),
            cell_config_loader: self.cell_configs_loader.dupe(),
            buck_out_dir: self.buck_out_dir.clone(),
            interpreter_platform,
            interpreter_architecture,
            interpreter_xcode_version,
            starlark_profiler_instrumentation_override: self
                .starlark_profiler_instrumentation_override
                .dupe(),
            disable_starlark_types: self.disable_starlark_types,
            skip_targets_with_duplicate_names: self.skip_targets_with_duplicate_names,
            record_target_call_stacks: self.record_target_call_stacks,
        })
    }

    pub fn get_re_connection(&self) -> ReConnectionHandle {
        self.base_context.re_client_manager.get_re_connection()
    }
}

struct CellConfigLoader {
    project_root: ProjectRoot,
    working_dir: ProjectRelativePathBuf,
    /// Reuses build config from the previous invocation if there is one
    reuse_current_config: bool,
    config_overrides: Vec<ConfigOverride>,
    loaded_cell_configs:
        AsyncOnceCell<SharedResult<(CellResolver, LegacyBuckConfigs, HashSet<AbsNormPathBuf>)>>,
}

impl CellConfigLoader {
    pub async fn cells_and_configs(
        &self,
        dice_ctx: &DiceComputations,
    ) -> SharedResult<(CellResolver, LegacyBuckConfigs, HashSet<AbsNormPathBuf>)> {
        self.loaded_cell_configs
            .get_or_init(async move {
                if self.reuse_current_config {
                    // If there is a previous command and --reuse-current-config is set, then the old config is used, ignoring any overrides.
                    if dice_ctx.is_cell_resolver_key_set().await?
                        && dice_ctx.is_legacy_configs_key_set().await?
                    {
                        if !self.config_overrides.is_empty() {
                            warn!(
                                "Found config overrides while using --reuse-current-config flag. Ignoring overrides [{}] and using current config instead",
                                truncate_container(self.config_overrides.iter().map(|e| &*e.config_override), 200),
                            );
                        }
                        return Ok::<(CellResolver, LegacyBuckConfigs, HashSet<AbsNormPathBuf>), anyhow::Error>((
                            dice_ctx.get_cell_resolver().await?,
                            dice_ctx.get_legacy_configs().await?,
                            HashSet::new(),
                        )).shared_error();
                    } else {
                        // If there is no previous command but the flag was set, then the flag is ignored, the command behaves as if there isn't the reuse config flag.
                        warn!(
                            "--reuse-current-config flag was set, but there was no previous invocation detected. Ignoring --reuse-current-config flag"
                        );
                    }
                }
                parse_legacy_cells(self.config_overrides.iter(), &self.working_dir, &self.project_root)
                    .shared_error()
            })
            .await
            .clone()
    }
}

struct DiceCommandDataProvider {
    cell_configs_loader: Arc<CellConfigLoader>,
    execution_strategy: ExecutionStrategy,
    events: EventDispatcher,
    concurrency: Option<Result<usize, SharedError>>,
    executor_config: Arc<CommandExecutorConfig>,
    blocking_executor: Arc<dyn BlockingExecutor>,
    materializer: Arc<dyn Materializer>,
    re_connection: Arc<ReConnectionHandle>,
    build_signals: BuildSignalsInstaller,
    forkserver: Option<ForkserverClient>,
    upload_all_actions: bool,
    run_action_knobs: RunActionKnobs,
    skip_cache_read: bool,
    skip_cache_write: bool,
    create_unhashed_symlink_lock: Arc<Mutex<()>>,
    starlark_debugger: Option<BuckStarlarkDebuggerHandle>,
    keep_going: bool,
    http_client: CountingHttpClient,
}

#[async_trait]
impl DiceDataProvider for DiceCommandDataProvider {
    async fn provide(&self, ctx: &DiceComputations) -> anyhow::Result<UserComputationData> {
        let (cell_resolver, legacy_configs, _): (CellResolver, LegacyBuckConfigs, _) =
            self.cell_configs_loader.cells_and_configs(ctx).await?;

        // TODO(cjhopman): The CellResolver and the legacy configs shouldn't be leaves on the graph. This should
        // just be setting the config overrides and host platform override as leaves on the graph.

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

        let config_threads = root_config.parse("build", "threads")?.unwrap_or(0);

        let concurrency = match self.concurrency.as_ref() {
            Some(v) => v.dupe()?,
            None => parse_concurrency(config_threads)?,
        };

        if let Some(max_lines) = root_config.parse("ui", "thread_line_limit")? {
            self.events
                .instant_event(buck2_data::ConsolePreferences { max_lines });
        }

        let enable_miniperf = root_config
            .parse::<RolloutPercentage>("buck2", "miniperf2")?
            .unwrap_or_else(RolloutPercentage::always)
            .roll();

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

        let executor_global_knobs = ExecutorGlobalKnobs {
            enable_miniperf,
            log_action_keys,
        };

        let host_sharing_broker =
            HostSharingBroker::new(HostSharingStrategy::SmallerTasksFirst, concurrency);

        // We use the job count for the low pass filter too. The low pass filter prevents sending
        // RE-eligile tasks to local if their concurrency is higher than our threshold. While it
        // doesn't *have* to be the same as the concurrency we give the actual executor, it's a
        // reasonable pick, because if we send more tasks than our concurrency limit allows, we
        // would expect to start losing out to RE in terms of perf.
        let low_pass_filter = LowPassFilter::new(concurrency);

        let mut data = DiceData::new();
        data.set(self.events.dupe());

        let cycle_detector = if root_config
            .parse::<bool>("build", "lazy_cycle_detector")?
            .unwrap_or(true)
        {
            Some(create_cycle_detector())
        } else {
            None
        };
        let has_cycle_detector = cycle_detector.is_some();

        let mut run_action_knobs = self.run_action_knobs.dupe();
        run_action_knobs.use_network_action_output_cache |= root_config
            .parse::<bool>("buck2", "use_network_action_output_cache")?
            .unwrap_or(false);

        run_action_knobs.enforce_re_timeouts = root_config
            .parse::<bool>("buck2", "enforce_re_timeouts")?
            .unwrap_or(true);

        let mut data = UserComputationData {
            data,
            tracker: Arc::new(BuckDiceTracker::new(self.events.dupe())),
            cycle_detector,
            activation_tracker: Some(self.build_signals.activation_tracker.dupe()),
            ..Default::default()
        };

        let worker_pool = Arc::new(WorkerPool::new());

        let critical_path_backend = root_config
            .parse("buck2", "critical_path_backend2")?
            .unwrap_or(CriticalPathBackendName::Default);

        set_fallback_executor_config(&mut data.data, self.executor_config.dupe());
        data.set_re_client(self.re_connection.get_client());
        data.set_command_executor(Box::new(CommandExecutorFactory::new(
            self.re_connection.dupe(),
            host_sharing_broker,
            low_pass_filter,
            self.materializer.dupe(),
            self.blocking_executor.dupe(),
            self.execution_strategy,
            executor_global_knobs,
            self.upload_all_actions,
            self.forkserver.dupe(),
            self.skip_cache_read,
            self.skip_cache_write,
            ctx.global_data()
                .get_io_provider()
                .project_root()
                .to_owned(),
            worker_pool,
        )));
        data.set_blocking_executor(self.blocking_executor.dupe());
        data.set_http_client(self.http_client.dupe());
        data.set_materializer(self.materializer.dupe());
        data.set_build_signals(self.build_signals.build_signals.dupe());
        data.set_run_action_knobs(run_action_knobs);
        data.set_create_unhashed_symlink_lock(self.create_unhashed_symlink_lock.dupe());
        data.set_starlark_debugger_handle(self.starlark_debugger.clone().map(|v| Box::new(v) as _));
        data.set_keep_going(self.keep_going);
        data.set_critical_path_backend(critical_path_backend);
        data.spawner = Arc::new(BuckSpawner);

        let tags = vec![
            format!("lazy-cycle-detector:{}", has_cycle_detector),
            format!("miniperf:{}", enable_miniperf),
        ];
        self.events.instant_event(buck2_data::TagEvent { tags });

        self.events.instant_event(buck2_data::CommandOptions {
            concurrency: concurrency as _,
        });

        Ok(data)
    }
}

fn create_cycle_detector() -> Arc<dyn UserCycleDetector> {
    Arc::new(PairDiceCycleDetector(
        CycleDetectorAdapter::<LoadCycleDescriptor>::new(),
        CycleDetectorAdapter::<ConfiguredGraphCycleDescriptor>::new(),
    ))
}

struct DiceCommandUpdater {
    file_watcher: Arc<dyn FileWatcher>,
    cell_config_loader: Arc<CellConfigLoader>,
    buck_out_dir: ProjectRelativePathBuf,
    interpreter_platform: InterpreterHostPlatform,
    interpreter_architecture: InterpreterHostArchitecture,
    interpreter_xcode_version: Option<XcodeVersionInfo>,
    starlark_profiler_instrumentation_override: StarlarkProfilerConfiguration,
    disable_starlark_types: bool,
    record_target_call_stacks: bool,
    skip_targets_with_duplicate_names: bool,
}

#[async_trait]
impl DiceUpdater for DiceCommandUpdater {
    async fn update(
        &self,
        ctx: DiceTransactionUpdater,
        user_data: &mut UserComputationData,
    ) -> anyhow::Result<DiceTransactionUpdater> {
        let (cell_resolver, legacy_configs, _): (CellResolver, LegacyBuckConfigs, _) = self
            .cell_config_loader
            .cells_and_configs(ctx.existing_state().await.deref())
            .await?;
        // TODO(cjhopman): The CellResolver and the legacy configs shouldn't be leaves on the graph. This should
        // just be setting the config overrides and host platform override as leaves on the graph.

        let configuror = BuildInterpreterConfiguror::new(
            Some(prelude_path(&cell_resolver)?),
            self.interpreter_platform,
            self.interpreter_architecture,
            self.interpreter_xcode_version.clone(),
            self.record_target_call_stacks,
            self.skip_targets_with_duplicate_names,
            configure_build_file_globals,
            configure_package_file_globals,
            configure_extension_file_globals,
            configure_bxl_file_globals,
            None,
        )?;

        let (mut ctx, mergebase) = self.file_watcher.sync(ctx).await?;
        user_data.set_mergebase(mergebase);

        ctx.set_buck_out_path(Some(self.buck_out_dir.clone()))?;

        setup_interpreter(
            &mut ctx,
            cell_resolver,
            configuror,
            legacy_configs,
            self.starlark_profiler_instrumentation_override.dupe(),
            self.disable_starlark_types,
        )?;

        Ok(ctx)
    }
}

impl<'a> Drop for ServerCommandContext<'a> {
    fn drop(&mut self) {
        // Ensure we cancel the heartbeat guard first.
        std::mem::drop(self.heartbeat_guard_handle.take());
    }
}

#[async_trait]
impl<'a> ServerCommandContextTrait for ServerCommandContext<'a> {
    fn working_dir(&self) -> &ProjectRelativePath {
        &self.working_dir
    }

    fn working_dir_abs(&self) -> &WorkingDir {
        &self.working_dir_abs
    }

    fn project_root(&self) -> &ProjectRoot {
        &self.base_context.project_root
    }

    fn materializer(&self) -> Arc<dyn Materializer> {
        self.base_context.materializer.dupe()
    }

    /// Provides a DiceTransaction, initialized on first use and shared after initialization.
    async fn dice_accessor(&self, _private: PrivateStruct) -> SharedResult<DiceAccessor> {
        let (build_signals_installer, deferred_build_signals) = create_build_signals();

        let is_nested_invocation = if let Some(uuid) = &self.daemon_uuid_from_client {
            uuid == &daemon_id::DAEMON_UUID.to_string()
        } else {
            false
        };

        Ok(DiceAccessor {
            dice_handler: self.base_context.dice_manager.dupe(),
            data: Box::new(self.dice_data_constructor(build_signals_installer).await),
            setup: Box::new(self.dice_updater().await?),
            is_nested_invocation,
            sanitized_argv: self.sanitized_argv.clone(),
            exit_when_different_state: self.exit_when_different_state,
            build_signals: deferred_build_signals,
        })
    }

    fn events(&self) -> &EventDispatcher {
        &self.base_context.events
    }

    fn stderr(&self) -> anyhow::Result<StderrOutputGuard<'_>> {
        Ok(StderrOutputGuard {
            _phantom: PhantomData,
            inner: BufWriter::with_capacity(
                // TODO(nga): no need to buffer here.
                4096,
                StderrOutputWriter::new(self)?,
            ),
        })
    }

    /// Gathers metadata to attach to events for when a command starts and stops.
    async fn request_metadata(&self) -> anyhow::Result<HashMap<String, String>> {
        // Facebook only: metadata collection for Scribe writes
        facebook_only();

        let mut metadata = metadata::collect();

        metadata.insert(
            "io_provider".to_owned(),
            self.base_context.io.name().to_owned(),
        );

        metadata.insert(
            "materializer".to_owned(),
            self.base_context.materializer.name().to_owned(),
        );

        if let Some(oncall) = &self.oncall {
            metadata.insert("oncall".to_owned(), oncall.clone());
        }

        Ok(metadata)
    }

    /// Gathers metadata from buckconfig to attach to events for when a command enters the critical
    /// section
    async fn config_metadata(
        &self,
        ctx: &DiceComputations,
    ) -> anyhow::Result<HashMap<String, String>> {
        // Facebook only: metadata collection for Scribe writes
        facebook_only();

        fn add_config(
            map: &mut HashMap<String, String>,
            cfg: &LegacyBuckConfig,
            section: &'static str,
            key: &'static str,
            field_name: &'static str,
        ) {
            if let Some(value) = cfg.get(section, key) {
                map.insert(field_name.to_owned(), value.to_owned());
            }
        }

        fn extract_scuba_defaults(
            config: Option<&LegacyBuckConfig>,
        ) -> Option<serde_json::Map<String, serde_json::Value>> {
            let config = config?.get("scuba", "defaults")?;
            let unescaped_config = shlex::split(config)?.join("");
            let sample_json: serde_json::Value = serde_json::from_str(&unescaped_config).ok()?;
            sample_json.get("normals")?.as_object().cloned()
        }

        let mut metadata = HashMap::new();
        // In the case of invalid configuration (e.g. something like buck2 build -c X), `dice_ctx_default` returns an
        // error. We won't be able to get configs to log in that case, but we shouldn't crash.
        let (cells, configs, paths): (CellResolver, LegacyBuckConfigs, HashSet<AbsNormPathBuf>) =
            self.cell_configs_loader.cells_and_configs(ctx).await?;

        // Add legacy config paths to I/O tracing (if enabled).
        if let Some(tracing_provider) = self
            .base_context
            .io
            .as_any()
            .downcast_ref::<TracingIoProvider>()
        {
            tracing_provider.add_config_paths(&self.base_context.project_root, paths);
        }

        let root_cell_config = configs.get(cells.root_cell());
        if let Ok(config) = root_cell_config {
            add_config(&mut metadata, config, "log", "repository", "repository");

            // Buck1 honors a configuration field, `scuba.defaults`, by drawing values from the configuration value and
            // inserting them verbatim into Scuba samples. Buck2 doesn't write to Scuba in the same way that Buck1
            // does, but metadata in this function indirectly makes its way to Scuba, so it makes sense to respect at
            // least some of the data within it.
            //
            // The configuration field is expected to be the canonical JSON representation for a Scuba sample, which is
            // to say something like this:
            // ```
            // {
            //   "normals": { "key": "value" },
            //   "ints": { "key": 0 },
            // }
            // ```
            //
            // TODO(swgillespie) - This only covers the normals since Buck2's event protocol only allows for string
            // metadata. Depending on what sort of things we're missing by dropping int default columns, we might want
            // to consider adding support to the protocol for integer metadata.

            if let Ok(cwd_cell_name) = cells.find(&self.working_dir) {
                let cwd_cell_config = configs.get(cwd_cell_name).ok();
                if let Some(normals_obj) = extract_scuba_defaults(cwd_cell_config) {
                    for (key, value) in normals_obj.iter() {
                        if let Some(value) = value.as_str() {
                            metadata.insert(key.clone(), value.to_owned());
                        }
                    }
                }

                // `client.id` is often set via the `-c` flag; `-c` configuration is assigned to the cwd cell and not
                // the root cell.
                if let Some(config) = cwd_cell_config {
                    add_config(&mut metadata, config, "client", "id", "client");
                    add_config(
                        &mut metadata,
                        config,
                        "cache",
                        "schedule_type",
                        "schedule_type",
                    );
                }
            }
        }

        Ok(metadata)
    }

    fn log_target_pattern(
        &self,
        providers_patterns: &[ParsedPattern<ConfiguredProvidersPatternExtra>],
    ) {
        let patterns = providers_patterns.map(|pat| buck2_data::TargetPattern {
            value: format!("{}", pat),
        });

        self.events()
            .instant_event(buck2_data::ParsedTargetPatterns {
                target_patterns: patterns,
            })
    }

    fn cancellation_context(&self) -> &ExplicitCancellationContext {
        self.cancellations
    }
}


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

use allocative::Allocative;
use buck2_data::*;
use buck2_events::dispatch::with_dispatcher_async;
use buck2_events::dispatch::EventDispatcher;
use dice::DiceEvent;
use dice::DiceEventListener;
use dupe::Dupe;
use futures::channel::mpsc;
use futures::channel::mpsc::UnboundedReceiver;
use futures::channel::mpsc::UnboundedSender;
use futures::StreamExt;

/// The BuckDiceTracker keeps track of the started/finished events for a dice computation and periodically sends a snapshot to the client.
///
/// There are too many events coming out of dice for us to forward them all to the client, so we need to aggregate
/// them in some way in the daemon.
///
/// The tracker will send a snapshot event every 500ms (only if there have been changes since the last snapshot).
///
/// A client won't necessarily get a final snapshot before a command returns.
#[derive(Allocative)]
pub struct BuckDiceTracker {
    #[allocative(skip)]
    event_forwarder: UnboundedSender<DiceEvent>,
}

const DICE_SNAPSHOT_INTERVAL: Duration = Duration::from_millis(500);

impl BuckDiceTracker {
    pub fn new(events: EventDispatcher) -> Self {
        let (event_forwarder, receiver) = mpsc::unbounded();

        std::thread::spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .unwrap();
            runtime.block_on(with_dispatcher_async(
                events.dupe(),
                Self::run_task(events, receiver),
            ))
        });

        Self { event_forwarder }
    }

    async fn run_task(events: EventDispatcher, mut receiver: UnboundedReceiver<DiceEvent>) {
        let mut needs_update = false;
        let mut states = HashMap::new();
        let mut interval = tokio::time::interval(DICE_SNAPSHOT_INTERVAL);
        interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
        // This will loop until the sender side of the channel is dropped.
        loop {
            tokio::select! {
                ev = receiver.next() => {
                    needs_update = true;
                    match ev {
                        Some(DiceEvent::Started{key_type}) => {
                            states.entry(key_type).or_insert_with(DiceKeyState::default).started += 1;
                        }
                        Some(DiceEvent::Finished{key_type}) => {
                            states.entry(key_type).or_insert_with(DiceKeyState::default).finished += 1;
                        }
                        Some(DiceEvent::CheckDepsStarted{key_type}) => {
                            states.entry(key_type).or_insert_with(DiceKeyState::default).check_deps_started += 1;
                        }
                        Some(DiceEvent::CheckDepsFinished{key_type}) => {
                            states.entry(key_type).or_insert_with(DiceKeyState::default).check_deps_finished += 1;
                        }
                        None => {
                            // This indicates that the sender side has been dropped and we can exit.
                            break;
                        }
                    }
                }
                _ = interval.tick() => {
                    if needs_update {
                        needs_update = false;
                        events.instant_event(DiceStateSnapshot {
                            key_states: states
                                .iter()
                                .map(|(k, v)| ((*k).to_owned(), v.clone()))
                                .collect(),
                        });
                    }
                }
            }
        }
    }
}

impl DiceEventListener for BuckDiceTracker {
    fn event(&self, event: DiceEvent) {
        let _ = self.event_forwarder.unbounded_send(event);
    }
}


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

use std::fmt;
use std::io::Write;
use std::path::Path;

use async_recursion::async_recursion;
use async_trait::async_trait;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::file_ops::FileOps;
use buck2_common::file_ops::RawPathMetadata;
use buck2_common::file_ops::RawSymlink;
use buck2_common::io::fs::FsIoProvider;
use buck2_common::io::IoProvider;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_execute::digest_config::HasDigestConfig;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::stdout_partial_output::StdoutPartialOutput;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use buck2_util::commas::commas;
use dice::DiceTransaction;
use dupe::Dupe;
use gazebo::variants::VariantName;

use crate::ctx::ServerCommandContext;

pub(crate) async fn file_status_command(
    ctx: &ServerCommandContext<'_>,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: buck2_cli_proto::FileStatusRequest,
) -> anyhow::Result<buck2_cli_proto::GenericResponse> {
    run_server_command(
        FileStatusServerCommand { req },
        ctx,
        partial_result_dispatcher,
    )
    .await
}
struct FileStatusServerCommand {
    req: buck2_cli_proto::FileStatusRequest,
}

struct FileStatusResult<'a> {
    /// Number of things we check
    checked: usize,
    /// Number of ones that were bad
    bad: usize,
    /// Whether to write matches
    verbose: bool,
    stdout: StdoutPartialOutput<'a>,
}

impl FileStatusResult<'_> {
    fn checking(&mut self) {
        self.checked += 1;
    }

    fn report<T>(
        &mut self,
        kind: &str,
        path: &ProjectRelativePath,
        fs: &T,
        dice: &T,
    ) -> anyhow::Result<()>
    where
        T: PartialEq + fmt::Display + ?Sized,
    {
        if fs != dice {
            writeln!(
                self.stdout,
                "MISMATCH: {} at {}: fs = {}, dice = {}",
                kind, path, fs, dice,
            )?;
            self.bad += 1;
        } else if self.verbose {
            writeln!(self.stdout, "Match: {} at {}: {}", kind, path, fs)?;
        }

        Ok(())
    }
}

#[derive(PartialEq)]
struct DirList(Vec<FileNameBuf>);

impl fmt::Display for DirList {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let mut comma = commas();
        for path in &self.0 {
            comma(f)?;
            write!(f, "{}", path)?;
        }
        Ok(())
    }
}

#[async_trait]
impl ServerCommandTemplate for FileStatusServerCommand {
    type StartEvent = buck2_data::FileStatusCommandStart;
    type EndEvent = buck2_data::FileStatusCommandEnd;
    type Response = buck2_cli_proto::GenericResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut stdout: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        let file_ops = ctx.file_ops();
        let cell_resolver = ctx.get_cell_resolver().await?;
        let project_root = server_ctx.project_root();
        let digest_config = ctx.global_data().get_digest_config();

        let io = FsIoProvider::new(project_root.dupe(), digest_config.cas_digest_config());
        let stdout = stdout.as_writer();

        let mut result = FileStatusResult {
            checked: 0,
            bad: 0,
            verbose: self.req.verbose,
            stdout,
        };

        let mut stderr = server_ctx.stderr()?;

        for path in &self.req.paths {
            let path = project_root.relativize_any(AbsPath::new(Path::new(path))?)?;
            writeln!(&mut stderr, "Check file status: {}", path)?;
            check_file_status(&file_ops, &cell_resolver, &io, &path, &mut result).await?;
        }
        if result.bad != 0 {
            Err(anyhow::anyhow!("Failed with {} mismatches", result.bad))
        } else {
            writeln!(
                &mut stderr,
                "No mismatches detected ({} entries checked)",
                result.checked
            )?;
            Ok(buck2_cli_proto::GenericResponse {})
        }
    }

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

#[async_recursion]
async fn check_file_status(
    file_ops: &dyn FileOps,
    cell_resolver: &CellResolver,
    io: &FsIoProvider,
    path: &ProjectRelativePath,
    result: &mut FileStatusResult,
) -> anyhow::Result<()> {
    result.checking();

    let cell_path = cell_resolver.get_cell_path(path)?;
    if file_ops.is_ignored(cell_path.as_ref()).await? {
        return Ok(());
    }

    let fs_metadata = io.read_path_metadata_if_exists(path.to_owned()).await?;

    let dice_metadata = file_ops
        .read_path_metadata_if_exists(cell_path.as_ref())
        .await?;

    let (fs_metadata, dice_metadata) = match (&fs_metadata, &dice_metadata) {
        (Some(fs), Some(dice)) => (fs, dice),
        (fs, dice) => {
            result.report("existence", path, &fs.is_some(), &dice.is_some())?;
            return Ok(());
        }
    };

    result.report(
        "entry kind",
        path,
        fs_metadata.variant_name(),
        dice_metadata.variant_name(),
    )?;

    match (fs_metadata, dice_metadata) {
        (
            RawPathMetadata::Symlink {
                at: fs_at,
                to: fs_to,
            },
            RawPathMetadata::Symlink {
                at: dice_at,
                to: dice_to,
            },
        ) => {
            let dice_at = cell_resolver.resolve_path(dice_at.as_ref().as_ref())?;
            result.report("symlink component location", path, fs_at, &dice_at)?;

            result.report(
                "symlink destination kind",
                path,
                fs_to.variant_name(),
                dice_to.variant_name(),
            )?;

            match (fs_to, dice_to) {
                (RawSymlink::Relative(fs_rel), RawSymlink::Relative(dice_rel)) => {
                    let dice_rel = cell_resolver.resolve_path(dice_rel.as_ref().as_ref())?;
                    result.report("relative symlink destination", path, fs_rel, &dice_rel)?;
                }
                (RawSymlink::External(fs_ext), RawSymlink::External(dice_ext)) => {
                    result.report("external symlink destination", path, fs_ext, dice_ext)?;
                }
                _ => {
                    // Reported earlier
                }
            };
        }
        (RawPathMetadata::File(fs_file), RawPathMetadata::File(dice_file)) => {
            result.report("file metadata", path, fs_file, dice_file)?;
        }
        (RawPathMetadata::Directory, RawPathMetadata::Directory) => {
            let fs_read_dir = io.read_dir(path.to_owned()).await?;
            let dice_read_dir = file_ops.read_dir(cell_path.as_ref()).await?;

            // No point checking file types here, we'll do that when we inspect them.
            let mut fs_names = fs_read_dir
                .iter()
                .map(|f| anyhow::Ok(FileName::new(&f.file_name)?.to_owned()))
                .collect::<Result<Vec<_>, _>>()?;

            let mut dice_names = dice_read_dir
                .included
                .iter()
                .map(|e| e.file_name.clone())
                .collect::<Vec<_>>();

            fs_names.sort();
            dice_names.sort();

            let fs_names = DirList(fs_names);
            let dice_names = DirList(dice_names);

            result.report("directory contents", path, &fs_names, &dice_names)?;

            for file_name in &fs_names.0 {
                check_file_status(file_ops, cell_resolver, io, &path.join(file_name), result)
                    .await?;
            }
        }
        (_, _) => {
            // Reported earlier
        }
    }

    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 std::sync::Mutex;
use std::time::Duration;

use buck2_events::dispatch::EventDispatcher;
use dupe::Dupe;
use tokio::task::JoinHandle;

use crate::ctx::BaseServerCommandContext;
use crate::snapshot;

// Spawns a thread to occasionally output snapshots of resource utilization.
pub struct HeartbeatGuard {
    handle: JoinHandle<()>,
    collector: snapshot::SnapshotCollector,
    events: Arc<Mutex<Option<EventDispatcher>>>,
}

impl HeartbeatGuard {
    pub fn new(ctx: &BaseServerCommandContext) -> Self {
        let events = Arc::new(Mutex::new(Some(ctx.events.dupe())));
        let collector = snapshot::SnapshotCollector::new(
            ctx.re_client_manager.dupe(),
            ctx.blocking_executor.dupe(),
            ctx.daemon_start_time,
            ctx.dice_manager.unsafe_dice().dupe(),
            ctx.materializer.dupe(),
            Some(ctx.events.sink().dupe()),
            ctx.http_client.dupe(),
        );

        // NOTE: This doesn't use the ambient dispatcher wrappers because we want to control the
        // exact lifetime of the dispatcher.
        let handle = tokio::spawn({
            let events = events.dupe();
            let collector = collector.clone();
            async move {
                let mut interval = tokio::time::interval(Duration::from_secs(1));
                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
                loop {
                    let snapshot = collector.create_snapshot();
                    match events.lock().expect("Poisoned lock").as_ref() {
                        Some(events) => events.instant_event(Box::new(snapshot)),
                        None => break,
                    }
                    interval.tick().await;
                }
            }
        });

        Self {
            handle,
            collector,
            events,
        }
    }
}

impl Drop for HeartbeatGuard {
    fn drop(&mut self) {
        let mut maybe_events = self.events.lock().expect("Poisoned lock");
        // Synchronously remove access for sending new heartbeats.
        if let Some(events) = maybe_events.take() {
            // Send one last snapshot.
            let collector = self.collector.dupe();
            events.instant_event(Box::new(collector.create_snapshot()));
        }
        // Cancel the task as well.
        self.handle.abort();
    }
}


/*
 * 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::client_context::HostArchOverride;
use buck2_cli_proto::client_context::HostPlatformOverride;
use buck2_core::soft_error;
use buck2_interpreter::extra::xcode::XcodeVersionInfo;
use buck2_interpreter::extra::InterpreterHostArchitecture;
use buck2_interpreter::extra::InterpreterHostPlatform;

pub fn get_host_info(
    host_platform: HostPlatformOverride,
    host_arch: HostArchOverride,
    host_xcode_override: &Option<String>,
) -> anyhow::Result<(
    InterpreterHostPlatform,
    InterpreterHostArchitecture,
    Option<XcodeVersionInfo>,
)> {
    let interpreter_platform = match host_platform {
        HostPlatformOverride::Linux => InterpreterHostPlatform::Linux,
        HostPlatformOverride::MacOs => InterpreterHostPlatform::MacOS,
        HostPlatformOverride::Windows => InterpreterHostPlatform::Windows,
        HostPlatformOverride::DefaultPlatform => match std::env::consts::OS {
            "linux" => InterpreterHostPlatform::Linux,
            "macos" => InterpreterHostPlatform::MacOS,
            "windows" => InterpreterHostPlatform::Windows,
            "freebsd" => InterpreterHostPlatform::FreeBsd,
            _ => InterpreterHostPlatform::Unknown,
        },
    };
    // This compiles in the target architecture, which should be sufficient, as
    // we currently only run e.g. x86_64 on x86_64.
    let interpreter_architecture = match host_arch {
        HostArchOverride::AArch64 => InterpreterHostArchitecture::AArch64,
        HostArchOverride::X8664 => InterpreterHostArchitecture::X86_64,
        HostArchOverride::DefaultArch => match std::env::consts::ARCH {
            "aarch64" => InterpreterHostArchitecture::AArch64,
            "x86_64" => InterpreterHostArchitecture::X86_64,
            "arm" => InterpreterHostArchitecture::Arm,
            "x86" => InterpreterHostArchitecture::X86,
            "mips" => InterpreterHostArchitecture::Mips,
            "mips64" => InterpreterHostArchitecture::Mips64,
            "powerpc" => InterpreterHostArchitecture::PowerPc,
            "powerpc64" => InterpreterHostArchitecture::PowerPc64,
            _ => InterpreterHostArchitecture::Unknown,
        },
    };

    let interpreter_xcode_version = match host_xcode_override {
        Some(s) => Some(
            XcodeVersionInfo::from_version_and_build(s.as_str())
                .context("Constructing `XcodeVersionInfo` from string.")?,
        ),
        None if interpreter_platform == InterpreterHostPlatform::MacOS => {
            match XcodeVersionInfo::new()
                .context("Constructing `XcodeVersionInfo` using host platform MacOS.")
            {
                Ok(v) => v,
                Err(e) => {
                    soft_error!("invalid_xcode_version", e)?;
                    None
                }
            }
        }
        _ => None,
    };

    Ok((
        interpreter_platform,
        interpreter_architecture,
        interpreter_xcode_version,
    ))
}

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

    #[test]
    fn test_host_info_create_xcode_override() {
        let linux = InterpreterHostPlatform::Linux;
        let linux_override = HostPlatformOverride::Linux;
        let mac = InterpreterHostPlatform::MacOS;
        let mac_override = HostPlatformOverride::MacOs;
        let aarch64 = InterpreterHostArchitecture::AArch64;
        let aarch64_override = HostArchOverride::AArch64;
        let xcode_str = "14.3.0-14C18";
        let xcode_string = Some(xcode_str.to_owned());

        let xcode_override_on_linux =
            get_host_info(linux_override, aarch64_override, &xcode_string).unwrap();
        let want1 = (
            linux,
            aarch64,
            XcodeVersionInfo::from_version_and_build(xcode_str).ok(),
        );
        assert_eq!(xcode_override_on_linux, want1);

        let xcode_override_on_mac =
            get_host_info(mac_override, aarch64_override, &xcode_string).unwrap();
        let want2 = (
            mac,
            aarch64,
            XcodeVersionInfo::from_version_and_build(xcode_str).ok(),
        );
        assert_eq!(xcode_override_on_mac, want2);
    }
}


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

// TODO(raulgarcia4): Consider moving out this file to buck2_common,
// similarly to what was done with buck2_common::process_stats.
pub struct AllocatorStats {
    pub bytes_active: Option<u64>,
    pub bytes_allocated: Option<u64>,
}

pub fn get_allocator_stats() -> anyhow::Result<AllocatorStats> {
    fn set(stats: &serde_json::Value, val: &str, to: &mut Option<u64>) -> anyhow::Result<()> {
        if let serde_json::Value::Number(stat) = &stats["jemalloc"]["stats"][val] {
            *to = stat.as_u64();
            Ok(())
        } else {
            Err(anyhow::anyhow!("Allocator stat '{}' not found.", val))
        }
    }

    // This options configuration flag string is passed to `malloc_stats_print()`.
    // Explanation: (<https://jemalloc.net/jemalloc.3.html>)
    // The default configuration prints minimal output, formatted as JSON.
    let alloc_stats_options = "Jmdablxg";
    let alloc_stats: serde_json::Value =
        serde_json::from_str(&allocator_stats(alloc_stats_options)?)?;

    let mut bytes_active = None;
    let mut bytes_allocated = None;
    set(&alloc_stats, "active", &mut bytes_active)?;
    set(&alloc_stats, "allocated", &mut bytes_allocated)?;

    Ok(AllocatorStats {
        bytes_active,
        bytes_allocated,
    })
}

#[cfg(test)]
mod tests {
    use buck2_common::memory::has_jemalloc_stats;

    use crate::jemalloc_stats::get_allocator_stats;

    #[test]
    fn test_get_allocator_stats() -> anyhow::Result<()> {
        if has_jemalloc_stats() {
            if let Ok(alloc_stats) = get_allocator_stats() {
                assert!(alloc_stats.bytes_active.is_some());
                assert!(alloc_stats.bytes_allocated.is_some());
                return Ok(());
            }
            return Err(anyhow::anyhow!("Allocator stats not found"));
        }
        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.
 */

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

pub mod active_commands;
pub mod builtin_docs;
mod clean_stale;
mod configs;
mod ctx;
pub mod daemon;
mod dice_tracker;
mod file_status;
mod heartbeat_guard;
mod host_info;
mod jemalloc_stats;
pub mod lsp;
mod materialize;
mod net_io;
pub mod profile;
mod snapshot;
mod subscription;
mod trace_io;


/*
 * 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::future::Future;
use std::io;
use std::io::ErrorKind;
use std::path::Path;
use std::thread;

use buck2_cli_proto::*;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::dice::file_ops::HasFileOps;
use buck2_common::file_ops::FileOps;
use buck2_common::package_listing::dice::HasPackageListingResolver;
use buck2_common::result::SharedResult;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::build_file_cell::BuildFileCell;
use buck2_core::cells::cell_path::CellPathRef;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::pattern::pattern_type::ProvidersPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::name::TargetName;
use buck2_events::dispatch::span_async;
use buck2_events::dispatch::with_dispatcher;
use buck2_events::dispatch::with_dispatcher_async;
use buck2_interpreter::path::BxlFilePath;
use buck2_interpreter::path::OwnedStarlarkModulePath;
use buck2_interpreter::path::StarlarkPath;
use buck2_interpreter_for_build::interpreter::dice_calculation_delegate::HasCalculationDelegate;
use buck2_interpreter_for_build::interpreter::global_interpreter_state::HasGlobalInterpreterState;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::ctx::ServerCommandDiceContext;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::streaming_request_handler::StreamingRequestHandler;
use dice::DiceEquality;
use dice::DiceTransaction;
use dupe::Dupe;
use futures::channel::mpsc::UnboundedSender;
use futures::FutureExt;
use futures::SinkExt;
use futures::StreamExt;
use gazebo::prelude::StrExt;
use itertools::Itertools;
use lsp_server::Connection;
use lsp_server::Message;
use lsp_types::Range;
use lsp_types::Url;
use starlark::docs::Doc;
use starlark::docs::DocItem;
use starlark::docs::Identifier;
use starlark::docs::Location;
use starlark::errors::EvalMessage;
use starlark::lsp::server::server_with_connection;
use starlark::lsp::server::LspContext;
use starlark::lsp::server::LspEvalResult;
use starlark::lsp::server::LspUrl;
use starlark::lsp::server::StringLiteralResult;
use starlark::syntax::AstModule;
use tokio::runtime::Handle;
use tokio::sync::Mutex;
use tokio::sync::MutexGuard;

use crate::builtin_docs::docs::get_builtin_docs;
use crate::builtin_docs::docs::get_prelude_docs;

static DOCS_DIRECTORY_KEY: &str = "directory";
static DOCS_BUILTIN_KEY: &str = "builtin";

/// Errors when [`LspContext::resolve_load()`] cannot resolve a given path.
#[derive(thiserror::Error, Debug)]
enum ResolveLoadError {
    /// The scheme provided was not correct or supported.
    #[error("Url `{}` was expected to be of type `{}`", .1, .0)]
    WrongScheme(String, LspUrl),
}

/// Manage refreshing DocsCache.
struct DocsCacheManager {
    docs_cache: Mutex<DocsCache>,
    fs: ProjectRoot,
    /// Used for checking if the DocsCache need refreshing. We need to refresh the DocsCache
    /// if the previous dice version does not match the current one.
    valid_at: DiceEquality,
}

impl DocsCacheManager {
    async fn new(fs: ProjectRoot, dice_ctx: DiceTransaction) -> anyhow::Result<Self> {
        Ok(Self {
            docs_cache: Mutex::new(Self::new_docs_cache(&fs, &dice_ctx).await?),
            fs,
            valid_at: dice_ctx.equality_token(),
        })
    }

    async fn get_cache(
        &self,
        current_dice_ctx: DiceTransaction,
    ) -> anyhow::Result<MutexGuard<DocsCache>> {
        let mut docs_cache = self.docs_cache.lock().await;

        let fs = &self.fs;
        match self.is_reusable(&current_dice_ctx) {
            true => (),
            false => {
                let new_docs_cache = Self::new_docs_cache(fs, &current_dice_ctx).await?;
                *docs_cache = new_docs_cache
            }
        };

        Ok(docs_cache)
    }

    fn is_reusable(&self, dice_ctx: &DiceTransaction) -> bool {
        dice_ctx.equivalent(&self.valid_at)
    }

    async fn new_docs_cache<'v>(
        fs: &ProjectRoot,
        dice_ctx: &DiceTransaction,
    ) -> anyhow::Result<DocsCache> {
        fn flatten(docs: Vec<Doc>) -> Vec<Doc> {
            docs.into_iter()
                .flat_map(|x| match x.item {
                    DocItem::Module(module) => module
                        .members
                        .into_iter()
                        .map(|(name, v)| Doc {
                            id: Identifier {
                                name,
                                location: x.id.location.clone(),
                            },
                            item: v.to_doc_item(),
                            custom_attrs: x.custom_attrs.clone(),
                        })
                        .collect(),
                    DocItem::Object(_) => vec![],
                    _ => {
                        vec![x]
                    }
                })
                .collect::<Vec<_>>()
        }

        let cell_resolver = dice_ctx.get_cell_resolver().await?;
        let global_interpreter_state = dice_ctx.get_global_interpreter_state().await?;
        let mut builtin_docs = flatten(get_builtin_docs(global_interpreter_state)?);
        let builtin_names = builtin_docs.iter().map(|d| d.id.name.as_str()).collect();
        let prelude_docs = flatten(get_prelude_docs(dice_ctx, &builtin_names).await?);
        builtin_docs.extend(prelude_docs);
        DocsCache::new(&builtin_docs, dice_ctx, fs, &cell_resolver).await
    }
}

/// Store rendered starlark representations of Doc objects for builtin symbols,
/// their names, and their real or virtual paths
struct DocsCache {
    /// Mapping of global names to URLs. These can either be files (for global symbols in the
    /// prelude), or `starlark:` urls for rust native types and functions.
    global_urls: HashMap<String, LspUrl>,
    /// Mapping of starlark: urls to a synthesized starlark representation.
    native_starlark_files: HashMap<LspUrl, String>,
}

#[derive(thiserror::Error, Debug)]
enum DocsCacheError {
    #[error("Duplicate global symbol `{}` detected. Existing URL was `{}`, new URL was `{}`", .name, .existing, .new)]
    DuplicateGlobalSymbol {
        name: String,
        existing: LspUrl,
        new: LspUrl,
    },
}

impl DocsCache {
    async fn get_prelude_uri(
        location: &Location,
        dice_ctx: &DiceTransaction,
        fs: &ProjectRoot,
        cell_resolver: &CellResolver,
    ) -> anyhow::Result<LspUrl> {
        // We strip the "@" off of the path normally to avoid having it in the docs,
        // in paths, etc, but it needs to be put back in to create an import path correctly.
        let path = if location.path.contains("//")
            && !location.path.starts_with('@')
            && !location.path.starts_with("//")
        {
            format!("@{}", location.path)
        } else {
            location.path.to_owned()
        };

        let bfc = BuildFileCell::new(cell_resolver.root_cell());
        let calculator = dice_ctx
            .get_interpreter_calculator(cell_resolver.root_cell(), bfc)
            .await?;

        let root_import_path = ImportPath::new_same_cell(
            cell_resolver.get_cell_path(ProjectRelativePath::new("non_existent.bzl")?)?,
        )?;
        let starlark_file = StarlarkPath::LoadFile(&root_import_path);
        let loaded_import_path = calculator.resolve_load(starlark_file, &path).await?;

        let relative_path =
            cell_resolver.resolve_path(loaded_import_path.borrow().path().as_ref())?;
        let abs_path = fs.resolve(&relative_path);
        Ok(Url::from_file_path(abs_path).unwrap().try_into()?)
    }

    async fn new(
        builtin_symbols: &[Doc],
        dice_ctx: &DiceTransaction,
        fs: &ProjectRoot,
        cell_resolver: &CellResolver,
    ) -> anyhow::Result<Self> {
        Self::new_with_lookup(builtin_symbols, |location| async {
            Self::get_prelude_uri(location, dice_ctx, fs, cell_resolver).await
        })
        .await
    }

    async fn new_with_lookup<
        'a,
        F: Fn(&'a Location) -> Fut + 'a,
        Fut: Future<Output = anyhow::Result<LspUrl>>,
    >(
        builtin_symbols: &'a [Doc],
        location_lookup: F,
    ) -> anyhow::Result<Self> {
        let mut global_urls = HashMap::with_capacity(builtin_symbols.len());
        let mut native_starlark_files = HashMap::new();
        for doc in builtin_symbols {
            let url = match &doc.id.location {
                Some(l) => location_lookup(l).await?,
                None => {
                    let filename = Path::new(&doc.id.name);
                    let filename = filename.with_extension(match filename.extension() {
                        None => "bzl".to_owned(),
                        Some(e) => format!("{}.bzl", e.to_str().expect("path is UTF-8")),
                    });
                    let path = Path::new("/native")
                        .join(output_subdir_for_doc(doc)?.as_path())
                        .join(filename);

                    let url =
                        LspUrl::try_from(Url::parse(&format!("starlark:{}", path.display()))?)?;

                    native_starlark_files.insert(url.clone(), doc.render_as_code());
                    url
                }
            };
            if let Some(existing) = global_urls.insert(doc.id.name.clone(), url.clone()) {
                return Err(DocsCacheError::DuplicateGlobalSymbol {
                    name: doc.id.name.clone(),
                    existing,
                    new: url,
                }
                .into());
            }
        }
        Ok(Self {
            global_urls,
            native_starlark_files,
        })
    }

    fn native_starlark_file(&self, url: &LspUrl) -> Option<&String> {
        // We only load starlark: urls into this hash map, so other types
        // will just get a "None" back.
        self.native_starlark_files.get(url)
    }

    fn url_for_symbol(&self, symbol: &str) -> Option<&LspUrl> {
        self.global_urls.get(symbol)
    }
}

#[derive(Debug, thiserror::Error)]
enum DocPathError {
    #[error("Directory traversal was found in documentation path `{}` provided for `{}`", .path, .name)]
    InvalidDirectory {
        name: String,
        path: String,
        source: anyhow::Error,
    },
    #[error("Invalid custom attributes were found on `{}`: {}", .name, format_custom_attr_error(.keys_and_values))]
    InvalidCustomAttributes {
        name: String,
        keys_and_values: Vec<(String, String)>,
    },
    #[error("Conflicting custom attributes were found on `{}`: {}", .name, format_custom_attr_error(.keys_and_values))]
    ConflictingCustomAttributes {
        name: String,
        keys_and_values: Vec<(String, String)>,
    },
}

fn format_custom_attr_error(keys_and_values: &[(String, String)]) -> String {
    let mut ret = "{".to_owned();
    ret.push_str(
        &keys_and_values
            .iter()
            .map(|(k, v)| format!("`{}` => `{}`", k, v))
            .join(", "),
    );
    ret.push('}');
    ret
}

/// Get the output subdirectory for a [`Doc`] based on the `directory` custom attr, if present.
pub fn output_subdir_for_doc(doc: &Doc) -> anyhow::Result<ForwardRelativePathBuf> {
    let unknown_keys: Vec<_> = doc
        .custom_attrs
        .iter()
        .filter(|(k, _)| *k != DOCS_DIRECTORY_KEY && *k != DOCS_BUILTIN_KEY)
        .map(|(k, v)| (k.clone(), v.clone()))
        .collect();
    if !unknown_keys.is_empty() {
        return Err(DocPathError::InvalidCustomAttributes {
            name: doc.id.name.to_owned(),
            keys_and_values: unknown_keys,
        }
        .into());
    }

    match (
        doc.custom_attrs.get(DOCS_DIRECTORY_KEY),
        doc.custom_attrs.get(DOCS_BUILTIN_KEY),
    ) {
        (Some(path), None) | (None, Some(path)) => {
            match ForwardRelativePathBuf::new(path.to_owned()) {
                Ok(fp) => Ok(fp),
                Err(e) => Err(DocPathError::InvalidDirectory {
                    name: doc.id.name.to_owned(),
                    path: path.to_owned(),
                    source: e,
                }
                .into()),
            }
        }
        (Some(dir), Some(builtin)) => Err(DocPathError::ConflictingCustomAttributes {
            name: doc.id.name.to_owned(),
            keys_and_values: vec![
                (DOCS_DIRECTORY_KEY.to_owned(), dir.clone()),
                (DOCS_BUILTIN_KEY.to_owned(), builtin.clone()),
            ],
        }
        .into()),
        (None, None) => Ok(ForwardRelativePathBuf::new(String::new())?),
    }
}

struct BuckLspContext<'a> {
    server_ctx: &'a dyn ServerCommandContextTrait,
    fs: ProjectRoot,
    docs_cache_manager: DocsCacheManager,
    runtime: Handle,
}

#[derive(Debug, thiserror::Error)]
enum BuckLspContextError {
    /// The scheme provided was not correct or supported.
    #[error("Url `{}` was expected to be of type `{}`", .1, .0)]
    WrongScheme(String, LspUrl),
}

impl<'a> BuckLspContext<'a> {
    async fn new(
        server_ctx: &'a dyn ServerCommandContextTrait,
    ) -> anyhow::Result<BuckLspContext<'a>> {
        let (fs, docs_cache_manager) = server_ctx
            .with_dice_ctx(async move |server_ctx, dice_ctx| {
                let fs = server_ctx.project_root().clone();

                let docs_cache_manager = DocsCacheManager::new(fs.clone(), dice_ctx).await?;

                Ok((fs, docs_cache_manager))
            })
            .await?;

        Ok(Self {
            server_ctx,
            fs,
            docs_cache_manager,
            runtime: Handle::current(),
        })
    }

    // We are purposefully only holding the DICE context briefly since the LSP logic lives in a
    // long-running thread, and we don't want it to block continuously while holding DICE since
    // it will block other concurrent commands.
    async fn with_dice_ctx<F, Fut, R>(&self, exec: F) -> anyhow::Result<R>
    where
        F: FnOnce(DiceTransaction) -> Fut + Send,
        Fut: Future<Output = anyhow::Result<R>> + Send,
        R: Send,
    {
        self.server_ctx
            .with_dice_ctx(|_, dice_ctx| async move { Ok(exec(dice_ctx).await) })
            .await?
    }

    async fn import_path(&self, path: &Path) -> anyhow::Result<OwnedStarlarkModulePath> {
        let abs_path = AbsPath::new(path)?;
        let relative_path = self.fs.relativize_any(abs_path)?;
        let cell_resolver = self
            .with_dice_ctx(|dice_ctx| async move { dice_ctx.get_cell_resolver().await })
            .await?;
        let cell_path = cell_resolver.get_cell_path(&relative_path)?;

        match path.extension() {
            Some(e) if e == "bxl" => Ok(OwnedStarlarkModulePath::BxlFile(BxlFilePath::new(
                cell_path,
            )?)),
            _ => {
                // Instantiating a BuildFileCell off of the path of the file being originally evaluated.
                // Unlike the build commands and such, there is no meaningful "build file cell" versus
                // the cell of the file being currently evaluated.
                let bfc = BuildFileCell::new(cell_path.cell());

                Ok(OwnedStarlarkModulePath::LoadFile(
                    ImportPath::new_hack_for_lsp(cell_path, bfc)?,
                ))
            }
        }
    }

    async fn starlark_import_path(&self, path: &Path) -> anyhow::Result<OwnedStarlarkModulePath> {
        // The "absolute" path from LSP urls doesn't work here, they have to be relative
        // to get a ProjectRelativePath. We already guaranteed that things start with a '/'
        // (rooted from `starlark:`, see LspUrl), so just drop it.
        let cell_resolver = self
            .with_dice_ctx(|dice_ctx| async move { dice_ctx.get_cell_resolver().await })
            .await?;

        let cell_path = cell_resolver.get_cell_path(&ProjectRelativePath::new(
            path.to_string_lossy().trim_start_match('/'),
        )?)?;

        match path.extension() {
            Some(e) if e == "bxl" => Ok(OwnedStarlarkModulePath::BxlFile(BxlFilePath::new(
                cell_path,
            )?)),
            _ => Ok(OwnedStarlarkModulePath::LoadFile(
                ImportPath::new_same_cell(cell_path)?,
            )),
        }
    }

    async fn parse_file_with_contents(
        &self,
        uri: &LspUrl,
        content: String,
    ) -> SharedResult<LspEvalResult> {
        let import_path: OwnedStarlarkModulePath = match uri {
            LspUrl::File(path) => self.import_path(path).await,
            LspUrl::Starlark(path) => self.starlark_import_path(path).await,
            LspUrl::Other(_) => Err(BuckLspContextError::WrongScheme(
                "file:// or starlark:".to_owned(),
                uri.clone(),
            )
            .into()),
        }?;

        Ok(self
            .with_dice_ctx(|dice_ctx| async move {
                let calculator = dice_ctx
                    .get_interpreter_calculator(
                        import_path.borrow().cell(),
                        import_path.borrow().build_file_cell(),
                    )
                    .await?;

                let module_path = import_path.borrow();
                let path = module_path.starlark_path();
                let ast = calculator.prepare_eval_with_content(path, content)?;
                Ok(LspEvalResult {
                    diagnostics: vec![],
                    ast: Some(ast),
                })
            })
            .await?)
    }

    async fn parse_file_from_string(
        &self,
        current_package: CellPathRef<'_>,
        literal: &str,
    ) -> anyhow::Result<Option<StringLiteralResult>> {
        match ForwardRelativePath::new(literal) {
            Ok(package_relative) => {
                let cell_resolver = self
                    .with_dice_ctx(|dice_ctx| async move { dice_ctx.get_cell_resolver().await })
                    .await?;
                let relative_path =
                    cell_resolver.resolve_path(current_package.join(package_relative).as_ref())?;

                let path = self.fs.resolve(&relative_path);
                let url = Url::from_file_path(path).unwrap().try_into()?;
                let string_literal = StringLiteralResult {
                    url,
                    location_finder: None,
                };
                Ok(Some(string_literal))
            }
            Err(_) => Ok(None),
        }
    }

    async fn parse_target_from_string(
        &self,
        current_package: CellPathRef<'_>,
        literal: &str,
    ) -> anyhow::Result<Option<StringLiteralResult>> {
        let cell_resolver = self
            .with_dice_ctx(|dice_ctx| async move { dice_ctx.get_cell_resolver().await })
            .await?;
        match ParsedPattern::<ProvidersPatternExtra>::parsed_opt_absolute(
            literal,
            Some(current_package),
            current_package.cell(),
            &cell_resolver,
        ) {
            Ok(ParsedPattern::Target(package, target, _)) => {
                let res = self
                    .with_dice_ctx(async move |dice_ctx| {
                        Ok(dice_ctx
                            .resolve_package_listing(package.dupe())
                            .await
                            .and_then(|listing| {
                                let relative_path = cell_resolver
                                    .resolve_package(package.dupe())?
                                    .join(listing.buildfile());
                                let path = self.fs.resolve(&relative_path);
                                match Url::from_file_path(path).unwrap().try_into() {
                                    Ok(url) => {
                                        let string_literal = StringLiteralResult {
                                            url,
                                            location_finder: Some(Box::new(|ast, _url| {
                                                Ok(Self::find_target(ast, target))
                                            })),
                                        };
                                        Ok(Some(string_literal))
                                    }
                                    Err(e) => Err(e.into()),
                                }
                            }))
                    })
                    .await?;
                Ok(res?)
            }
            _ => Ok(None),
        }
    }

    fn find_target(ast: &AstModule, target: TargetName) -> Option<Range> {
        ast.find_function_call_with_name(target.as_str())
            .map(Range::from)
    }
}

impl<'a> LspContext for BuckLspContext<'a> {
    fn parse_file_with_contents(&self, uri: &LspUrl, content: String) -> LspEvalResult {
        let dispatcher = self.server_ctx.events().dupe();
        self.runtime
            .block_on(with_dispatcher_async(dispatcher, async {
                match uri {
                    LspUrl::File(_) | LspUrl::Starlark(_) => {
                        match self.parse_file_with_contents(uri, content).await {
                            Ok(result) => result,
                            Err(e) => {
                                let message = EvalMessage::from_anyhow(uri.path(), e.inner());
                                LspEvalResult {
                                    diagnostics: vec![message.into()],
                                    ast: None,
                                }
                            }
                        }
                    }
                    _ => LspEvalResult::default(),
                }
            }))
    }

    fn resolve_load(&self, path: &str, current_file: &LspUrl) -> anyhow::Result<LspUrl> {
        let dispatcher = self.server_ctx.events().dupe();
        self.runtime
            .block_on(with_dispatcher_async(dispatcher, async {
                match current_file {
                    LspUrl::File(current_file) => {
                        let current_import_path = self.import_path(current_file).await?;
                        let borrowed_current_import_path = current_import_path.borrow();
                        let url = self
                            .with_dice_ctx(async move |dice_ctx| {
                                let calculator = dice_ctx
                                    .get_interpreter_calculator(
                                        borrowed_current_import_path.cell(),
                                        borrowed_current_import_path.build_file_cell(),
                                    )
                                    .await?;

                                let starlark_file = borrowed_current_import_path.starlark_path();
                                let loaded_import_path =
                                    calculator.resolve_load(starlark_file, path).await?;
                                let relative_path = dice_ctx
                                    .get_cell_resolver()
                                    .await?
                                    .resolve_path(loaded_import_path.borrow().path().as_ref())?;
                                let abs_path = self.fs.resolve(&relative_path);
                                Ok(Url::from_file_path(abs_path).unwrap().try_into()?)
                            })
                            .await?;

                        Ok(url)
                    }
                    _ => Err(ResolveLoadError::WrongScheme(
                        "file://".to_owned(),
                        current_file.clone(),
                    )
                    .into()),
                }
            }))
    }

    fn resolve_string_literal(
        &self,
        literal: &str,
        current_file: &LspUrl,
    ) -> anyhow::Result<Option<StringLiteralResult>> {
        let dispatcher = self.server_ctx.events().dupe();
        self.runtime
            .block_on(with_dispatcher_async(dispatcher, async {
                let import_path = match current_file {
                    LspUrl::File(current_file) => {
                        Ok(self.import_path(current_file.parent().unwrap()).await?)
                    }
                    _ => Err(ResolveLoadError::WrongScheme(
                        "file://".to_owned(),
                        current_file.clone(),
                    )),
                }?;
                let current_package = import_path.path();

                // Right now we swallow the errors up as they can happen for a lot of reasons that are
                // perfectly recoverable (e.g. an invalid cell is specified, we can't list an invalid
                // package path, an absolute path is requested, etc).
                //
                // anyhow::Error sort of obscures the root causes, so we just have
                // to assume if it failed, it's fine, and we can maybe try the next thing.
                if let Ok(Some(string_literal)) = self
                    .parse_target_from_string(current_package, literal)
                    .await
                {
                    Ok(Some(string_literal))
                } else if let Ok(Some(string_literal)) =
                    self.parse_file_from_string(current_package, literal).await
                {
                    Ok(Some(string_literal))
                } else {
                    Ok(None)
                }
            }))
    }

    fn get_load_contents(&self, uri: &LspUrl) -> anyhow::Result<Option<String>> {
        let dispatcher = self.server_ctx.events().dupe();
        self.runtime
            .block_on(with_dispatcher_async(dispatcher, async {
                match uri {
                    LspUrl::File(path) => {
                        let path = self.import_path(path).await?;

                        self.with_dice_ctx(async move |dice_ctx| {
                            match <dyn FileOps>::read_file(
                                &dice_ctx.file_ops(),
                                path.borrow().path().as_ref(),
                            )
                            .await
                            {
                                Ok(s) => Ok(Some(s)),
                                Err(e) => match e.downcast_ref::<io::Error>() {
                                    Some(inner_e) if inner_e.kind() == ErrorKind::NotFound => {
                                        Ok(None)
                                    }
                                    _ => Err(e),
                                },
                            }
                        })
                        .await
                    }
                    LspUrl::Starlark(_) => {
                        let docs_cache = self
                            .with_dice_ctx(async move |dice_ctx| {
                                self.docs_cache_manager.get_cache(dice_ctx).await
                            })
                            .await?;

                        Ok(docs_cache.native_starlark_file(uri).cloned())
                    }
                    _ => Err(
                        BuckLspContextError::WrongScheme("file://".to_owned(), uri.clone()).into(),
                    ),
                }
            }))
    }

    fn get_url_for_global_symbol(
        &self,
        _current_file: &LspUrl,
        symbol: &str,
    ) -> anyhow::Result<Option<LspUrl>> {
        let dispatcher = self.server_ctx.events().dupe();
        self.runtime
            .block_on(with_dispatcher_async(dispatcher, async {
                let docs_cache = self
                    .with_dice_ctx(|dice_ctx| async {
                        self.docs_cache_manager.get_cache(dice_ctx).await
                    })
                    .await?;
                Ok(docs_cache.url_for_symbol(symbol).cloned())
            }))
    }
}

pub(crate) async fn run_lsp_server_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::LspMessage>,
    req: StreamingRequestHandler<buck2_cli_proto::LspRequest>,
) -> anyhow::Result<buck2_cli_proto::LspResponse> {
    let metadata = ctx.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::LspCommandStart {}.into()),
    };
    span_async(start_event, async move {
        let result = run_lsp_server(ctx, partial_result_dispatcher, req).await;
        let end_event = command_end(metadata, &result, buck2_data::LspCommandEnd {});
        (result, end_event)
    })
    .await
}

/// Run an LSP server for a given client.
async fn run_lsp_server(
    ctx: &dyn ServerCommandContextTrait,
    mut partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::LspMessage>,
    mut req: StreamingRequestHandler<LspRequest>,
) -> anyhow::Result<LspResponse> {
    // This gets a bit messy because the various frameworks don't quite work the same way.
    // - tonic has async streams (which we pull from with .message().await)
    // - The LSP server uses crossbeam channels, which are synchronous, so we need to run
    //   recv from that channel it its own thread.
    // - We then need an async aware channel that takes from the LSP recv thread, and sends
    //   to the client. It needs to be async aware so that we can use select! and process
    //   client and server communication at the same time
    //
    // We also start up a thread to run the actual LSP, and make sure that we close those
    // threads down when we either fail to send/recv on the client channel or the server one.
    //
    // tl;dr; This creates two threads, and a few plumbing channels to get the client sending
    //        receiving to/from the LSP server.

    let (send_to_server, server_receiver) = crossbeam_channel::unbounded();
    let (send_to_client, client_receiver) = crossbeam_channel::unbounded();
    let (events_from_server, mut events_to_client) = futures::channel::mpsc::unbounded();

    let connection = Connection {
        sender: send_to_client,
        receiver: server_receiver,
    };

    let dispatcher = ctx.events().dupe();
    let buck_lsp_ctx = BuckLspContext::new(ctx).await?;

    tokio::task::block_in_place(|| {
        thread::scope(|scope| {
            let recv_thread = scope.spawn(move || {
                let runtime = tokio::runtime::Builder::new_current_thread()
                    .build()
                    .unwrap();
                runtime.block_on(recv_from_lsp(client_receiver, events_from_server))
            });

            let server_thread = scope.spawn(with_dispatcher(dispatcher, || {
                move || server_with_connection(connection, buck_lsp_ctx)
            }));

            let res = {
                let runtime = tokio::runtime::Builder::new_current_thread()
                    .build()
                    .unwrap();

                runtime.block_on(async move {
                    loop {
                        let message_handler_res = tokio::select! {
                            m = req.message().fuse() => {
                                handle_incoming_lsp_message(&send_to_server, m)
                            },
                            m = events_to_client.next() => {
                                Ok(handle_outgoing_lsp_message(&mut partial_result_dispatcher, m))
                            },
                        };
                        match message_handler_res {
                            Ok(Some(res)) => break Ok(res),
                            Ok(None) => {}
                            Err(e) => break Err(e),
                        }
                    }
                })
            };

            let _ignored = recv_thread.join();
            let _ignored = server_thread.join();
            res
        })
    })
}

/// Receive messages from the LSP's channel, and pass them to the client after encapsulating them.
///
/// This returns `Ok(())` when the other end of the connection disconnects.
async fn recv_from_lsp(
    to_client: crossbeam_channel::Receiver<Message>,
    mut event_sender: UnboundedSender<buck2_cli_proto::LspMessage>,
) -> anyhow::Result<()> {
    loop {
        let msg = to_client.recv()?;

        let lsp_json = serde_json::to_string(&msg).unwrap();
        let res = buck2_cli_proto::LspMessage { lsp_json };
        match event_sender.send(res).await {
            Ok(_) => {}
            Err(e) if e.is_disconnected() => break Ok(()),
            Err(e) => break Err(e.into()),
        }
    }
}

/// Unpacks and validates messages from the client, sending them to the LSP.
///
/// Returns:
///     - `Ok(None)` if the message could be passed to the LSP
///     - `Ok(Some(LspResponse))` if the message was actually an error or the LSP could not accept
///                               the provided message. The stream should be terminated after this.
///     - `Err` if a payload could not be decoded.
fn handle_incoming_lsp_message(
    from_client: &crossbeam_channel::Sender<Message>,
    message: anyhow::Result<LspRequest>,
) -> anyhow::Result<Option<LspResponse>> {
    if let Ok(m) = message {
        let message = serde_json::from_str(&m.lsp_json)?;
        if from_client.send(message).is_ok() {
            return Ok(None);
        }
    }
    Ok(Some(LspResponse {}))
}

/// Handles an outgoing message from the LSP
///
/// Returns:
///     - `None` if the message could be passed to the client / event dispatcher.
///     - `Some(LspResponse)` if there was no message. This happens when the server is done
///                           sending messages, and the stream should be disconnected.
fn handle_outgoing_lsp_message(
    partial_result_dispatcher: &mut PartialResultDispatcher<buck2_cli_proto::LspMessage>,
    event: Option<buck2_cli_proto::LspMessage>,
) -> Option<LspResponse> {
    match event {
        Some(event) => {
            partial_result_dispatcher.emit(event);
            None
        }
        None => Some(LspResponse {}),
    }
}

#[cfg(test)]
mod test {
    use lsp_types::Url;
    use maplit::hashmap;
    use starlark::docs::Doc;
    use starlark::docs::DocFunction;
    use starlark::docs::DocItem;
    use starlark::docs::Identifier;
    use starlark::docs::Location;
    use starlark::lsp::server::LspUrl;

    use crate::lsp::DocsCache;
    use crate::lsp::DOCS_DIRECTORY_KEY;

    #[test]
    fn cache_builds() -> anyhow::Result<()> {
        let docs = vec![
            Doc {
                id: Identifier {
                    name: "native_function1".to_owned(),
                    location: None,
                },
                item: DocItem::Function(DocFunction::default()),
                custom_attrs: Default::default(),
            },
            Doc {
                id: Identifier {
                    name: "native_function2".to_owned(),
                    location: None,
                },
                item: DocItem::Function(DocFunction::default()),
                custom_attrs: hashmap! { DOCS_DIRECTORY_KEY.to_owned() => "subdir".to_owned() },
            },
            Doc {
                id: Identifier {
                    name: "prelude_function".to_owned(),
                    location: Some(Location {
                        path: "//dir/prelude.bzl".to_owned(),
                        position: None,
                    }),
                },
                item: DocItem::Function(DocFunction::default()),
                custom_attrs: Default::default(),
            },
        ];

        let runtime = tokio::runtime::Builder::new_current_thread()
            .build()
            .unwrap();
        async fn lookup_function(location: &Location) -> anyhow::Result<LspUrl> {
            match location.path.as_str() {
                "//dir/prelude.bzl" => Ok(LspUrl::try_from(Url::parse(
                    "file:///usr/local/dir/prelude.bzl",
                )?)?),
                p => Err(anyhow::anyhow!("Unknown path {}", p)),
            }
        }

        let cache =
            runtime.block_on(async { DocsCache::new_with_lookup(&docs, lookup_function).await })?;

        assert_eq!(
            &LspUrl::try_from(Url::parse("starlark:/native/native_function1.bzl")?)?,
            cache.url_for_symbol("native_function1").unwrap()
        );
        assert_eq!(
            &LspUrl::try_from(Url::parse("starlark:/native/subdir/native_function2.bzl")?)?,
            cache.url_for_symbol("native_function2").unwrap()
        );
        assert_eq!(
            &LspUrl::try_from(Url::parse("file:/usr/local/dir/prelude.bzl")?)?,
            cache.url_for_symbol("prelude_function").unwrap()
        );

        Ok(())
    }
}