apollo-language-server 0.6.0

use apollo_compiler::{
    ast::{
        Definition, Directive, DirectiveDefinition, DirectiveLocation, Document,
        EnumTypeDefinition, InputObjectTypeDefinition, InputValueDefinition, NamedType,
        ScalarTypeDefinition, Type, Value,
    },
    InvalidNameError, Name, Node,
};
use ropey::Rope;
use std::collections::HashMap;
use url::Url;

use crate::{
    specs::{link::LINK_DIRECTIVE, SpecDirective, SpecType, KNOWN_SPECS},
    utils::lsp_range_from_ast_sourcespan::lsp_range_from_ast_sourcespan,
};

#[derive(Debug)]
// A link looks something like this:
// @link(url: "https://specs.apollo.dev/federation/v2.7", as: "fed", import: ["@key", { name: "@override", as: "ovr" }])
// In order to consume this in a meaningful way, we need to parse out the name and version from the slug,
// the alias `as` (if it exists), and the imports.
pub struct ParsedLink {
    pub spec_name: String,
    pub version: semver::Version,
    pub imported_as: String,
    pub imports: Option<Vec<ParsedImport>>,
    pub node: Node<Directive>,
}

impl ParsedLink {
    // Parse a @link directive node into a `ParsedLink` (if able)
    pub(crate) fn from_link_directive(link_directive: &Node<Directive>) -> Option<ParsedLink> {
        let url_string = link_directive.specified_argument_by_name("url")?.as_str()?;
        let parsed_url = Url::parse(url_string).ok()?;

        // get spec identifier from the path segments (e.g. federation/v2.0 ->
        // federation)
        let mut segments = parsed_url.path_segments()?;
        let spec_name = segments.next()?.to_string();

        // get version from the path segments (e.g. federation/v2.0 -> 2.0)
        let version_string = segments.next()?.strip_prefix('v')?;
        // parse version string into a semver::Version (we have to hack in a
        // patch version of 0 because semver requires it)
        let parsed_version =
            semver::Version::parse(format!("{}.0", &version_string).as_str()).ok()?;

        // if the link directive has an `as` argument, we'll use this to prefix
        // all of the directives and types imported by the spec
        let imported_as = link_directive
            .specified_argument_by_name("as")
            .map(|as_arg| as_arg.as_str().unwrap_or_default().to_string())
            .unwrap_or(spec_name.clone());

        // imports can be either a string or an object of { name: "<name>", as: "<alias>" }
        // this is normalized into a list of ParsedImport
        let imports =
            ParsedImport::imports_from_arg(link_directive.specified_argument_by_name("import"));

        Some(ParsedLink {
            spec_name,
            version: parsed_version,
            imported_as,
            imports,
            node: link_directive.clone(),
        })
    }

    pub(crate) fn unimported_directives(&self) -> HashMap<&String, &SpecDirective> {
        let spec = KNOWN_SPECS
            .get(&self.spec_name)
            .unwrap()
            .get(format!("{}.{}", &self.version.major, &self.version.minor).as_str())
            .unwrap();

        if let Some(imports) = &self.imports {
            spec.directives
                .iter()
                .filter(|(name, _directive)| !imports.iter().any(|import| import.name == **name))
                .collect::<HashMap<_, _>>()
        } else {
            spec.directives.iter().by_ref().collect()
        }
    }

    fn position_for_next_import(&self, source_text: &Rope) -> Option<ImportPosition> {
        if let Some(import_arg) = self.node.specified_argument_by_name("import") {
            let args_list = import_arg.as_list()?;

            if let Some(last_arg) = args_list.last() {
                Some(ImportPosition::ImportsNonEmpty(
                    lsp_range_from_ast_sourcespan(last_arg.location()?, source_text)?.end,
                ))
            } else {
                let position =
                    lsp_range_from_ast_sourcespan(import_arg.location()?, source_text)?.end;

                // offset it back by 1 so the position is inside the brackets of
                // the empty import list
                Some(ImportPosition::ImportsEmpty(lsp::Position {
                    line: position.line,
                    character: position.character - 1,
                }))
            }
        } else {
            Some(ImportPosition::NoImport(
                lsp_range_from_ast_sourcespan(
                    self.node.specified_argument_by_name("url")?.location()?,
                    source_text,
                )?
                .end,
            ))
        }
    }

    pub(crate) fn text_edit_for_import(
        &self,
        source_text: &Rope,
        directive: &DirectiveDefinition,
    ) -> lsp::TextEdit {
        match self.position_for_next_import(source_text).unwrap() {
            ImportPosition::ImportsNonEmpty(position) => lsp::TextEdit {
                range: lsp::Range {
                    start: position,
                    end: position,
                },
                new_text: format!(", \"@{}\"", directive.name),
            },
            ImportPosition::ImportsEmpty(position) => lsp::TextEdit {
                range: lsp::Range {
                    start: position,
                    end: position,
                },
                new_text: format!("\"@{}\"", directive.name),
            },
            ImportPosition::NoImport(position) => lsp::TextEdit {
                range: lsp::Range {
                    start: position,
                    end: position,
                },
                new_text: format!(", import: [\"@{}\"]", directive.name),
            },
        }
    }
}

enum ImportPosition {
    NoImport(lsp::Position),
    ImportsEmpty(lsp::Position),
    ImportsNonEmpty(lsp::Position),
}

#[derive(Debug, Clone)]
// A ParsedImport is a normalized representation of an import argument in a
// @link directive. The import list items can be either a string or an object.
// If it's a string, we'll use it as the name and leave the alias as None. If
// it's an object, we'll use the name and alias provided.
pub struct ParsedImport {
    pub name: String,
    pub import_as: Option<String>,
}

impl ParsedImport {
    // parse a @link's import argument into a list of `ParsedImport`s
    fn imports_from_arg(arg: Option<&Node<Value>>) -> Option<Vec<ParsedImport>> {
        let import_list = arg?.as_list()?;

        Some(
            import_list
                .iter()
                .filter_map(|import| {
                    // if the import is a string, we'll use it as the name and
                    // leave the alias as None
                    if let Some(import_name) = import.as_str() {
                        Some(ParsedImport {
                            name: import_name
                                .strip_prefix('@')
                                .unwrap_or(import_name)
                                .to_string(),
                            import_as: None,
                        })
                    } else if let Some(import_obj) = import.as_object() {
                        let lookup = import_obj
                            .iter()
                            .filter_map(|(name, value)| {
                                Some((name.to_string(), value.as_str()?.to_string()))
                            })
                            .collect::<HashMap<_, _>>();
                        let raw_name = lookup.get("name")?;
                        let name = raw_name.strip_prefix('@').unwrap_or(raw_name).to_string();

                        let import_as = lookup
                            .get("as")
                            .map(|a| a.strip_prefix('@').unwrap_or(a).to_string());

                        Some(ParsedImport { name, import_as })
                    } else {
                        // TODO: error unexpected type in import list, couldn't parse as string or object
                        None
                    }
                })
                .collect(),
        )
    }
}

// Given a spec name and a schema, find the first link directive with a matching
// spec name
pub fn find_link_for_spec(spec_name: &str, document: &Document) -> Option<ParsedLink> {
    document.definitions.iter().find_map(|def| match def {
        Definition::SchemaExtension(schema_ext) => {
            schema_ext.directives.iter().find_map(|directive| {
                if directive.name == LINK_DIRECTIVE {
                    let parsed_link = ParsedLink::from_link_directive(directive)?;
                    if parsed_link.spec_name == spec_name {
                        Some(parsed_link)
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
        }
        Definition::SchemaDefinition(schema_def) => {
            schema_def.directives.iter().find_map(|directive| {
                if directive.name == LINK_DIRECTIVE {
                    let parsed_link = ParsedLink::from_link_directive(directive)?;
                    if parsed_link.spec_name == spec_name {
                        Some(parsed_link)
                    } else {
                        None
                    }
                } else {
                    None
                }
            })
        }
        _ => None,
    })
}

const SPEC_SCALARS: [&str; 5] = ["String", "Int", "Float", "Boolean", "ID"];

pub fn update_name_with_link(name: &Name, link: &ParsedLink) -> Result<Name, InvalidNameError> {
    let binding = link.imports.clone().unwrap_or_default();
    let import = binding.iter().find(|i| i.name == name.as_str());

    if let Some(import) = import {
        if let Some(import_as) = &import.import_as {
            Name::new(import_as)
        } else {
            Ok(name.clone())
        }
    } else {
        // special case where directive name matches spec name and doesn't get a prefix
        if name.as_str() == link.spec_name {
            Ok(name.clone())
        } else {
            Name::new(&format!("{}__{}", link.imported_as, name.as_str()))
        }
    }
}

fn update_type_with_link(ty: &Type, link: &ParsedLink) -> Type {
    let inner_named_type = ty.inner_named_type().as_str();
    if SPEC_SCALARS.contains(&inner_named_type) {
        return ty.clone();
    }

    let binding = link.imports.clone().unwrap_or_default();
    let import = binding.iter().find(|i| i.name == inner_named_type);

    if let Some(import) = import {
        if let Some(import_as) = &import.import_as {
            update_inner_named_type(ty, import_as)
        } else {
            ty.clone()
        }
    } else {
        // special case where directive name matches spec name and doesn't get a prefix
        if ty.inner_named_type().as_str() == link.spec_name {
            ty.clone()
        } else {
            update_inner_named_type(ty, &format!("{}__{}", link.imported_as, inner_named_type))
        }
    }
}

fn update_inner_named_type(ty: &Type, new_name: &str) -> Type {
    match ty {
        Type::Named(_) => Type::Named(Name::new(new_name).unwrap()),
        Type::NonNullNamed(_) => Type::NonNullNamed(Name::new(new_name).unwrap()),
        Type::List(inner) => Type::List(Box::new(update_inner_named_type(inner, new_name))),
        Type::NonNullList(inner) => {
            Type::NonNullList(Box::new(update_inner_named_type(inner, new_name)))
        }
    }
}

impl SpecDirective {
    pub fn update_with_link(&self, link: &ParsedLink) -> Result<SpecDirective, InvalidNameError> {
        // TODO: special case where directive name matches spec name and doesn't
        // get a prefix

        // can we just mutate the args in the node?
        Ok(SpecDirective {
            node: Node::new(DirectiveDefinition {
                name: Name::new(&update_name_with_link(&self.node.name, link)?)?,
                arguments: self
                    .node
                    .arguments
                    .iter()
                    .map(|argument| {
                        Node::new(InputValueDefinition {
                            ty: Node::new(update_type_with_link(&argument.ty, link)),
                            ..argument.as_ref().clone()
                        })
                    })
                    .collect(),
                ..self.node.as_ref().clone()
            }),
        })
    }
}

impl SpecType<InputObjectTypeDefinition> {
    pub fn update_with_link(
        &self,
        link: &ParsedLink,
    ) -> Result<SpecType<InputObjectTypeDefinition>, InvalidNameError> {
        Ok(SpecType {
            node: Node::new(InputObjectTypeDefinition {
                name: update_name_with_link(&self.node.name, link)?,
                fields: self
                    .node
                    .fields
                    .iter()
                    .map(|field| {
                        Node::new(InputValueDefinition {
                            ty: Node::new(update_type_with_link(&field.ty, link)),
                            ..field.as_ref().clone()
                        })
                    })
                    .collect(),
                ..self.node.as_ref().clone()
            }),
        })
    }
}

impl SpecType<ScalarTypeDefinition> {
    pub fn update_with_link(
        &self,
        link: &ParsedLink,
    ) -> Result<SpecType<ScalarTypeDefinition>, InvalidNameError> {
        Ok(SpecType {
            node: Node::new(ScalarTypeDefinition {
                name: update_name_with_link(&self.node.name, link)?,
                ..self.node.as_ref().clone()
            }),
        })
    }
}

impl SpecType<EnumTypeDefinition> {
    pub fn update_with_link(
        &self,
        link: &ParsedLink,
    ) -> Result<SpecType<EnumTypeDefinition>, InvalidNameError> {
        Ok(SpecType {
            node: Node::new(EnumTypeDefinition {
                name: update_name_with_link(&self.node.name, link)?,
                ..self.node.as_ref().clone()
            }),
        })
    }
}

// Given a `DirectiveDefinition`, determine if it's compatible with our `@link`
// directive A "compatible" signature only needs the required `url` argument and
// at least the `SCHEMA` location directive @link(url: String!) on SCHEMA.
// Unfortunately, the above constraints make our `@contact` directive
// "compatible", so we should also enforce that other unexpected args aren't
// present.
pub fn is_link_compatible(directive_definition: &DirectiveDefinition) -> bool {
    let expected_args = ["url", "import", "as", "for"];
    let only_has_expected_args = directive_definition
        .arguments
        .iter()
        .all(|arg| expected_args.contains(&arg.name.as_str()));

    if !only_has_expected_args {
        return false;
    }

    let has_compatible_url_arg = directive_definition.arguments.iter().any(|arg| {
        arg.name == "url"
            && arg.ty == Node::new(Type::NonNullNamed(NamedType::new("String").unwrap()))
    });

    if !has_compatible_url_arg {
        return false;
    }

    directive_definition
        .locations
        .contains(&DirectiveLocation::Schema)
}