apollo-language-server 0.6.0

use super::subgraph::Subgraph;
use crate::{
    diagnostics::{lsp_range_from_line_column_range, reformat_satisfiability_error},
    goto_definition::goto_definition_at_position,
    hover::get_hover_at_position,
};
use apollo_federation_types::{
    composition, composition::Issue, config::SchemaSource, javascript::SubgraphDefinition,
};
use std::{collections::HashMap, sync::Arc};

#[derive(Default, Debug, Clone)]
pub(crate) struct KnownSubgraphs {
    pub(crate) root_uri: String,
    pub(crate) by_name: HashMap<String, SchemaSource>,
    pub(crate) by_uri: HashMap<lsp::Url, String>,
}

impl KnownSubgraphs {
    #[cfg(any(not(feature = "wasm"), test))]
    pub(crate) fn uri_for_name(&self, name: &str) -> Option<lsp::Url> {
        let source = self.by_name.get(name)?;
        match source {
            SchemaSource::File { file } => {
                if file.is_relative() {
                    // build a uri from config.root_uri and file path
                    let uri =
                        lsp::Url::from_directory_path(&self.root_uri).expect("Failed to parse URL");
                    Some(
                        uri.join(file.to_str().expect("Failed to convert path to string"))
                            .expect("Failed to join URL"),
                    )
                } else {
                    Some(lsp::Url::from_file_path(file).expect("Failed to convert path to URL"))
                }
            }
            _ => None,
        }
    }

    #[cfg(all(feature = "wasm", not(test)))]
    pub(crate) fn uri_for_name(&self, _: &str) -> Option<lsp::Url> {
        None
    }
}

#[derive(Debug, Default)]
pub struct Supergraph {
    subgraphs_by_uri: HashMap<lsp::Url, Arc<Subgraph>>,
    subgraphs_by_name: HashMap<String, Arc<Subgraph>>,
    pub(crate) known_subgraphs: KnownSubgraphs,
}

type DiagnosticsBySubgraph = HashMap<lsp::Url, (Vec<lsp::Diagnostic>, i32)>;

impl Supergraph {
    pub fn new(known_subgraphs: KnownSubgraphs) -> Supergraph {
        Supergraph {
            subgraphs_by_uri: HashMap::new(),
            subgraphs_by_name: HashMap::new(),
            known_subgraphs,
        }
    }

    pub(crate) fn insert(&mut self, subgraph: Subgraph) {
        let subgraph = Arc::new(subgraph);
        if self.subgraphs_by_name.contains_key(&subgraph.name) {
            panic!("Subgraph `{}` already exists", subgraph.name);
        }
        self.subgraphs_by_name
            .insert(subgraph.name.clone(), subgraph.clone());
        self.subgraphs_by_uri.insert(subgraph.uri.clone(), subgraph);
    }

    pub(crate) fn update(&mut self, subgraph: Subgraph) {
        let subgraph = Arc::new(subgraph);
        self.subgraphs_by_name
            .insert(subgraph.name.clone(), subgraph.clone());
        self.subgraphs_by_uri.insert(subgraph.uri.clone(), subgraph);
    }

    pub(crate) fn remove(&mut self, uri: &lsp::Url) -> Option<Subgraph> {
        if let Some(subgraph) = self.subgraphs_by_uri.remove(uri) {
            self.subgraphs_by_name.remove(&subgraph.name);
            Some(Arc::try_unwrap(subgraph).unwrap())
        } else {
            None
        }
    }

    #[cfg(any(not(feature = "wasm"), test))]
    pub(crate) async fn add_known_subgraph(&mut self, name: String, source: SchemaSource) {
        match source {
            SchemaSource::File { file } => {
                let uri = if file.is_relative() {
                    // build a uri from config.root_uri and file path
                    let uri = lsp::Url::from_directory_path(&self.known_subgraphs.root_uri)
                        .expect("Failed to parse URL");
                    uri.join(file.to_str().expect("Failed to convert path to string"))
                        .expect("Failed to join URL")
                } else {
                    lsp::Url::from_file_path(&file).expect("Failed to convert path to URL")
                };

                self.known_subgraphs
                    .by_name
                    .insert(name.clone(), SchemaSource::File { file: file.clone() });
                self.known_subgraphs.by_uri.insert(uri, name);
            }
            _ => {
                self.known_subgraphs
                    .by_name
                    .insert(name.clone(), source.clone());
            }
        };
    }

    #[cfg(any(not(feature = "wasm"), test))]
    pub(crate) async fn remove_known_subgraph(&mut self, name: &str) {
        let removed = self.known_subgraphs.by_name.remove(name);
        if let Some(SchemaSource::File { file }) = removed {
            let uri = if file.is_relative() {
                // build a uri from config.root_uri and file path
                let uri = lsp::Url::from_directory_path(&self.known_subgraphs.root_uri)
                    .expect("Failed to parse URL");
                uri.join(file.to_str().expect("Failed to convert path to string"))
                    .expect("Failed to join URL")
            } else {
                lsp::Url::from_file_path(&file).expect("Failed to convert path to URL")
            };

            self.known_subgraphs.by_uri.remove(&uri);
        }
    }

    pub(crate) fn subgraph_by_uri(&self, uri: &lsp::Url) -> Option<&Arc<Subgraph>> {
        self.subgraphs_by_uri.get(uri)
    }

    // We won't necessarily have a URI for every subgraph. i.e., when rover
    // loads a supergraph.yaml with subgraphs from graphref or introspection, we
    // won't have a URI to point to in the workspace. This only matters for
    // reporting diagnostics - in which case we can report to the root URI.
    pub(crate) fn uri_for_name(&self, name: &str) -> Option<lsp::Url> {
        (self
            .subgraphs_by_name
            .get(name)
            .map(|subgraph| subgraph.uri.clone()))
        .or_else(|| self.known_subgraphs.uri_for_name(name))
    }

    pub fn subgraph_definitions(&self) -> Vec<SubgraphDefinition> {
        self.subgraphs_by_name
            .iter()
            .map(|(uri, subgraph)| SubgraphDefinition {
                name: subgraph.name.clone(),
                url: uri.to_string(),
                sdl: subgraph.source_text.to_string(),
            })
            .collect()
    }

    pub fn get_invalid_subgraph_uris(&self) -> Vec<String> {
        self.subgraphs_by_name
            .values()
            .filter(|subgraph| subgraph.has_diagnostics())
            .map(|subgraph| subgraph.uri.to_string())
            .collect()
    }

    pub fn subgraphs_are_invalid(&self) -> bool {
        self.subgraphs_by_name
            .values()
            .any(|subgraph| subgraph.has_diagnostics())
    }

    pub fn diagnostics_for_subgraph(&self, uri: &lsp::Url) -> (Vec<lsp::Diagnostic>, i32) {
        let subgraph = self
            .subgraph_by_uri(uri)
            .unwrap_or_else(|| panic!("No subgraph found at uri: {}", uri));
        (subgraph.diagnostics(), subgraph.version)
    }

    pub fn version_for_subgraph(&self, uri: &lsp::Url) -> Option<i32> {
        self.subgraph_by_uri(uri).map(|s| s.version)
    }

    pub fn diagnostics_for_composition(
        &self,
        issues: Vec<Issue>,
    ) -> (DiagnosticsBySubgraph, Vec<lsp::Diagnostic>) {
        let mut diagnostics_by_subgraph = self
            .subgraphs_by_uri
            .iter()
            .map(|(uri, subgraph)| (uri.clone(), (vec![], subgraph.version)))
            .collect::<DiagnosticsBySubgraph>();

        for uri in self.known_subgraphs.by_uri.keys() {
            // Don't overwrite an entry from above which might have included
            // version information for a subgraph that went from "known" to an
            // actually open file. It's worth mentioning that we don't want to
            // remove the open file's subgraph from the known subgraphs map
            // since the file could eventually be closed (and we'd still need to
            // know about it).
            if diagnostics_by_subgraph.contains_key(uri) {
                continue;
            }
            diagnostics_by_subgraph.insert(uri.clone(), (vec![], 0));
        }

        let mut unattributed_diagnostics = vec![];
        let first_character_range = lsp::Range {
            start: lsp::Position {
                line: 0,
                character: 0,
            },
            end: lsp::Position {
                line: 0,
                character: 1,
            },
        };
        for issue in issues {
            if issue.locations.is_empty() {
                let message = if issue.code == "SATISFIABILITY_ERROR" {
                    reformat_satisfiability_error(&issue)
                } else {
                    issue.message.clone()
                };
                // If a location is not associated with a subgraph, we'll add it
                // to the diagnostics of every subgraph. We could eventually
                // attach these to each subgraph's schema definition or
                // extension, but also not all subgraphs have one.
                unattributed_diagnostics.push(lsp::Diagnostic {
                    range: first_character_range,
                    severity: Some(match issue.severity {
                        composition::Severity::Error => lsp::DiagnosticSeverity::ERROR,
                        composition::Severity::Warning => lsp::DiagnosticSeverity::WARNING,
                    }),
                    message,
                    code: Some(lsp::NumberOrString::String(issue.code.to_string())),
                    ..Default::default()
                });
                continue;
            }
            for location in issue.locations.iter() {
                if let Some(uri) = location
                    .subgraph
                    .as_ref()
                    .and_then(|subgraph| self.uri_for_name(subgraph))
                {
                    diagnostics_by_subgraph
                        .get_mut(&uri)
                        .unwrap_or_else(|| {
                            panic!("Unexpectedly missing subgraph with uri `{}`", uri)
                        })
                        .0
                        .push(lsp::Diagnostic {
                            range: location
                                .range
                                .as_ref()
                                .map(|range| lsp_range_from_line_column_range(range.clone()))
                                .unwrap_or(first_character_range),
                            severity: Some(match issue.severity {
                                composition::Severity::Error => lsp::DiagnosticSeverity::ERROR,
                                composition::Severity::Warning => lsp::DiagnosticSeverity::WARNING,
                            }),
                            message: issue.message.clone(),
                            code: Some(lsp::NumberOrString::String(issue.code.to_string())),
                            ..Default::default()
                        });
                } else {
                    unattributed_diagnostics.push(lsp::Diagnostic {
                        range: first_character_range,
                        severity: Some(match issue.severity {
                            composition::Severity::Error => lsp::DiagnosticSeverity::ERROR,
                            composition::Severity::Warning => lsp::DiagnosticSeverity::WARNING,
                        }),
                        message: issue.message.clone(),
                        code: Some(lsp::NumberOrString::String(issue.code.to_string())),
                        ..Default::default()
                    });
                }
            }
        }
        (diagnostics_by_subgraph, unattributed_diagnostics)
    }

    pub fn on_hover(&self, uri: &lsp::Url, position: &lsp::Position) -> Option<lsp::Hover> {
        let subgraph = self
            .subgraph_by_uri(uri)
            .unwrap_or_else(|| panic!("No subgraph found at uri: {}", uri));

        let doc = subgraph.cst.document();

        get_hover_at_position(&doc, subgraph.schema(), &subgraph.source_text, position)
    }

    pub fn goto_definition(
        &self,
        uri: &lsp::Url,
        position: &lsp::Position,
    ) -> Option<Vec<lsp::LocationLink>> {
        let subgraph = self
            .subgraph_by_uri(uri)
            .unwrap_or_else(|| panic!("No subgraph found at uri: {}", uri));

        let doc = subgraph.cst.document();

        goto_definition_at_position(
            uri,
            &doc,
            subgraph.schema(),
            &subgraph.source_text,
            position,
            Some(&self.subgraphs_by_uri),
        )
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use crate::graph::{supergraph::KnownSubgraphs, Graph};
    use itertools::Itertools;

    #[test]
    fn captures_subgraph_names_in_uri_from_query_params_or_defaults() {
        let mut graph = Graph::new(
            lsp::Url::parse("file:///path/to/subgraph1.graphql?subgraph_name=MyGraph").unwrap(),
            r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.10")
            type Query { foo: String }
            "#
            .to_string(),
            0,
            KnownSubgraphs::default(),
            Default::default(),
        );

        // add a second subgraph with no name set by query param
        graph.update(
            lsp::Url::parse("file:///path/to/subgraph2.graphql").unwrap(),
            r#"
            extend schema @link(url: "https://specs.apollo.dev/federation/v2.10")
            type Query { bar: String }
            "#
            .to_string(),
            1,
            Default::default(),
        );

        let supergraph = graph.supergraph().unwrap();

        assert_eq!(
            supergraph
                .subgraph_definitions()
                .iter()
                .map(|d| d.name.as_str())
                .sorted()
                .collect::<Vec<_>>(),
            // name is derived from query param (`MyGraph`) or defaults to full URI (`file:///path/to/subgraph2.graphql`)
            vec!["MyGraph", "file:///path/to/subgraph2.graphql"]
        );
    }

    #[test]
    fn test_unimported_spec_directives() {
        let subgraph_uri = lsp::Url::parse("file:///path/to/subgraph.graphql").unwrap();
        let graph = Graph::new(
            subgraph_uri.clone(),
            r#"
            extend schema 
                @link(url: "https://specs.apollo.dev/federation/v2.10", import: ["@key"])
                @link(url: "https://specs.apollo.dev/connect/v0.1")
            type Query { foo: String }
            "#
            .to_string(),
            0,
            KnownSubgraphs::default(),
            Default::default(),
        );

        let unimported = graph
            .supergraph()
            .unwrap()
            .subgraph_by_uri(&subgraph_uri)
            .unwrap()
            .links()
            .values()
            .flat_map(|parsed_link| parsed_link.unimported_directives())
            .collect::<HashMap<_, _>>();

        assert!(!unimported.contains_key(&"key".to_string()));
        assert!(unimported.contains_key(&"source".to_string()));
        assert!(unimported.contains_key(&"connect".to_string()));
    }
}