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

use allocative::Allocative;
use buck2_util::collections::ordered_set::OrderedSet;
use dupe::Clone_;
use dupe::Dupe;
use starlark_map::small_set;
use starlark_map::Equivalent;

use crate::query::environment::LabeledNode;

#[derive(Debug, Clone, Dupe, Allocative)]
pub struct LabelIndexed<T: LabeledNode>(pub T);

impl<T: LabeledNode> PartialEq for LabelIndexed<T> {
    fn eq(&self, other: &Self) -> bool {
        self.0.node_ref() == other.0.node_ref()
    }
}
impl<T: LabeledNode> Eq for LabelIndexed<T> {}
impl<T: LabeledNode> Hash for LabelIndexed<T> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.node_ref().hash(state)
    }
}

impl<T: LabeledNode> Ord for LabelIndexed<T>
where
    T::NodeRef: Ord,
{
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        self.0.node_ref().cmp(other.0.node_ref())
    }
}

impl<T: LabeledNode> PartialOrd for LabelIndexed<T>
where
    T::NodeRef: PartialOrd,
{
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        self.0.node_ref().partial_cmp(other.0.node_ref())
    }
}

struct LabelIndexer<'a, T: LabeledNode>(&'a T::NodeRef);

impl<'a, T: LabeledNode> Equivalent<LabelIndexed<T>> for LabelIndexer<'a, T> {
    fn equivalent(&self, key: &LabelIndexed<T>) -> bool {
        self.0.eq(key.0.node_ref())
    }
}
impl<'a, T: LabeledNode> Hash for LabelIndexer<'a, T> {
    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
        self.0.hash(state)
    }
}

/// This is a Set that contains `LabeledNode`s, but for all trait operation like `Eq`, `Ord` all
/// operate only on the `NodeRef` of the `LabeledNode`.
#[derive(Debug, Clone, Eq, PartialEq, Hash, Allocative)]
pub struct LabelIndexedSet<T: LabeledNode> {
    nodes: OrderedSet<LabelIndexed<T>>,
}

impl<T: LabeledNode> LabelIndexedSet<T> {
    pub fn new() -> Self {
        Self {
            nodes: OrderedSet::new(),
        }
    }

    pub fn with_capacity(n: usize) -> Self {
        Self {
            nodes: OrderedSet::with_capacity(n),
        }
    }

    pub fn len(&self) -> usize {
        self.nodes.len()
    }

    pub fn get(&self, value: &T::NodeRef) -> Option<&T> {
        self.nodes.get(&LabelIndexer(value)).map(|e| &e.0)
    }

    pub fn take(&mut self, value: &T::NodeRef) -> Option<T> {
        self.nodes.take(&LabelIndexer(value)).map(|e| e.0)
    }

    pub fn iter(&self) -> Iter<T> {
        Iter {
            iter: self.nodes.iter(),
        }
    }

    #[allow(clippy::should_implement_trait)] // the std trait requires concrete or boxed iterator type
    pub fn into_iter(self) -> impl ExactSizeIterator<Item = T> {
        self.nodes.into_iter().map(|e| e.0)
    }

    pub fn insert(&mut self, value: T) -> bool {
        self.nodes.insert(LabelIndexed(value))
    }

    pub fn contains(&self, value: &T::NodeRef) -> bool {
        self.nodes.contains(&LabelIndexer(value))
    }

    pub fn get_index(&self, index: usize) -> Option<&T> {
        self.nodes.get_index(index).map(|e| &e.0)
    }

    pub fn get_index_of(&self, value: &T::NodeRef) -> Option<usize> {
        self.nodes.get_index_of(&LabelIndexer(value))
    }

    pub fn last(&self) -> Option<&T> {
        self.nodes.last().map(|e| &e.0)
    }
}

#[derive(Clone_)]
pub struct Iter<'a, T: LabeledNode> {
    iter: small_set::Iter<'a, LabelIndexed<T>>,
}

impl<'a, T: LabeledNode> Iterator for Iter<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<Self::Item> {
        self.iter.next().map(|e| &e.0)
    }
}

impl<'a, T: LabeledNode> ExactSizeIterator for Iter<'a, T> {
    fn len(&self) -> usize {
        self.iter.len()
    }
}

impl<'a, T: LabeledNode> Dupe for Iter<'a, T> {}

impl<'a, T: LabeledNode> IntoIterator for &'a LabelIndexedSet<T> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        Iter {
            iter: self.nodes.iter(),
        }
    }
}

impl<T: LabeledNode> FromIterator<T> for LabelIndexedSet<T> {
    fn from_iter<Iter: IntoIterator<Item = T>>(iter: Iter) -> Self {
        let iter = iter.into_iter();
        let mut res = Self::with_capacity(iter.size_hint().0);
        for item in iter {
            res.insert(item);
        }
        res
    }
}


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

//! Implementation of the cli and query_* attr query language.

use buck2_query_parser::parse_expr;
use buck2_query_parser::placeholder::QUERY_PERCENT_S_PLACEHOLDER;
use starlark_map::small_set::SmallSet;

use crate::query::syntax::simple::eval::values::QueryResultExt;
use crate::query::syntax::simple::functions::QueryFunctions;
use crate::query::syntax::simple::functions::QueryFunctionsVisitLiterals;
use crate::query::syntax::simple::functions::QueryLiteralVisitor;

/// Look through the expression to find all the target literals.
/// Adds those that are found to `result` set.
pub fn extract_target_literals<F: QueryFunctions>(
    functions: &F,
    query: &str,
    result: &mut SmallSet<String>,
) -> anyhow::Result<()> {
    let parsed = parse_expr(query)?;
    struct LiteralExtractor<'a> {
        literals: &'a mut SmallSet<String>,
    }
    impl QueryLiteralVisitor for LiteralExtractor<'_> {
        fn target_pattern(&mut self, pattern: &str) -> anyhow::Result<()> {
            if pattern != QUERY_PERCENT_S_PLACEHOLDER {
                self.literals.get_or_insert_owned(pattern);
            }
            Ok(())
        }
    }
    let mut visitor = LiteralExtractor { literals: result };
    functions
        .visit_literals(&mut visitor, &parsed)
        .into_anyhow(query)?;
    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.
 */

pub mod error;
pub mod evaluator;
pub mod file_set;
pub mod label_indexed;
pub mod literals;
pub mod multi_query;
pub mod set;
pub mod tests;
pub mod values;


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

//! Implementation of the cli and query_* attr query language.
use buck2_query_parser::placeholder::QUERY_PERCENT_S_PLACEHOLDER;
use futures::stream::FuturesOrdered;
use futures::Future;
use futures::StreamExt;
use indexmap::IndexMap;

use crate::query::environment::QueryTarget;
use crate::query::syntax::simple::eval::set::TargetSet;
use crate::query::syntax::simple::eval::values::QueryEvaluationValue;

/// Used to represent the results for a "multi-query" (one that contains a "%s" and potentially is applied against multiple literals).
pub struct MultiQueryResult<T: QueryTarget>(
    pub IndexMap<String, anyhow::Result<QueryEvaluationValue<T>>>,
);

impl<T: QueryTarget> MultiQueryResult<T> {
    pub fn merged(self) -> anyhow::Result<QueryEvaluationValue<T>> {
        let mut iter = self.0.into_iter();
        let (first_literal, mut results) = match iter.next() {
            Some((literal, value)) => (literal, value?),
            None => {
                return Ok(QueryEvaluationValue::TargetSet(TargetSet::new()));
            }
        };
        for (name, value) in iter {
            let value = value?;
            match (value, &mut results) {
                (
                    QueryEvaluationValue::TargetSet(value),
                    QueryEvaluationValue::TargetSet(results),
                ) => results.extend(&value),
                (QueryEvaluationValue::FileSet(value), QueryEvaluationValue::FileSet(results)) => {
                    results.insert_all(&value)
                }
                _ => unreachable!(
                    "no queries should return different types for different literals, but somehow that happened for `{}` and `{}`",
                    first_literal, name
                ),
            }
        }
        Ok(results)
    }
}

pub async fn process_multi_query<T, Fut, F, A: AsRef<str>>(
    query: &str,
    query_args: &[A],
    func: F,
) -> MultiQueryResult<T>
where
    T: QueryTarget,
    Fut: Future<Output = (String, anyhow::Result<QueryEvaluationValue<T>>)>,
    F: Fn(String, String) -> Fut,
{
    let mut queue: FuturesOrdered<_> = query_args
        .iter()
        .map(|input| {
            let input = input.as_ref();
            let query = query.replace(QUERY_PERCENT_S_PLACEHOLDER, input);
            let input = input.to_owned();
            func(input, query)
        })
        .collect();

    let mut results = IndexMap::new();
    while let Some((query, result)) = queue.next().await {
        results.insert(query, result);
    }
    MultiQueryResult(results)
}


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

use allocative::Allocative;
use buck2_query::query::environment::LabeledNode;
use display_container::fmt_container;
use dupe::IterDupedExt;
use fancy_regex::Regex;
use indexmap::IndexSet;

use crate::query::environment::QueryTarget;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::file_set::FileNode;
use crate::query::syntax::simple::eval::file_set::FileSet;
use crate::query::syntax::simple::eval::label_indexed;
use crate::query::syntax::simple::eval::label_indexed::LabelIndexedSet;

#[derive(Debug, Eq, PartialEq, Clone, Allocative)]
pub struct TargetSet<T: QueryTarget> {
    targets: LabelIndexedSet<T>,
}

impl<T: QueryTarget> TargetSet<T> {
    pub fn new() -> Self {
        Self {
            targets: LabelIndexedSet::new(),
        }
    }

    pub fn with_capacity(n: usize) -> Self {
        Self {
            targets: LabelIndexedSet::with_capacity(n),
        }
    }

    pub fn insert(&mut self, value: T) -> bool {
        self.targets.insert(value)
    }

    pub fn is_empty(&self) -> bool {
        self.targets.len() == 0
    }

    pub fn len(&self) -> usize {
        self.targets.len()
    }

    fn filter<F: Fn(&T) -> anyhow::Result<bool>>(&self, filter: F) -> anyhow::Result<TargetSet<T>> {
        let mut targets = LabelIndexedSet::new();
        for target in self.targets.iter() {
            if filter(target)? {
                targets.insert(target.dupe());
            }
        }
        Ok(Self { targets })
    }

    pub fn buildfile(&self) -> FileSet {
        let mut files = IndexSet::new();
        for target in self.targets.iter() {
            files.insert(FileNode(target.buildfile_path().path()));
        }
        FileSet::new(files)
    }

    pub fn inputs(&self) -> anyhow::Result<FileSet> {
        let mut files = IndexSet::new();
        for target in self.targets.iter() {
            target.inputs_for_each(|file| {
                files.insert(FileNode(file));
                anyhow::Ok(())
            })?;
        }
        Ok(FileSet::new(files))
    }

    // TODO(cjhopman): Does this even make sense?
    // TODO(cjhopman): I think this needs a heap to allocate values
    pub fn labels(&self, _attr: &str) -> anyhow::Result<()> {
        Err(anyhow::anyhow!(QueryError::FunctionUnimplemented(
            "labels()"
        )))
    }

    pub fn union(&self, right: &TargetSet<T>) -> TargetSet<T> {
        let mut targets = LabelIndexedSet::new();
        for target in self.targets.iter() {
            targets.insert(target.dupe());
        }
        for target in right.targets.iter() {
            targets.insert(target.dupe());
        }
        Self { targets }
    }

    pub fn iter_names(&self) -> impl Iterator<Item = &T::NodeRef> + Clone {
        self.targets.iter().map(|e| e.node_ref())
    }

    pub fn iter(&self) -> Iter<T> {
        self.targets.iter()
    }

    #[allow(clippy::should_implement_trait)] // the std trait requires concrete or boxed iterator type
    pub fn into_iter(self) -> impl Iterator<Item = T> {
        self.targets.into_iter()
    }

    pub fn contains(&self, item: &T::NodeRef) -> bool {
        self.targets.contains(item)
    }

    pub fn get(&self, item: &T::NodeRef) -> Option<&T> {
        self.targets.get(item)
    }

    pub fn get_index(&self, index: usize) -> Option<&T> {
        self.targets.get_index(index)
    }

    pub fn get_index_of(&self, item: &T::NodeRef) -> Option<usize> {
        self.targets.get_index_of(item)
    }

    pub fn last(&self) -> Option<&T> {
        self.targets.last()
    }
}

pub type Iter<'a, T> = label_indexed::Iter<'a, T>;

impl<'a, T: QueryTarget> IntoIterator for &'a TargetSet<T> {
    type Item = &'a T;
    type IntoIter = Iter<'a, T>;

    fn into_iter(self) -> Self::IntoIter {
        self.targets.iter()
    }
}

impl<T: QueryTarget> Extend<T> for TargetSet<T> {
    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
        for target in iter {
            self.targets.insert(target);
        }
    }
}

impl<'a, T: QueryTarget> Extend<&'a T> for TargetSet<T> {
    fn extend<I: IntoIterator<Item = &'a T>>(&mut self, iter: I) {
        self.extend(iter.into_iter().duped())
    }
}

impl<T: QueryTarget> FromIterator<T> for TargetSet<T> {
    fn from_iter<Iter: IntoIterator<Item = T>>(iter: Iter) -> Self {
        let targets = LabelIndexedSet::from_iter(iter);
        Self { targets }
    }
}

/// This contains additional TargetSet functions implemented via the core
/// functions on TargetSet itself.
pub trait TargetSetExt {
    type T: QueryTarget;

    fn filter<F: Fn(&Self::T) -> anyhow::Result<bool>>(
        &self,
        filter: F,
    ) -> anyhow::Result<TargetSet<Self::T>>;

    fn attrfilter(
        &self,
        attribute: &str,
        filter: &dyn Fn(&str) -> anyhow::Result<bool>,
    ) -> anyhow::Result<TargetSet<Self::T>> {
        self.filter(move |node| {
            node.map_attr(attribute, |val| match val {
                None => Ok(false),
                Some(v) => Self::T::attr_any_matches(v, &filter),
            })
        })
    }

    fn nattrfilter(
        &self,
        attribute: &str,
        filter: &dyn Fn(&str) -> anyhow::Result<bool>,
    ) -> anyhow::Result<TargetSet<Self::T>> {
        self.filter(move |node| {
            node.map_attr(attribute, |val| match val {
                None => Ok(false),
                Some(v) => Ok(!Self::T::attr_any_matches(v, &filter)?),
            })
        })
    }

    fn attrregexfilter(&self, attribute: &str, value: &str) -> anyhow::Result<TargetSet<Self::T>> {
        let regex = Regex::new(value)?;
        let filter = move |s: &'_ str| -> anyhow::Result<bool> { Ok(regex.is_match(s)?) };
        self.attrfilter(attribute, &filter)
    }

    fn filter_name(&self, regex: &str) -> anyhow::Result<TargetSet<Self::T>> {
        let re = Regex::new(regex)?;
        self.filter(|node| Ok(re.is_match(&node.node_ref().to_string())?))
    }

    fn kind(&self, regex: &str) -> anyhow::Result<TargetSet<Self::T>> {
        let re = Regex::new(regex)?;
        self.filter(|node| Ok(re.is_match(&node.rule_type())?))
    }

    fn intersect(&self, right: &TargetSet<Self::T>) -> anyhow::Result<TargetSet<Self::T>> {
        self.filter(|node| Ok(right.contains(node.node_ref())))
    }

    fn difference(&self, right: &TargetSet<Self::T>) -> anyhow::Result<TargetSet<Self::T>> {
        self.filter(|node| Ok(!right.contains(node.node_ref())))
    }
}

impl<T: QueryTarget> TargetSetExt for TargetSet<T> {
    type T = T;

    fn filter<F: Fn(&T) -> anyhow::Result<bool>>(&self, filter: F) -> anyhow::Result<TargetSet<T>> {
        TargetSet::filter(self, filter)
    }
}

impl<T: QueryTarget> Display for TargetSet<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt_container(f, "[", "]", self.targets.iter().map(|t| t.node_ref()))
    }
}


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

//! Implementation of the cli and query_* attr query language.

#![cfg(test)]

use std::borrow::Cow;

use async_trait::async_trait;
use buck2_core::build_file_path::BuildFilePath;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_query::query::environment::LabeledNode;
use buck2_query_parser::parse_expr;
use derive_more::Display;
use dupe::Dupe;
use serde::Serialize;
use serde::Serializer;

use crate::query::environment::QueryEnvironment;
use crate::query::environment::QueryTarget;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::evaluator::QueryEvaluator;
use crate::query::syntax::simple::eval::file_set::FileSet;
use crate::query::syntax::simple::eval::set::TargetSet;
use crate::query::syntax::simple::functions::DefaultQueryFunctionsModule;
use crate::query::traversal::AsyncTraversalDelegate;

#[derive(Clone, Hash, PartialEq, Eq, Debug, Display)]
struct TargetRef(String);

#[derive(Debug, Display, Serialize)]
struct TargetAttr(String);

#[derive(Debug, Clone, Dupe, Eq, PartialEq)]
struct Target {}

impl LabeledNode for Target {
    type NodeRef = TargetRef;

    fn node_ref(&self) -> &Self::NodeRef {
        unimplemented!()
    }
}

impl QueryTarget for Target {
    type Attr<'a> = TargetAttr;

    fn inputs_for_each<E, F: FnMut(CellPath) -> Result<(), E>>(&self, _func: F) -> Result<(), E> {
        unimplemented!()
    }

    fn rule_type(&self) -> Cow<str> {
        unimplemented!()
    }

    fn buildfile_path(&self) -> &BuildFilePath {
        unimplemented!()
    }

    fn deps<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Self::NodeRef> + Send + 'a> {
        unimplemented!()
    }

    fn exec_deps<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Self::NodeRef> + Send + 'a> {
        unimplemented!()
    }

    fn target_deps<'a>(&'a self) -> Box<dyn Iterator<Item = &'a Self::NodeRef> + Send + 'a> {
        unimplemented!()
    }

    fn special_attrs_for_each<E, F: FnMut(&str, &Self::Attr<'_>) -> Result<(), E>>(
        &self,
        _func: F,
    ) -> Result<(), E> {
        unimplemented!()
    }

    fn attr_any_matches(
        _attr: &Self::Attr<'_>,
        _filter: &dyn Fn(&str) -> anyhow::Result<bool>,
    ) -> anyhow::Result<bool> {
        unimplemented!()
    }

    fn attrs_for_each<E, F: FnMut(&str, &Self::Attr<'_>) -> Result<(), E>>(
        &self,
        _func: F,
    ) -> Result<(), E> {
        unimplemented!()
    }

    fn map_attr<R, F: FnMut(Option<&Self::Attr<'_>>) -> R>(&self, _key: &str, _func: F) -> R {
        unimplemented!()
    }

    fn call_stack(&self) -> Option<String> {
        None
    }

    fn attr_to_string_alternate(&self, _attr: &Self::Attr<'_>) -> String {
        unimplemented!("not needed for tests")
    }

    fn attr_serialize<S: Serializer>(
        &self,
        _attr: &Self::Attr<'_>,
        _serializer: S,
    ) -> Result<S::Ok, S::Error> {
        unimplemented!("not needed for tests")
    }
}

struct Env;
#[async_trait]
impl QueryEnvironment for Env {
    type Target = Target;

    async fn get_node(&self, _node_ref: &TargetRef) -> anyhow::Result<Self::Target> {
        unimplemented!()
    }

    async fn get_node_for_default_configured_target(
        &self,
        _node_ref: &TargetRef,
    ) -> anyhow::Result<MaybeCompatible<Self::Target>> {
        unimplemented!()
    }

    async fn eval_literals(&self, _literal: &[&str]) -> anyhow::Result<TargetSet<Self::Target>> {
        unimplemented!()
    }

    async fn eval_file_literal(&self, _literal: &str) -> anyhow::Result<FileSet> {
        unimplemented!()
    }

    async fn dfs_postorder(
        &self,
        _root: &TargetSet<Self::Target>,
        _delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
    ) -> anyhow::Result<()> {
        unimplemented!()
    }

    async fn depth_limited_traversal(
        &self,
        _root: &TargetSet<Self::Target>,
        _delegate: &mut dyn AsyncTraversalDelegate<Self::Target>,
        _depth: u32,
    ) -> anyhow::Result<()> {
        unimplemented!()
    }

    async fn owner(&self, _paths: &FileSet) -> anyhow::Result<TargetSet<Self::Target>> {
        unimplemented!()
    }
}

#[tokio::test]
pub async fn test_missing_arg() -> anyhow::Result<()> {
    let input = "kind(a, kind(a, kind()))";
    let parsed = parse_expr(input)?;
    match QueryEvaluator::new(&Env, &DefaultQueryFunctionsModule::new())
        .eval(&parsed)
        .await
    {
        Ok(_) => panic!(),
        Err(e) => {
            let err = QueryError::convert_error(e, input);
            let msg = format!("{:#}", err);
            let expected = "too few args. function `kind` requires at least ";
            if !msg.contains(expected) {
                return Err(err.context(format!("Expected error to contain `{}`", expected)));
            }
        }
    }
    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.
 */

//! Implementation of the cli and query_* attr query language.

use buck2_query_parser::spanned::Spanned;
use gazebo::variants::VariantName;

use crate::query::environment::QueryTarget;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::file_set::FileSet;
use crate::query::syntax::simple::eval::multi_query::MultiQueryResult;
use crate::query::syntax::simple::eval::set::TargetSet;

pub enum QueryEvaluationResult<T: QueryTarget> {
    Single(QueryEvaluationValue<T>),
    Multiple(MultiQueryResult<T>),
}

/// Used as a value in query evaluation, may appear in arguments to functions, results of functions etc.
#[derive(Debug, VariantName, Eq, PartialEq)]
pub enum QueryValue<T: QueryTarget> {
    String(String),
    Integer(u64),
    TargetSet(TargetSet<T>),
    FileSet(FileSet),
}

/// Used as a value in query evaluation where sets are valid, may appear in arguments to functions, results of functions etc.
#[derive(Debug)]
pub enum QueryValueSet<T: QueryTarget> {
    TargetSet(TargetSet<T>),
    FileSet(FileSet),
}

/// Used as the final result of evaluating a query. A literal at the top-level is treated specially and so this has
/// a more limited set of possibilities than a general QueryValue (for example `//foo/...` becomes a TargetSet in
/// `buck query //foo/...` rather than being a String).
#[derive(Debug, VariantName)]
pub enum QueryEvaluationValue<T: QueryTarget> {
    TargetSet(TargetSet<T>),
    FileSet(FileSet),
}

impl<T: QueryTarget> QueryEvaluationValue<T> {
    pub fn try_into_targets(self) -> anyhow::Result<TargetSet<T>> {
        match self {
            QueryEvaluationValue::TargetSet(targets) => Ok(targets),
            v => {
                return Err(QueryError::InvalidType {
                    expected: "targets",
                    actual: v.variant_name(),
                }
                .into());
            }
        }
    }
}

pub type QueryResult<T> = Result<Spanned<T>, Spanned<QueryError>>;

pub trait QueryResultExt<T> {
    fn into_anyhow(self, query: &str) -> anyhow::Result<T>;
}

impl<T> QueryResultExt<T> for QueryResult<T> {
    fn into_anyhow(self, query: &str) -> anyhow::Result<T> {
        match self {
            Ok(v) => Ok(v.value),
            Err(err) => Err(QueryError::convert_error(err, query)),
        }
    }
}

impl<T: QueryTarget> From<TargetSet<T>> for QueryValue<T> {
    fn from(v: TargetSet<T>) -> Self {
        QueryValue::TargetSet(v)
    }
}

impl<T: QueryTarget> From<FileSet> for QueryValue<T> {
    fn from(v: FileSet) -> Self {
        QueryValue::FileSet(v)
    }
}


/*
 * 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::fmt::Debug;
use std::marker::PhantomData;

use async_trait::async_trait;
use buck2_query_derive::query_module;
use gazebo::variants::VariantName;

use crate::query::environment::QueryEnvironment;
use crate::query::environment::QueryTarget;
use crate::query::environment::TraversalFilter;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::evaluator::QueryEvaluator;
use crate::query::syntax::simple::eval::set::TargetSet;
use crate::query::syntax::simple::eval::values::QueryEvaluationValue;
use crate::query::syntax::simple::eval::values::QueryValue;
use crate::query::syntax::simple::functions::helpers::CapturedExpr;
use crate::query::syntax::simple::functions::AugmentedQueryFunctions;
use crate::query::syntax::simple::functions::QueryFunctions;

pub(crate) struct DepsContextFunctions<'a, Env: QueryEnvironment> {
    target: &'a Env::Target,
}

impl<'a, Env: QueryEnvironment> Debug for DepsContextFunctions<'a, Env> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("DepsContextFunctions")
            .finish_non_exhaustive()
    }
}

#[query_module(Env)]
impl<'a, Env: QueryEnvironment> DepsContextFunctions<'a, Env> {
    async fn first_order_deps(&self, env: &Env) -> Result<QueryValue<Env::Target>, QueryError> {
        let mut deps = TargetSet::new();
        for dep in self.target.deps() {
            deps.insert(env.get_node(dep).await?);
        }
        Ok(QueryValue::TargetSet(deps))
    }

    async fn exec_deps(&self, env: &Env) -> Result<QueryValue<Env::Target>, QueryError> {
        let mut deps = TargetSet::new();
        for dep in self.target.exec_deps() {
            deps.insert(env.get_node(dep).await?);
        }
        Ok(QueryValue::TargetSet(deps))
    }

    async fn target_deps(&self, env: &Env) -> Result<QueryValue<Env::Target>, QueryError> {
        let mut deps = TargetSet::new();
        for dep in self.target.target_deps() {
            deps.insert(env.get_node(dep).await?);
        }
        Ok(QueryValue::TargetSet(deps))
    }
}

pub(crate) struct DepsFunction<Env: QueryEnvironment> {
    pub(crate) _marker: PhantomData<Env>,
}

impl<Env: QueryEnvironment> DepsFunction<Env> {
    pub(crate) async fn invoke_deps(
        &self,
        env: &Env,
        functions: &dyn QueryFunctions<Env = Env>,
        targets: &TargetSet<Env::Target>,
        depth: Option<i32>,
        captured_expr: Option<&CapturedExpr<'_>>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        let filter = match captured_expr {
            Some(expr) => {
                struct Filter<'a, Env: QueryEnvironment> {
                    inner_env: &'a Env,
                    functions: &'a dyn QueryFunctions<Env = Env>,
                    expr: &'a CapturedExpr<'a>,
                }

                #[async_trait]
                impl<'a, T: QueryTarget, Env: QueryEnvironment<Target = T>> TraversalFilter<T> for Filter<'a, Env> {
                    async fn get_children(&self, target: &T) -> anyhow::Result<TargetSet<T>> {
                        let augmented_functions = AugmentedQueryFunctions::augment(
                            self.functions,
                            Box::new(DepsContextFunctions { target }),
                        );
                        let evaluator = QueryEvaluator::new(self.inner_env, &augmented_functions);
                        match evaluator.eval_parsed_query(self.expr.expr).await {
                            Ok(v) => match v.value {
                                QueryEvaluationValue::TargetSet(v) => Ok(v),
                                v => Err(QueryError::InvalidType {
                                    expected: "targets",
                                    actual: v.variant_name(),
                                }
                                .into()),
                            },
                            Err(e) => Err(QueryError::drop_spans(e)),
                        }
                    }
                }

                Some(Filter {
                    inner_env: env,
                    functions,
                    expr,
                })
            }
            None => None,
        };

        let filter_ref = filter
            .as_ref()
            .map(|v| v as &dyn TraversalFilter<Env::Target>);

        env.deps(targets, depth, filter_ref).await
    }
}


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

use std::fmt::Write;

use enum_iterator::IntoEnumIterator;
use indexmap::IndexMap;
use itertools::Itertools;

use crate::query::syntax::simple::functions::helpers::QueryArgType;

pub struct MarkdownOptions {
    /// Whether to include alt text (used for expression type links).
    pub include_alt_text: bool,
}

// Instances created by #[query_module]
pub struct ArgDescription {
    pub name: String,
    pub repr_format: String,
    pub arg_type: QueryArgType,
}

impl ArgDescription {
    pub fn render_markdown(&self) -> String {
        format!(
            "{}: {}",
            &self.name,
            self.repr_format
                .replace("{}", &format!("*{}*", self.arg_type.repr()))
        )
    }
}

// Instances created by #[query_module]
pub struct FunctionDescription {
    pub name: &'static str,
    pub args: Vec<ArgDescription>,

    pub short_help: Option<String>,
    pub details: Option<String>,
}

impl FunctionDescription {
    pub fn render_short_markdown(&self) -> String {
        format!(
            " - {}({}){}",
            self.name,
            self.args.iter().map(|v| v.render_markdown()).join(", "),
            &match &self.short_help {
                None => "".to_owned(),
                Some(v) => format!(": {}", v),
            }
        )
    }

    pub fn render_markdown(&self) -> String {
        let mut rendered = format!(
            "### {}({})\n\n",
            self.name,
            self.args.iter().map(|v| v.render_markdown()).join(", ")
        );
        if let Some(v) = &self.short_help {
            writeln!(rendered, "{}\n", v).unwrap();
        }

        if let Some(v) = &self.details {
            writeln!(rendered, "{}\n", v).unwrap();
        }
        rendered
    }
}

// Instances created by #[query_module]
pub struct ModuleDescription {
    pub functions: IndexMap<&'static str, FunctionDescription>,

    pub short_help: Option<String>,
    pub details: Option<String>,
}

impl ModuleDescription {
    pub fn render_markdown(&self) -> String {
        let mut rendered = format!(
            "## {}\n\n",
            match &self.short_help {
                Some(v) => v,
                None => "Query functions",
            }
        );
        if let Some(v) = &self.details {
            writeln!(rendered, "{}\n", v).unwrap();
        }

        for (_, func) in &self.functions {
            writeln!(rendered, "{}\n", func.render_short_markdown()).unwrap();
        }

        for (_, func) in &self.functions {
            writeln!(rendered, "{}\n", func.render_markdown()).unwrap();
        }

        rendered
    }
}

pub struct QueryEnvironmentDescription {
    pub name: String,
    pub mods: Vec<ModuleDescription>,
}

impl QueryEnvironmentDescription {
    pub fn render_markdown(&self, options: &MarkdownOptions) -> String {
        format!(
            indoc::indoc! {r#"
            # {}

            ## Query Value Types
            {}

            {}
            "#},
            &self.name,
            QueryArgType::into_enum_iter()
                .map(|v| render_arg_type_markdown(v, options))
                .join("\n\n"),
            self.mods.iter().map(|v| v.render_markdown()).join("\n\n")
        )
    }
}

fn render_arg_type_markdown(v: QueryArgType, options: &MarkdownOptions) -> String {
    let mut rendered = if options.include_alt_text {
        format!("*[{}]: {}\n", v.repr(), v.short_description(),)
    } else {
        String::new()
    };
    rendered.push_str(&format!("- *{}*: {}", v.repr(), v.description()));
    rendered
}


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

//! Provides utilities for defining query functions.
//!
//! To allow a new type to be used as an argument, implement the [QueryFunctionArg] trait for that type.

use async_trait::async_trait;
use buck2_query_parser::spanned::Spanned;
use buck2_query_parser::Expr;
use buck2_query_parser::SpannedExpr;
use enum_iterator::IntoEnumIterator;
use gazebo::variants::VariantName;

use crate::query::environment::QueryEnvironment;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::evaluator::QueryEvaluator;
use crate::query::syntax::simple::eval::file_set::FileSet;
use crate::query::syntax::simple::eval::set::TargetSet;
use crate::query::syntax::simple::eval::values::QueryValue;
use crate::query::syntax::simple::eval::values::QueryValueSet;

// IntoEnumIterator triggers this lint :/
#[cfg_attr(feature = "gazebo_lint", allow(gazebo_lint_impl_dupe))]
mod _scoped_allow {
    use dupe::Dupe;

    use super::*;

    #[derive(Debug, Clone, Dupe, Copy, IntoEnumIterator)]
    pub enum QueryArgType {
        String,
        Integer,
        TargetSet,
        FileSet,
        Set,
        Expression,
        Value,
    }
}

pub use _scoped_allow::*;

impl QueryArgType {
    pub fn repr(self) -> &'static str {
        match self {
            QueryArgType::String => "string",
            QueryArgType::Integer => "integer",
            QueryArgType::TargetSet => "target expression",
            QueryArgType::FileSet => "file expression",
            QueryArgType::Set => "target or file expression",
            QueryArgType::Expression => "query expression",
            QueryArgType::Value => "any value",
        }
    }

    pub fn short_description(self) -> &'static str {
        match self {
            QueryArgType::TargetSet => {
                "a target expression, either a literal or the return value of a function"
            }
            QueryArgType::FileSet => {
                "a file expression, either a literal or the return value of a function"
            }
            QueryArgType::Set => {
                "a file or target expression, either a literal or the return value of a function"
            }
            QueryArgType::Expression => {
                "a valid query expression, evaluated in a function-specific context"
            }
            _ => self.description(),
        }
    }

    pub fn description(self) -> &'static str {
        match self {
            QueryArgType::String => {
                "a string. For example `non_quoted_string` or `\"quoted string\"`"
            }
            QueryArgType::Integer => {
                "an integer. query integers must be positive and fit in a `u32`"
            }
            QueryArgType::TargetSet => {
                "A target set expression. This could be a literal build target (`\"cell//some:target\"`), a literal build target pattern \
                (`\"cell//package:\"` or `\"cell//recursive/...\"`) or the result of another function that returns a target set. For \
                queries in cli commands (like `buck2 query`), literals can be relative to the current working dir (like `some:target` \
                or `...`)."
            }
            QueryArgType::FileSet => {
                "A file set expression. This could be a file literal like `path/to/a.file` or the return value of a function that \
                returns files (for example, the `buildfile()` function)."
            }
            QueryArgType::Set => {
                "A file set or target set expression. This could be a literal like `path/to/a.file` or `\"cell//some:target\"`,
                or the return value of a function that returns files or targets."
            }
            QueryArgType::Expression => {
                "A query expression. This is used for functions that capture an expression and evaluate it in another context. \
                For example, the `deps()` function can accept an expression that it uses to find the children of a node to \
                customize the deps traversal."
            }
            QueryArgType::Value => "Any query value.",
        }
    }
}

#[async_trait]
pub trait QueryFunction<Env: QueryEnvironment>: Send + Sync {
    fn name(&self) -> &'static str;

    async fn invoke(
        &self,
        evaluator: &QueryEvaluator<Env>,
        args: &[SpannedExpr<'_>],
    ) -> Result<QueryValue<Env::Target>, QueryError>;

    fn arg_type(&self, idx: usize) -> Result<QueryArgType, QueryError>;
}

#[async_trait]
pub trait QueryBinaryOp<Env: QueryEnvironment>: Send + Sync {
    fn name(&self) -> &'static str;

    async fn invoke(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError>;
}

#[async_trait]
pub trait QueryFunctionArg<'a, Env: QueryEnvironment>: Sized + Send + Sync {
    const ARG_TYPE: QueryArgType;

    fn describe_format() -> String {
        "{}".to_owned()
    }

    /// accept_none() is called when no arg is provided. This is used for optional args.
    fn accept_none() -> Option<Self> {
        None
    }

    async fn eval(
        evaluator: &QueryEvaluator<'a, Env>,
        expr: &'a Spanned<Expr<'a>>,
    ) -> Result<Self, QueryError> {
        let result = evaluator.eval(expr).await?;
        Self::accept(evaluator.env(), result.value).await
    }

    /// accept() will be called with the value of the evaluated expression for an arg. It
    /// should perform type checking and return the typed value or an error.
    async fn accept(env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError>;
}

pub struct CapturedExpr<'a> {
    pub expr: &'a Spanned<Expr<'a>>,
}

/// Allows a "captured expression" as a function argument. This is used for a function like
/// `deps()` that uses an argument similar to a lambda where it will evaluate the expression
/// itself in some other context.
#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for CapturedExpr<'a> {
    const ARG_TYPE: QueryArgType = QueryArgType::Expression;

    async fn eval(
        _evaluator: &QueryEvaluator<'a, Env>,
        expr: &'a Spanned<Expr<'a>>,
    ) -> Result<Self, QueryError> {
        Ok(CapturedExpr { expr })
    }

    async fn accept(_env: &Env, _val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        unimplemented!("accept shouldn't be called since eval doesn't use it")
    }
}

#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for QueryValue<Env::Target> {
    const ARG_TYPE: QueryArgType = QueryArgType::Value;

    async fn accept(_env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        Ok(val)
    }
}

#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for QueryValueSet<Env::Target> {
    const ARG_TYPE: QueryArgType = QueryArgType::Set;

    async fn accept(env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        match val {
            QueryValue::String(s) => Ok(QueryValueSet::TargetSet(env.eval_literals(&[&s]).await?)),
            QueryValue::TargetSet(x) => Ok(QueryValueSet::TargetSet(x)),
            QueryValue::FileSet(x) => Ok(QueryValueSet::FileSet(x)),
            _ => Err(QueryError::InvalidType {
                expected: "file or target set",
                actual: val.variant_name(),
            }),
        }
    }
}

/// Straightforward implementation for String.
#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for String {
    const ARG_TYPE: QueryArgType = QueryArgType::String;

    async fn accept(_env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        match val {
            QueryValue::String(s) => Ok(s),
            _ => Err(QueryError::InvalidType {
                expected: "string",
                actual: val.variant_name(),
            }),
        }
    }
}

/// Straightforward implementation for TargetSet.
#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for TargetSet<Env::Target> {
    const ARG_TYPE: QueryArgType = QueryArgType::TargetSet;

    async fn accept(env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        match val {
            QueryValue::String(s) => Ok(env.eval_literals(&[&s]).await?),
            QueryValue::TargetSet(t) => Ok(t),
            _ => Err(QueryError::InvalidType {
                expected: "target_set",
                actual: val.variant_name(),
            }),
        }
    }
}

/// Straightforward implementation for FileSet.
#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for FileSet {
    const ARG_TYPE: QueryArgType = QueryArgType::FileSet;

    async fn accept(env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        match val {
            QueryValue::String(v) => Ok(env.eval_file_literal(&v).await?),
            QueryValue::FileSet(t) => Ok(t),
            _ => Err(QueryError::InvalidType {
                expected: "target_set",
                actual: val.variant_name(),
            }),
        }
    }
}

/// Straightforward implementation for u64.
#[async_trait]
impl<'a, Env: QueryEnvironment> QueryFunctionArg<'a, Env> for u64 {
    const ARG_TYPE: QueryArgType = QueryArgType::Integer;

    async fn accept(_env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        match val {
            QueryValue::Integer(v) => Ok(v),
            _ => Err(QueryError::InvalidType {
                expected: "uint",
                actual: val.variant_name(),
            }),
        }
    }
}

/// `Option<T>` is used for optional args and so this is the (only?) implementation that supports a "none" arg.
#[async_trait]
impl<'a, Env: QueryEnvironment, A: QueryFunctionArg<'a, Env>> QueryFunctionArg<'a, Env>
    for Option<A>
{
    const ARG_TYPE: QueryArgType = A::ARG_TYPE;

    fn accept_none() -> Option<Self> {
        Some(None)
    }

    async fn accept(env: &Env, val: QueryValue<Env::Target>) -> Result<Self, QueryError> {
        Ok(Some(A::accept(env, val).await?))
    }

    async fn eval(
        evaluator: &QueryEvaluator<'a, Env>,
        expr: &'a Spanned<Expr<'a>>,
    ) -> Result<Self, QueryError> {
        Ok(Some(A::eval(evaluator, expr).await?))
    }

    fn describe_format() -> String {
        "?{}".to_owned()
    }
}

// Helper for buck_query_proc_macro implementations. Evaluates an arg at an index and tries to convert it to the QueryFunctionArg type, providing decent errors on failures.
pub async fn eval_arg<'a, Env: QueryEnvironment, A: QueryFunctionArg<'a, Env>>(
    func_name: &str,
    evaluator: &QueryEvaluator<'a, Env>,
    args: &'a [Spanned<Expr<'a>>],
    idx: usize,
) -> Result<A, QueryError> {
    match args.get(idx) {
        Some(v) => A::eval(evaluator, v).await,
        None => match A::accept_none() {
            Some(v) => Ok(v),
            None => Err(QueryError::TooFewArgs {
                function: func_name.to_owned(),
                min: idx + 1,
                actual: args.len(),
            }),
        },
    }
}


/*
 * 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::fmt::Debug;
use std::marker::PhantomData;

use allocative::Allocative;
use async_trait::async_trait;
use buck2_core::configuration::compatibility::MaybeCompatible;
use buck2_query_derive::query_module;
use buck2_query_parser::spanned::Spanned;
use buck2_query_parser::BinaryOp;
use buck2_query_parser::Expr;
use gazebo::variants::VariantName;

use crate::query::environment::QueryEnvironment;
use crate::query::syntax::simple::eval::error::QueryError;
use crate::query::syntax::simple::eval::evaluator::QueryEvaluator;
use crate::query::syntax::simple::eval::file_set::FileSet;
use crate::query::syntax::simple::eval::set::TargetSet;
use crate::query::syntax::simple::eval::set::TargetSetExt;
use crate::query::syntax::simple::eval::values::QueryResult;
use crate::query::syntax::simple::eval::values::QueryValue;
use crate::query::syntax::simple::eval::values::QueryValueSet;
use crate::query::syntax::simple::functions::deps::DepsFunction;
use crate::query::syntax::simple::functions::docs::ModuleDescription;
use crate::query::syntax::simple::functions::helpers::CapturedExpr;
use crate::query::syntax::simple::functions::helpers::QueryArgType;
use crate::query::syntax::simple::functions::helpers::QueryBinaryOp;
use crate::query::syntax::simple::functions::helpers::QueryFunction;

pub mod deps;
pub mod docs;
pub mod helpers;

pub trait QueryLiteralVisitor {
    fn target_pattern(&mut self, pattern: &str) -> anyhow::Result<()>;
}

pub trait HasModuleDescription {
    fn describe() -> ModuleDescription;
}

#[async_trait]
pub trait QueryFunctions: Debug + Send + Sync {
    type Env: QueryEnvironment;

    fn get(&self, name: &str) -> Option<&dyn QueryFunction<Self::Env>>;

    fn get_op(&self, op: BinaryOp) -> Option<&dyn QueryBinaryOp<Self::Env>>;
}

pub trait QueryFunctionsVisitLiterals: Debug + Send + Sync {
    fn visit_literals(
        &self,
        visitor: &mut dyn QueryLiteralVisitor,
        expr: &Spanned<Expr>,
    ) -> QueryResult<()>;
}

impl<F: QueryFunctions> QueryFunctionsVisitLiterals for F {
    fn visit_literals(
        &self,
        visitor: &mut dyn QueryLiteralVisitor,
        expr: &Spanned<Expr>,
    ) -> QueryResult<()> {
        fn visit_literals_recurse<F: QueryFunctions>(
            this: &F,
            visitor: &mut dyn QueryLiteralVisitor,
            expr: &Expr,
        ) -> Result<(), QueryError> {
            match expr {
                Expr::Function {
                    function_name,
                    args,
                } => match this.get(function_name) {
                    Some(func) => {
                        for (i, arg) in args.iter().enumerate() {
                            visit_literals_item(
                                this,
                                visitor,
                                arg,
                                matches!(
                                    func.arg_type(i)?,
                                    QueryArgType::TargetSet
                                        | QueryArgType::Set
                                        | QueryArgType::Value
                                ),
                            )?;
                        }
                        Ok(())
                    }
                    None => Err(QueryError::UnknownFunction(
                        (*function_name.fragment()).to_owned(),
                    )),
                },
                Expr::BinaryOpSequence(left, exprs) => {
                    visit_literals_item(this, visitor, left, true)?;
                    // All binary ops are on targetsets currently.
                    for (_, right) in exprs {
                        visit_literals_item(this, visitor, right, true)?;
                    }
                    Ok(())
                }
                Expr::Set(args) => {
                    for arg in args {
                        visitor.target_pattern(arg)?;
                    }
                    Ok(())
                }
                Expr::FileSet(_args) => Ok(()),
                Expr::String(..) | Expr::Integer(..) => {
                    panic!(
                        "This shouldn't be called with literals, they should be handled in the caller"
                    )
                }
            }
        }

        fn visit_literals_item<F: QueryFunctions>(
            this: &F,
            visitor: &mut dyn QueryLiteralVisitor,
            expr: &Spanned<Expr>,
            is_target_expr: bool,
        ) -> QueryResult<()> {
            expr.map_res(|value| -> Result<(), QueryError> {
                match value {
                    Expr::String(val) => {
                        if is_target_expr {
                            visitor.target_pattern(val)?;
                        }
                    }
                    Expr::Integer(..) => {
                        // ignored
                    }
                    _ => visit_literals_recurse(this, visitor, value)?,
                }
                Ok(())
            })
        }

        visit_literals_item(self, visitor, expr, true)
    }
}

#[derive(Allocative)]
#[allocative(bound = "")]
pub struct DefaultQueryFunctionsModule<Env: QueryEnvironment> {
    implementation: DefaultQueryFunctions<Env>,
}

impl<Env: QueryEnvironment> Debug for DefaultQueryFunctionsModule<Env> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("DefaultQueryFunctionsModule")
            .finish_non_exhaustive()
    }
}

impl<Env: QueryEnvironment> DefaultQueryFunctionsModule<Env> {
    pub fn new() -> Self {
        Self {
            implementation: DefaultQueryFunctions::new(),
        }
    }
}

type QueryFuncResult<Env> =
    std::result::Result<QueryValue<<Env as QueryEnvironment>::Target>, QueryError>;

async fn accept_target_set<Env: QueryEnvironment>(
    env: &Env,
    val: QueryValue<Env::Target>,
) -> Result<TargetSet<Env::Target>, QueryError> {
    match val {
        QueryValue::TargetSet(x) => Ok(x),
        QueryValue::String(literal) => Ok(env.eval_literals(&[&literal]).await?),
        _ => Err(QueryError::InvalidType {
            expected: "target_set",
            actual: val.variant_name(),
        }),
    }
}

/// Common query functions
#[query_module(Env)]
impl<Env: QueryEnvironment> DefaultQueryFunctionsModule<Env> {
    /// Computes all dependency paths.
    ///
    /// The `allpaths(from, to)` function evaluates to the graph formed by paths between the target expressions from and to, following the dependencies between nodes. For example, the value of
    /// `buck query "allpaths('//foo:bar', '//foo/bar/lib:baz')"`
    /// is the dependency graph rooted at the single target node `//foo:bar`, that includes all target nodes that depend (transitively) on `//foo/bar/lib:baz`.
    ///
    /// The two arguments to `allpaths()` can themselves be expressions. For example, the command:
    /// `buck query "allpaths(kind(java_library, '//...'), '//foo:bar')"`
    /// shows all the paths between any java_library in the repository and the target `//foo:bar`.
    ///
    /// We recommend using `allpaths()` with the `--output-format=dot` parameter to generate a graphviz file that can then be rendered as an image. For example:
    ///
    /// ```ignore
    /// $ buck query "allpaths('//foo:bar', '//foo/bar/lib:baz')" --output-format=dot --output-file=result.dot
    /// $ dot -Tpng result.dot -o image.png
    /// ```
    ///
    /// Graphviz is an open-source graph-visualization software tool. Graphviz uses the dot language to describe graphs.
    async fn allpaths(
        &self,
        env: &Env,
        from: TargetSet<Env::Target>,
        to: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self.implementation.allpaths(env, &from, &to).await?.into())
    }

    async fn somepath(
        &self,
        env: &Env,
        from: TargetSet<Env::Target>,
        to: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self.implementation.somepath(env, &from, &to).await?.into())
    }

    async fn attrfilter(
        &self,
        attr: String,
        value: String,
        targets: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .attrfilter(&attr, &value, &targets)?
            .into())
    }

    async fn nattrfilter(
        &self,
        attr: String,
        value: String,
        targets: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .nattrfilter(&attr, &value, &targets)?
            .into())
    }

    async fn attrregexfilter(
        &self,
        attr: String,
        value: String,
        targets: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .attrregexfilter(&attr, &value, &targets)?
            .into())
    }

    async fn buildfile(&self, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
        Ok(self.implementation.buildfile(&targets).into())
    }

    async fn rbuildfiles(
        &self,
        env: &Env,
        universe: FileSet,
        argset: FileSet,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .rbuildfiles(env, &universe, &argset)
            .await?
            .into())
    }

    async fn allbuildfiles(
        &self,
        env: &Env,
        universe: TargetSet<Env::Target>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .allbuildfiles(env, &universe)
            .await?
            .into())
    }

    async fn deps(
        &self,
        evaluator: &QueryEvaluator<'_, Env>,
        targets: TargetSet<Env::Target>,
        depth: Option<u64>,
        captured_expr: Option<CapturedExpr<'_>>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .deps(
                evaluator.env(),
                evaluator.functions(),
                &targets,
                depth.map(|v| v as i32),
                captured_expr.as_ref(),
            )
            .await?
            .into())
    }

    async fn filter(&self, regex: String, set: QueryValueSet<Env::Target>) -> QueryFuncResult<Env> {
        match set {
            QueryValueSet::TargetSet(targets) => Ok(self
                .implementation
                .filter_target_set(&regex, &targets)?
                .into()),
            QueryValueSet::FileSet(files) => {
                Ok(self.implementation.filter_file_set(&regex, &files)?.into())
            }
        }
    }

    async fn inputs(&self, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
        Ok(self.implementation.inputs(&targets)?.into())
    }

    async fn kind(&self, regex: String, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
        Ok(targets.kind(&regex)?.into())
    }

    async fn labels(&self, attr: String, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
        self.implementation.labels(&attr, &targets)
    }

    async fn owner(&self, env: &Env, files: FileSet) -> QueryFuncResult<Env> {
        Ok(self.implementation.owner(env, &files).await?.into())
    }

    async fn rdeps(
        &self,
        env: &Env,
        universe: TargetSet<Env::Target>,
        targets: TargetSet<Env::Target>,
        depth: Option<u64>,
    ) -> QueryFuncResult<Env> {
        Ok(self
            .implementation
            .rdeps(env, &universe, &targets, depth.map(|v| v as i32))
            .await?
            .into())
    }

    async fn testsof(&self, env: &Env, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
        Ok(self.implementation.testsof(env, &targets).await?.into())
    }

    // These three functions are intentionally implemented as errors. They are only available within the context
    // of a deps functions 3rd parameter expr. When used in that context, the QueryFunctions will be augmented to
    // have non-erroring implementations.
    async fn first_order_deps(&self) -> QueryFuncResult<Env> {
        self.implementation.first_order_deps()
    }
    async fn target_deps(&self) -> QueryFuncResult<Env> {
        Err(QueryError::NotAvailableInContext("target_deps"))
    }
    async fn exec_deps(&self) -> QueryFuncResult<Env> {
        Err(QueryError::NotAvailableInContext("exec_deps"))
    }

    #[binary_op(BinaryOp::Intersect)]
    async fn intersect(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        self.implementation.intersect(env, left, right).await
    }

    #[binary_op(BinaryOp::Except)]
    async fn except(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        self.implementation.except(env, left, right).await
    }

    #[binary_op(BinaryOp::Union)]
    async fn union(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        self.implementation.union(env, left, right).await
    }
}

#[derive(Allocative)]
#[allocative(bound = "")]
pub struct DefaultQueryFunctions<Env: QueryEnvironment> {
    _marker: std::marker::PhantomData<Env>,
}

impl<Env: QueryEnvironment> DefaultQueryFunctions<Env> {
    pub fn new() -> Self {
        Self {
            _marker: PhantomData,
        }
    }
}

impl<Env: QueryEnvironment> DefaultQueryFunctions<Env> {
    pub async fn allpaths(
        &self,
        env: &Env,
        from: &TargetSet<Env::Target>,
        to: &TargetSet<Env::Target>,
    ) -> Result<TargetSet<Env::Target>, QueryError> {
        Ok(env.allpaths(from, to).await?)
    }

    pub async fn somepath(
        &self,
        env: &Env,
        from: &TargetSet<Env::Target>,
        to: &TargetSet<Env::Target>,
    ) -> Result<TargetSet<Env::Target>, QueryError> {
        Ok(env.somepath(from, to).await?)
    }

    pub fn attrfilter(
        &self,
        attr: &str,
        value: &str,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        targets.attrfilter(attr, &|v| Ok(v == value))
    }

    pub fn nattrfilter(
        &self,
        attr: &str,
        value: &str,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        targets.nattrfilter(attr, &|v| Ok(v == value))
    }

    pub fn attrregexfilter(
        &self,
        attr: &str,
        value: &str,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        targets.attrregexfilter(attr, value)
    }

    pub fn buildfile(&self, targets: &TargetSet<Env::Target>) -> FileSet {
        targets.buildfile()
    }

    pub async fn allbuildfiles(
        &self,
        env: &Env,
        universe: &TargetSet<Env::Target>,
    ) -> anyhow::Result<FileSet> {
        env.allbuildfiles(universe).await
    }

    pub async fn rbuildfiles(
        &self,
        env: &Env,
        universe: &FileSet,
        argset: &FileSet,
    ) -> anyhow::Result<FileSet> {
        env.rbuildfiles(universe, argset).await
    }

    pub async fn deps(
        &self,
        env: &Env,
        functions: &dyn QueryFunctions<Env = Env>,
        targets: &TargetSet<Env::Target>,
        depth: Option<i32>,
        captured_expr: Option<&CapturedExpr<'_>>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        DepsFunction::<Env> {
            _marker: PhantomData,
        }
        .invoke_deps(env, functions, targets, depth, captured_expr)
        .await
    }

    pub fn filter_target_set(
        &self,
        regex: &str,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        targets.filter_name(regex)
    }

    pub fn filter_file_set(&self, regex: &str, files: &FileSet) -> anyhow::Result<FileSet> {
        files.filter_name(regex)
    }

    pub fn inputs(&self, targets: &TargetSet<Env::Target>) -> anyhow::Result<FileSet> {
        targets.inputs()
    }

    pub fn labels(
        &self,
        _attr: &str,
        _targets: &TargetSet<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        Err(QueryError::FunctionUnimplemented("labels"))
    }

    pub async fn owner(
        &self,
        env: &Env,
        files: &FileSet,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        env.owner(files).await
    }

    pub async fn rdeps(
        &self,
        env: &Env,
        universe: &TargetSet<Env::Target>,
        targets: &TargetSet<Env::Target>,
        depth: Option<i32>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        env.rdeps(universe, targets, depth).await
    }

    pub async fn testsof(
        &self,
        env: &Env,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<TargetSet<Env::Target>> {
        env.testsof(targets).await
    }

    pub async fn testsof_with_default_target_platform(
        &self,
        env: &Env,
        targets: &TargetSet<Env::Target>,
    ) -> anyhow::Result<Vec<MaybeCompatible<Env::Target>>> {
        env.testsof_with_default_target_platform(targets).await
    }

    // These three functions are intentionally implemented as errors. They are only available within the context
    // of a deps functions 3rd parameter expr. When used in that context, the QueryFunctions will be augmented to
    // have non-erroring implementations.
    pub fn first_order_deps(&self) -> QueryFuncResult<Env> {
        Err(QueryError::NotAvailableInContext("first_order_deps"))
    }
    pub fn target_deps(&self) -> QueryFuncResult<Env> {
        Err(QueryError::NotAvailableInContext("target_deps"))
    }
    pub fn exec_deps(&self) -> QueryFuncResult<Env> {
        Err(QueryError::NotAvailableInContext("exec_deps"))
    }

    pub async fn intersect(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        let left = accept_target_set(env, left).await?;
        let right = accept_target_set(env, right).await?;
        Ok(QueryValue::TargetSet(left.intersect(&right)?))
    }

    pub async fn except(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        let left = accept_target_set(env, left).await?;
        let right = accept_target_set(env, right).await?;
        Ok(QueryValue::TargetSet(left.difference(&right)?))
    }

    pub async fn union(
        &self,
        env: &Env,
        left: QueryValue<Env::Target>,
        right: QueryValue<Env::Target>,
    ) -> Result<QueryValue<Env::Target>, QueryError> {
        // If the operations are of the same type, which + join them.
        // If one is a string, and the other a FileSet or TargetSet, we can promote the string
        match (left, right) {
            (QueryValue::TargetSet(l), QueryValue::TargetSet(r)) => {
                Ok(QueryValue::TargetSet(l.union(&r)))
            }
            (QueryValue::String(l), QueryValue::TargetSet(r)) => {
                let l = env.eval_literals(&[&l]).await?;
                Ok(QueryValue::TargetSet(l.union(&r)))
            }
            (QueryValue::TargetSet(l), QueryValue::String(r)) => {
                let r = env.eval_literals(&[&r]).await?;
                Ok(QueryValue::TargetSet(l.union(&r)))
            }
            (QueryValue::String(l), QueryValue::String(r)) => {
                // Important that String + treats both as target literals, since that's what
                // buck1 does - we blur the lines between string and targetset
                Ok(QueryValue::TargetSet(env.eval_literals(&[&l, &r]).await?))
            }
            (QueryValue::FileSet(l), QueryValue::FileSet(r)) => {
                Ok(QueryValue::FileSet(l.union(&r)))
            }
            (QueryValue::String(l), QueryValue::FileSet(r)) => {
                let l = env.eval_file_literal(&l).await?;
                Ok(QueryValue::FileSet(l.union(&r)))
            }
            (QueryValue::FileSet(l), QueryValue::String(r)) => {
                let r = env.eval_file_literal(&r).await?;
                Ok(QueryValue::FileSet(l.union(&r)))
            }
            (left, right) => Err(QueryError::UnionIncompatibleTypes(
                left.variant_name(),
                right.variant_name(),
            )),
        }
    }
}

pub struct AugmentedQueryFunctions<'a, Env: QueryEnvironment> {
    inner: &'a dyn QueryFunctions<Env = Env>,
    extra: Box<dyn QueryFunctions<Env = Env> + 'a>,
}

impl<'a, Env: QueryEnvironment> Debug for AugmentedQueryFunctions<'a, Env> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("AugmentedQueryFunctions")
            .finish_non_exhaustive()
    }
}

impl<'a, Env: QueryEnvironment> AugmentedQueryFunctions<'a, Env> {
    pub fn augment(
        inner: &'a dyn QueryFunctions<Env = Env>,
        extra: Box<dyn QueryFunctions<Env = Env> + 'a>,
    ) -> Self {
        Self { inner, extra }
    }
}

impl<'a, Env: QueryEnvironment> QueryFunctions for AugmentedQueryFunctions<'a, Env> {
    type Env = Env;
    fn get(&self, name: &str) -> Option<&dyn QueryFunction<Env>> {
        match self.extra.get(name) {
            None => self.inner.get(name),
            Some(v) => Some(v),
        }
    }

    fn get_op(&self, op: BinaryOp) -> Option<&dyn QueryBinaryOp<Env>> {
        match self.extra.get_op(op) {
            None => self.inner.get_op(op),
            Some(v) => Some(v),
        }
    }
}


/*
 * 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 proc_macro2::Span;
use proc_macro2::TokenStream;
use quote::quote;
use quote::quote_spanned;
use syn::Ident;
use syn::Path;
use syn::Type;
use syn::TypePath;
use syn::TypeReference;

use crate::parse::Argument;
use crate::parse::DocString;
use crate::parse::Method;
use crate::parse::Parsed;

struct MethodCodegen {
    method_def: TokenStream,
    method_describe: TokenStream,
    method_dispatch: Option<TokenStream>,
    op_dispatch: Option<TokenStream>,
}

enum ContextArgument {
    None,
    Evaluator,
    Environment,
}

fn is_type_name(x: &Type, name: &str) -> bool {
    if let Type::Path(TypePath {
        path: Path { segments, .. },
        ..
    }) = x
    {
        if let Some(seg1) = segments.last() {
            return seg1.ident == name;
        }
    }
    false
}

fn is_ref_ctx_type_name(x: &Type, name: &str) -> syn::Result<bool> {
    match x {
        Type::Reference(TypeReference {
            mutability: Some(_),
            ..
        }) => Err(syn::Error::new_spanned(
            x,
            "#[query_module] context args (QueryEvaluator/QueryEnvironment) should not be mutable",
        )),
        Type::Reference(TypeReference {
            mutability: None,
            elem: x,
            ..
        }) => Ok(is_type_name(x, name)),
        _ if is_type_name(x, name) => Err(syn::Error::new_spanned(
            x,
            "#[query_module] context args should be taken by reference, not value",
        )),
        _ => Ok(false),
    }
}

fn is_evaluator_arg(ty: &Type) -> syn::Result<bool> {
    is_ref_ctx_type_name(ty, "QueryEvaluator")
}

fn is_env_arg(ty: &Type, env_ty: &Type) -> syn::Result<bool> {
    match ty {
        Type::Reference(TypeReference {
            mutability: Some(_),
            ..
        }) => Err(syn::Error::new_spanned(
            ty,
            "#[query_module] context args (QueryEvaluator/QueryEnvironment) should not be mutable",
        )),
        Type::Reference(TypeReference {
            mutability: None,
            elem: x,
            ..
        }) => Ok(env_ty == &**x),
        _ if env_ty == ty => Err(syn::Error::new_spanned(
            ty,
            "#[query_module] context args should be taken by reference, not value",
        )),
        _ => Ok(false),
    }
}

fn ctx_arg_for(arg: &Argument, env_ty: &Type) -> syn::Result<ContextArgument> {
    Ok(match arg {
        Argument { ty, .. } if is_evaluator_arg(ty)? => ContextArgument::Evaluator,
        Argument { ty, .. } if is_env_arg(ty, env_ty)? => ContextArgument::Environment,
        _ => ContextArgument::None,
    })
}

fn analyze_args<'a>(
    method: &'a Method,
    env_ty: &Type,
) -> syn::Result<(ContextArgument, Vec<&'a Argument>)> {
    let mut args = method.args.iter().peekable();
    let ctx_arg = match args.peek() {
        Some(v) => ctx_arg_for(v, env_ty)?,
        None => ContextArgument::None,
    };
    match ctx_arg {
        ContextArgument::None => {}
        _ => {
            args.next();
        }
    };

    Ok((ctx_arg, args.map(|a| match ctx_arg_for(a, env_ty)? {
        ContextArgument::None => {Ok(a)},
        _ => {
            Err(syn::Error::new(a.span, "context arguments (QueryEvaluator/QueryEnvironment) can only be the first non-receiver arg"))
        }
    }).collect::<syn::Result<Vec<_>>>()?))
}

fn gen_for_docstring(span: Span, docs: Option<&DocString>) -> TokenStream {
    match &docs {
        Some(DocString {
            short_help,
            details,
        }) => quote_spanned!(span =>
            short_help: Some(#short_help.to_owned()),
            details: Some(#details.to_owned())
        ),
        None => quote_spanned!(span =>
            short_help: None,
            details: None
        ),
    }
}

fn gen_for_method(parsed: &Parsed, method: &Method) -> syn::Result<MethodCodegen> {
    let func_ty = Ident::new(
        &format!("__{}QueryFunction", &method.name),
        method.name.span(),
    );

    let func_ident = &method.name;
    let func_ident_str = func_ident.to_string();
    let method_doc = gen_for_docstring(method.name.span(), method.docs.as_ref());
    let env_ident = &parsed.args.env_ty;
    let env_target = quote! {<#env_ident as QueryEnvironment>::Target};
    let self_ty = &parsed.module.self_ty;
    let (impl_generics, ty_generics, where_clause) = parsed.module.generics.split_for_impl();

    let (ctx_arg, value_args) = analyze_args(method, &parsed.args.env_ty)?;

    let pass_ctx = match ctx_arg {
        ContextArgument::None => quote! {},
        ContextArgument::Evaluator => quote! {evaluator,},
        ContextArgument::Environment => quote! {evaluator.env(),},
    };

    let mut describe_args = Vec::new();
    let mut pass_args = Vec::new();
    let mut arg_type_match = Vec::new();
    for (i, arg) in value_args.iter().enumerate() {
        let arg_type = &arg.ty;
        let as_arg_type = quote! {<#arg_type as QueryFunctionArg<'_, #env_ident>>};
        let arg_name = arg.name.to_string();
        arg_type_match.push(quote_spanned!(arg.span => #i => Ok(#as_arg_type::ARG_TYPE)));
        pass_args
            .push(quote_spanned!(arg.span => eval_arg(self.name(), evaluator, args, #i).await?));
        describe_args.push(quote_spanned!(arg.span => ArgDescription {
            name: #arg_name.to_owned(),
            repr_format: #as_arg_type::describe_format(),
            arg_type: #as_arg_type::ARG_TYPE
        }))
    }

    let method_describe = quote! {
        #func_ident_str => FunctionDescription {
            name: #func_ident_str,
            args: vec![
                #(#describe_args,)*
            ],
            #method_doc,
        }
    };

    match &method.binary_op {
        Some(op) => {
            if value_args.len() != 2 {
                return Err(syn::Error::new(
                    method.name.span(),
                    "binary ops must take exactly two arguments",
                ));
            }

            let op_dispatch = Some(quote_spanned!(method.name.span() =>
                #op => Some(#func_ty::ref_cast(self) as &dyn QueryBinaryOp<#env_ident>)
            ));

            let method_def = quote_spanned!(method.name.span() =>
                #[derive(RefCast)]
                #[repr(transparent)]
                struct #func_ty #impl_generics #where_clause(#self_ty);

                #[async_trait]
                impl #impl_generics QueryBinaryOp<#env_ident> for #func_ty #ty_generics #where_clause {
                    fn name(&self) -> &'static str { #func_ident_str }

                    async fn invoke(
                        &self,
                        env: &Env,
                        left: QueryValue<#env_target>,
                        right: QueryValue<#env_target>,
                    ) -> Result<QueryValue<#env_target>, QueryError> {
                        self.0.#func_ident(
                            env,
                            left,
                            right
                        ).await
                    }
                }
            );

            Ok(MethodCodegen {
                method_def,
                method_describe,
                method_dispatch: None,
                op_dispatch,
            })
        }
        None => {
            let max_args = value_args.len();

            let method_dispatch = Some(quote_spanned!(method.name.span() =>
                stringify!(#func_ident) => Some(#func_ty::ref_cast(self) as &dyn QueryFunction<#env_ident>)
            ));

            let method_def = quote_spanned!(method.name.span() =>
                #[derive(RefCast)]
                #[repr(transparent)]
                struct #func_ty #impl_generics #where_clause(#self_ty);

                #[async_trait]
                impl #impl_generics QueryFunction<#env_ident> for #func_ty #ty_generics #where_clause {
                    fn name(&self) -> &'static str { stringify!(#func_ident) }

                    async fn invoke(
                        &self,
                        evaluator: &QueryEvaluator<#env_ident>,
                        args: &[SpannedExpr<'_>],
                    ) -> Result<QueryValue<#env_target>, QueryError> {
                        self.0.#func_ident(
                            #pass_ctx
                            #(#pass_args,)*
                        ).await
                    }

                    fn arg_type(&self, idx: usize) -> Result<QueryArgType, QueryError> {
                        match idx {
                            #(#arg_type_match,)*
                            v => Err(QueryError::TooManyArgs {
                                function: self.name().to_owned(),
                                max: #max_args,
                                actual: idx + 1,
                            })
                        }
                    }
                }
            );

            Ok(MethodCodegen {
                method_def,
                method_describe,
                method_dispatch,
                op_dispatch: None,
            })
        }
    }
}

pub(crate) fn codegen(parsed: Parsed) -> syn::Result<TokenStream> {
    let mod_name = quote!(__buck_query_gen);
    let mut methods = Vec::new();
    for method in &parsed.module.methods {
        methods.push(gen_for_method(&parsed, method)?);
    }

    let method_defs = methods.iter().map(|v| &v.method_def);
    let method_describe = methods.iter().map(|v| &v.method_describe);
    let method_dispatch = methods
        .iter()
        .map(|v| &v.method_dispatch)
        .filter_map(|v| v.as_ref());
    let op_dispatch = methods
        .iter()
        .map(|v| &v.op_dispatch)
        .filter_map(|v| v.as_ref());

    let (impl_generics, _ty_generics, where_clause) = parsed.module.generics.split_for_impl();
    let self_ty = &parsed.module.self_ty;
    let env_ty = &parsed.args.env_ty;

    let modified_original = &parsed.module.original_impl;
    let module_docs = gen_for_docstring(Span::call_site(), parsed.module.docs.as_ref());

    let result = quote! {
        #modified_original

        mod #mod_name {
            #![allow(non_camel_case_types)]
            #![allow(unused)]
            #![allow(clippy::match_single_binding)]

            use ::buck2_query::__derive_refs::ref_cast::RefCast;
            use ::buck2_query::__derive_refs::buck2_query_parser::{BinaryOp, SpannedExpr};
            use ::buck2_query::__derive_refs::indexmap::indexmap;

            use super::*;
            use ::buck2_query::query::{
                environment::{QueryEnvironment, QueryTarget},
                syntax::simple::{
                    eval::{error::QueryError, evaluator::QueryEvaluator, values::QueryValue},
                    functions::docs::{ModuleDescription, ArgDescription, FunctionDescription},
                    functions::{QueryFunctions, HasModuleDescription},
                    functions::helpers::{
                        eval_arg, QueryArgType, QueryBinaryOp, QueryFunction, QueryFunctionArg,
                    },
                },
            };

            #(#method_defs)*

            impl #impl_generics QueryFunctions for #self_ty #where_clause {
                type Env = #env_ty;
                fn get(&self, name: &str) -> Option<&dyn QueryFunction<#env_ty>> {
                    match name {
                        #(#method_dispatch,)*
                        _ => None
                    }
                }

                fn get_op(&self, op: BinaryOp) -> Option<&dyn QueryBinaryOp<#env_ty>> {
                    match op {
                        #(#op_dispatch,)*
                        _ => None
                    }
                }
            }

            impl #impl_generics HasModuleDescription for #self_ty #where_clause {
                fn describe() -> ModuleDescription {
                    ModuleDescription {
                        functions: indexmap! {
                                #(#method_describe,)*
                        },
                        #module_docs,
                    }
                }
            }

        }

    };

    Ok(result)
}


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

#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]
#![feature(try_blocks)]

#[allow(unused_extern_crates)] // proc_macro is very special
extern crate proc_macro;

mod codegen;
mod parse;

/// Generates a QueryFunctions and HasModuleDescription impls for a type.
///
/// The HasModuleDescriptions will be based on the rustdoc comments on the
/// module and functions (and on the extract function names/types).
///
/// Example usage:
///
/// ```ignore
/// struct MyFunctionsContext {
///   <fill with any data>
/// }
///
/// // The query_module attribute must be passed the QueryEnvironment type to be used, it can be a type parameter
/// // as in this example or can be a concrete type.
/// #[query_module(Env)]
/// impl <Env: QueryEnvironment + OtherTrait> MyFunctionsContext {
///     // functions can be either async or not.
///     async fn some_func(
///         // functions must take a &self receiver
///         &self,
///         // the first argument can optionally be a &Env or &QueryEvaluator.
///         env: &Env,
///         // remaining arguments must all be types that implement QueryFunctionArg
///         roots: TargetSet<Env::Target>
///     ) -> Result<QueryValue<Env>, QueryError> {
///         <...>
///     }
///
///     // binary_ops can also be implemented, the binary_op type needs to be passed in an attribute
///     #[binary_op(BinaryOp::Intersect)]
///     fn intersect(&self, env: &Env, left: TargetSet<Env::Target>, right: TargetSet<Env::Target>) -> Result<TargetSet<Env::Target>, QueryError> {
///         <...>
///     }
/// }
/// ```
///
/// Each function and binary op will be available on the QueryFunctions impl, each one uses a
/// unique opaque struct that implements QueryFunction/BinaryOp. For the `some_func` in the example above, one
/// of those structs looks like:
///
/// ```ignore
/// #[derive(RefCast)]
/// #[repr(transparent)]
/// struct __OpaqueSomeFunc(MyFunctionsContext);
///
/// impl QueryFunction<Env> for __OpaqueSomeFunc {
///   <...>
///   fn invoke(&self, ...) {
///     self.0.some_func(<converted args>)
///   }
/// }
/// ```
///
/// and the generated QueryFunctions::get() will look something like:
///
/// ```ignore
///   fn get(&self, name: &str) -> Option<&dyn QueryFunction<Env>> {
///     match name {
///         "some_func" => Some(__OpaqueSomeFunc::ref_cast(self) as &dyn QueryFunction<Env>),
///         ...
///     }
///   }
/// ```
#[proc_macro_attribute]
pub fn query_module(
    attr: proc_macro::TokenStream,
    input: proc_macro::TokenStream,
) -> proc_macro::TokenStream {
    let result: syn::Result<_> = try {
        let parsed = parse::parse(attr.into(), input.into())?;
        codegen::codegen(parsed)?
    };

    match result {
        Ok(v) => v.into(),
        Err(e) => e.to_compile_error().into(),
    }
}


/*
 * 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 itertools::Itertools;
use proc_macro2::Span;
use proc_macro2::TokenStream;
use syn::parse::Result;
use syn::spanned::Spanned;
use syn::Error;
use syn::Expr;
use syn::ExprLit;
use syn::FnArg;
use syn::Generics;
use syn::Ident;
use syn::ImplItem;
use syn::Item;
use syn::ItemImpl;
use syn::Lit;
use syn::Meta;
use syn::MetaNameValue;
use syn::Pat;
use syn::PatIdent;
use syn::PatType;
use syn::Path;
use syn::Type;

/// Validates and parses the macro input into structured data.
pub(crate) fn parse(args: TokenStream, item: TokenStream) -> Result<Parsed> {
    // Parse input first so that any error with it takes precedence.
    let module = syn::parse2::<Module>(item)?;
    let args = syn::parse2::<QueryModuleArgs>(args)?;

    Ok(Parsed { module, args })
}

#[derive(Debug, Eq, PartialEq)]
pub(crate) struct DocString {
    pub short_help: String,
    pub details: String,
}

impl DocString {
    fn parse(attrs: &[syn::Attribute]) -> Option<Self> {
        let mut docs = Vec::new();
        for attr in attrs {
            if attr.path().is_ident("doc") {
                if let Meta::NameValue(MetaNameValue {
                    value:
                        Expr::Lit(ExprLit {
                            lit: Lit::Str(s), ..
                        }),
                    ..
                }) = &attr.meta
                {
                    docs.push(s.value());
                }
            }
        }
        if docs.is_empty() {
            None
        } else {
            let mut lines = docs.into_iter().peekable();
            let short_help = lines.next().unwrap().trim().to_owned();
            let mut lines = lines.skip_while(|v| v.is_empty());

            let details = lines.join("\n");
            let details = textwrap::dedent(&details);
            Some(Self {
                short_help,
                details,
            })
        }
    }
}

#[derive(Debug)]
pub(crate) struct Parsed {
    pub module: Module,
    pub args: QueryModuleArgs,
}

#[derive(Debug)]
pub(crate) struct QueryModuleArgs {
    pub env_ty: Type,
}

#[derive(Debug)]
pub(crate) struct Argument {
    pub name: Ident,
    pub span: Span,
    pub ty: Type,
}

#[derive(Debug)]
pub(crate) struct Method {
    pub docs: Option<DocString>,
    pub name: Ident,
    pub args: Vec<Argument>,
    pub binary_op: Option<Path>,
}

#[derive(Debug)]
pub(crate) struct Module {
    pub docs: Option<DocString>,
    pub generics: Generics,
    /// The Self type of the impl.
    pub self_ty: Box<Type>,
    pub methods: Vec<Method>,

    /// Holds the original annotated impl, with any attributes we consumed removed.
    pub original_impl: ItemImpl,
}

impl syn::parse::Parse for Module {
    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
        let span = input.span();
        let mut item = match input.parse::<Item>() {
            Ok(Item::Impl(item)) => item,
            Ok(item) => {
                return Err(Error::new_spanned(
                    item,
                    "`#[query_module(..)]` can only be used on impls",
                ));
            }
            Err(e) => return Err(Error::new(span, format!("{}", e))),
        };

        let ItemImpl {
            // Should we care if there are other attributes?
            attrs,
            defaultness,
            unsafety,
            impl_token: _impl_token,
            generics,
            trait_,
            self_ty,
            brace_token: _brace_token,
            items,
        } = &mut item;

        if let Some(v) = defaultness {
            return Err(Error::new_spanned(
                v,
                "`#[query_module]` can't handle `default` in impl item",
            ));
        }
        if let Some(v) = unsafety {
            return Err(Error::new_spanned(
                v,
                "`#[query_module]` can't handle `unsafe` in impl item",
            ));
        }
        if trait_.is_some() {
            return Err(Error::new_spanned(
                self_ty,
                "`#[query_module]` should only be used on an inherent impl, not a trait impl",
            ));
        }

        let docs = DocString::parse(attrs);

        let mut methods = Vec::new();
        let binary_op_path = Path::from(Ident::new("binary_op", Span::call_site()));
        for item in items {
            match item {
                ImplItem::Fn(method) => {
                    let sig = &method.sig;
                    let span = sig.span();

                    if let Some(v) = sig.unsafety {
                        return Err(Error::new_spanned(
                            v,
                            "`#[query_module]` methods can't be unsafe.",
                        ));
                    }

                    if sig.asyncness.is_none() {
                        return Err(Error::new_spanned(
                            sig,
                            "`#[query_module]` methods must be async.",
                        ));
                    }

                    let ident = &sig.ident;

                    let mut args = sig.inputs.iter();
                    match args.next() {
                        Some(FnArg::Receiver(_)) => {}
                        _ => {
                            return Err(Error::new(
                                span,
                                "`#[query_module]` methods must have a `&self` parameter",
                            ));
                        }
                    }

                    let mut parsed_args = Vec::new();
                    for arg in args {
                        let span = arg.span();
                        match arg {
                            FnArg::Receiver(_) => {
                                unreachable!("a method can't have multiple receiver args")
                            }
                            FnArg::Typed(PatType { pat, ty, .. }) => {
                                let name = match &**pat {
                                    Pat::Ident(PatIdent { ident, .. }) => ident,
                                    _ => {
                                        return Err(Error::new(
                                            span,
                                            "`#[query_module]` method arguments can only have simple identity captures. You can destructure them within the method impl.",
                                        ));
                                    }
                                };
                                parsed_args.push(Argument {
                                    span,
                                    ty: (**ty).clone(),
                                    name: name.clone(),
                                })
                            }
                        }
                    }

                    // We look for a #[binary_op(BinaryOp::Foo)] attribute to identify functions that implement binary ops. We need to also remove that attribute from the original code.
                    let mut binary_op = None;
                    let original_attrs = std::mem::take(&mut method.attrs);

                    let docs = DocString::parse(&original_attrs);

                    for attr in original_attrs {
                        if *attr.path() == binary_op_path {
                            if let Ok(path) = attr.parse_args::<Path>() {
                                binary_op = Some(path);
                            } else {
                                return Err(Error::new_spanned(
                                    &attr.meta,
                                    "#[query_module] attribute `binary_op` should receive a single argument identifying the binary op like `#[binary_op(BinaryOp::Intersect)]`",
                                ));
                            }
                        } else {
                            method.attrs.push(attr)
                        }
                    }

                    methods.push(Method {
                        name: ident.clone(),
                        args: parsed_args,
                        binary_op,
                        docs,
                    });
                }
                _ => {
                    return Err(Error::new_spanned(
                        item,
                        "`#[query_module]` only supports methods within the annotated impl",
                    ));
                }
            };
        }

        Ok(Self {
            docs,
            generics: generics.clone(),
            self_ty: self_ty.clone(),
            methods,
            original_impl: item,
        })
    }
}

impl syn::parse::Parse for QueryModuleArgs {
    fn parse(input: syn::parse::ParseStream) -> Result<Self> {
        const ARGS_ERROR: &str = "this attribute takes a single argument, the identifier for the QueryEnvironment type. use `#[query_module(Env)]`, for example";
        let span = input.span();

        let env_ty: Type = match input.parse() {
            Ok(v) => v,
            Err(_) => return Err(Error::new(span, ARGS_ERROR)),
        };

        Ok(QueryModuleArgs { env_ty })
    }
}

#[cfg(test)]
mod test {
    use quote::quote;
    use syn::parse_quote;
    use syn::Generics;
    use syn::Type;

    use crate::parse::DocString;
    use crate::parse::Parsed;

    #[test]
    fn test() {
        let args = quote! {Env};

        let item = quote! {
            /// module-level docs line 1
            ///
            /// module-level docs line 2
            impl <Env: QueryEnvironment> MySimpleFunctions<Env> {
                /// buildfile docs
                async fn buildfile(&self, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
                    Ok(targets.buildfile()?.into())
                }

                /// deps docs
                async fn deps(
                    &self,
                    evaluator: &QueryEvaluator<'_, Env>,
                    targets: TargetSet<Env::Target>,
                    depth: Option<u64>,
                    captured_expr: Option<CapturedExpr<'_>>,
                ) -> QueryFuncResult<Env> {
                    DepsFunction::<Env> {
                        _marker: PhantomData,
                    }
                    .invoke_deps(evaluator, targets, depth, captured_expr)
                    .await
                }

                /// filter docs
                async fn filter(&self, regex: String, targets: TargetSet<Env::Target>) -> QueryFuncResult<Env> {
                    Ok(targets.filter_name(&regex)?.into())
                }
            }
        };

        let Parsed { module, .. } = super::parse(args, item).unwrap();

        assert_eq!(
            &Some(DocString {
                short_help: "module-level docs line 1".to_owned(),
                details: "module-level docs line 2".to_owned()
            }),
            &module.docs
        );

        let expected: Generics = parse_quote!(<Env: QueryEnvironment>);
        assert_eq!(&expected, &module.generics);
        let expected: Type = parse_quote!(MySimpleFunctions<Env>);
        assert_eq!(&expected, &*module.self_ty);

        assert_eq!(3, module.methods.len());

        let method = module.methods.get(0).unwrap();
        assert_eq!("buildfile", method.name.to_string());
        let args = &method.args;
        assert_eq!(1, args.len());
        let expected: Type = parse_quote!(TargetSet<Env::Target>);
        assert_eq!(&expected, &args.get(0).unwrap().ty);

        let method = module.methods.get(1).unwrap();
        assert_eq!("deps", method.name.to_string());
        let args = &method.args;
        assert_eq!(4, args.len());
        let expected: Type = parse_quote!(&QueryEvaluator<'_, Env>);
        assert_eq!(&expected, &args.get(0).unwrap().ty);
        let expected: Type = parse_quote!(TargetSet<Env::Target>);
        assert_eq!(&expected, &args.get(1).unwrap().ty);
        let expected: Type = parse_quote!(Option<u64>);
        assert_eq!(&expected, &args.get(2).unwrap().ty);
        let expected: Type = parse_quote!(Option<CapturedExpr<'_>>);
        assert_eq!(&expected, &args.get(3).unwrap().ty);

        let method = module.methods.get(2).unwrap();
        assert_eq!("filter", method.name.to_string());
        let args = &method.args;
        assert_eq!(2, args.len());
        let expected: Type = parse_quote!(String);
        assert_eq!(&expected, &args.get(0).unwrap().ty);
        let expected: Type = parse_quote!(TargetSet<Env::Target>);
        assert_eq!(&expected, &args.get(1).unwrap().ty);
    }
}


/*
 * 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 async_trait::async_trait;
use buck2_build_api::actions::query::ActionQueryNode;
use buck2_build_api::query::oneshot::CqueryOwnerBehavior;
use buck2_build_api::query::oneshot::QueryFrontend;
use buck2_build_api::query::oneshot::QUERY_FRONTEND;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::target::label::TargetLabel;
use buck2_node::configured_universe::CqueryUniverse;
use buck2_node::nodes::configured::ConfiguredTargetNode;
use buck2_node::nodes::unconfigured::TargetNode;
use buck2_query::query::syntax::simple::eval::values::QueryEvaluationResult;
use dice::DiceComputations;

use crate::aquery::evaluator::get_aquery_evaluator;
use crate::cquery::evaluator::get_cquery_evaluator;
use crate::cquery::evaluator::preresolve_literals_and_build_universe;
use crate::dice::get_dice_query_delegate;
use crate::uquery::evaluator::get_uquery_evaluator;

struct QueryFrontendImpl;

pub(crate) fn init_query_frontend() {
    QUERY_FRONTEND.init(&QueryFrontendImpl);
}

#[async_trait]
impl QueryFrontend for QueryFrontendImpl {
    async fn eval_uquery(
        &self,
        ctx: &DiceComputations,
        working_dir: &ProjectRelativePath,
        query: &str,
        query_args: &[String],
        global_target_platform: Option<TargetLabel>,
    ) -> anyhow::Result<QueryEvaluationResult<TargetNode>> {
        let evaluator = get_uquery_evaluator(ctx, working_dir, global_target_platform).await?;

        evaluator.eval_query(query, query_args).await
    }

    async fn eval_cquery(
        &self,
        ctx: &DiceComputations,
        working_dir: &ProjectRelativePath,
        owner_behavior: CqueryOwnerBehavior,
        query: &str,
        query_args: &[String],
        global_target_platform: Option<TargetLabel>,
        target_universe: Option<&[String]>,
    ) -> anyhow::Result<QueryEvaluationResult<ConfiguredTargetNode>> {
        let evaluator =
            get_cquery_evaluator(ctx, working_dir, global_target_platform, owner_behavior).await?;

        // TODO(nga): this should support configured target patterns
        //   similarly to what we do for `build` command.
        //   Something like this should work:
        //   ```
        //   buck2 cquery --target-universe android//:binary 'deps("some//:lib (<arm32>)")'
        //   ```
        evaluator
            .eval_query(query, query_args, target_universe.as_ref().map(|v| &v[..]))
            .await
    }

    async fn eval_aquery(
        &self,
        ctx: &DiceComputations,
        working_dir: &ProjectRelativePath,
        query: &str,
        query_args: &[String],
        global_target_platform: Option<TargetLabel>,
    ) -> anyhow::Result<QueryEvaluationResult<ActionQueryNode>> {
        let evaluator = get_aquery_evaluator(ctx, working_dir, global_target_platform).await?;

        evaluator.eval_query(query, query_args).await
    }

    async fn universe_from_literals(
        &self,
        ctx: &DiceComputations,
        cwd: &ProjectRelativePath,
        literals: &[String],
        global_target_platform: Option<TargetLabel>,
    ) -> anyhow::Result<CqueryUniverse> {
        let query_delegate = get_dice_query_delegate(ctx, cwd, global_target_platform).await?;
        Ok(
            preresolve_literals_and_build_universe(&query_delegate, literals)
                .await?
                .0,
        )
    }
}