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

//! Utilities for providing the completion language feature
use std::iter::once;

use common::ArgumentName;
use common::DirectiveName;
use common::Named;
use common::NamedItem;
use common::Span;
use fnv::FnvHashSet;
use graphql_ir::FragmentDefinitionName;
use graphql_ir::OperationDefinitionName;
use graphql_ir::Program;
use graphql_ir::VariableDefinition;
use graphql_ir::VariableName;
use graphql_ir::DIRECTIVE_ARGUMENTS;
use graphql_syntax::Argument;
use graphql_syntax::ConstantValue;
use graphql_syntax::Directive;
use graphql_syntax::DirectiveLocation;
use graphql_syntax::ExecutableDefinition;
use graphql_syntax::ExecutableDocument;
use graphql_syntax::FragmentSpread;
use graphql_syntax::InlineFragment;
use graphql_syntax::LinkedField;
use graphql_syntax::List;
use graphql_syntax::OperationDefinition;
use graphql_syntax::ScalarField;
use graphql_syntax::Selection;
use graphql_syntax::TokenKind;
use graphql_syntax::Value;
use intern::string_key::StringKey;
use intern::Lookup;
use log::debug;
use lsp_types::request::Completion;
use lsp_types::request::Request;
use lsp_types::request::ResolveCompletionItem;
use lsp_types::CompletionItem;
use lsp_types::CompletionItemKind;
use lsp_types::CompletionResponse;
use lsp_types::Documentation;
use lsp_types::InsertTextFormat;
use lsp_types::MarkupContent;
use lsp_types::MarkupKind;
use schema::Argument as SchemaArgument;
use schema::Directive as SchemaDirective;
use schema::InputObject;
use schema::SDLSchema;
use schema::Schema;
use schema::Type;
use schema::TypeReference;
use schema::TypeWithFields;

use crate::lsp_runtime_error::LSPRuntimeResult;
use crate::node_resolution_info::TypePath;
use crate::node_resolution_info::TypePathItem;
use crate::server::GlobalState;
use crate::LSPRuntimeError;
use crate::SchemaDocumentation;

#[derive(Debug, Clone)]
pub enum CompletionKind {
    FieldName {
        existing_linked_field: bool,
    },
    FragmentSpread,
    DirectiveName {
        location: DirectiveLocation,
    },
    ArgumentName {
        has_colon: bool,
        existing_names: FnvHashSet<StringKey>,
        kind: ArgumentKind,
    },
    ArgumentValue {
        executable_name: ExecutableName,
        argument_name: StringKey,
        kind: ArgumentKind,
    },
    InlineFragmentType {
        existing_inline_fragment: bool,
    },
    InputObjectFieldName {
        name: StringKey,
        existing_names: FnvHashSet<StringKey>,
        input_field_path: Vec<StringKey>,
    },
}

#[derive(Debug, Clone)]
pub enum ArgumentKind {
    Field,
    Directive(DirectiveName),
    ArgumentsDirective(StringKey),
}

#[derive(Debug)]
pub struct CompletionRequest {
    /// The type of the completion request we're responding to
    kind: CompletionKind,
    /// A list of type metadata that we can use to resolve the leaf
    /// type the request is being made against
    type_path: TypePath,
    /// The project the request belongs to
    pub project_name: StringKey,
}

impl CompletionRequest {
    fn new(project_name: StringKey, kind: CompletionKind, type_path: TypePath) -> Self {
        Self {
            kind,
            type_path,
            project_name,
        }
    }
}

#[derive(Debug, Copy, Clone)]
pub enum ExecutableName {
    Operation(StringKey),
    Fragment(FragmentDefinitionName),
}

trait ArgumentLike {
    fn name(&self) -> StringKey;
    fn type_(&self) -> &TypeReference<Type>;
}

impl ArgumentLike for &SchemaArgument {
    fn name(&self) -> StringKey {
        self.name.0
    }
    fn type_(&self) -> &TypeReference<Type> {
        &self.type_
    }
}

impl ArgumentLike for &VariableDefinition {
    fn name(&self) -> StringKey {
        self.name.item.0
    }
    fn type_(&self) -> &TypeReference<Type> {
        &self.type_
    }
}

struct CompletionRequestBuilder {
    project_name: StringKey,
    current_executable_name: Option<ExecutableName>,
}

impl CompletionRequestBuilder {
    fn new(project_name: StringKey) -> Self {
        Self {
            project_name,
            current_executable_name: None,
        }
    }

    fn new_request(&self, kind: CompletionKind, type_path: Vec<TypePathItem>) -> CompletionRequest {
        CompletionRequest::new(self.project_name, kind, type_path.into())
    }

    fn create_completion_request(
        &mut self,
        document: ExecutableDocument,
        position_span: Span,
    ) -> Option<CompletionRequest> {
        for definition in document.definitions {
            match &definition {
                ExecutableDefinition::Operation(operation) => {
                    if operation.location.contains(position_span) {
                        self.current_executable_name = operation
                            .name
                            .as_ref()
                            .map(|name| ExecutableName::Operation(name.value));
                        let (_, kind) = operation.operation.clone()?;
                        let type_path = vec![TypePathItem::Operation(kind)];

                        debug!(
                            "Completion request is within operation: {:?}",
                            operation.name
                        );
                        let OperationDefinition {
                            selections,
                            directives,
                            ..
                        } = operation;

                        let directive_location = kind.into();

                        if let Some(req) = self.build_request_from_selection_or_directives(
                            selections,
                            directives,
                            directive_location,
                            position_span,
                            type_path,
                        ) {
                            return Some(req);
                        }
                    }
                    // Check if the position span is within this operation's span
                }
                ExecutableDefinition::Fragment(fragment) => {
                    if fragment.location.contains(position_span) {
                        self.current_executable_name = Some(ExecutableName::Fragment(
                            FragmentDefinitionName(fragment.name.value),
                        ));
                        let type_name = fragment.type_condition.type_.value;
                        let type_path = vec![TypePathItem::FragmentDefinition { type_name }];
                        if let Some(req) = self.build_request_from_selection_or_directives(
                            &fragment.selections,
                            &fragment.directives,
                            DirectiveLocation::FragmentDefinition,
                            position_span,
                            type_path,
                        ) {
                            return Some(req);
                        }
                    }
                }
            }
        }
        None
    }

    fn build_request_from_selections(
        &self,
        selections: &List<Selection>,
        position_span: Span,
        mut type_path: Vec<TypePathItem>,
    ) -> Option<CompletionRequest> {
        for item in &selections.items {
            if item.span().contains(position_span) {
                return match item {
                    Selection::LinkedField(node) => {
                        if node.name.span.contains(position_span) {
                            return Some(self.new_request(
                                CompletionKind::FieldName {
                                    existing_linked_field: true,
                                },
                                type_path,
                            ));
                        }
                        let LinkedField {
                            name,
                            selections,
                            directives,
                            arguments,
                            ..
                        } = node;
                        type_path.push(TypePathItem::LinkedField { name: name.value });
                        if let Some(arguments) = arguments {
                            if arguments.span.contains(position_span) {
                                return self.build_request_from_arguments(
                                    arguments,
                                    position_span,
                                    type_path,
                                    ArgumentKind::Field,
                                );
                            }
                        }
                        self.build_request_from_selection_or_directives(
                            selections,
                            directives,
                            DirectiveLocation::Field,
                            position_span,
                            type_path,
                        )
                    }
                    Selection::FragmentSpread(spread) => {
                        let FragmentSpread {
                            name, directives, ..
                        } = spread;
                        if name.span.contains(position_span) {
                            Some(self.new_request(CompletionKind::FragmentSpread, type_path))
                        } else {
                            self.build_request_from_directives(
                                directives,
                                DirectiveLocation::FragmentSpread,
                                position_span,
                                type_path,
                                Some(name.value),
                            )
                        }
                    }
                    Selection::InlineFragment(node) => {
                        let InlineFragment {
                            selections,
                            directives,
                            type_condition,
                            ..
                        } = node;
                        if let Some(type_condition) = type_condition {
                            let type_name = type_condition.type_.value;
                            if type_condition.span.contains(position_span) {
                                return Some(self.new_request(
                                    CompletionKind::InlineFragmentType {
                                        existing_inline_fragment: selections.start.kind
                                            != TokenKind::Empty,
                                    },
                                    type_path,
                                ));
                            }
                            type_path.push(TypePathItem::InlineFragment { type_name });
                        }
                        self.build_request_from_selection_or_directives(
                            selections,
                            directives,
                            DirectiveLocation::InlineFragment,
                            position_span,
                            type_path,
                        )
                    }
                    Selection::ScalarField(node) => {
                        if node.name.span.contains(position_span) {
                            return Some(self.new_request(
                                CompletionKind::FieldName {
                                    existing_linked_field: false,
                                },
                                type_path,
                            ));
                        }
                        let ScalarField {
                            directives,
                            name,
                            arguments,
                            ..
                        } = node;
                        type_path.push(TypePathItem::ScalarField { name: name.value });
                        if let Some(arguments) = arguments {
                            if arguments.span.contains(position_span) {
                                return self.build_request_from_arguments(
                                    arguments,
                                    position_span,
                                    type_path,
                                    ArgumentKind::Field,
                                );
                            }
                        }
                        self.build_request_from_directives(
                            directives,
                            DirectiveLocation::Field,
                            position_span,
                            type_path,
                            None,
                        )
                    }
                };
            }
        }
        // The selection list is empty or the current cursor is out of any of the selection
        Some(self.new_request(
            CompletionKind::FieldName {
                existing_linked_field: false,
            },
            type_path,
        ))
    }

    fn build_request_from_constant_input_value(
        &self,
        position_span: Span,
        type_path: Vec<TypePathItem>,
        mut input_field_path: Vec<StringKey>,
        constant_value: &ConstantValue,
        name: StringKey,
    ) -> Option<CompletionRequest> {
        match constant_value {
            ConstantValue::List(list) => list
                .items
                .iter()
                .find(|arg| arg.span().contains(position_span))
                .and_then(|constant_value| {
                    self.build_request_from_constant_input_value(
                        position_span,
                        type_path,
                        input_field_path,
                        constant_value,
                        name,
                    )
                }),
            ConstantValue::Object(arguments) => {
                if let Some(constant_argument) = arguments
                    .items
                    .iter()
                    .find(|arg| arg.span.contains(position_span))
                {
                    input_field_path.push(constant_argument.name());
                    self.build_request_from_constant_input_value(
                        position_span,
                        type_path,
                        input_field_path,
                        &constant_argument.value,
                        name,
                    )
                } else {
                    Some(self.new_request(
                        CompletionKind::InputObjectFieldName {
                            name,
                            existing_names:
                                arguments.items.iter().map(|item| item.name()).collect(),
                            input_field_path,
                        },
                        type_path,
                    ))
                }
            }
            _ => None,
        }
    }

    fn build_request_from_input_value(
        &self,
        position_span: Span,
        type_path: Vec<TypePathItem>,
        mut input_field_path: Vec<StringKey>,
        value: &Value,
        name: StringKey,
    ) -> Option<CompletionRequest> {
        match value {
            Value::List(list) => list
                .items
                .iter()
                .find(|arg| arg.span().contains(position_span))
                .and_then(|value| {
                    self.build_request_from_input_value(
                        position_span,
                        type_path,
                        input_field_path,
                        value,
                        name,
                    )
                }),
            Value::Object(arguments) => {
                if let Some(position_argument) = arguments
                    .items
                    .iter()
                    .find(|arg| arg.span.contains(position_span))
                {
                    input_field_path.push(position_argument.name());
                    self.build_request_from_input_value(
                        position_span,
                        type_path,
                        input_field_path,
                        &position_argument.value,
                        name,
                    )
                } else {
                    Some(self.new_request(
                        CompletionKind::InputObjectFieldName {
                            name,
                            existing_names:
                                arguments.items.iter().map(|item| item.name()).collect(),
                            input_field_path,
                        },
                        type_path,
                    ))
                }
            }
            Value::Constant(constant_value) => self.build_request_from_constant_input_value(
                position_span,
                type_path,
                input_field_path,
                constant_value,
                name,
            ),
            _ => None,
        }
    }

    fn build_request_from_arguments(
        &self,
        arguments: &List<Argument>,
        position_span: Span,
        type_path: Vec<TypePathItem>,
        kind: ArgumentKind,
    ) -> Option<CompletionRequest> {
        for (
            i,
            Argument {
                name,
                value,
                colon,
                span,
                ..
            },
        ) in arguments.items.iter().enumerate()
        {
            if span.contains(position_span) {
                return if name.span.contains(position_span) {
                    Some(self.new_request(
                        CompletionKind::ArgumentName {
                            has_colon: colon.kind != TokenKind::Empty,
                            existing_names:
                                arguments.items.iter().map(|arg| arg.name.value).collect(),
                            kind,
                        },
                        type_path,
                    ))
                } else if let Some(executable_name) = self.current_executable_name {
                    match value {
                        Value::Constant(ConstantValue::Null(token))
                            if token.kind == TokenKind::Empty =>
                        {
                            Some(self.new_request(
                                CompletionKind::ArgumentValue {
                                    argument_name: name.value,
                                    executable_name,
                                    kind,
                                },
                                type_path,
                            ))
                        }
                        Value::Constant(constant_value) => self
                            .build_request_from_constant_input_value(
                                position_span,
                                type_path,
                                Default::default(),
                                constant_value,
                                name.value,
                            ),
                        Value::Variable(_) => Some(self.new_request(
                            CompletionKind::ArgumentValue {
                                argument_name: name.value,
                                executable_name,
                                kind,
                            },
                            type_path,
                        )),
                        value => self.build_request_from_input_value(
                            position_span,
                            type_path,
                            Default::default(),
                            value,
                            name.value,
                        ),
                    }
                } else {
                    None
                };
            } else if span.end <= position_span.start {
                let is_cursor_in_next_white_space = {
                    if let Some(next_argument) = arguments.items.get(i + 1) {
                        position_span.start < next_argument.span.start
                    } else {
                        position_span.start < arguments.span.end
                    }
                };
                if is_cursor_in_next_white_space {
                    // Handles the following special case
                    // (args1:  | args2:$var)
                    //          ^ cursor here
                    // The cursor is on the white space between args1 and args2.
                    // We want to autocomplete the value if it's empty.
                    return if let Some(executable_name) = self.current_executable_name {
                        match value {
                            Value::Constant(ConstantValue::Null(token))
                                if token.kind == TokenKind::Empty =>
                            {
                                Some(self.new_request(
                                    CompletionKind::ArgumentValue {
                                        argument_name: name.value,
                                        executable_name,
                                        kind,
                                    },
                                    type_path,
                                ))
                            }
                            _ => Some(self.new_request(
                                CompletionKind::ArgumentName {
                                    has_colon: false,
                                    existing_names:
                                        arguments.items.iter().map(|arg| arg.name.value).collect(),
                                    kind,
                                },
                                type_path,
                            )),
                        }
                    } else {
                        None
                    };
                }
            }
        }
        // The argument list is empty or the cursor is not on any of the argument
        Some(self.new_request(
            CompletionKind::ArgumentName {
                has_colon: false,
                existing_names: arguments.items.iter().map(|arg| arg.name.value).collect(),
                kind,
            },
            type_path,
        ))
    }

    fn build_request_from_directives(
        &self,
        directives: &[Directive],
        location: DirectiveLocation,
        position_span: Span,
        type_path: Vec<TypePathItem>,
        fragment_spread_name: Option<StringKey>,
    ) -> Option<CompletionRequest> {
        for directive in directives {
            if !directive.span.contains(position_span) {
                continue;
            };
            return if directive.name.span.contains(position_span) {
                Some(self.new_request(CompletionKind::DirectiveName { location }, type_path))
            } else if let Some(arguments) = &directive.arguments {
                if arguments.span.contains(position_span) {
                    self.build_request_from_arguments(
                        arguments,
                        position_span,
                        type_path,
                        if let Some(fragment_spread_name) = fragment_spread_name {
                            if directive.name.value == *DIRECTIVE_ARGUMENTS {
                                ArgumentKind::ArgumentsDirective(fragment_spread_name)
                            } else {
                                ArgumentKind::Directive(DirectiveName(directive.name.value))
                            }
                        } else {
                            ArgumentKind::Directive(DirectiveName(directive.name.value))
                        },
                    )
                } else {
                    None
                }
            } else {
                // The directive doesn't have a name `@|`
                Some(self.new_request(CompletionKind::DirectiveName { location }, type_path))
            };
        }
        None
    }

    fn build_request_from_selection_or_directives(
        &self,
        selections: &List<Selection>,
        directives: &[Directive],
        directive_location: DirectiveLocation,
        position_span: Span,
        type_path: Vec<TypePathItem>,
    ) -> Option<CompletionRequest> {
        if selections.span.contains(position_span) {
            // TODO(brandondail) handle when the completion occurs at/within the start token
            self.build_request_from_selections(selections, position_span, type_path)
        } else {
            self.build_request_from_directives(
                directives,
                directive_location,
                position_span,
                type_path,
                None,
            )
        }
    }
}

fn completion_items_for_request(
    request: CompletionRequest,
    schema: &SDLSchema,
    schema_documentation: impl SchemaDocumentation,
    program: &Program,
) -> Option<Vec<CompletionItem>> {
    let kind = request.kind;
    debug!("completion_items_for_request: {:?}", kind);
    match kind {
        CompletionKind::FragmentSpread => {
            let leaf_type = request.type_path.resolve_leaf_type(schema)?;
            Some(resolve_completion_items_for_fragment_spread(
                leaf_type, program, schema, true,
            ))
        }
        CompletionKind::FieldName {
            existing_linked_field,
        } => match request.type_path.resolve_leaf_type(schema)? {
            Type::Interface(interface_id) => {
                let interface = schema.interface(interface_id);
                Some(merge_completion_items_ordered([
                    resolve_completion_items_for_fields(
                        interface,
                        schema,
                        schema_documentation,
                        existing_linked_field,
                    ),
                    resolve_completion_items_typename(Type::Interface(interface_id), schema),
                    resolve_completion_items_for_inline_fragment(
                        Type::Interface(interface_id),
                        schema,
                        false,
                    ),
                    resolve_completion_items_for_fragment_spread(
                        Type::Interface(interface_id),
                        program,
                        schema,
                        false,
                    ),
                ]))
            }
            Type::Object(object_id) => Some(merge_completion_items_ordered([
                resolve_completion_items_for_fields(
                    schema.object(object_id),
                    schema,
                    schema_documentation,
                    existing_linked_field,
                ),
                resolve_completion_items_typename(Type::Object(object_id), schema),
                resolve_completion_items_for_fragment_spread(
                    Type::Object(object_id),
                    program,
                    schema,
                    false,
                ),
            ])),
            Type::Union(union_id) => Some(merge_completion_items_ordered([
                resolve_completion_items_typename(Type::Union(union_id), schema),
                resolve_completion_items_for_inline_fragment(Type::Union(union_id), schema, false),
                resolve_completion_items_for_fragment_spread(
                    Type::Union(union_id),
                    program,
                    schema,
                    false,
                ),
            ])),
            Type::Enum(_) | Type::InputObject(_) | Type::Scalar(_) => None,
        },
        CompletionKind::DirectiveName { location } => {
            let directives = schema.directives_for_location(location);
            let items = directives
                .iter()
                .map(|directive| completion_item_from_directive(directive, schema))
                .collect();
            Some(items)
        }
        CompletionKind::ArgumentName {
            has_colon,
            existing_names,
            kind,
        } => match kind {
            ArgumentKind::Field => {
                let (_, field) = request.type_path.resolve_current_field(schema)?;
                Some(resolve_completion_items_for_argument_name(
                    field.arguments.iter(),
                    schema,
                    existing_names,
                    has_colon,
                ))
            }
            ArgumentKind::ArgumentsDirective(fragment_spread_name) => {
                let fragment = program.fragment(FragmentDefinitionName(fragment_spread_name))?;
                Some(resolve_completion_items_for_argument_name(
                    fragment.variable_definitions.iter(),
                    schema,
                    existing_names,
                    has_colon,
                ))
            }
            ArgumentKind::Directive(directive_name) => {
                Some(resolve_completion_items_for_argument_name(
                    schema.get_directive(directive_name)?.arguments.iter(),
                    schema,
                    existing_names,
                    has_colon,
                ))
            }
        },
        CompletionKind::ArgumentValue {
            executable_name,
            argument_name,
            kind,
        } => {
            let argument_type = match kind {
                ArgumentKind::Field => {
                    let (_, field) = request.type_path.resolve_current_field(schema)?;
                    &field.arguments.named(ArgumentName(argument_name))?.type_
                }
                ArgumentKind::ArgumentsDirective(fragment_spread_name) => {
                    let fragment =
                        program.fragment(FragmentDefinitionName(fragment_spread_name))?;
                    &fragment
                        .variable_definitions
                        .named(VariableName(argument_name))?
                        .type_
                }
                ArgumentKind::Directive(directive_name) => {
                    &schema
                        .get_directive(directive_name)?
                        .arguments
                        .named(ArgumentName(argument_name))?
                        .type_
                }
            };
            Some(resolve_completion_items_for_argument_value(
                schema,
                argument_type,
                program,
                executable_name,
            ))
        }
        CompletionKind::InlineFragmentType {
            existing_inline_fragment,
        } => {
            let type_ = request.type_path.resolve_leaf_type(schema)?;
            Some(resolve_completion_items_for_inline_fragment(
                type_,
                schema,
                existing_inline_fragment,
            ))
        }
        CompletionKind::InputObjectFieldName {
            name,
            existing_names,
            input_field_path,
        } => {
            let (_, field) = request.type_path.resolve_current_field(schema)?;

            fn resolve_root_input_field<'a>(
                schema: &'a SDLSchema,
                input_object: &'a TypeReference<Type>,
            ) -> Option<&'a InputObject> {
                match input_object {
                    TypeReference::Named(Type::InputObject(input_object_id)) => {
                        Some(schema.input_object(*input_object_id))
                    }
                    TypeReference::Named(_) => None,
                    TypeReference::NonNull(inner) => resolve_root_input_field(schema, inner),
                    TypeReference::List(inner) => resolve_root_input_field(schema, inner),
                }
            }

            fn resolve_input_field<'a>(
                schema: &'a SDLSchema,
                input_object: &'a InputObject,
                field_name: &StringKey,
            ) -> Option<&'a InputObject> {
                input_object
                    .fields
                    .iter()
                    .find(|field| field.name.0 == *field_name)
                    .and_then(|field| resolve_root_input_field(schema, &field.type_))
            }

            let field_argument = field
                .arguments
                .iter()
                .find(|argument| argument.name() == name)?;

            let mut input_object = resolve_root_input_field(schema, &field_argument.type_)?;

            for input_field_name in input_field_path.iter() {
                input_object = resolve_input_field(schema, input_object, input_field_name)?;
            }

            Some(resolve_completion_items_for_input_object(
                input_object,
                schema,
                existing_names,
            ))
        }
    }
}

fn resolve_completion_items_typename(type_: Type, schema: &SDLSchema) -> Vec<CompletionItem> {
    if type_.is_root_type(schema) {
        vec![]
    } else {
        let mut item = CompletionItem::new_simple("__typename".to_owned(), "String!".to_owned());
        item.kind = Some(CompletionItemKind::FIELD);
        vec![item]
    }
}

fn resolve_completion_items_for_input_object(
    input_object: &InputObject,
    schema: &SDLSchema,
    existing_names: FnvHashSet<StringKey>,
) -> Vec<CompletionItem> {
    input_object
        .fields
        .iter()
        .filter(|arg| !existing_names.contains(&arg.name()))
        .map(|arg| {
            let label = arg.name().lookup().to_string();
            let detail = schema.get_type_string(arg.type_());
            let kind = match arg.type_().inner() {
                Type::InputObject(_) => Some(CompletionItemKind::STRUCT),
                Type::Scalar(_) => Some(CompletionItemKind::FIELD),
                _ => None,
            };

            CompletionItem {
                label: label.clone(),
                kind,
                detail: Some(detail),
                documentation: None,
                deprecated: None,
                preselect: None,
                sort_text: None,
                filter_text: None,
                insert_text: Some(format!("{}: $1", label)),
                insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
                text_edit: None,
                additional_text_edits: None,
                command: Some(lsp_types::Command::new(
                    "Suggest".into(),
                    "editor.action.triggerSuggest".into(),
                    None,
                )),
                data: None,
                tags: None,
                ..Default::default()
            }
        })
        .collect()
}

fn resolve_completion_items_for_argument_name<T: ArgumentLike>(
    arguments: impl Iterator<Item = T>,
    schema: &SDLSchema,
    existing_names: FnvHashSet<StringKey>,
    has_colon: bool,
) -> Vec<CompletionItem> {
    arguments
        .filter(|arg| !existing_names.contains(&arg.name()))
        .map(|arg| {
            let label = arg.name().lookup().into();
            let detail = schema.get_type_string(arg.type_());
            if has_colon {
                CompletionItem::new_simple(label, detail)
            } else {
                CompletionItem {
                    label: label.clone(),
                    kind: None,
                    detail: Some(detail),
                    documentation: None,
                    deprecated: None,
                    preselect: None,
                    sort_text: None,
                    filter_text: None,
                    insert_text: Some(format!("{}: $1", label)),
                    insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
                    text_edit: None,
                    additional_text_edits: None,
                    command: Some(lsp_types::Command::new(
                        "Suggest".into(),
                        "editor.action.triggerSuggest".into(),
                        None,
                    )),
                    data: None,
                    tags: None,
                    ..Default::default()
                }
            }
        })
        .collect()
}

fn resolve_completion_items_for_inline_fragment(
    type_: Type,
    schema: &SDLSchema,
    existing_inline_fragment: bool,
) -> Vec<CompletionItem> {
    match type_ {
        Type::Interface(id) => {
            let interface = schema.interface(id);
            once(type_)
                .chain(
                    interface
                        .implementing_objects
                        .iter()
                        .filter_map(|id| schema.get_type(schema.object(*id).name.item.0)),
                )
                .collect()
        }
        Type::Union(id) => {
            let union = schema.union(id);
            once(type_)
                .chain(
                    union
                        .members
                        .iter()
                        .filter_map(|id| schema.get_type(schema.object(*id).name.item.0)),
                )
                .collect()
        }
        Type::Enum(_) | Type::Object(_) | Type::InputObject(_) | Type::Scalar(_) => vec![],
    }
    .into_iter()
    .map(|type_| {
        let type_name = schema.get_type_name(type_).lookup();
        if existing_inline_fragment {
            CompletionItem::new_simple(type_name.to_owned(), "".into())
        } else {
            CompletionItem {
                label: format!("... on {type_name}"),
                kind: None,
                detail: None,
                documentation: None,
                deprecated: None,
                preselect: None,
                sort_text: None,
                filter_text: None,
                insert_text: Some(format!("... on {type_name} {{\n\t$1\n}}")),
                insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
                text_edit: None,
                additional_text_edits: None,
                command: Some(lsp_types::Command::new(
                    "Suggest".into(),
                    "editor.action.triggerSuggest".into(),
                    None,
                )),
                data: None,
                tags: None,
                ..Default::default()
            }
        }
    })
    .collect()
}

fn resolve_completion_items_for_argument_value(
    schema: &SDLSchema,
    type_: &TypeReference<Type>,
    program: &Program,
    executable_name: ExecutableName,
) -> Vec<CompletionItem> {
    let mut completion_items = match executable_name {
        ExecutableName::Fragment(name) => {
            if let Some(fragment) = program.fragment(name) {
                fragment
                    .used_global_variables
                    .iter()
                    .chain(fragment.variable_definitions.iter())
                    .filter(|variable| variable.type_.eq(type_))
                    .map(|variable| {
                        CompletionItem::new_simple(format!("${}", variable.name.item,), "".into())
                    })
                    .collect()
            } else {
                vec![]
            }
        }
        ExecutableName::Operation(name) => {
            if let Some(operation) = program.operation(OperationDefinitionName(name)) {
                operation
                    .variable_definitions
                    .iter()
                    .filter(|variable| variable.type_.eq(type_))
                    .map(|variable| {
                        CompletionItem::new_simple(format!("${}", variable.name.item,), "".into())
                    })
                    .collect()
            } else {
                vec![]
            }
        }
    };

    if !type_.is_list() {
        if let Type::Enum(id) = type_.inner() {
            let enum_ = schema.enum_(id);
            completion_items.extend(
                enum_
                    .values
                    .iter()
                    .map(|value| CompletionItem::new_simple(value.value.to_string(), "".into())),
            )
        }
    }

    completion_items
}

fn resolve_completion_items_for_fields<T: TypeWithFields + Named>(
    type_: &T,
    schema: &SDLSchema,
    schema_documentation: impl SchemaDocumentation,
    existing_linked_field: bool,
) -> Vec<CompletionItem> {
    type_
        .fields()
        .iter()
        .map(|field_id| {
            let field = schema.field(*field_id);
            let field_name = field.name.item.to_string();
            let deprecated = field.deprecated();
            let is_deprecated = deprecated.is_some();
            let deprecated_reason = deprecated
                .and_then(|deprecated| deprecated.reason)
                .map(|reason| format!("Deprecated: {}", reason));
            let args = create_arguments_snippets(field.arguments.iter(), schema);
            let insert_text = match (
                existing_linked_field
                    || matches!(field.type_.inner(), Type::Scalar(_) | Type::Enum(_)), // don't insert { }
                args.is_empty(), // don't insert arguments
            ) {
                (true, true) => None,
                (true, false) => Some(format!("{}({})", field_name, args.join(", "))),
                (false, true) => Some(format!("{} {{\n\t$1\n}}", field_name)),
                (false, false) => Some(format!(
                    "{}({}) {{\n\t${}\n}}",
                    field_name,
                    args.join(", "),
                    args.len() + 1
                )),
            };
            let (insert_text_format, command) = if insert_text.is_some() {
                (
                    Some(lsp_types::InsertTextFormat::SNIPPET),
                    Some(lsp_types::Command::new(
                        "Suggest".into(),
                        "editor.action.triggerSuggest".into(),
                        None,
                    )),
                )
            } else {
                (None, None)
            };

            let type_description = schema_documentation
                .get_type_description(schema.get_type_name(field.type_.inner()).lookup());

            let field_description = schema_documentation
                .get_field_description(type_.name().lookup(), field.name.item.lookup());

            let type_name = schema.get_type_string(&field.type_);
            let documentation = make_markdown_table_documentation(
                field.name.item.lookup(),
                &type_name,
                field_description.unwrap_or(""),
                type_description.unwrap_or(""),
            );

            let kind = match field.type_.inner() {
                Type::Enum(_) => Some(CompletionItemKind::ENUM),
                // There is no Kind for union, so we'll use interface
                Type::Interface(_) | Type::Union(_) => Some(CompletionItemKind::INTERFACE),
                Type::Object(_) | Type::InputObject(_) => Some(CompletionItemKind::STRUCT),
                Type::Scalar(_) => Some(CompletionItemKind::FIELD),
            };

            CompletionItem {
                label: field_name,
                kind,
                detail: deprecated_reason.or(Some(type_name)),
                documentation: Some(documentation),
                deprecated: Some(is_deprecated),
                preselect: None,
                sort_text: None,
                filter_text: None,
                insert_text,
                insert_text_format,
                text_edit: None,
                additional_text_edits: None,
                command,
                data: None,
                tags: None,
                ..Default::default()
            }
        })
        .collect()
}

fn resolve_completion_items_for_fragment_spread(
    type_: Type,
    source_program: &Program,
    schema: &SDLSchema,
    existing_fragment_spread: bool,
) -> Vec<CompletionItem> {
    source_program
        .fragments()
        .filter(|fragment| schema.are_overlapping_types(fragment.type_condition, type_))
        .map(|fragment| {
            let label = if existing_fragment_spread {
                fragment.name.item.to_string()
            } else {
                format!("...{}", fragment.name.item)
            };
            let detail = schema
                .get_type_name(fragment.type_condition)
                .lookup()
                .to_string();
            if fragment.variable_definitions.is_empty() {
                return CompletionItem::new_simple(label, detail);
            }
            // Create a snippet if the fragment has required argumentDefinition with no default values
            let args = create_arguments_snippets(fragment.variable_definitions.iter(), schema);
            if args.is_empty() {
                return CompletionItem::new_simple(label, detail);
            }
            let insert_text = format!("{} @arguments({})", label, args.join(", "));
            CompletionItem {
                label,
                kind: None,
                detail: Some(detail),
                documentation: None,
                deprecated: None,
                preselect: None,
                sort_text: None,
                filter_text: None,
                insert_text: Some(insert_text),
                insert_text_format: Some(lsp_types::InsertTextFormat::SNIPPET),
                text_edit: None,
                additional_text_edits: None,
                command: Some(lsp_types::Command::new(
                    "Suggest".into(),
                    "editor.action.triggerSuggest".into(),
                    None,
                )),
                data: None,
                tags: None,
                ..Default::default()
            }
        })
        .collect()
}

fn merge_completion_items_ordered<I: IntoIterator<Item = Vec<CompletionItem>>>(
    completion_item_groups: I,
) -> Vec<CompletionItem> {
    completion_item_groups
        .into_iter()
        .enumerate()
        .flat_map(|(index, mut items)| {
            items.iter_mut().for_each(|item| {
                item.sort_text = Some(format!(
                    "{}{}",
                    index,
                    item.sort_text.clone().unwrap_or_else(|| item.label.clone())
                ));
            });
            items
        })
        .collect()
}

fn completion_item_from_directive(
    directive: &SchemaDirective,
    schema: &SDLSchema,
) -> CompletionItem {
    let SchemaDirective {
        name, arguments, ..
    } = directive;

    // Always use the name of the directive as the label
    let label = name.to_string();

    // We can return a snippet with the expected arguments of the directive
    let (insert_text, insert_text_format) = if arguments.is_empty() {
        (label.clone(), InsertTextFormat::PLAIN_TEXT)
    } else {
        let args = create_arguments_snippets(arguments.iter(), schema);
        if args.is_empty() {
            (label.clone(), InsertTextFormat::PLAIN_TEXT)
        } else {
            let insert_text = format!("{}({})", label, args.join(", "));
            (insert_text, InsertTextFormat::SNIPPET)
        }
    };

    let documentation = directive.description.map(|desc| {
        Documentation::MarkupContent(MarkupContent {
            kind: MarkupKind::Markdown,
            value: desc.to_string(),
        })
    });

    CompletionItem {
        label,
        kind: None,
        detail: None,
        documentation,
        deprecated: None,
        preselect: None,
        sort_text: None,
        filter_text: None,
        insert_text: Some(insert_text),
        insert_text_format: Some(insert_text_format),
        text_edit: None,
        additional_text_edits: None,
        command: Some(lsp_types::Command::new(
            "Suggest".into(),
            "editor.action.triggerSuggest".into(),
            None,
        )),
        data: None,
        tags: None,
        ..Default::default()
    }
}

fn create_arguments_snippets<T: ArgumentLike>(
    arguments: impl Iterator<Item = T>,
    schema: &SDLSchema,
) -> Vec<String> {
    let mut cursor_location = 1;
    let mut args = vec![];

    for arg in arguments {
        if let TypeReference::NonNull(type_) = arg.type_() {
            let value_snippet = match type_ {
                t if t.is_list() => format!("[${}]", cursor_location),
                t if schema.is_string(t.inner()) => format!("\"${}\"", cursor_location),
                _ => format!("${}", cursor_location),
            };
            let str = format!("{}: {}", arg.name(), value_snippet);
            args.push(str);
            cursor_location += 1;
        }
    }
    args
}

pub fn on_completion(
    state: &impl GlobalState,
    params: <Completion as Request>::Params,
) -> LSPRuntimeResult<<Completion as Request>::Result> {
    match state.extract_executable_document_from_text(&params.text_document_position, 0) {
        Ok((document, position_span)) => {
            let project_name = state
                .extract_project_name_from_url(&params.text_document_position.text_document.uri)?;
            let schema = &state.get_schema(&project_name)?;
            let items = resolve_completion_items(
                document,
                position_span,
                project_name,
                schema,
                state.get_schema_documentation(project_name.lookup()),
                &state.get_program(&project_name)?,
            )
            .unwrap_or_else(Vec::new);
            Ok(Some(CompletionResponse::Array(items)))
        }
        Err(graphql_err) => {
            let js_server = state.get_js_language_sever().ok_or_else(|| {
                if matches!(graphql_err, LSPRuntimeError::ExpectedError) {
                    LSPRuntimeError::ExpectedError
                } else {
                    LSPRuntimeError::UnexpectedError(format!(
                        "Unable to get completion {:?}",
                        &graphql_err,
                    ))
                }
            })?;

            if let Ok(response) = js_server.on_complete(&params, state) {
                Ok(response)
            } else {
                Err(graphql_err)
            }
        }
    }
}

pub(crate) fn on_resolve_completion_item(
    _state: &impl GlobalState,
    params: <ResolveCompletionItem as Request>::Params,
) -> LSPRuntimeResult<<ResolveCompletionItem as Request>::Result> {
    // We currently don't do anything with the selected item
    // and we just return an input

    Ok(params)
}

fn resolve_completion_items(
    document: ExecutableDocument,
    position_span: Span,
    project_name: StringKey,
    schema: &SDLSchema,
    schema_documentation: impl SchemaDocumentation,
    progam: &Program,
) -> Option<Vec<CompletionItem>> {
    let completion_request = CompletionRequestBuilder::new(project_name)
        .create_completion_request(document, position_span);
    completion_request.and_then(|completion_request| {
        completion_items_for_request(completion_request, schema, schema_documentation, progam)
    })
}

fn make_markdown_table_documentation(
    field_name: &str,
    type_name: &str,
    field_description: &str,
    type_description: &str,
) -> Documentation {
    Documentation::MarkupContent(MarkupContent {
        kind: MarkupKind::Markdown,
        value: [
            format!("| **Field: {}** |", field_name),
            "| :--- |".to_string(),
            format!("| {} |", field_description),
            format!("| **Type: {}** |", type_name),
            format!("| {} |", type_description),
        ]
        .join("\n"),
    })
}

#[cfg(test)]
mod test;


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

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

use common::SourceLocationKey;
use common::Span;
use graphql_ir::build;
use graphql_ir::Program;
use graphql_syntax::parse_executable;
use graphql_syntax::parse_executable_with_error_recovery;
use intern::string_key::Intern;
use lsp_types::CompletionItem;
use lsp_types::Documentation;
use relay_test_schema::get_test_schema;

use super::make_markdown_table_documentation;
use super::resolve_completion_items;

fn parse_and_resolve_completion_items(
    source: &str,
    program: Option<Program>,
) -> Option<Vec<CompletionItem>> {
    let pos = source.find('|').unwrap() - 1;
    let next_source = source.replace("|", "");
    let document = parse_executable_with_error_recovery(
        &next_source,
        SourceLocationKey::standalone("/test/file"),
    )
    .item;

    let position_span = Span {
        start: pos as u32,
        end: pos as u32,
    };
    let test_schema = get_test_schema();

    resolve_completion_items(
        document,
        position_span,
        "test_project".intern(),
        &Arc::clone(&test_schema),
        Arc::clone(&test_schema),
        &program.unwrap_or_else(|| Program::new(Arc::clone(&test_schema))),
    )
}

fn build_test_program(source: &str) -> Program {
    let document = parse_executable(source, SourceLocationKey::Generated).unwrap();
    let ir = build(&get_test_schema(), &document.definitions).unwrap();
    Program::from_definitions(get_test_schema(), ir)
}

fn assert_labels(items: Vec<CompletionItem>, labels: Vec<&str>) {
    let mut completion_labels = items
        .into_iter()
        .map(|item| item.label)
        .collect::<HashSet<_>>();

    assert_eq!(
        completion_labels.len(),
        labels.len(),
        "Provided labels {:?} do not match completion items {:?}",
        &completion_labels,
        &labels
    );
    for label in labels {
        assert!(
            completion_labels.remove(label),
            "Expected to have {} in the set",
            label
        );
    }
    assert!(completion_labels.is_empty());
}

fn assert_no_typename_label(items: Option<Vec<CompletionItem>>) {
    assert!(
        items
            .unwrap()
            .into_iter()
            .all(|item| item.label != "__typename"),
        "__typename is invalid on root types"
    );
}

#[test]
fn scalar_field() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on User {
                profile_picture {
                    uri|
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["uri", "width", "height", "test_enums", "__typename"],
    );
}

#[test]
fn linked_field() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
                marketplace_settings {
                    locatio| {
                    }
                }
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec!["location", "categories", "__typename"]);
}

#[test]
fn whitespace_in_linked_field() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on User {
                profile_picture {
                    uri
                    |
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["uri", "test_enums", "width", "height", "__typename"],
    );
}

#[test]
fn whitespace_in_fragment() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Image {
                |
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["uri", "width", "height", "test_enums", "__typename"],
    );
}

#[test]
fn whitespace_in_inline_fragment() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Node {
                ... on Image {
                    |
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["uri", "width", "height", "test_enums", "__typename"],
    );
}

#[test]
fn whitespace_in_object_type() {
    let items = parse_and_resolve_completion_items(
        r#"
          fragment Test on CommentsEdge {
            |
          }
        "#,
        Some(build_test_program(
            r#"
            fragment ObjectTypeFragment on CommentsEdge {
              __typename
            }

            fragment InterfaceFragment on CommentsEdgeInterface {
              __typename
            }

            fragment UnrelatedFragment on Task {
              __typename
            }
        "#,
        )),
    );
    assert_labels(
        items.unwrap(),
        vec![
            "cursor",
            "node",
            "__typename",
            "...ObjectTypeFragment",
            "...InterfaceFragment",
        ],
    )
}

#[test]
fn whitespace_in_interface() {
    let items = parse_and_resolve_completion_items(
        r#"
          fragment Test on CommentsEdgeInterface {
            |
          }
        "#,
        Some(build_test_program(
            r#"
            fragment ImplementingFragment on CommentsEdge {
              __typename
            }

            fragment InterfaceFragment on CommentsEdgeInterface {
              __typename
            }

            fragment UnrelatedFragment on Task {
              __typename
            }
        "#,
        )),
    );
    assert_labels(
        items.unwrap(),
        vec![
            "cursor",
            "source",
            "node",
            "__typename",
            "... on CommentsEdgeInterface",
            "... on CommentsEdge",
            "...ImplementingFragment",
            "...InterfaceFragment",
        ],
    )
}

#[test]
fn whitespace_in_union() {
    let items = parse_and_resolve_completion_items(
        r#"
          fragment Test on CommentBody {
            |
          }
        "#,
        Some(build_test_program(
            r#"
                fragment UnionFragment on CommentBody {
                  __typename
                }

                fragment UnionVariantFragment on PlainCommentBody {
                  __typename
                }

                fragment UnrelatedFragment on Task {
                  __typename
                }
            "#,
        )),
    );
    assert_labels(
        items.unwrap(),
        vec![
            "__typename",
            "... on CommentBody",
            "... on PlainCommentBody",
            "... on MarkdownCommentBody",
            "...UnionFragment",
            "...UnionVariantFragment",
        ],
    )
}

#[test]
fn inline_fragment_on_object() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
                ... on |a
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec![]);
}

#[test]
fn inline_fragment_on_interface() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Named {
                ... on |a
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["... on Named", "... on User", "... on SimpleNamed"],
    );
}

#[test]
fn inline_fragment_on_interface_with_existing_inline_fragment() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Named {
                ... on |a {}
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec!["Named", "User", "SimpleNamed"]);
}

#[test]
fn inline_fragment_on_union() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on MaybeNode {
                ... on |a
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec![
            "... on MaybeNode",
            "... on Story",
            "... on FakeNode",
            "... on NonNode",
        ],
    );
}

#[test]
fn inline_fragment_on_union_with_existing_inline_fragment() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on MaybeNode {
                ... on |a {}
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["MaybeNode", "Story", "FakeNode", "NonNode"],
    );
}

#[test]
fn directive() {
    let items = parse_and_resolve_completion_items(
        r#"
        fragment Test on User {
            profile_picture @a| {
                uri
            }
        }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec![
            "prependEdge",
            "deleteRecord",
            "appendNode",
            "deleteEdge",
            "__clientField",
            "appendEdge",
            "required",
            "stream_connection",
            "match",
            "customDirective",
            "prependNode",
            "fixme_fat_interface",
            "stream",
            "include",
            "connection",
            "skip",
            "fb_actor_change",
            "waterfall",
        ],
    );
}

#[test]
fn directive_on_scalar_field() {
    let items = parse_and_resolve_completion_items(
        r#"
        fragment Test on User {
            profile_picture {
                uri @|
            }
        }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec![
            "prependEdge",
            "deleteRecord",
            "appendNode",
            "deleteEdge",
            "__clientField",
            "appendEdge",
            "required",
            "stream_connection",
            "match",
            "customDirective",
            "prependNode",
            "fixme_fat_interface",
            "stream",
            "include",
            "connection",
            "skip",
            "fb_actor_change",
            "waterfall",
        ],
    );
}

#[test]
fn empty_argument_list() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
                profile_picture @stream(|) {
                    uri
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["label", "initial_count", "if", "use_customized_batch"],
    );
}

#[test]
fn argument_name_without_value() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
                profile_picture @stream(i|) {
                    uri
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["label", "initial_count", "if", "use_customized_batch"],
    );
}

#[test]
fn argument_name_with_existing_name() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
                profile_picture @stream(
                    if: false
                    |
                ) {
                    uri
                }
            }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec!["label", "initial_count", "use_customized_batch"],
    );
}

#[test]
fn fragment_spread_on_the_same_type() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Viewer {
               ...T|
            }
        "#,
        Some(build_test_program(
            r#"
        fragment TestFragment on Viewer {
           __typename
        }

        fragment TestFragment2 on Viewer {
            __typename
         }
      "#,
        )),
    );
    assert_labels(items.unwrap(), vec!["TestFragment", "TestFragment2"]);
}

#[test]
fn fragment_spread_on_interface() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Actor {
               ...T|
            }
        "#,
        Some(build_test_program(
            r#"
        fragment TestFragment on Page {
           __typename
        }

        fragment TestFragment2 on Node {
            __typename
        }
    "#,
        )),
    );
    assert_labels(items.unwrap(), vec!["TestFragment", "TestFragment2"]);
}

#[test]
fn argument_value_object() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Mutation {
                commentCreate(inpu|) {
                    __typename
                }
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec!["input"]);
}

#[test]
fn argument_value_constant_object() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Mutation {
                commentCreate(input: {
                    feedbackId: "some-id"
                    |
                }) {
                    __typename
                }
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec!["client_mutation_id", "feedback"]);
}

#[test]
fn argument_value_constant_object_nested() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Mutation {
                commentCreate(input: {
                    feedbackId: "some-id"
                    feedback: {
                        |
                    }
                }) {
                    __typename
                }
            }
        "#,
        None,
    );
    assert_labels(items.unwrap(), vec!["comment"]);
}

#[test]
fn argument_value() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on User {
                profilePicture(size:|) {
                    uri
                }
            }
        "#,
        Some(build_test_program(
            r#"
            fragment Test on User
                @argumentDefinitions(
                    pictureSize: {type: "[Int]"},
                    pictureSize2: {type: "[Int]"}
                ) {
                profilePicture(size: $pictureSize) {
                    uri
                }
            }
        "#,
        )),
    );
    assert_labels(items.unwrap(), vec!["$pictureSize", "$pictureSize2"]);
}

#[test]
fn argument_value_between_names() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on User {
                profilePicture(size: |, preset: SMALL) {
                    uri
                }
            }
        "#,
        Some(build_test_program(
            r#"
            fragment Test on User
                @argumentDefinitions(
                    pictureSize: {type: "[Int]"},
                    pictureSize2: {type: "[Int]"}
                ) {
                profilePicture(size: $pictureSize) {
                    uri
                }
            }
        "#,
        )),
    );
    assert_labels(items.unwrap(), vec!["$pictureSize", "$pictureSize2"]);
}

#[test]
fn no_typename_on_query() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Query {
                |
            }
        "#,
        None,
    );
    assert_no_typename_label(items);
}

#[test]
fn no_typename_on_mutation() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Mutation {
                |
            }
        "#,
        None,
    );
    assert_no_typename_label(items);
}

#[test]
fn no_typename_on_subscription() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on Subscription {
                |
            }
        "#,
        None,
    );
    assert_no_typename_label(items);
}

#[test]
fn empty_directive() {
    let items = parse_and_resolve_completion_items(
        r#"
        fragment Test on User {
            profile_picture @| {
                uri
            }
        }
        "#,
        None,
    );
    assert_labels(
        items.unwrap(),
        vec![
            "prependEdge",
            "deleteRecord",
            "appendNode",
            "deleteEdge",
            "__clientField",
            "appendEdge",
            "required",
            "stream_connection",
            "match",
            "customDirective",
            "prependNode",
            "fixme_fat_interface",
            "stream",
            "include",
            "connection",
            "skip",
            "fb_actor_change",
            "waterfall",
        ],
    );
}

#[test]
fn field_documentation() {
    let items = parse_and_resolve_completion_items(
        r#"
            fragment Test on User {
                profile_picture {
                    uri|
                }
            }
        "#,
        None,
    )
    .unwrap();

    let docs = items
        .into_iter()
        .map(|item| (item.label, item.documentation))
        .collect::<HashMap<String, Option<Documentation>>>();

    assert_eq!(docs.len(), 5);
    assert_eq!(
        *docs.get("uri").unwrap(),
        Some(make_markdown_table_documentation(
            "uri",
            "String",
            "URI where the image can be found",
            "",
        ))
    );

    assert_eq!(
        *docs.get("width").unwrap(),
        Some(make_markdown_table_documentation(
            "width",
            "Int",
            "Width in pixels",
            "",
        ))
    );

    assert_eq!(
        *docs.get("height").unwrap(),
        Some(make_markdown_table_documentation(
            "height",
            "Int",
            "Height in pixels",
            "",
        ))
    );

    assert_eq!(
        *docs.get("test_enums").unwrap(),
        Some(make_markdown_table_documentation(
            "test_enums",
            "TestEnums",
            "",
            "",
        ))
    );
}


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

use intern::string_key::Intern;
use lsp_types::request::Request;
use schema::Schema;
use serde::Deserialize;
use serde::Serialize;

use crate::lsp_runtime_error::LSPRuntimeError;
use crate::lsp_runtime_error::LSPRuntimeResult;
use crate::server::GlobalState;

mod types;

pub(crate) struct ExploreSchemaForType {}

#[derive(Deserialize, Serialize)]
pub(crate) struct ExploreSchemaForTypeParams {
    pub item: String,
    pub filter: Option<String>,
    pub schema_name: String,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct ExploreSchemaForTypeResult {
    schema_item: types::SchemaExplorerTypeReference<types::SchemaExplorerSchemaType>,
}

impl Request for ExploreSchemaForType {
    type Params = ExploreSchemaForTypeParams;
    type Result = ExploreSchemaForTypeResult;
    const METHOD: &'static str = "relay/exploreSchemaForType";
}

pub(crate) fn on_explore_schema_for_type(
    state: &impl GlobalState,
    params: <ExploreSchemaForType as Request>::Params,
) -> LSPRuntimeResult<<ExploreSchemaForType as Request>::Result> {
    // TODO these should not be expected errors
    let schema_name: &str = &params.schema_name;
    let schema = state.get_schema(&schema_name.intern())?;

    let type_ = if params.item == "Query" {
        schema.query_type()
    } else if params.item == "Subscription" {
        schema.subscription_type()
    } else if params.item == "Mutation" {
        schema.mutation_type()
    } else {
        let item: &str = &params.item;
        schema.get_type(item.intern())
    };

    // TODO these should not be expected errors
    let schema_item = type_
        .map(|type_| {
            types::get_full_schema_explorer_type_reference(
                type_,
                &params.item,
                &schema,
                &state.get_schema_documentation(&params.schema_name),
                &params.filter.as_ref().map(|x| x.to_lowercase()),
                None,
            )
        })
        .ok_or(LSPRuntimeError::ExpectedError)?;

    Ok(ExploreSchemaForTypeResult { schema_item })
}


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

use intern::string_key::StringKey;
use intern::Lookup;
use schema::DirectiveValue;
use schema::EnumID;
use schema::FieldID;
use schema::InputObjectID;
use schema::InterfaceID;
use schema::ObjectID;
use schema::SDLSchema;
use schema::ScalarID;
use schema::Schema;
use schema::Type;
use schema::UnionID;
use schema_documentation::SchemaDocumentation;
use serde::Deserialize;
use serde::Serialize;

const DEFAULT_COUNT: usize = 100;

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerTypeReference<T> {
    name: String,
    description: Option<String>,
    kind: String,
    #[serde(flatten)]
    inner: T,
}

#[derive(Deserialize, Serialize)]
#[serde(tag = "kind")]
pub(crate) enum SchemaExplorerSchemaType {
    Object(SchemaExplorerObject),
    Interface(SchemaExplorerInterface),
    Enum(SchemaExplorerEnum),
    Union(SchemaExplorerUnion),
    Scalar(SchemaExplorerScalar),
    InputObject(SchemaExplorerInputObject),
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerInterface {
    implementing_objects: Vec<SchemaExplorerTypeReference<()>>,
    fields: Vec<SchemaExplorerField>,
    interfaces: Vec<SchemaExplorerTypeReference<()>>,
    is_extension: bool,
    directives: Vec<SchemaExplorerDirective>,
}
#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerObject {
    fields: Vec<SchemaExplorerField>,
    interfaces: Vec<SchemaExplorerTypeReference<()>>,
    is_extension: bool,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerInputObject {
    // This is named differently than in definitions.rs
    arguments: Vec<SchemaExplorerFieldArgument>,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerEnum {
    values: Vec<String>,
    is_extension: bool,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerUnion {
    is_extension: bool,
    members: Vec<SchemaExplorerTypeReference<()>>,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerField {
    field_name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    field_description: Option<String>,
    arguments: Vec<SchemaExplorerFieldArgument>,
    rendered_type_name: String,
    type_reference: SchemaExplorerTypeReference<()>,
    is_extension: bool,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerFieldArgument {
    argument_name: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    argument_description: Option<String>,
    default_value: Option<String>,
    rendered_type_name: String,
    type_reference: SchemaExplorerTypeReference<()>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerScalar {
    is_extension: bool,
    directives: Vec<SchemaExplorerDirective>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerDirective {
    directive_name: String,
    arguments: Vec<SchemaExplorerDirectiveArgument>,
}

#[derive(Deserialize, Serialize)]
pub(crate) struct SchemaExplorerDirectiveArgument {
    argument_name: String,
    value: String,
}

pub(crate) fn get_full_schema_explorer_type_reference(
    type_: Type,
    type_name: &str,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
    filter: &Option<String>,
    count: Option<usize>,
) -> SchemaExplorerTypeReference<SchemaExplorerSchemaType> {
    let description = documentation
        .get_type_description(type_name)
        .map(|description| description.to_string());

    let schema_explorer_schema_type = match type_ {
        Type::Object(object_id) => {
            let object =
                get_schema_explorer_object(object_id, schema, documentation, filter, count);
            SchemaExplorerSchemaType::Object(object)
        }
        Type::Interface(interface_id) => {
            let interface =
                get_schema_explorer_interface(interface_id, schema, documentation, filter);
            SchemaExplorerSchemaType::Interface(interface)
        }
        Type::Enum(enum_id) => {
            let enum_ = get_schema_explorer_enum(enum_id, schema);
            SchemaExplorerSchemaType::Enum(enum_)
        }
        Type::InputObject(input_object_id) => {
            let input_object =
                get_schema_explorer_input_object(input_object_id, schema, documentation);
            SchemaExplorerSchemaType::InputObject(input_object)
        }
        Type::Union(union_id) => {
            let union_ = get_schema_explorer_union(union_id, schema, documentation);
            SchemaExplorerSchemaType::Union(union_)
        }
        Type::Scalar(scalar_id) => {
            let scalar = get_schema_explorer_scalar(scalar_id, schema);
            SchemaExplorerSchemaType::Scalar(scalar)
        }
    };

    SchemaExplorerTypeReference {
        kind: get_kind(type_),
        name: type_name.to_string(),
        description,
        inner: schema_explorer_schema_type,
    }
}

fn get_empty_schema_explorer_type_reference(
    type_: Type,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
) -> SchemaExplorerTypeReference<()> {
    let name = schema.get_type_name(type_).to_string();
    let description = documentation
        .get_type_description(&name)
        .map(|description| description.to_string());
    SchemaExplorerTypeReference {
        kind: get_kind(type_),
        name,
        description,
        inner: (),
    }
}

fn get_schema_explorer_enum(enum_id: EnumID, schema: &SDLSchema) -> SchemaExplorerEnum {
    let enum_ = schema.enum_(enum_id);
    let values = enum_
        .values
        .iter()
        .map(|value| value.value.to_string())
        .collect::<Vec<_>>();
    SchemaExplorerEnum {
        values,
        is_extension: enum_.is_extension,
        directives: get_schema_explorer_directives(&enum_.directives),
    }
}

fn get_schema_explorer_interface(
    interface_id: InterfaceID,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
    filter: &Option<String>,
) -> SchemaExplorerInterface {
    let interface = schema.interface(interface_id);
    let implementing_objects = interface
        .implementing_objects
        .iter()
        .map(|object_id| {
            get_empty_schema_explorer_type_reference(
                Type::Object(*object_id),
                schema,
                documentation,
            )
        })
        .collect::<Vec<_>>();

    let interface_type_name = schema.get_type_name(Type::Interface(interface_id));
    let fields = interface
        .fields
        .iter()
        .filter_map(|field_id| {
            get_schema_explorer_field(
                *field_id,
                interface_type_name,
                schema,
                documentation,
                filter,
            )
        })
        .collect::<Vec<_>>();

    let subinterfaces = interface
        .interfaces
        .iter()
        .map(|interface_id| {
            get_empty_schema_explorer_type_reference(
                Type::Interface(*interface_id),
                schema,
                documentation,
            )
        })
        .collect::<Vec<_>>();

    SchemaExplorerInterface {
        implementing_objects,
        fields,
        interfaces: subinterfaces,
        is_extension: interface.is_extension,
        directives: get_schema_explorer_directives(&interface.directives),
    }
}

fn get_schema_explorer_input_object(
    input_object_id: InputObjectID,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
) -> SchemaExplorerInputObject {
    let input_object = schema.input_object(input_object_id);
    let arguments = input_object
        .fields
        .iter()
        .map(|arg| {
            let type_reference =
                get_empty_schema_explorer_type_reference(arg.type_.inner(), schema, documentation);

            SchemaExplorerFieldArgument {
                argument_name: arg.name.to_string(),
                argument_description: None,
                default_value: arg.default_value.as_ref().map(|value| value.to_string()),
                rendered_type_name: schema.get_type_string(&arg.type_),
                type_reference,
            }
        })
        .collect::<Vec<_>>();

    SchemaExplorerInputObject {
        arguments,
        directives: get_schema_explorer_directives(&input_object.directives),
    }
}

fn get_schema_explorer_object(
    object_id: ObjectID,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
    filter: &Option<String>,
    count: Option<usize>,
) -> SchemaExplorerObject {
    let object = schema.object(object_id);
    let object_type_name = schema.get_type_name(Type::Object(object_id));
    let fields = object
        .fields
        .iter()
        .filter_map(|field_id| {
            get_schema_explorer_field(*field_id, object_type_name, schema, documentation, filter)
        })
        .take(count.unwrap_or(DEFAULT_COUNT))
        .collect::<Vec<_>>();

    let interfaces = object
        .interfaces
        .iter()
        .map(|interface_id| {
            get_empty_schema_explorer_type_reference(
                Type::Interface(*interface_id),
                schema,
                documentation,
            )
        })
        .collect::<Vec<_>>();

    SchemaExplorerObject {
        fields,
        interfaces,
        is_extension: object.is_extension,
        directives: get_schema_explorer_directives(&object.directives),
    }
}

fn get_schema_explorer_union(
    union_id: UnionID,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
) -> SchemaExplorerUnion {
    let union_ = schema.union(union_id);

    let members = union_
        .members
        .iter()
        .map(|id| {
            get_empty_schema_explorer_type_reference(Type::Object(*id), schema, documentation)
        })
        .collect::<Vec<_>>();

    SchemaExplorerUnion {
        members,
        is_extension: union_.is_extension,
        directives: get_schema_explorer_directives(&union_.directives),
    }
}

fn get_schema_explorer_scalar(scalar_id: ScalarID, schema: &SDLSchema) -> SchemaExplorerScalar {
    let scalar = schema.scalar(scalar_id);

    SchemaExplorerScalar {
        is_extension: scalar.is_extension,
        directives: get_schema_explorer_directives(&scalar.directives),
    }
}

fn get_schema_explorer_field(
    field_id: FieldID,
    parent_type_name: StringKey,
    schema: &SDLSchema,
    documentation: &impl SchemaDocumentation,
    filter: &Option<String>,
) -> Option<SchemaExplorerField> {
    let field = schema.field(field_id);
    let field_name = field.name.item.to_string();

    let field_type_name = schema.get_type_name(field.type_.inner()).to_string();

    let field_type_description = documentation
        .get_type_description(&field_type_name)
        .map(|field_type_description| field_type_description.to_string());
    let field_description = documentation
        .get_field_description(parent_type_name.lookup(), &field_name)
        .map(|field_description| field_description.to_string());

    if let Some(filter) = filter {
        if !field_passes_filter(
            filter,
            field_name.to_lowercase(),
            field_type_name.to_lowercase(),
            field_type_description.as_ref().map(|s| s.to_lowercase()),
        ) {
            return None;
        }
    }

    let type_reference =
        get_empty_schema_explorer_type_reference(field.type_.inner(), schema, documentation);

    let arguments = field
        .arguments
        .iter()
        .map(|arg| {
            let type_reference =
                get_empty_schema_explorer_type_reference(arg.type_.inner(), schema, documentation);

            let argument_description = documentation
                .get_field_argument_description(
                    parent_type_name.lookup(),
                    &field_name,
                    arg.name.0.lookup(),
                )
                .map(|field_argument_description| field_argument_description.to_string());

            SchemaExplorerFieldArgument {
                argument_name: arg.name.to_string(),
                argument_description,
                rendered_type_name: schema.get_type_string(&arg.type_),
                type_reference,
                default_value: arg.default_value.as_ref().map(|value| value.to_string()),
            }
        })
        .collect::<Vec<_>>();

    Some(SchemaExplorerField {
        field_name,
        field_description,
        type_reference,
        rendered_type_name: schema.get_type_string(&field.type_),
        arguments,
        is_extension: field.is_extension,
        directives: get_schema_explorer_directives(&field.directives),
    })
}

fn field_passes_filter(
    filter: &str,
    field_name: String,
    field_type_name: String,
    field_type_description: Option<String>,
) -> bool {
    if field_name.contains(filter) || field_type_name.contains(filter) {
        true
    } else if let Some(field_type_description) = field_type_description {
        field_type_description.contains(filter)
    } else {
        false
    }
}

fn get_schema_explorer_directives(directives: &[DirectiveValue]) -> Vec<SchemaExplorerDirective> {
    directives
        .iter()
        .map(|directive| -> SchemaExplorerDirective {
            SchemaExplorerDirective {
                directive_name: directive.name.to_string(),
                arguments: directive
                    .arguments
                    .iter()
                    .map(|arg| SchemaExplorerDirectiveArgument {
                        argument_name: arg.name.to_string(),
                        value: arg.value.to_string(),
                    })
                    .collect::<Vec<_>>(),
            }
        })
        .collect::<Vec<_>>()
}

fn get_kind(type_: Type) -> String {
    match type_ {
        Type::Enum(_) => "Enum".to_string(),
        Type::Union(_) => "Union".to_string(),
        Type::InputObject(_) => "InputObject".to_string(),
        Type::Object(_) => "Object".to_string(),
        Type::Scalar(_) => "Scalar".to_string(),
        Type::Interface(_) => "Interface".to_string(),
    }
}


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

use common::Location as IRLocation;
use graphql_ir::Program;
use graphql_ir::Visitor;
use intern::string_key::StringKey;
use schema::Schema;

use crate::find_field_usages::FieldUsageFinder;

pub(crate) fn find_field_locations(
    program: &Program,
    field_name: StringKey,
    type_name: StringKey,
) -> Option<Vec<IRLocation>> {
    let type_ = program.schema.get_type(type_name)?;
    let mut usage_finder = FieldUsageFinder::new(&program.schema, type_, field_name);
    usage_finder.visit_program(program);

    Some(usage_finder.get_locations())
}


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

mod find_field_locations;

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

use common::Location as IRLocation;
use common::WithLocation;
pub(crate) use find_field_locations::find_field_locations;
use graphql_ir::FragmentDefinition;
use graphql_ir::InlineFragment;
use graphql_ir::LinkedField;
use graphql_ir::OperationDefinition;
use graphql_ir::Program;
use graphql_ir::ScalarField;
use graphql_ir::Visitor;
use intern::string_key::Intern;
use intern::string_key::StringKey;
use itertools::Itertools;
use lsp_types::request::Request;
use schema::FieldID;
use schema::SDLSchema;
use schema::Schema;
use schema::Type;
use serde::Deserialize;
use serde::Serialize;

use crate::location::transform_relay_location_to_lsp_location;
use crate::server::GlobalState;
use crate::LSPRuntimeError;
use crate::LSPRuntimeResult;

// This implementation of FindFieldUsages find matching fields in:
//   - exact type matches
//   - subtypes (of the input type)
//   - (not supertypes, since this could introduce many false positives)
// It currently does a shallow traversal of Operations and Fragments,
//  not following fragment spreads to named fragments
//   - this could result in false negatives, but saves on memory / time
pub struct FindFieldUsages {}

#[derive(Deserialize, Serialize)]
pub struct FindFieldUsagesParams {
    pub schema_name: String,
    pub type_name: String,
    pub field_name: String,
}

#[derive(Deserialize, Serialize)]
struct FindFieldUsageResultItem {
    location_uri: String,
    location_range: lsp_types::Range,
    label: String,
}

#[derive(Deserialize, Serialize)]
pub struct FindFieldUsagesResult {
    usages: Vec<FindFieldUsageResultItem>,
}

impl Request for FindFieldUsages {
    type Params = FindFieldUsagesParams;
    type Result = FindFieldUsagesResult;
    const METHOD: &'static str = "relay/findFieldUsages";
}

pub fn on_find_field_usages(
    state: &impl GlobalState,
    params: <FindFieldUsages as Request>::Params,
) -> LSPRuntimeResult<<FindFieldUsages as Request>::Result> {
    let schema_name = params.schema_name.intern();
    let type_name = params.type_name.intern();
    let field_name = params.field_name.intern();

    let schema = state.get_schema(&schema_name)?;
    let program = state.get_program(&schema_name)?;
    let root_dir = &state.root_dir();

    let ir_locations = get_usages(&program, &schema, type_name, field_name)?;
    let lsp_locations = ir_locations
        .into_iter()
        .map(|(label, ir_location)| {
            let lsp_location = transform_relay_location_to_lsp_location(root_dir, ir_location)?;
            Ok(FindFieldUsageResultItem {
                location_uri: lsp_location.uri.to_string(),
                location_range: lsp_location.range,
                label,
            })
        })
        .collect::<Result<Vec<_>, LSPRuntimeError>>()?;
    Ok(FindFieldUsagesResult {
        usages: lsp_locations,
    })
}

pub fn get_usages(
    program: &Program,
    schema: &Arc<SDLSchema>,
    type_name: StringKey,
    field_name: StringKey,
) -> LSPRuntimeResult<Vec<(String, IRLocation)>> {
    let type_ = schema.get_type(type_name).ok_or_else(|| {
        LSPRuntimeError::UnexpectedError(format!("Type {} not found!", type_name))
    })?;
    let mut usage_finder = FieldUsageFinder::new(schema, type_, field_name);
    usage_finder.visit_program(program);

    let mut result = Vec::with_capacity(usage_finder.usages.len());
    for (label, locations) in usage_finder.usages.into_iter().sorted() {
        if let [location] = locations.as_slice() {
            // exactly 1 location, so no need to enumerate
            result.push((label.to_string(), *location));
        } else {
            for (idx, location) in locations.into_iter().enumerate() {
                result.push((format!("{} - {}", label, idx), location));
            }
        }
    }
    Ok(result)
}

#[derive(Default)]
struct FieldUsageFinderScope {
    // Types that a visited field acts on
    //  - reset when we see a linked field, a fragment, or an operation
    //  - pushed when we see type refinement (inline fragments)
    types: Vec<StringKey>,
    // name of the enclosing Fragment or Operation
    label: Option<StringKey>,
}

pub(crate) struct FieldUsageFinder<'schema> {
    usages: HashMap<StringKey, Vec<IRLocation>>,
    schema: &'schema Arc<SDLSchema>,
    type_: Type,
    field_name: StringKey,
    current_scope: FieldUsageFinderScope,
}

impl<'schema> FieldUsageFinder<'schema> {
    pub(crate) fn new(
        schema: &'schema Arc<SDLSchema>,
        type_: Type,
        field_name: StringKey,
    ) -> FieldUsageFinder<'schema> {
        FieldUsageFinder {
            usages: Default::default(),
            schema,
            type_,
            field_name,
            current_scope: Default::default(),
        }
    }

    fn match_field(&mut self, field: &WithLocation<FieldID>) -> bool {
        // check field name match
        if self.schema.field(field.item).name.item == self.field_name {
            // check for
            // - exact type match
            // - inferred types (spread on concrete should match all abstract)
            for curr_typename in &self.current_scope.types {
                if self.schema.is_named_type_subtype_of(
                    self.schema.get_type(*curr_typename).unwrap(),
                    self.type_,
                ) {
                    return true;
                }
            }
        }
        false
    }

    fn add_field(&mut self, field: &WithLocation<FieldID>) {
        let current_label = self
            .current_scope
            .label
            .expect("Expected label in find_field_usages");
        self.usages
            .entry(current_label)
            .or_default()
            .push(field.location);
    }

    pub(crate) fn get_locations(&self) -> Vec<IRLocation> {
        self.usages
            .values()
            .into_iter()
            .flatten()
            .copied()
            .collect_vec()
    }
}

impl Visitor for FieldUsageFinder<'_> {
    const NAME: &'static str = "FieldUsageFinder";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn visit_operation(&mut self, operation: &OperationDefinition) {
        // Check that scope is empty + populate the name + type (Query/Mutation/Subscription)
        //  before recursively visiting the operation's selections
        assert!(self.current_scope.label.is_none());
        assert!(self.current_scope.types.is_empty());
        self.current_scope.label = Some(operation.name.item.0);
        self.current_scope
            .types
            .push(self.schema.get_type_name(operation.type_));

        self.default_visit_operation(operation);

        self.current_scope.types.pop();
        self.current_scope.label = None;
        assert!(self.current_scope.types.is_empty());
    }

    fn visit_fragment(&mut self, fragment: &FragmentDefinition) {
        assert!(self.current_scope.label.is_none());
        assert!(self.current_scope.types.is_empty());

        self.current_scope.label = Some(fragment.name.item.0);
        self.current_scope
            .types
            .push(self.schema.get_type_name(fragment.type_condition));

        self.default_visit_fragment(fragment);

        self.current_scope.types.pop();
        self.current_scope.label = None;
        assert!(self.current_scope.types.is_empty());
    }

    fn visit_inline_fragment(&mut self, fragment: &InlineFragment) {
        // Inline fragments might not have a type condition
        // fragment Foo on User {
        //   ... {
        //     name
        //   }
        // }
        let should_pop = if let Some(type_) = fragment.type_condition {
            self.current_scope
                .types
                .push(self.schema.get_type_name(type_));
            true
        } else {
            false
        };
        self.default_visit_inline_fragment(fragment);

        if should_pop {
            self.current_scope.types.pop();
        }
    }

    fn visit_linked_field(&mut self, field: &LinkedField) {
        if self.match_field(&field.definition) {
            self.add_field(&field.definition);
        }

        // save all enclosing types
        let prev_types = std::mem::take(&mut self.current_scope.types);

        // remember linked type
        let linked_type = self.schema.field(field.definition.item).type_.inner();
        self.current_scope
            .types
            .push(self.schema.get_type_name(linked_type));

        self.default_visit_linked_field(field);

        self.current_scope.types.pop();
        assert!(self.current_scope.types.is_empty());
        self.current_scope.types = prev_types;
    }

    fn visit_scalar_field(&mut self, field: &ScalarField) {
        if self.match_field(&field.definition) {
            self.add_field(&field.definition);
        }
    }
}


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

use common::Span;
use relay_docblock::DocblockIr;

use super::DefinitionDescription;
use crate::docblock_resolution_info::create_docblock_resolution_info;
use crate::docblock_resolution_info::DocblockResolutionInfo;
use crate::lsp_runtime_error::LSPRuntimeError;
use crate::LSPRuntimeResult;

pub fn get_docblock_definition_description(
    docblock_ir: &DocblockIr,
    position_span: Span,
) -> LSPRuntimeResult<DefinitionDescription> {
    let resolution = create_docblock_resolution_info(docblock_ir, position_span)
        .ok_or(LSPRuntimeError::ExpectedError)?;
    match resolution {
        DocblockResolutionInfo::Type(type_name) => Ok(DefinitionDescription::Type { type_name }),
        DocblockResolutionInfo::RootFragment(fragment_name) => {
            Ok(DefinitionDescription::Fragment { fragment_name })
        }
        DocblockResolutionInfo::FieldName(_) => {
            // The field name _id_ the definition of the field.
            Err(LSPRuntimeError::ExpectedError)
        }
        DocblockResolutionInfo::Deprecated => {
            // We don't currently have any go to definition support for
            // schema directives.
            Err(LSPRuntimeError::ExpectedError)
        }
    }
}