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

use std::collections::HashSet;
use std::path;

use anyhow::Context;
use buck2_artifact::artifact::artifact_type::BaseArtifactKind;
use buck2_build_api::build::BuildProviderType;
use buck2_build_api::build::ProviderArtifacts;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_query::__derive_refs::indexmap::IndexMap;
use itertools::Itertools;
use tracing::info;

pub(crate) fn create_unhashed_outputs(
    provider_artifacts: Vec<ProviderArtifacts>,
    artifact_fs: &ArtifactFs,
    fs: &ProjectRoot,
) -> anyhow::Result<u64> {
    let buck_out_root = fs.resolve(artifact_fs.buck_out_path_resolver().root());

    let start = std::time::Instant::now();
    // The following IndexMap will contain a key of the unhashed/symlink path and values of all the hashed locations that map to the unhashed location.
    let mut unhashed_to_hashed: IndexMap<AbsNormPathBuf, HashSet<AbsNormPathBuf>> = IndexMap::new();
    for provider_artifact in provider_artifacts {
        if !matches!(provider_artifact.provider_type, BuildProviderType::Default) {
            continue;
        }

        match provider_artifact.values.iter().exactly_one() {
            Ok((artifact, _)) => match artifact.as_parts() {
                (BaseArtifactKind::Build(build), _projected_path) => {
                    if let Some(unhashed_path) =
                        artifact_fs.retrieve_unhashed_location(build.get_path())
                    {
                        let path = artifact_fs.resolve_build(build.get_path());
                        let abs_unhashed_path = fs.resolve(&unhashed_path);
                        let entry = unhashed_to_hashed
                            .entry(abs_unhashed_path)
                            .or_insert_with(HashSet::new);
                        entry.insert(fs.resolve(&path));
                    }
                }
                _ => {}
            },
            Err(_) => {}
        };
    }
    // The IndexMap is used now to determine if and what conflicts exist where multiple hashed artifact locations
    // all want a symlink to the same unhashed artifact location and deal with them accordingly.
    let mut num_unhashed_links_made = 0;
    for (unhashed, hashed_set) in unhashed_to_hashed {
        if hashed_set.len() == 1 {
            create_unhashed_link(&unhashed, hashed_set.iter().next().unwrap(), &buck_out_root)?;
            num_unhashed_links_made += 1;
        } else {
            info!(
                "The following outputs have a conflicting unhashed path at {}: {:?}",
                unhashed, hashed_set
            );
        }
    }
    let duration = start.elapsed();
    info!(
        "Creating {} output compatibility symlinks in {:3}s",
        num_unhashed_links_made,
        duration.as_secs_f64()
    );
    Ok(num_unhashed_links_made)
}

fn create_unhashed_link(
    unhashed_path: &AbsNormPathBuf,
    original_path: &AbsNormPathBuf,
    buck_out_root: &AbsNormPathBuf,
) -> anyhow::Result<()> {
    // Remove the final path separator if it exists so that the path looks like a file and not a directory or else symlink() fails.
    tracing::debug!("Creating link: `{}` -> `{}`", unhashed_path, original_path);

    let mut abs_unhashed_path = unhashed_path.to_owned();
    if let Some(path) = unhashed_path
        .to_str()
        .unwrap()
        .strip_suffix(path::is_separator)
    {
        abs_unhashed_path = AbsNormPathBuf::from(path.to_owned())?;
    }

    // We are going to need to clear the path between buck-out and the symlink we want to create.
    // To do this, we need to traverse forward out of buck_out_root and towards our symlink, and
    // delete any files or symlinks we find along the way. As soon as we find one, we can stop.

    if let Some(parent) = abs_unhashed_path.parent() {
        for prefix in iter_reverse_ancestors(parent, buck_out_root.as_ref()) {
            let meta = match fs_util::symlink_metadata_if_exists(prefix)? {
                Some(meta) => meta,
                None => continue,
            };

            if meta.is_file() || meta.is_symlink() {
                fs_util::remove_file(prefix)
                    .with_context(|| "was not able to remove file while cleaning up prefixes")?;
            }
        }

        fs_util::create_dir_all(parent)
            .with_context(|| "while creating unhashed directory for symlink")?;
    }

    match fs_util::symlink_metadata(&abs_unhashed_path) {
        Ok(metadata) => {
            if metadata.is_dir() {
                fs_util::remove_dir_all(&abs_unhashed_path)
                    .with_context(|| "was not able to remove absolute unhashed path (directory)")?
            } else {
                fs_util::remove_file(&abs_unhashed_path)
                    .with_context(|| "was not able to remove absolute unhashed path (file)")?
            }
        }
        Err(_) => {}
    }
    fs_util::symlink(original_path, abs_unhashed_path)
        .with_context(|| "was not able to symlink original path to absolute unhashed path")?;
    Ok(())
}

/// Iterate over the path components between stop_at and path.
fn iter_reverse_ancestors<'a>(
    path: &'a AbsPath,
    stop_at: &'_ AbsPath,
) -> impl Iterator<Item = &'a AbsPath> {
    let ancestors = path
        .ancestors()
        .take_while(|a| *a != stop_at)
        .collect::<Vec<_>>();

    ancestors.into_iter().rev()
}

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

    #[test]
    fn test_iter_reverse_ancestors() {
        let prefix = if cfg!(windows) { "C:" } else { "" };
        let root = AbsNormPathBuf::try_from(format!("{prefix}/repo/buck-out/v2")).unwrap();
        let path =
            AbsNormPathBuf::try_from(format!("{prefix}/repo/buck-out/v2/foo/bar/some")).unwrap();

        let mut iter = iter_reverse_ancestors(&path, &root);
        assert_eq!(
            iter.next().unwrap().to_str().unwrap(),
            &format!("{prefix}/repo/buck-out/v2/foo"),
        );
        assert_eq!(
            iter.next().unwrap().to_str().unwrap(),
            &format!("{prefix}/repo/buck-out/v2/foo/bar"),
        );
        assert_eq!(
            iter.next().unwrap().to_str().unwrap(),
            &format!("{prefix}/repo/buck-out/v2/foo/bar/some"),
        );
        assert_eq!(iter.next(), None);
    }
}


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

use std::io::Write;

use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::query::oneshot::QUERY_FRONTEND;
use buck2_common::dice::cells::HasCellResolver;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceTransaction;

use crate::commands::query::printer::QueryResultPrinter;
use crate::commands::query::printer::ShouldPrintProviders;

pub async fn aquery_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: buck2_cli_proto::AqueryRequest,
) -> anyhow::Result<buck2_cli_proto::AqueryResponse> {
    run_server_command(AqueryServerCommand { req }, ctx, partial_result_dispatcher).await
}

struct AqueryServerCommand {
    req: buck2_cli_proto::AqueryRequest,
}

#[async_trait]
impl ServerCommandTemplate for AqueryServerCommand {
    type StartEvent = buck2_data::AqueryCommandStart;
    type EndEvent = buck2_data::AqueryCommandEnd;
    type Response = buck2_cli_proto::AqueryResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

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

    fn is_success(&self, response: &Self::Response) -> bool {
        response.error_messages.is_empty()
    }
}

async fn aquery(
    server_ctx: &dyn ServerCommandContextTrait,
    mut stdout: impl Write,
    ctx: DiceTransaction,
    request: &buck2_cli_proto::AqueryRequest,
) -> anyhow::Result<buck2_cli_proto::AqueryResponse> {
    let cell_resolver = ctx.get_cell_resolver().await?;

    let output_configuration = QueryResultPrinter::from_request_options(
        &cell_resolver,
        &request.output_attributes,
        request.unstable_output_format,
    )?;

    let buck2_cli_proto::AqueryRequest {
        query,
        query_args,
        context,
        ..
    } = request;

    let client_ctx = context
        .as_ref()
        .context("No client context (internal error)")?;
    let global_target_platform =
        target_platform_from_client_context(client_ctx, server_ctx, &ctx).await?;

    let query_result = QUERY_FRONTEND
        .get()?
        .eval_aquery(
            &ctx,
            server_ctx.working_dir(),
            query,
            query_args,
            global_target_platform,
        )
        .await?;

    let result = match query_result {
        QueryEvaluationResult::Single(targets) => {
            output_configuration
                .print_single_output(&mut stdout, targets, false, ShouldPrintProviders::No)
                .await
        }
        QueryEvaluationResult::Multiple(results) => {
            output_configuration
                .print_multi_output(&mut stdout, results, false, ShouldPrintProviders::No)
                .await
        }
    };
    let error_messages = match result {
        Ok(_) => vec![],
        Err(e) => vec![format!("{:#}", e)],
    };
    Ok(buck2_cli_proto::AqueryResponse { error_messages })
}


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

use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::analysis::calculation::RuleAnalysisCalculation;
use buck2_build_api::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;
use buck2_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_build_api::query::oneshot::QUERY_FRONTEND;
use buck2_cli_proto::CqueryRequest;
use buck2_cli_proto::CqueryResponse;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_core::provider::label::ConfiguredProvidersLabel;
use buck2_core::provider::label::ProvidersName;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use buck2_util::truncate::truncate;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;

use crate::commands::query::printer::ProviderLookUp;
use crate::commands::query::printer::QueryResultPrinter;
use crate::commands::query::printer::ShouldPrintProviders;

pub async fn cquery_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: CqueryRequest,
) -> anyhow::Result<CqueryResponse> {
    run_server_command(CqueryServerCommand { req }, ctx, partial_result_dispatcher).await
}

struct CqueryServerCommand {
    req: CqueryRequest,
}

#[async_trait]
impl ServerCommandTemplate for CqueryServerCommand {
    type StartEvent = buck2_data::CQueryCommandStart;
    type EndEvent = buck2_data::CQueryCommandEnd;
    type Response = CqueryResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

    fn start_event(&self) -> buck2_data::CQueryCommandStart {
        buck2_data::CQueryCommandStart {
            query: truncate(&self.req.query, 50000),
            query_args: truncate(&self.req.query_args.join(","), 1000),
            target_universe: truncate(&self.req.target_universe.join(","), 1000),
        }
    }

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

    fn is_success(&self, response: &Self::Response) -> bool {
        response.error_messages.is_empty()
    }
}

async fn cquery(
    server_ctx: &dyn ServerCommandContextTrait,
    mut stdout: impl Write,
    ctx: DiceTransaction,
    request: &CqueryRequest,
) -> anyhow::Result<CqueryResponse> {
    let cell_resolver = ctx.get_cell_resolver().await?;
    let output_configuration = QueryResultPrinter::from_request_options(
        &cell_resolver,
        &request.output_attributes,
        request.unstable_output_format,
    )?;

    let CqueryRequest {
        query,
        query_args,
        target_universe,
        context,
        show_providers,
        correct_owner,
        ..
    } = request;
    // The request will always have a universe value, an empty one indicates the user didn't provide a universe.
    let target_universe: Option<&[String]> = if target_universe.is_empty() {
        None
    } else {
        Some(target_universe)
    };
    let client_ctx = context
        .as_ref()
        .context("No client context (internal error)")?;

    let target_call_stacks = client_ctx.target_call_stacks;

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

    let owner_behavior = match correct_owner {
        true => CqueryOwnerBehavior::Correct,
        false => CqueryOwnerBehavior::Deprecated,
    };

    let query_result = QUERY_FRONTEND
        .get()?
        .eval_cquery(
            &ctx,
            server_ctx.working_dir(),
            owner_behavior,
            query,
            query_args,
            global_target_platform,
            target_universe,
        )
        .await?;

    let should_print_providers = if *show_providers {
        ShouldPrintProviders::Yes(&*ctx as &dyn ProviderLookUp<ConfiguredTargetNode>)
    } else {
        ShouldPrintProviders::No
    };

    let result = match query_result {
        QueryEvaluationResult::Single(targets) => {
            output_configuration
                .print_single_output(
                    &mut stdout,
                    targets,
                    target_call_stacks,
                    should_print_providers,
                )
                .await
        }
        QueryEvaluationResult::Multiple(results) => {
            output_configuration
                .print_multi_output(
                    &mut stdout,
                    results,
                    target_call_stacks,
                    should_print_providers,
                )
                .await
        }
    };

    let error_messages = match result {
        Ok(_) => vec![],
        Err(e) => vec![format!("{:#}", e)],
    };

    Ok(CqueryResponse { error_messages })
}

#[async_trait]
impl ProviderLookUp<ConfiguredTargetNode> for DiceComputations {
    async fn lookup(
        &self,
        t: &ConfiguredTargetNode,
    ) -> anyhow::Result<MaybeCompatible<FrozenProviderCollectionValue>> {
        self.get_providers(&ConfiguredProvidersLabel::new(
            t.label().dupe(),
            ProvidersName::Default,
        ))
        .await
    }
}


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

use thiserror::Error;

pub mod aquery;
pub mod cquery;
pub mod printer;
pub mod uquery;

#[derive(Debug, Error)]
enum QueryCommandError {
    #[error(
        "query result was a set of files and one or more --output-attribute was requested, but files have not attributes"
    )]
    FileSetHasNoAttributes,
}


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

#![allow(clippy::drop_non_drop)] // FIXME?

use std::fmt::Display;
use std::fmt::Formatter;
use std::io::Write;

use async_trait::async_trait;
use buck2_build_api::actions::query::ActionQueryNode;
use buck2_build_api::actions::query::PRINT_ACTION_NODE;
use buck2_build_api::interpreter::rule_defs::provider::collection::FrozenProviderCollectionValue;
use buck2_cli_proto::QueryOutputFormat;
use buck2_core::cells::CellResolver;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_query::query::environment::QueryTarget;
use buck2_query::query::environment::QueryTargets;
use buck2_query::query::syntax::simple::eval::file_set::FileSet;
use buck2_query::query::syntax::simple::eval::multi_query::MultiQueryResult;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationValue;
use buck2_util::indent::indent;
use dupe::Clone_;
use dupe::Copy_;
use dupe::Dupe_;
use gazebo::variants::UnpackVariants;
use indent_write::fmt::IndentWriter;
use regex::RegexSet;
use serde::ser::SerializeMap;
use serde::ser::SerializeSeq;
use serde::Serialize;
use serde::Serializer;

use crate::commands::query::QueryCommandError;
use crate::dot::targets::DotTargetGraph;
use crate::dot::Dot;
use crate::dot::DotCompact;

#[derive(Copy_, Dupe_, Clone_, UnpackVariants)]
pub enum ShouldPrintProviders<'a, T> {
    No,
    Yes(&'a dyn ProviderLookUp<T>),
}

#[async_trait]
pub trait ProviderLookUp<T: QueryTarget>: Send + Sync {
    async fn lookup(&self, t: &T)
    -> anyhow::Result<MaybeCompatible<FrozenProviderCollectionValue>>;
}

#[derive(Debug)]
pub struct QueryResultPrinter<'a> {
    resolver: &'a CellResolver,
    attributes: Option<RegexSet>,
    output_format: QueryOutputFormat,
}

struct TargetSetJsonPrinter<'a, T: QueryTarget> {
    value: Vec<PrintableQueryTarget<'a, T>>,
    is_complex: bool,
}

impl<'a, T: QueryTarget> TargetSetJsonPrinter<'a, T> {
    async fn new(
        target_call_stacks: bool,
        print_providers: ShouldPrintProviders<'a, T>,
        attributes: &'a Option<RegexSet>,
        targets: &'a TargetSet<T>,
    ) -> anyhow::Result<TargetSetJsonPrinter<'a, T>> {
        Ok(TargetSetJsonPrinter {
            value: printable_targets(targets, print_providers, attributes, target_call_stacks)
                .await?,
            is_complex: attributes.is_some()
                || target_call_stacks
                || print_providers.unpack_yes().is_some(),
        })
    }
}

struct PrintableQueryTarget<'a, T: QueryTarget> {
    value: &'a T,
    attributes: &'a Option<RegexSet>,
    providers: Option<FrozenProviderCollectionValue>,
    target_call_stacks: bool,
}

impl<'a, T: QueryTarget> PrintableQueryTarget<'a, T> {
    fn label(&self) -> String {
        self.value.node_ref().to_string()
    }
}

impl<'a, T: QueryTarget> Display for PrintableQueryTarget<'a, T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}", self.value.node_ref())?;

        if self.target_call_stacks || self.providers.is_some() {
            writeln!(f)?;
        }

        if self.target_call_stacks {
            match self.value.call_stack() {
                Some(call_stack) => {
                    write!(f, "{}", indent("  ", &call_stack))?;
                }
                None => {
                    // This is `aquery`.
                }
            }
        }

        if let Some(providers) = &self.providers {
            use std::fmt::Write;
            write!(
                IndentWriter::new("  ", f),
                "{:#}",
                providers.provider_collection()
            )?;
        }

        Ok(())
    }
}

impl<'a, T: QueryTarget> Serialize for PrintableQueryTarget<'a, T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut map = serializer.serialize_map(None)?;

        QueryTargets::for_all_attrs(self.value, |attr_name, attr_value| {
            if let Some(attr_regex) = self.attributes {
                if attr_regex.is_match(attr_name) {
                    struct AttrValueSerialize<'a, 'b, T: QueryTarget> {
                        target: &'a T,
                        attr: &'a T::Attr<'b>,
                    }

                    impl<'a, 'b, T: QueryTarget> Serialize for AttrValueSerialize<'a, 'b, T> {
                        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
                        where
                            S: Serializer,
                        {
                            self.target.attr_serialize(self.attr, serializer)
                        }
                    }

                    map.serialize_entry(
                        attr_name,
                        &AttrValueSerialize {
                            target: self.value,
                            attr: attr_value,
                        },
                    )?;
                }
            }
            Ok(())
        })?;

        if self.target_call_stacks {
            map.serialize_entry("buck.target_call_stack", &self.value.call_stack())?;
        }

        if let Some(providers) = &self.providers {
            map.serialize_entry("buck.providers", providers)?;
        }

        map.end()
    }
}

impl<'a, T: QueryTarget> Serialize for TargetSetJsonPrinter<'a, T> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        if self.is_complex {
            serializer.collect_map(self.value.iter().map(|target| (target.label(), target)))
        } else {
            serializer.collect_seq(self.value.iter().map(|target| target.label()))
        }
    }
}

struct FileSetJsonPrinter<'a> {
    value: &'a FileSet,
    resolver: &'a CellResolver,
}

impl<'a> Serialize for FileSetJsonPrinter<'a> {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut seq = serializer.serialize_seq(None)?;
        for file in self.value.iter() {
            seq.serialize_element(
                &self
                    .resolver
                    .resolve_path(file.as_ref())
                    .map_err(serde::ser::Error::custom)?
                    .to_string(),
            )?;
        }
        SerializeSeq::end(seq)
    }
}

impl<'a> QueryResultPrinter<'a> {
    /// Utility for creating from the options in their protobuf form.
    pub fn from_request_options(
        resolver: &'a CellResolver,
        attributes: &[String],
        output_format: i32,
    ) -> anyhow::Result<Self> {
        Self::from_options(
            resolver,
            attributes,
            QueryOutputFormat::from_i32(output_format)
                .expect("cli should send a valid output_format enum"),
        )
    }

    pub fn from_options(
        resolver: &'a CellResolver,
        attributes: &[String],
        output_format: QueryOutputFormat,
    ) -> anyhow::Result<Self> {
        let output_format = match (output_format, attributes.is_empty()) {
            // following buck1's behavior, if any attributes are requested we use json output instead of list output
            (QueryOutputFormat::Default, false) => QueryOutputFormat::Json,
            (v, _) => v,
        };

        let attributes = if attributes.is_empty() {
            None
        } else {
            Some(RegexSet::new(attributes)?)
        };

        Ok(Self {
            resolver,
            attributes,
            output_format,
        })
    }

    pub async fn print_multi_output<'b, T: QueryTarget, W: std::io::Write>(
        &self,
        mut output: W,
        multi_result: MultiQueryResult<T>,
        target_call_stacks: bool,
        print_providers: ShouldPrintProviders<'b, T>,
    ) -> anyhow::Result<()> {
        match (self.output_format, &self.attributes) {
            // A multi-query only has interesting output with --json output. For non-json output it gets merged together.
            // TODO(cjhopman): buck1 does this really odd thing that a multi-query that requests any attributes
            // gets the entire result merged together rather than printed as a multi-query. We match that behavior, but
            // it really doesn't make sense and we should migrate off of that.
            (QueryOutputFormat::Json, None) => {
                let multi_result = multi_result.0;
                let mut captured_error = Ok(());

                let mut ser = serde_json::Serializer::pretty(&mut output);
                let mut seq = ser.serialize_map(Some(multi_result.len()))?;
                for (arg, result) in multi_result {
                    match result {
                        Ok(v) => match v {
                            QueryEvaluationValue::TargetSet(targets) => seq.serialize_entry(
                                &arg,
                                &TargetSetJsonPrinter::new(
                                    target_call_stacks,
                                    print_providers,
                                    &self.attributes,
                                    &targets,
                                )
                                .await?,
                            )?,
                            QueryEvaluationValue::FileSet(files) => seq.serialize_entry(
                                &arg,
                                &FileSetJsonPrinter {
                                    resolver: self.resolver,
                                    value: &files,
                                },
                            )?,
                        },
                        Err(e) => {
                            seq.serialize_entry(
                                &arg,
                                &serde_json::json!({ "$error": format!("{:#}", e) }),
                            )?;
                            captured_error = Err(e);
                        }
                    }
                }
                SerializeMap::end(seq)?;
                std::mem::drop(ser);
                // need to add a newline to flush the output.
                writeln!(&mut output)?;
                captured_error
            }
            _ => {
                self.print_single_output(
                    output,
                    multi_result.merged()?,
                    target_call_stacks,
                    print_providers,
                )
                .await
            }
        }
    }

    pub async fn print_single_output<'b, T: QueryTarget, W: std::io::Write>(
        &self,
        mut output: W,
        result: QueryEvaluationValue<T>,
        call_stack: bool,
        print_providers: ShouldPrintProviders<'b, T>,
    ) -> anyhow::Result<()> {
        match result {
            QueryEvaluationValue::TargetSet(targets) => match self.output_format {
                QueryOutputFormat::Default => {
                    for target in
                        printable_targets(&targets, print_providers, &self.attributes, call_stack)
                            .await?
                    {
                        writeln!(&mut output, "{}", target)?;
                    }
                }
                QueryOutputFormat::Json => {
                    let mut ser = serde_json::Serializer::pretty(&mut output);
                    TargetSetJsonPrinter::new(
                        call_stack,
                        print_providers,
                        &self.attributes,
                        &targets,
                    )
                    .await?
                    .serialize(&mut ser)?;
                    std::mem::drop(ser);
                    // need to add a newline to flush the output.
                    writeln!(&mut output)?
                }
                QueryOutputFormat::Dot => {
                    Dot::render(
                        &DotTargetGraph {
                            targets,
                            attributes: self.attributes.clone(),
                        },
                        &mut output,
                    )?;
                }
                QueryOutputFormat::DotCompact => {
                    DotCompact::render(
                        &DotTargetGraph {
                            targets,
                            attributes: self.attributes.clone(),
                        },
                        &mut output,
                    )?;
                }
            },
            QueryEvaluationValue::FileSet(files) => {
                if self.attributes.is_some() {
                    return Err(QueryCommandError::FileSetHasNoAttributes.into());
                }
                match self.output_format {
                    QueryOutputFormat::Default => {
                        for file in files.iter() {
                            writeln!(
                                &mut output,
                                "{}",
                                self.resolver.resolve_path(file.as_ref())?
                            )?;
                        }
                    }
                    QueryOutputFormat::Json => {
                        let mut ser = serde_json::Serializer::pretty(&mut output);
                        FileSetJsonPrinter {
                            resolver: self.resolver,
                            value: &files,
                        }
                        .serialize(&mut ser)?;
                        std::mem::drop(ser);
                        // need to add a newline to flush the output.
                        writeln!(&mut output)?;
                    }
                    QueryOutputFormat::Dot => {
                        unimplemented!("dot output for files not implemented yet")
                    }
                    QueryOutputFormat::DotCompact => {
                        unimplemented!("dot_compact output for files not implemented yet")
                    }
                }
            }
        }

        Ok(())
    }
}

async fn printable_targets<'a, T: QueryTarget>(
    targets: &'a TargetSet<T>,
    print_providers: ShouldPrintProviders<'a, T>,
    attributes: &'a Option<RegexSet>,
    target_call_stacks: bool,
) -> anyhow::Result<Vec<PrintableQueryTarget<'a, T>>> {
    futures::future::join_all(targets.iter().map(|t| {
        let print_providers = &print_providers;
        async move {
            Ok(PrintableQueryTarget {
                value: t,
                attributes,
                target_call_stacks,
                providers: match print_providers {
                    ShouldPrintProviders::No => None,
                    ShouldPrintProviders::Yes(lookup) => {
                        Some(lookup.lookup(t).await?.require_compatible()?)
                    }
                },
            })
        }
    }))
    .await
    .into_iter()
    .collect::<anyhow::Result<_>>()
}

async fn print_action_node(
    stdout: &mut (dyn Write + Send),
    action: ActionQueryNode,
    json: bool,
    output_attributes: &[String],
    cell_resolver: &CellResolver,
) -> anyhow::Result<()> {
    // Dot/DotCompact output format don't make sense here.
    let unstable_output_format = if json {
        QueryOutputFormat::Json
    } else {
        QueryOutputFormat::Default
    } as i32;

    let query_result_printer = QueryResultPrinter::from_request_options(
        cell_resolver,
        output_attributes,
        unstable_output_format,
    )?;

    let mut result = TargetSet::new();
    result.insert(action);

    query_result_printer
        .print_single_output(
            stdout,
            QueryEvaluationValue::TargetSet(result),
            false,
            ShouldPrintProviders::No,
        )
        .await
}

pub(crate) fn init_print_action_node() {
    PRINT_ACTION_NODE.init(|stdout, action, json, output_attributes, cell_resolver| {
        Box::pin(print_action_node(
            stdout,
            action,
            json,
            output_attributes,
            cell_resolver,
        ))
    });
}


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

use anyhow::Context;
use async_trait::async_trait;
use buck2_build_api::query::oneshot::QUERY_FRONTEND;
use buck2_cli_proto::UqueryRequest;
use buck2_cli_proto::UqueryResponse;
use buck2_common::dice::cells::HasCellResolver;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceTransaction;

use crate::commands::query::printer::QueryResultPrinter;
use crate::commands::query::printer::ShouldPrintProviders;

pub async fn uquery_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: UqueryRequest,
) -> anyhow::Result<UqueryResponse> {
    run_server_command(UqueryServerCommand { req }, ctx, partial_result_dispatcher).await
}

struct UqueryServerCommand {
    req: UqueryRequest,
}

#[async_trait]
impl ServerCommandTemplate for UqueryServerCommand {
    type StartEvent = buck2_data::QueryCommandStart;
    type EndEvent = buck2_data::QueryCommandEnd;
    type Response = UqueryResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

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

    fn is_success(&self, response: &Self::Response) -> bool {
        response.error_messages.is_empty()
    }
}

async fn uquery(
    server_ctx: &dyn ServerCommandContextTrait,
    mut stdout: impl Write,
    ctx: DiceTransaction,
    request: &UqueryRequest,
) -> anyhow::Result<UqueryResponse> {
    let cell_resolver = ctx.get_cell_resolver().await?;
    let output_configuration = QueryResultPrinter::from_request_options(
        &cell_resolver,
        &request.output_attributes,
        request.unstable_output_format,
    )?;

    let UqueryRequest {
        query,
        query_args,
        context,
        ..
    } = request;

    let client_ctx = context
        .as_ref()
        .context("No client context (internal error)")?;

    let target_call_stacks = client_ctx.target_call_stacks;

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

    let query_result = QUERY_FRONTEND
        .get()?
        .eval_uquery(
            &ctx,
            server_ctx.working_dir(),
            query,
            query_args,
            global_target_platform,
        )
        .await?;

    let result = match query_result {
        QueryEvaluationResult::Single(targets) => {
            output_configuration
                .print_single_output(
                    &mut stdout,
                    targets,
                    target_call_stacks,
                    ShouldPrintProviders::No,
                )
                .await
        }
        QueryEvaluationResult::Multiple(results) => {
            output_configuration
                .print_multi_output(
                    &mut stdout,
                    results,
                    target_call_stacks,
                    ShouldPrintProviders::No,
                )
                .await
        }
    };

    let error_messages = match result {
        Ok(_) => vec![],
        Err(e) => vec![format!("{:#}", e)],
    };

    Ok(UqueryResponse { error_messages })
}


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

//! Server-side implementation of `buck2 targets` command
//! without `--streaming` or `--resolve-alias` arguments.

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

use buck2_cli_proto::targets_request;
use buck2_cli_proto::targets_request::TargetHashFileMode;
use buck2_cli_proto::targets_request::TargetHashGraphType;
use buck2_cli_proto::TargetsResponse;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_path::AbsPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::TargetLabel;
use buck2_node::load_patterns::load_patterns;
use buck2_node::load_patterns::MissingTargetBehavior;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::lookup::ConfiguredTargetNodeLookup;
use buck2_node::nodes::lookup::TargetNodeLookup;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use dice::DiceTransaction;
use dupe::Dupe;
use dupe::OptionDupedExt;

use crate::commands::targets::fmt::Stats;
use crate::commands::targets::fmt::TargetFormatter;
use crate::commands::targets::fmt::TargetInfo;
use crate::commands::targets::mk_error;
use crate::target_hash::TargetHashes;
use crate::target_hash::TargetHashesFileMode;

pub(crate) struct TargetHashOptions {
    file_mode: TargetHashesFileMode,
    fast_hash: bool,
    graph_type: TargetHashGraphType,
    recursive: bool,
}

impl TargetHashOptions {
    pub(crate) fn new(
        request: &targets_request::Other,
        cell_resolver: &CellResolver,
        fs: &ProjectRoot,
    ) -> anyhow::Result<Self> {
        let file_mode = TargetHashFileMode::from_i32(request.target_hash_file_mode)
            .expect("buck cli should send valid target hash file mode");
        let file_mode = match file_mode {
            TargetHashFileMode::PathsOnly => {
                let modified_paths = request
                    .target_hash_modified_paths
                    .iter()
                    .map(|path| {
                        let path = AbsPath::new(Path::new(&path))?;
                        cell_resolver.get_cell_path_from_abs_path(path, fs)
                    })
                    .collect::<anyhow::Result<_>>()?;
                TargetHashesFileMode::PathsOnly(modified_paths)
            }
            TargetHashFileMode::PathsAndContents => TargetHashesFileMode::PathsAndContents,
            TargetHashFileMode::NoFiles => TargetHashesFileMode::None,
        };

        Ok(Self {
            file_mode,
            fast_hash: request.target_hash_use_fast_hash,
            graph_type: TargetHashGraphType::from_i32(request.target_hash_graph_type)
                .expect("buck cli should send valid target hash graph type"),
            recursive: request.target_hash_recursive,
        })
    }
}

pub(crate) async fn targets_batch(
    server_ctx: &dyn ServerCommandContextTrait,
    dice: DiceTransaction,
    formatter: &dyn TargetFormatter,
    parsed_patterns: Vec<ParsedPattern<TargetPatternExtra>>,
    target_platform: Option<TargetLabel>,
    hash_options: TargetHashOptions,
    keep_going: bool,
) -> anyhow::Result<TargetsResponse> {
    let results = load_patterns(&dice, parsed_patterns, MissingTargetBehavior::Fail).await?;

    let target_hashes = match hash_options.graph_type {
        TargetHashGraphType::Configured => Some(
            TargetHashes::compute::<ConfiguredTargetNode, _>(
                dice.dupe(),
                ConfiguredTargetNodeLookup(&dice),
                results.iter_loaded_targets_by_package().collect(),
                target_platform,
                hash_options.file_mode,
                hash_options.fast_hash,
                hash_options.recursive,
            )
            .await?,
        ),
        TargetHashGraphType::Unconfigured => Some(
            TargetHashes::compute::<TargetNode, _>(
                dice.dupe(),
                TargetNodeLookup(&dice),
                results.iter_loaded_targets_by_package().collect(),
                target_platform,
                hash_options.file_mode,
                hash_options.fast_hash,
                hash_options.recursive,
            )
            .await?,
        ),
        _ => None,
    };

    let mut buffer = String::new();
    formatter.begin(&mut buffer);
    let mut stats = Stats::default();
    let mut needs_separator = false;
    for (package, result) in results.iter() {
        match result {
            Ok(res) => {
                stats.success += 1;
                for (_, node) in res.iter() {
                    stats.targets += 1;
                    let target_hash = target_hashes
                        .as_ref()
                        .and_then(|hashes| hashes.get(node.label()))
                        .duped()
                        .transpose()?;
                    if needs_separator {
                        formatter.separator(&mut buffer);
                    }
                    needs_separator = true;
                    formatter.target(TargetInfo { node, target_hash }, &mut buffer)
                }
            }
            Err(e) => {
                stats.errors += 1;
                let mut stderr = String::new();

                if needs_separator {
                    formatter.separator(&mut buffer);
                }
                needs_separator = true;
                formatter.package_error(package.dupe(), e.inner(), &mut buffer, &mut stderr);

                server_ctx.stderr()?.write_all(stderr.as_bytes())?;

                if !keep_going {
                    break;
                }
            }
        }
    }
    formatter.end(&stats, &mut buffer);
    if !keep_going && stats.errors != 0 {
        Err(mk_error(stats.errors))
    } else {
        Ok(TargetsResponse {
            error_count: stats.errors,
            serialized_targets_output: buffer,
        })
    }
}


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

use std::fmt::Write;
use std::sync::Arc;

use anyhow::Context;
use buck2_cli_proto::targets_request;
use buck2_cli_proto::targets_request::OutputFormat;
use buck2_cli_proto::targets_request::TargetHashGraphType;
use buck2_cli_proto::HasClientContext;
use buck2_cli_proto::TargetsRequest;
use buck2_core::bzl::ImportPath;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::package::PackageLabel;
use buck2_node::attrs::hacks::value_to_json;
use buck2_node::attrs::inspect_options::AttrInspectOptions;
use buck2_node::nodes::attributes::DEPS;
use buck2_node::nodes::attributes::INPUTS;
use buck2_node::nodes::attributes::ONCALL;
use buck2_node::nodes::attributes::PACKAGE;
use buck2_node::nodes::attributes::TARGET_CALL_STACK;
use buck2_node::nodes::attributes::TARGET_HASH;
use buck2_node::nodes::attributes::TYPE;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_util::indent::indent;
use gazebo::prelude::SliceExt;
use regex::RegexSet;

use crate::json::QuotedJson;
use crate::target_hash::BuckTargetHash;

#[derive(Debug, thiserror::Error)]
enum FormatterError {
    #[error("Attributes can only be specified when output format is JSON (internal error)")]
    AttrsOnlyWithJson,
    #[error("`output_format` is not set (internal error)")]
    OutputFormatNotSet,
}

pub(crate) struct TargetInfo<'a> {
    pub(crate) node: &'a TargetNode,
    pub(crate) target_hash: Option<BuckTargetHash>,
}

fn package_error_to_stderr(package: &PackageLabel, error: &anyhow::Error, stderr: &mut String) {
    writeln!(stderr, "Error parsing {package}\n{error:?}").unwrap();
}

#[allow(unused_variables)]
pub(crate) trait TargetFormatter: Send + Sync {
    fn begin(&self, buffer: &mut String) {}
    fn end(&self, stats: &Stats, buffer: &mut String) {}
    /// Called between each target/imports/package_error
    fn separator(&self, buffer: &mut String) {}
    fn target(&self, target_info: TargetInfo<'_>, buffer: &mut String) {}
    fn imports(
        &self,
        source: &CellPath,
        imports: &[ImportPath],
        package: Option<PackageLabel>,
        buffer: &mut String,
    ) {
    }
    fn package_error(
        &self,
        package: PackageLabel,
        error: &anyhow::Error,
        stdout: &mut String,
        stderr: &mut String,
    ) {
        package_error_to_stderr(&package, error, stderr);
    }
}

pub(crate) struct JsonWriter {
    pub(crate) json_lines: bool,
}

impl JsonWriter {
    pub(crate) fn begin(&self, buffer: &mut String) {
        if !self.json_lines {
            buffer.push_str("[\n");
        }
    }

    pub(crate) fn end(&self, buffer: &mut String) {
        if !self.json_lines {
            buffer.push_str("\n]\n");
        }
    }

    pub(crate) fn separator(&self, buffer: &mut String) {
        if !self.json_lines {
            buffer.push_str(",\n");
        }
    }

    pub(crate) fn entry_start(&self, buffer: &mut String) {
        if self.json_lines {
            buffer.push('{');
        } else {
            buffer.push_str("  {\n");
        }
    }

    pub(crate) fn entry_end(&self, buffer: &mut String, first: bool) {
        if self.json_lines {
            buffer.push_str("}\n");
        } else {
            if !first {
                buffer.push('\n');
            }
            buffer.push_str("  }");
        }
    }

    pub(crate) fn entry_item(
        &self,
        buffer: &mut String,
        first: &mut bool,
        key: &str,
        value: QuotedJson,
    ) {
        if *first {
            *first = false;
        } else if self.json_lines {
            buffer.push(',');
        } else {
            buffer.push_str(",\n");
        }
        if !self.json_lines {
            buffer.push_str("    ");
        }
        write!(
            buffer,
            "{}:{}",
            serde_json::to_string(key).unwrap(),
            value.as_str()
        )
        .unwrap();
    }
}

struct JsonFormat {
    attributes: Option<RegexSet>,
    attr_inspect_opts: AttrInspectOptions,
    target_call_stacks: bool,
    writer: JsonWriter,
}

impl TargetFormatter for JsonFormat {
    fn begin(&self, buffer: &mut String) {
        self.writer.begin(buffer)
    }

    fn end(&self, _stats: &Stats, buffer: &mut String) {
        self.writer.end(buffer)
    }

    fn separator(&self, buffer: &mut String) {
        self.writer.separator(buffer)
    }

    fn target(&self, target_info: TargetInfo<'_>, buffer: &mut String) {
        self.writer.entry_start(buffer);
        let mut first = true;

        fn print_attr(
            this: &JsonFormat,
            buffer: &mut String,
            first: &mut bool,
            k: &str,
            v: impl FnOnce() -> QuotedJson,
        ) {
            if let Some(filter) = &this.attributes {
                if !filter.is_match(k) {
                    return;
                }
            }
            this.writer.entry_item(buffer, first, k, v());
        }

        print_attr(self, buffer, &mut first, TYPE, || {
            QuotedJson::quote_str(&target_info.node.rule_type().to_string())
        });
        print_attr(self, buffer, &mut first, DEPS, || {
            QuotedJson::list(target_info.node.deps().map(QuotedJson::quote_display))
        });

        print_attr(self, buffer, &mut first, INPUTS, || {
            QuotedJson::list(target_info.node.inputs().map(QuotedJson::quote_display))
        });

        if let Some(hash) = target_info.target_hash {
            print_attr(self, buffer, &mut first, TARGET_HASH, || {
                QuotedJson::quote_display(hash)
            });
        }
        print_attr(self, buffer, &mut first, PACKAGE, || {
            QuotedJson::quote_display(target_info.node.label().pkg())
        });
        if let Some(oncall) = target_info.node.oncall() {
            print_attr(self, buffer, &mut first, ONCALL, || {
                QuotedJson::quote_display(oncall)
            });
        }

        for a in target_info.node.attrs(self.attr_inspect_opts) {
            print_attr(self, buffer, &mut first, a.name, || {
                QuotedJson::from_serde_json_value(
                    value_to_json(a.value, target_info.node.label().pkg()).unwrap(),
                )
            });
        }

        if self.target_call_stacks {
            match target_info.node.call_stack() {
                Some(call_stack) => {
                    print_attr(self, buffer, &mut first, TARGET_CALL_STACK, || {
                        QuotedJson::quote_str(&call_stack)
                    });
                }
                None => {
                    // Should not happen.
                }
            }
        }
        self.writer.entry_end(buffer, first);
    }

    fn imports(
        &self,
        source: &CellPath,
        imports: &[ImportPath],
        package: Option<PackageLabel>,
        buffer: &mut String,
    ) {
        self.writer.entry_start(buffer);
        let mut first = true;
        if let Some(package) = package {
            self.writer.entry_item(
                buffer,
                &mut first,
                PACKAGE,
                QuotedJson::quote_str(&package.to_string()),
            );
        }
        self.writer.entry_item(
            buffer,
            &mut first,
            "buck.file",
            QuotedJson::quote_str(&source.to_string()),
        );
        self.writer.entry_item(
            buffer,
            &mut first,
            "buck.imports",
            QuotedJson::list(imports.map(|d| QuotedJson::quote_display(d.path()))),
        );
        self.writer.entry_end(buffer, first);
    }

    fn package_error(
        &self,
        package: PackageLabel,
        error: &anyhow::Error,
        stdout: &mut String,
        stderr: &mut String,
    ) {
        // When an error happens we print it to stdout (as a JSON entry) and to stderr (as a human message).
        // If the user has keep-going turned on, they'll get the JSON on stdout, but also have the error message appear on stderr.
        // If the user has keep-going turned off, they'll only see one error message and then abort.
        package_error_to_stderr(&package, error, stderr);
        self.writer.entry_start(stdout);
        let mut first = true;
        self.writer.entry_item(
            stdout,
            &mut first,
            PACKAGE,
            QuotedJson::quote_display(package),
        );
        self.writer.entry_item(
            stdout,
            &mut first,
            "buck.error",
            QuotedJson::quote_str(&format!("{:?}", error)),
        );
        self.writer.entry_end(stdout, first);
    }
}

#[derive(Debug, Default)]
pub(crate) struct Stats {
    pub(crate) errors: u64,
    pub(crate) success: u64,
    pub(crate) targets: u64,
}

impl Stats {
    pub(crate) fn merge(&mut self, stats: &Stats) {
        self.errors += stats.errors;
        self.success += stats.success;
        self.targets += stats.targets;
    }
}

struct StatsFormat;

impl TargetFormatter for StatsFormat {
    fn end(&self, stats: &Stats, buffer: &mut String) {
        writeln!(buffer, "{:?}", stats).unwrap()
    }
}

struct TargetNameFormat {
    target_call_stacks: bool,
    target_hash_graph_type: TargetHashGraphType,
}
impl TargetFormatter for TargetNameFormat {
    fn target(&self, target_info: TargetInfo<'_>, buffer: &mut String) {
        if self.target_hash_graph_type != TargetHashGraphType::None {
            match target_info.target_hash {
                Some(hash) => {
                    writeln!(buffer, "{label} {hash}", label = target_info.node.label()).unwrap()
                }
                None => {} // print nothing if there is no hash and show_target_hash is specified.
            };
        } else {
            writeln!(buffer, "{}", target_info.node.label()).unwrap();
        }
        if self.target_call_stacks {
            print_target_call_stack_after_target(buffer, target_info.node.call_stack().as_deref());
        }
    }
}

pub(crate) fn print_target_call_stack_after_target(out: &mut String, call_stack: Option<&str>) {
    if let Some(call_stack) = call_stack {
        write!(out, "{}", indent("  ", call_stack)).unwrap();
    }
}

pub(crate) fn create_formatter(
    request: &TargetsRequest,
    other: &targets_request::Other,
) -> anyhow::Result<Arc<dyn TargetFormatter>> {
    let output_format = OutputFormat::from_i32(request.output_format)
        .context("Invalid value of `output_format` (internal error)")?;

    let target_call_stacks = request.client_context()?.target_call_stacks;

    match output_format {
        OutputFormat::Json | OutputFormat::JsonLines => {}
        _ => {
            // Self-check.
            if !other.output_attributes.is_empty() {
                return Err(FormatterError::AttrsOnlyWithJson.into());
            }
        }
    }

    match output_format {
        OutputFormat::Unknown => Err(FormatterError::OutputFormatNotSet.into()),
        OutputFormat::Stats => Ok(Arc::new(StatsFormat)),
        OutputFormat::Text => Ok(Arc::new(TargetNameFormat {
            target_call_stacks,
            target_hash_graph_type: TargetHashGraphType::from_i32(other.target_hash_graph_type)
                .expect("buck cli should send valid target hash graph type"),
        })),
        OutputFormat::Json | OutputFormat::JsonLines => Ok(Arc::new(JsonFormat {
            attributes: if other.output_attributes.is_empty() {
                None
            } else {
                Some(RegexSet::new(&other.output_attributes)?)
            },
            attr_inspect_opts: if other.include_default_attributes {
                AttrInspectOptions::All
            } else {
                AttrInspectOptions::DefinedOnly
            },
            target_call_stacks,
            writer: JsonWriter {
                json_lines: output_format == OutputFormat::JsonLines,
            },
        })),
    }
}


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

mod default;
pub(crate) mod fmt;
mod resolve_alias;
mod streaming;

use std::fs::File;
use std::io::BufWriter;
use std::io::Write;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_cli_proto::targets_request;
use buck2_cli_proto::targets_request::TargetHashGraphType;
use buck2_cli_proto::HasClientContext;
use buck2_cli_proto::TargetsRequest;
use buck2_cli_proto::TargetsResponse;
use buck2_common::dice::cells::HasCellResolver;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::pattern::parse_patterns_from_cli_args;
use buck2_server_ctx::pattern::target_platform_from_client_context;
use buck2_server_ctx::template::run_server_command;
use buck2_server_ctx::template::ServerCommandTemplate;
use dice::DiceTransaction;

use crate::commands::targets::default::targets_batch;
use crate::commands::targets::default::TargetHashOptions;
use crate::commands::targets::fmt::create_formatter;
use crate::commands::targets::resolve_alias::targets_resolve_aliases;
use crate::commands::targets::streaming::targets_streaming;

#[derive(Debug, thiserror::Error)]
enum TargetsCommandError {
    #[error("Missing field in proto request (internal error)")]
    MissingField,
}

pub(crate) enum Outputter {
    Stdout,
    File(BufWriter<File>),
}

impl Outputter {
    fn new(request: &TargetsRequest) -> anyhow::Result<Self> {
        match &request.output {
            None => Ok(Self::Stdout),
            Some(file) => Ok(Self::File(BufWriter::new(
                File::create(file).with_context(|| {
                    format!("Failed to open file `{}` for `targets` output ", file)
                })?,
            ))),
        }
    }

    fn write1(&mut self, stdout: &mut impl Write, x: &str) -> anyhow::Result<()> {
        match self {
            Self::Stdout => stdout.write_all(x.as_bytes())?,
            Self::File(f) => f.write_all(x.as_bytes())?,
        }
        Ok(())
    }

    fn write2(&mut self, stdout: &mut impl Write, x: &str, y: &str) -> anyhow::Result<()> {
        match self {
            Self::Stdout => {
                stdout.write_all(x.as_bytes())?;
                stdout.write_all(y.as_bytes())?;
            }
            Self::File(f) => {
                f.write_all(x.as_bytes())?;
                f.write_all(y.as_bytes())?;
            }
        }
        Ok(())
    }

    /// If this outputter should write anything to a file, do so, and return whatever buffer is left over.
    fn write_to_file(&mut self, buffer: String) -> anyhow::Result<String> {
        match self {
            Self::Stdout => Ok(buffer),
            Self::File(f) => {
                f.write_all(buffer.as_bytes())?;
                Ok(String::new())
            }
        }
    }

    fn flush(&mut self) -> anyhow::Result<()> {
        match self {
            Self::Stdout => Ok(()),
            Self::File(f) => Ok(f.flush()?),
        }
    }
}

pub async fn targets_command(
    server_ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
    req: TargetsRequest,
) -> anyhow::Result<TargetsResponse> {
    run_server_command(
        TargetsServerCommand { req },
        server_ctx,
        partial_result_dispatcher,
    )
    .await
}

struct TargetsServerCommand {
    req: TargetsRequest,
}

#[async_trait]
impl ServerCommandTemplate for TargetsServerCommand {
    type StartEvent = buck2_data::TargetsCommandStart;
    type EndEvent = buck2_data::TargetsCommandEnd;
    type Response = TargetsResponse;
    type PartialResult = buck2_cli_proto::StdoutBytes;

    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        mut partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        dice: DiceTransaction,
    ) -> anyhow::Result<Self::Response> {
        targets(
            server_ctx,
            &mut partial_result_dispatcher.as_writer(),
            dice,
            &self.req,
        )
        .await
    }

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

    fn end_event(&self, _response: &anyhow::Result<Self::Response>) -> Self::EndEvent {
        buck2_data::TargetsCommandEnd {
            unresolved_target_patterns: self.req.target_patterns.clone(),
        }
    }
}

async fn targets(
    server_ctx: &dyn ServerCommandContextTrait,
    stdout: &mut impl Write,
    mut dice: DiceTransaction,
    request: &TargetsRequest,
) -> anyhow::Result<TargetsResponse> {
    // TODO(nmj): Rather than returning fully formatted data in the TargetsResponse, we should
    //            instead return structured data, and return *that* to the CLI. The CLI should
    //            then handle printing. The current approach is just a temporary hack to fix some
    //            issues with printing to stdout.

    let cwd = server_ctx.working_dir();
    let cell_resolver = dice.get_cell_resolver().await?;
    let parsed_target_patterns = parse_patterns_from_cli_args::<TargetPatternExtra>(
        &mut dice,
        &request.target_patterns,
        cwd,
    )
    .await?;

    let mut outputter = Outputter::new(request)?;

    let response = match &request.targets {
        Some(targets_request::Targets::ResolveAlias(_)) => {
            targets_resolve_aliases(dice, request, parsed_target_patterns).await?
        }
        Some(targets_request::Targets::Other(other)) => {
            if other.streaming {
                let formatter = create_formatter(request, other)?;
                let hashing = match TargetHashGraphType::from_i32(other.target_hash_graph_type)
                    .expect("buck cli should send valid target hash graph type")
                {
                    TargetHashGraphType::None => None,
                    _ => Some(other.target_hash_use_fast_hash),
                };

                let res = targets_streaming(
                    server_ctx,
                    stdout,
                    dice,
                    formatter,
                    &mut outputter,
                    parsed_target_patterns,
                    other.keep_going,
                    other.cached,
                    other.imports,
                    hashing,
                    request.concurrency.as_ref().map(|x| x.concurrency as usize),
                )
                .await;
                // Make sure we always flush the outputter, even on failure, as we may have partially written to it
                outputter.flush()?;
                res?
            } else {
                let formatter = create_formatter(request, other)?;
                let client_ctx = request.client_context()?;
                let target_platform =
                    target_platform_from_client_context(client_ctx, server_ctx, &dice).await?;
                let fs = server_ctx.project_root();
                targets_batch(
                    server_ctx,
                    dice,
                    &*formatter,
                    parsed_target_patterns,
                    target_platform,
                    TargetHashOptions::new(other, &cell_resolver, fs)?,
                    other.keep_going,
                )
                .await?
            }
        }
        None => return Err(TargetsCommandError::MissingField.into()),
    };

    let response = TargetsResponse {
        error_count: response.error_count,
        serialized_targets_output: outputter.write_to_file(response.serialized_targets_output)?,
    };
    outputter.flush()?;
    Ok(response)
}

fn mk_error(errors: u64) -> anyhow::Error {
    // Simpler error so that we don't print long errors twice (when exiting buck2)
    let package_str = if errors == 1 { "package" } else { "packages" };
    anyhow::anyhow!("Failed to parse {} {}", errors, package_str)
}


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

//! Server-side implementation of `buck2 targets --resolve-alias` command.

#[derive(Debug, thiserror::Error)]
enum ResolveAliasError {
    #[error("`output_format` not set (internal error)")]
    OutputFormatNotSet,
    #[error("`--stat` format is not supported by `--resolve-alias`")]
    StatFormatNotSupported,
}

use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt::Write;
use std::sync::Arc;

use anyhow::Context;
use buck2_cli_proto::targets_request::OutputFormat;
use buck2_cli_proto::TargetsRequest;
use buck2_cli_proto::TargetsResponse;
use buck2_common::result::SharedResult;
use buck2_common::result::ToSharedResultExt;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::TargetLabel;
use buck2_node::nodes::attributes::PACKAGE;
use buck2_node::nodes::eval_result::EvaluationResult;
use buck2_node::nodes::frontend::TargetGraphCalculation;
use dice::DiceTransaction;
use dupe::Dupe;
use futures::stream::FuturesUnordered;
use futures::StreamExt;

use crate::commands::targets::fmt::JsonWriter;
use crate::json::QuotedJson;

trait ResolveAliasFormatter {
    /// Before writing anything.
    fn begin(&self, buffer: &mut String);

    /// After writing everything.
    fn end(&self, buffer: &mut String);

    /// Between items
    fn separator(&self, buffer: &mut String);

    /// Emit an alias
    fn emit(&self, alias: &str, label: &TargetLabel, buffer: &mut String);
}

impl ResolveAliasFormatter for JsonWriter {
    fn begin(&self, buffer: &mut String) {
        self.begin(buffer);
    }

    fn end(&self, buffer: &mut String) {
        self.end(buffer);
    }

    fn separator(&self, buffer: &mut String) {
        self.separator(buffer);
    }

    fn emit(&self, alias: &str, label: &TargetLabel, buffer: &mut String) {
        let mut first = true;
        self.entry_start(buffer);
        self.entry_item(buffer, &mut first, "alias", QuotedJson::quote_str(alias));
        // Using a format consistent wit hthe output of `buck2 targets`
        self.entry_item(
            buffer,
            &mut first,
            PACKAGE,
            QuotedJson::quote_display(label.pkg()),
        );
        self.entry_item(
            buffer,
            &mut first,
            "name",
            QuotedJson::quote_display(label.name()),
        );
        self.entry_end(buffer, first);
    }
}

struct LinesWriter;

impl ResolveAliasFormatter for LinesWriter {
    fn begin(&self, _buffer: &mut String) {}

    fn end(&self, _buffer: &mut String) {}

    fn separator(&self, buffer: &mut String) {
        buffer.push('\n');
    }

    fn emit(&self, _alias: &str, label: &TargetLabel, buffer: &mut String) {
        write!(buffer, "{}", label).unwrap();
    }
}

pub(crate) async fn targets_resolve_aliases(
    dice: DiceTransaction,
    request: &TargetsRequest,
    parsed_target_patterns: Vec<ParsedPattern<TargetPatternExtra>>,
) -> anyhow::Result<TargetsResponse> {
    // If we are only asked to resolve aliases, then don't expand any of the patterns, and just
    // print them out. This expects the aliases to resolve to individual targets.
    let parsed_target_patterns = std::iter::zip(&request.target_patterns, parsed_target_patterns)
        .map(|(alias, pattern)| match pattern {
            ParsedPattern::Target(package, target_name, TargetPatternExtra) => {
                Ok((package, target_name))
            }
            _ => Err(anyhow::anyhow!(
                "Invalid alias (does not expand to a single target): `{}`",
                alias.value
            )),
        })
        .collect::<Result<Vec<_>, _>>()?;

    let packages = parsed_target_patterns
        .iter()
        .map(|(package, _name)| package.dupe())
        .collect::<HashSet<_>>();

    let packages = packages
        .into_iter()
        .map(|package| {
            let dice = &dice;
            async move {
                (
                    package.dupe(),
                    dice.get_interpreter_results(package.dupe())
                        .await
                        .shared_error(),
                )
            }
        })
        .collect::<FuturesUnordered<_>>()
        .collect::<HashMap<PackageLabel, SharedResult<Arc<EvaluationResult>>>>()
        .await;

    let mut buffer = String::new();

    let output_format = OutputFormat::from_i32(request.output_format)
        .context("Invalid value of `output_format` (internal error)")?;

    let json_writer;

    let formatter = match output_format {
        OutputFormat::Unknown => return Err(ResolveAliasError::OutputFormatNotSet.into()),
        OutputFormat::Text => &LinesWriter as &dyn ResolveAliasFormatter,
        OutputFormat::Json => {
            json_writer = JsonWriter { json_lines: false };
            &json_writer as &dyn ResolveAliasFormatter
        }
        OutputFormat::JsonLines => {
            json_writer = JsonWriter { json_lines: true };
            &json_writer as &dyn ResolveAliasFormatter
        }
        OutputFormat::Stats => return Err(ResolveAliasError::StatFormatNotSupported.into()),
    };

    let mut needs_separator = false;

    formatter.begin(&mut buffer);

    for (alias, (package, target_name)) in
        std::iter::zip(&request.target_patterns, &parsed_target_patterns)
    {
        // NOTE: We don't technically need the node to get the label, but we need the node to
        // validate it exists.
        let node = packages
            .get(package)
            .with_context(|| format!("Package does not exist: `{}`", package))
            .and_then(|package_data| {
                package_data
                    .as_ref()
                    .map_err(|e| e.dupe())
                    .with_context(|| format!("Package cannot be evaluated: `{}`", package))?
                    .resolve_target(target_name)
                    .with_context(|| {
                        format!(
                            "Target does not exist in package `{}`: `{}`",
                            package, target_name,
                        )
                    })
            })
            .with_context(|| format!("Invalid alias: `{}`", alias.value))?;

        if needs_separator {
            formatter.separator(&mut buffer);
        }
        needs_separator = true;
        formatter.emit(&alias.value, node.label(), &mut buffer);
    }

    formatter.end(&mut buffer);

    Ok(TargetsResponse {
        error_count: 0,
        serialized_targets_output: buffer,
    })
}


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

//! Server-side implementation of `buck2 targets --streaming` command.

use std::collections::HashSet;
use std::io::Write;
use std::mem;
use std::sync::Arc;
use std::sync::Mutex;

use buck2_cli_proto::TargetsResponse;
use buck2_common::pattern::package_roots::find_package_roots_stream;
use buck2_common::pattern::resolve::ResolvedPattern;
use buck2_core::bzl::ImportPath;
use buck2_core::package::PackageLabel;
use buck2_core::pattern::pattern_type::PatternType;
use buck2_core::pattern::pattern_type::TargetPatternExtra;
use buck2_core::pattern::PackageSpec;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::name::TargetName;
use buck2_interpreter::load_module::InterpreterCalculation;
use buck2_interpreter::load_module::INTERPRETER_CALCULATION_IMPL;
use buck2_interpreter::path::PackageFilePath;
use buck2_node::nodes::eval_result::EvaluationResult;
use buck2_node::nodes::frontend::TargetGraphCalculation;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;
use dupe::IterDupedExt;
use futures::future::FutureExt;
use futures::Stream;
use futures::StreamExt;
use gazebo::prelude::VecExt;
use itertools::Either;
use itertools::Itertools;
use more_futures::spawn::spawn_cancellable;
use starlark_map::small_set::SmallSet;
use thiserror::Error;
use tokio::sync::Semaphore;

use crate::commands::targets::fmt::Stats;
use crate::commands::targets::fmt::TargetFormatter;
use crate::commands::targets::fmt::TargetInfo;
use crate::commands::targets::mk_error;
use crate::commands::targets::Outputter;
use crate::target_hash::TargetHashes;

pub(crate) async fn targets_streaming(
    server_ctx: &dyn ServerCommandContextTrait,
    stdout: &mut impl Write,
    dice: DiceTransaction,
    formatter: Arc<dyn TargetFormatter>,
    outputter: &mut Outputter,
    parsed_patterns: Vec<ParsedPattern<TargetPatternExtra>>,
    keep_going: bool,
    cached: bool,
    imports: bool,
    fast_hash: Option<bool>, // None = no hashing
    threads: Option<usize>,
) -> anyhow::Result<TargetsResponse> {
    struct Res {
        stats: Stats,           // Stats to merge in
        package: PackageLabel,  // The package I was operating on
        stderr: Option<String>, // Print to stderr (and break unless keep_going is set)
        stdout: String,         // Print to stdout
    }

    let imported = Arc::new(Mutex::new(SmallSet::new()));
    let threads = Arc::new(Semaphore::new(threads.unwrap_or(Semaphore::MAX_PERMITS)));

    let mut packages = stream_packages(&dice, parsed_patterns)
        .map(|x| {
            let formatter = formatter.dupe();
            let imported = imported.dupe();
            let threads = threads.dupe();
            let ctx = dice.dupe();

            spawn_cancellable(
                |_cancellation| {
                    {
                        async move {
                            let (package, spec) = x?;
                            let mut res = Res {
                                stats: Stats::default(),
                                package: package.dupe(),
                                stderr: None,
                                stdout: String::new(),
                            };
                            let targets = {
                                // This bit of code is the heavy CPU stuff, so guard it with the threads
                                let _permit = threads.acquire().await.unwrap();
                                load_targets(&ctx, package.dupe(), spec, cached, keep_going).await
                            };
                            let mut show_err = |err| {
                                res.stats.errors += 1;
                                let mut stderr = String::new();
                                formatter.package_error(
                                    package.dupe(),
                                    err,
                                    &mut res.stdout,
                                    &mut stderr,
                                );
                                res.stderr = Some(stderr);
                            };
                            match targets {
                                Ok((eval_result, targets, err)) => {
                                    if let Some(err) = err {
                                        show_err(&err);
                                        formatter.separator(&mut res.stdout);
                                    }
                                    res.stats.success += 1;
                                    if imports {
                                        let eval_imports = eval_result.imports();
                                        formatter.imports(
                                            &eval_result.buildfile_path().path(),
                                            eval_imports,
                                            Some(package.dupe()),
                                            &mut res.stdout,
                                        );
                                        imported
                                            .lock()
                                            .unwrap()
                                            .extend(eval_imports.iter().cloned());
                                    }
                                    for (i, node) in targets.iter().enumerate() {
                                        res.stats.targets += 1;
                                        if imports || i != 0 {
                                            formatter.separator(&mut res.stdout);
                                        }
                                        formatter.target(
                                            TargetInfo {
                                                node,
                                                target_hash: fast_hash.map(|fast| {
                                                    TargetHashes::compute_immediate_one(node, fast)
                                                }),
                                            },
                                            &mut res.stdout,
                                        )
                                    }
                                }
                                Err(err) => {
                                    show_err(&err);
                                }
                            }
                            anyhow::Ok(res)
                        }
                    }
                    .boxed()
                },
                &*dice.per_transaction_data().spawner,
                dice.per_transaction_data(),
            )
            .into_drop_cancel()
        })
        // Use unlimited parallelism - tokio will restrict us anyway
        .buffer_unordered(1000000);

    let mut buffer = String::new();
    formatter.begin(&mut buffer);
    let mut stats = Stats::default();
    let mut needs_separator = false;
    let mut package_files_seen = SmallSet::new();
    while let Some(res) = packages.next().await {
        let res = res?;
        stats.merge(&res.stats);
        if let Some(stderr) = &res.stderr {
            server_ctx.stderr()?.write_all(stderr.as_bytes())?;
            if !keep_going {
                return Err(mk_error(stats.errors));
            }
        }
        if !res.stdout.is_empty() {
            if needs_separator {
                formatter.separator(&mut buffer);
            }
            needs_separator = true;
            outputter.write2(stdout, &buffer, &res.stdout)?;
            buffer.clear();
        }
        if imports {
            // Need to also find imports from PACKAGE files
            let mut path = Some(PackageFilePath::for_dir(res.package.as_cell_path()));
            while let Some(x) = path {
                if package_files_seen.contains(&x) {
                    break;
                }
                package_files_seen.insert(x.clone());
                // These aren't cached, but the cost is relatively low (Starlark parsing),
                // and there aren't many, so we just do it on the main thread.
                // We ignore errors as these will bubble up as BUCK file errors already.
                if let Ok(Some(imports)) = package_imports(&dice, &x).await {
                    if needs_separator {
                        formatter.separator(&mut buffer);
                    }
                    needs_separator = true;
                    formatter.imports(x.path(), &imports, None, &mut buffer);
                    outputter.write1(stdout, &buffer)?;
                    buffer.clear();
                    imported.lock().unwrap().extend(imports.into_iter());
                }
                path = x.parent_package_file();
            }
        }
    }

    // Recursively chase down all imported paths
    let mut todo = mem::take(&mut *imported.lock().unwrap());
    let mut seen_imported = HashSet::new();
    while let Some(path) = todo.pop() {
        if seen_imported.insert(path.path().clone()) {
            // If these lead to an error, that's surpsing (we had a working module with it loaded)
            // so we should always propagate the error here (even with keep_going)
            if needs_separator {
                formatter.separator(&mut buffer);
            }
            needs_separator = true;
            // No need to parallelise these this step because it will already be on the DICE graph
            let loaded = dice.get_loaded_module_from_import_path(&path).await?;
            let imports = loaded.imports().cloned().collect::<Vec<_>>();
            formatter.imports(path.path(), &imports, None, &mut buffer);
            todo.extend(imports);
            outputter.write1(stdout, &buffer)?;
            buffer.clear();
        }
    }

    formatter.end(&stats, &mut buffer);
    Ok(TargetsResponse {
        error_count: stats.errors,
        serialized_targets_output: buffer,
    })
}

/// Given the patterns, separate into those which have an explicit package, and those which are recursive
fn stream_packages<T: PatternType>(
    dice: &DiceTransaction,
    patterns: Vec<ParsedPattern<T>>,
) -> impl Stream<Item = anyhow::Result<(PackageLabel, PackageSpec<T>)>> {
    let mut spec = ResolvedPattern::<T>::new();
    let mut recursive_paths = Vec::new();

    for pattern in patterns {
        match pattern {
            ParsedPattern::Target(package, target_name, extra) => {
                spec.add_target(package.dupe(), target_name, extra);
            }
            ParsedPattern::Package(package) => {
                spec.add_package(package.dupe());
            }
            ParsedPattern::Recursive(package) => {
                recursive_paths.push(package);
            }
        }
    }

    futures::stream::iter(spec.specs.into_iter().map(Ok))
        .chain(find_package_roots_stream(dice, recursive_paths).map(|x| Ok((x?, PackageSpec::All))))
}

#[derive(Error, Debug)]
enum TargetsError {
    #[error(
        "Unknown targets {} from package `{0}`.",
        _1.iter().map(|x| format!("`{}`", x)).join(", ")
    )]
    MissingTargets(PackageLabel, Vec<TargetName>),
}

/// Load the targets from a package. If `keep_going` is specified then it may return a `Some` error in the triple.
async fn load_targets(
    dice: &DiceComputations,
    package: PackageLabel,
    spec: PackageSpec<TargetPatternExtra>,
    cached: bool,
    keep_going: bool,
) -> anyhow::Result<(
    Arc<EvaluationResult>,
    Vec<TargetNode>,
    Option<anyhow::Error>,
)> {
    let result = if cached {
        dice.get_interpreter_results(package.dupe()).await?
    } else {
        dice.get_interpreter_results_uncached(package.dupe())
            .await?
    };

    match spec {
        PackageSpec::Targets(targets) => {
            if keep_going {
                let (miss, targets): (Vec<_>, Vec<_>) =
                    targets
                        .into_iter()
                        .partition_map(|(target, TargetPatternExtra)| {
                            match result.targets().get(target.as_ref()) {
                                None => Either::Left(target),
                                Some(x) => Either::Right(x.dupe()),
                            }
                        });
                let err = if miss.is_empty() {
                    None
                } else {
                    Some(TargetsError::MissingTargets(package.dupe(), miss).into())
                };
                Ok((result, targets, err))
            } else {
                let targets = targets.into_try_map(|(target, TargetPatternExtra)| {
                    anyhow::Ok(result.resolve_target(target.as_ref())?.dupe())
                })?;
                Ok((result, targets, None))
            }
        }
        PackageSpec::All => {
            let targets = result.targets().values().duped().collect();
            Ok((result, targets, None))
        }
    }
}

/// Return `None` if the PACKAGE file doesn't exist
async fn package_imports(
    dice: &DiceComputations,
    path: &PackageFilePath,
) -> anyhow::Result<Option<Vec<ImportPath>>> {
    INTERPRETER_CALCULATION_IMPL
        .get()?
        .get_package_file_deps(dice, path)
        .await
}


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

use buck2_query::query::environment::QueryTarget;
use buck2_query::query::environment::QueryTargets;
use buck2_query::query::syntax::simple::eval::set::TargetSet;
use regex::RegexSet;
use starlark_map::small_map::SmallMap;

use crate::dot::DotDigraph;
use crate::dot::DotEdge;
use crate::dot::DotNode;
use crate::dot::DotNodeAttrs;

pub struct DotTargetGraphNode<'a, T: QueryTarget>(&'a T, &'a DotTargetGraph<T>);

/// A simple adapter for creating a DotDiGraph for a TargetSet.
pub struct DotTargetGraph<T: QueryTarget> {
    pub targets: TargetSet<T>,
    pub attributes: Option<RegexSet>,
}

impl<'a, T: QueryTarget> DotDigraph<'a> for DotTargetGraph<T> {
    type Node = DotTargetGraphNode<'a, T>;

    fn name(&self) -> &str {
        "result_graph"
    }

    fn for_each_node<F: FnMut(&Self::Node) -> anyhow::Result<()>>(
        &'a self,
        mut f: F,
    ) -> anyhow::Result<()> {
        for node in self.targets.iter() {
            f(&DotTargetGraphNode(node, self))?;
        }
        Ok(())
    }

    fn for_each_edge<F: FnMut(&DotEdge) -> anyhow::Result<()>>(
        &'a self,
        node: &Self::Node,
        mut f: F,
    ) -> anyhow::Result<()> {
        for dep in node.0.deps() {
            // Only include edges to other nodes within the subgraph.
            if self.targets.contains(dep) {
                f(&DotEdge {
                    from: &node.0.node_ref().to_string(),
                    to: &dep.to_string(),
                })?;
            }
        }
        Ok(())
    }
}

impl<'a, T: QueryTarget> DotNode for DotTargetGraphNode<'a, T> {
    fn attrs(&self) -> anyhow::Result<DotNodeAttrs> {
        let extra = match &self.1.attributes {
            Some(attr_regex) => {
                let mut extra = SmallMap::new();
                QueryTargets::for_all_attrs::<anyhow::Error, _, _>(
                    self.0,
                    |attr_name, attr_value| {
                        if attr_regex.is_match(attr_name) {
                            extra.insert(
                                format!("buck_{}", attr_name),
                                self.0.attr_to_string_alternate(attr_value),
                            );
                        }
                        Ok(())
                    },
                )?;
                extra
            }
            None => SmallMap::new(),
        };
        Ok(DotNodeAttrs {
            style: Some("filled".to_owned()),
            color: Some("#DFECDF".to_owned()),
            extra,
            ..DotNodeAttrs::default()
        })
    }

    fn id(&self) -> String {
        self.0.node_ref().to_string()
    }
}


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

/// Common code executed in the end of command to produce `CommandEnd`.
pub fn command_end<R, D>(
    metadata: HashMap<String, String>,
    result: &anyhow::Result<R>,
    data: D,
) -> buck2_data::CommandEnd
where
    D: Into<buck2_data::command_end::Data>,
{
    command_end_ext(metadata, result, data.into(), |_| true)
}

pub fn command_end_ext<R, D, F>(
    metadata: HashMap<String, String>,
    result: &anyhow::Result<R>,
    data: D,
    is_success: F,
) -> buck2_data::CommandEnd
where
    F: FnOnce(&R) -> bool,
    D: Into<buck2_data::command_end::Data>,
{
    let (is_success, error_messages) = match result {
        Ok(r) => (is_success(r), Vec::new()),
        Err(e) => (false, vec![format!("{:#}", e)]),
    };
    buck2_data::CommandEnd {
        is_success,
        error_messages,
        metadata,
        data: Some(data.into()),
    }
}