apollo-language-server 0.7.0

use apollo_compiler::ast::DirectiveDefinition;
use itertools::Itertools;
use once_cell::sync::Lazy;
use std::{collections::HashMap, vec};

use crate::{
    federation::link::ParsedLink,
    server::MaxSpecVersions,
    specs::{
        connect::{
            CONNECT_DIRECTIVE, CONNECT_SPECS_BY_VERSION, CONNECT_SPEC_NAME, SOURCE_DIRECTIVE,
        },
        federation::{
            CACHE_TAG_DIRECTIVE, FEDERATION_SPECS_BY_VERSION, FEDERATION_SPEC_NAME, KEY_DIRECTIVE,
            LINK_DIRECTIVE, PROVIDES_DIRECTIVE, REQUIRES_DIRECTIVE,
        },
        SpecStatus,
    },
};

#[derive(Default, Clone)]
struct CompletionItemUpdate {
    insert_text_snippet_portion: String,
    label_details: Option<lsp::CompletionItemLabelDetails>,
    sort_text: Option<String>,
    // This is metadata used to filter out completions that are already present
    // in the schema and shouldn't be applied to an actual CompletionItem
    spec_name: Option<String>,
    spec_version: Option<semver::Version>,
}

static CUSTOM_SNIPPETS_BY_DIRECTIVE_NAME: Lazy<HashMap<String, Vec<CompletionItemUpdate>>> = {
    Lazy::new(|| {
        let supported_connect_specs = CONNECT_SPECS_BY_VERSION
            .iter()
            .filter(|(_, spec)| spec.status == SpecStatus::Supported)
            .map(|(version_string, spec)| {
                (
                    CONNECT_SPEC_NAME,
                    semver_version(version_string),
                    vec!["@connect", "@source"],
                    &spec.status,
                )
            })
            .sorted_by_key(|(_, version, ..)| version.clone());

        let latest_connect_spec = CONNECT_SPECS_BY_VERSION
            .iter()
            .find(|(_, spec)| spec.status == SpecStatus::Latest)
            .map(|(version_string, spec)| {
                (
                    CONNECT_SPEC_NAME,
                    semver_version(version_string),
                    vec!["@connect", "@source"],
                    &spec.status,
                )
            })
            .unwrap();

        let supported_federation_specs = FEDERATION_SPECS_BY_VERSION
            .iter()
            .filter(|(_, spec)| spec.status == SpecStatus::Supported)
            .map(|(version_string, spec)| {
                (
                    FEDERATION_SPEC_NAME,
                    semver_version(version_string),
                    vec!["@key"],
                    &spec.status,
                )
            })
            .sorted_by_key(|(_, version, ..)| version.clone());

        let latest_federation_spec = FEDERATION_SPECS_BY_VERSION
            .iter()
            .find(|(_, spec)| spec.status == SpecStatus::Latest)
            .map(|(version_string, spec)| {
                (
                    FEDERATION_SPEC_NAME,
                    semver_version(version_string),
                    vec!["@key"],
                    &spec.status,
                )
            })
            .unwrap();

        let experimental_federation_specs = FEDERATION_SPECS_BY_VERSION
            .iter()
            .filter(|(_, spec)| spec.status == SpecStatus::Experimental)
            .map(|(version_string, spec)| {
                (
                    FEDERATION_SPEC_NAME,
                    semver_version(version_string),
                    vec!["@key"],
                    &spec.status,
                )
            })
            .sorted_by_key(|(_, version, ..)| version.clone());

        let links = std::iter::once(latest_federation_spec)
            .chain(std::iter::once(latest_connect_spec))
            .chain(supported_federation_specs.rev())
            .chain(supported_connect_specs.rev())
            .chain(experimental_federation_specs.rev());

        let fieldset_snippet = "(fields: \"$1\")";
        HashMap::from([
            (
                KEY_DIRECTIVE.to_string(),
                vec![CompletionItemUpdate {
                    insert_text_snippet_portion: fieldset_snippet.to_string(),
                    ..Default::default()
                },]
            ),
            (
                REQUIRES_DIRECTIVE.to_string(),
                vec![CompletionItemUpdate {
                    insert_text_snippet_portion: fieldset_snippet.to_string(),
                    ..Default::default()
                },]
            ),
            (
                PROVIDES_DIRECTIVE.to_string(),
                vec![CompletionItemUpdate {
                    insert_text_snippet_portion: fieldset_snippet.to_string(),
                    ..Default::default()
                },]
            ),
            (
                CONNECT_DIRECTIVE.to_string(),
                vec![CompletionItemUpdate {
                    insert_text_snippet_portion: "(\n  source: \"$1\"\n  http: { ${2|GET,POST,PUT,PATCH,DELETE|}: \"$3\" }\n  selection: \"\"\"\n    $4\n  \"\"\"\n)"
                    .to_string(),
                    ..Default::default()
                },]
            ),
            (SOURCE_DIRECTIVE.to_string(), vec![CompletionItemUpdate {
                insert_text_snippet_portion: "(name: \"$1\", http: { baseURL: \"$2\" })".to_string(),
                ..Default::default()
            }]),
            (CACHE_TAG_DIRECTIVE.to_string(), vec![CompletionItemUpdate {
                insert_text_snippet_portion: "(format: \"$1\")".to_string(),
                ..Default::default()
            }]),
            (LINK_DIRECTIVE.to_string(),
            links.enumerate().map(|(i, (spec_name, version, default_imports, status))| {
                CompletionItemUpdate {
                    insert_text_snippet_portion: format!(
                        "(url: \"https://specs.apollo.dev/{spec_name}/v{}.{}\", import: [{}])",
                        version.major,
                        version.minor,
                        default_imports.iter().map(|import| format!("\"{}\"", import)).join(", "),
                    ),
                    label_details: Some(lsp::CompletionItemLabelDetails {
                        detail: None,
                        description: Some(
                            format!("{} v{}.{} {}", spec_name, version.major, version.minor, match status {
                                SpecStatus::Latest => "🚀",
                                SpecStatus::Supported => "✅",
                                SpecStatus::Experimental => "🧪",
                                _ => "",
                            })
                        ),
                    }),
                    sort_text: Some(format!("@link-{:03}", i)),
                    spec_name: Some(spec_name.to_string()),
                    spec_version: Some(version)
                }
            }).collect::<Vec<_>>())
        ])
    })
};

pub(super) fn apply_custom_completions_for_spec_directive(
    completion_item: lsp::CompletionItem,
    directive: &DirectiveDefinition,
    alias_to_spec_map: &HashMap<&String, &String>,
    links_in_schema: Option<&HashMap<String, ParsedLink>>,
    should_include_at_prefix: bool,
    max_spec_versions: &MaxSpecVersions,
    should_suggest_snippet: bool,
) -> Vec<lsp::CompletionItem> {
    let Some(unaliased_spec_directive) = alias_to_spec_map.get(&directive.name.to_string()) else {
        return vec![completion_item];
    };

    let Some(custom_completion_updates) = CUSTOM_SNIPPETS_BY_DIRECTIVE_NAME
        .get(*unaliased_spec_directive)
        .cloned()
    else {
        return vec![completion_item];
    };

    custom_completion_updates
        .into_iter()
        .filter(|completion_update| {
            // if we have a spec_name, this is a `@link` completion
            let (Some(links_in_schema), Some(spec_name), Some(spec_version)) = (
                links_in_schema,
                completion_update.spec_name.as_ref(),
                completion_update.spec_version.as_ref(),
            ) else {
                return true;
            };

            if let (Some(max_federation_version), Some(max_connect_version)) = (
                max_spec_versions.federation.as_ref(),
                max_spec_versions.connect.as_ref(),
            ) {
                if (spec_name == FEDERATION_SPEC_NAME && spec_version > max_federation_version)
                    || (spec_name == CONNECT_SPEC_NAME && spec_version > max_connect_version)
                {
                    return false;
                }
            }

            !links_in_schema.keys().any(|link_in_schema| {
                completion_update.spec_name.as_ref().unwrap() == link_in_schema
            })
        })
        .map(|completion_update| {
            let mut completion_item = completion_item.clone();
            if should_suggest_snippet {
                completion_item.insert_text = Some(format!(
                    "{}{}{}",
                    should_include_at_prefix.then_some("@").unwrap_or_default(),
                    directive.name,
                    completion_update.insert_text_snippet_portion
                ));
            }

            if let Some(label_details) = completion_update.label_details {
                completion_item.label_details = Some(label_details);
            }

            if let Some(sort_text) = completion_update.sort_text {
                completion_item.sort_text = Some(sort_text);
            }

            completion_item
        })
        .collect()
}

fn semver_version(version: &str) -> semver::Version {
    semver::Version::parse(&format!("{}.0", version)).unwrap()
}