hive-router 0.2.0

use ahash::{HashMap, HashSet};

use crate::query_planner::{
    consumer_schema::ConsumerSchema, state::supergraph_state::OperationKind,
};
use graphql_tools::parser::{
    query::Type,
    schema::{Definition, TypeDefinition},
};

#[derive(Debug)]
pub struct FieldTypeInfo {
    pub output_type_name: String,
    pub nullability: FieldNullability,
}

/// The nullability shape of a field's type. Similar to the GraphQL type wrappers.
#[derive(Debug, Clone)]
pub enum FieldNullability {
    Leaf {
        non_null: bool,
    },
    List {
        non_null: bool,
        item: Box<FieldNullability>,
    },
}

impl FieldNullability {
    // __typename: String! (https://spec.graphql.org/September2025/#sec-Type-Name-Introspection)
    #[inline]
    pub fn type_name() -> Self {
        FieldNullability::Leaf { non_null: true }
    }

    fn leaf(non_null: bool) -> Self {
        FieldNullability::Leaf { non_null }
    }

    fn list(non_null: bool, item: FieldNullability) -> Self {
        FieldNullability::List {
            non_null,
            item: Box::new(item),
        }
    }

    /// Whether this position's value is Non-Null, so a `null` here must bubble to its parent.
    #[inline]
    pub fn is_non_null(&self) -> bool {
        match self {
            FieldNullability::Leaf { non_null } | FieldNullability::List { non_null, .. } => {
                *non_null
            }
        }
    }

    /// The element's nullability shape, when this position is a list.
    #[inline]
    pub fn list_item(&self) -> Option<&FieldNullability> {
        match self {
            FieldNullability::List { item, .. } => Some(item),
            FieldNullability::Leaf { .. } => None,
        }
    }
}

#[derive(Debug, Default)]
pub struct SchemaMetadata {
    pub possible_types: PossibleTypes,
    pub enum_values: HashMap<String, HashSet<String>>,
    pub type_fields: HashMap<String, HashMap<String, FieldTypeInfo>>,
    pub object_types: HashSet<String>,
    pub scalar_types: HashSet<String>,
    pub union_types: HashSet<String>,
    pub interface_types: HashSet<String>,
    pub query_type_name: Option<String>,
    pub mutation_type_name: Option<String>,
    pub subscription_type_name: Option<String>,
}

impl SchemaMetadata {
    pub fn is_object_type(&self, name: &str) -> bool {
        self.object_types.contains(name)
    }

    pub fn is_scalar_type(&self, name: &str) -> bool {
        self.scalar_types.contains(name)
    }

    pub fn is_union_type(&self, name: &str) -> bool {
        self.union_types.contains(name)
    }

    pub fn is_interface_type(&self, name: &str) -> bool {
        self.interface_types.contains(name)
    }

    pub fn get_type_fields(&self, type_name: &str) -> Option<&HashMap<String, FieldTypeInfo>> {
        self.type_fields.get(type_name)
    }

    /// Gets the list of types that implement an interface or are members of a union.
    /// Returns None if the type is not an interface or has no implementors.
    pub fn get_possible_types(&self, interface_name: &str) -> Option<&HashSet<String>> {
        self.possible_types.map.get(interface_name)
    }

    /// panics if the root type is not found for the given operation kind.
    /// Use only when you are sure the root type is present,
    /// like after the operation was validated to have matching root type.
    pub fn expect_root_type_name(&self, operation_kind: Option<&OperationKind>) -> &str {
        let root_type_name = match operation_kind {
            None | Some(OperationKind::Query) => &self.query_type_name,
            Some(OperationKind::Mutation) => &self.mutation_type_name,
            Some(OperationKind::Subscription) => &self.subscription_type_name,
        };

        root_type_name
            .as_deref()
            // SAFETY: The root type is guaranteed to exist for the given operation kind,
            // as the operation was validated.
            .unwrap_or_else(|| {
                panic!(
                    "Root type not found for operation kind: {:?}",
                    operation_kind
                )
            })
    }
}

#[derive(Debug, Default)]
pub struct PossibleTypes {
    map: HashMap<String, HashSet<String>>,
}

impl PossibleTypes {
    pub fn entity_satisfies_type_condition(&self, type_name: &str, type_condition: &str) -> bool {
        if type_name == type_condition {
            true
        } else if let Some(possible_types_of_type) = self.map.get(type_condition) {
            possible_types_of_type.contains(type_name)
        } else {
            false
        }
    }

    pub fn get_possible_types(&self, type_name: &str) -> HashSet<String> {
        let mut possible_types = self.map.get(type_name).cloned().unwrap_or_default();
        possible_types.insert(type_name.to_string());
        possible_types
    }
}

pub trait SchemaWithMetadata {
    fn schema_metadata(&self) -> SchemaMetadata;
}

impl SchemaWithMetadata for ConsumerSchema {
    fn schema_metadata(&self) -> SchemaMetadata {
        let mut first_possible_types: HashMap<String, Vec<String>> = HashMap::default();
        let mut type_fields: HashMap<String, HashMap<String, FieldTypeInfo>> = HashMap::default();
        let mut enum_values: HashMap<String, HashSet<String>> = HashMap::default();
        let mut scalar_types: HashSet<String> = HashSet::from_iter(vec![
            "Boolean".to_string(),
            "Float".to_string(),
            "Int".to_string(),
            "ID".to_string(),
            "String".to_string(),
        ]);
        let mut object_types: HashSet<String> = HashSet::default();
        let mut union_types: HashSet<String> = HashSet::default();
        let mut interface_types: HashSet<String> = HashSet::default();

        for definition in &self.document.definitions {
            match definition {
                Definition::TypeDefinition(TypeDefinition::Enum(enum_type)) => {
                    let name = enum_type.name.to_string();

                    let mut values = HashSet::default();
                    for enum_value in &enum_type.values {
                        values.insert(enum_value.name.to_string());
                    }
                    enum_values.insert(name, values);
                }
                Definition::TypeDefinition(TypeDefinition::Object(object_type)) => {
                    let name = object_type.name.to_string();
                    object_types.insert(name.clone());
                    let fields = type_fields.entry(name).or_default();
                    for field in &object_type.fields {
                        let field_type_name = field.field_type.type_name();
                        fields.insert(
                            field.name.to_string(),
                            FieldTypeInfo {
                                output_type_name: field_type_name,
                                nullability: field.field_type.field_nullability(),
                            },
                        );
                    }

                    for interface in &object_type.implements_interfaces {
                        let interface_name = interface.to_string();
                        let possible_types_entry =
                            first_possible_types.entry(interface_name).or_default();
                        possible_types_entry.push(object_type.name.to_string());
                    }
                }
                Definition::TypeDefinition(TypeDefinition::Interface(interface_type)) => {
                    let name = interface_type.name.to_string();
                    interface_types.insert(name.clone());
                    let mut fields = HashMap::default();
                    for field in &interface_type.fields {
                        let field_type_name = field.field_type.type_name();
                        fields.insert(
                            field.name.to_string(),
                            FieldTypeInfo {
                                output_type_name: field_type_name,
                                nullability: field.field_type.field_nullability(),
                            },
                        );
                    }
                    type_fields.insert(name, fields);
                    for interface_name in &interface_type.implements_interfaces {
                        let interface_name = interface_name.to_string();
                        let possible_types_entry =
                            first_possible_types.entry(interface_name).or_default();
                        possible_types_entry.push(interface_type.name.to_string());
                    }
                }
                Definition::TypeDefinition(TypeDefinition::Union(union_type)) => {
                    let name = union_type.name.to_string();
                    union_types.insert(name.clone());
                    let mut types = vec![];
                    for member in &union_type.types {
                        types.push(member.to_string());
                    }
                    first_possible_types.insert(name, types);
                }
                Definition::TypeDefinition(TypeDefinition::Scalar(scalar_type)) => {
                    scalar_types.insert(scalar_type.name.to_string());
                }
                _ => {}
            }
        }

        let query_type_name = self.document.query_type_name().cloned();
        let mutation_type_name = self.document.mutation_type_name().cloned();
        let subscription_type_name = self.document.subscription_type_name().cloned();

        let mut final_possible_types: HashMap<String, HashSet<String>> = HashMap::default();
        // Re-iterate over the possible_types
        for (definition_name_of_x, first_possible_types_of_x) in &first_possible_types {
            let mut possible_types_of_x: HashSet<String> = HashSet::default();
            for definition_name_of_y in first_possible_types_of_x {
                possible_types_of_x.insert(definition_name_of_y.to_string());
                let possible_types_of_y = first_possible_types.get(definition_name_of_y);
                if let Some(possible_types_of_y) = possible_types_of_y {
                    for definition_name_of_z in possible_types_of_y {
                        possible_types_of_x.insert(definition_name_of_z.to_string());
                    }
                }
            }
            final_possible_types.insert(definition_name_of_x.to_string(), possible_types_of_x);
        }

        SchemaMetadata {
            possible_types: PossibleTypes {
                map: final_possible_types,
            },
            enum_values,
            type_fields,
            object_types,
            scalar_types,
            union_types,
            interface_types,
            query_type_name,
            mutation_type_name,
            subscription_type_name,
        }
    }
}

trait TypeName {
    fn type_name(&self) -> String;
    fn field_nullability(&self) -> FieldNullability;
}

impl TypeName for Type<'_, String> {
    fn type_name(&self) -> String {
        match self {
            graphql_tools::parser::schema::Type::NamedType(named_type) => named_type.to_string(),
            graphql_tools::parser::schema::Type::NonNullType(non_null_type) => {
                non_null_type.type_name()
            }
            graphql_tools::parser::schema::Type::ListType(list_type) => list_type.type_name(),
        }
    }

    fn field_nullability(&self) -> FieldNullability {
        use graphql_tools::parser::schema::Type;

        match self {
            Type::NonNullType(inner) => match inner.as_ref() {
                Type::ListType(item) => FieldNullability::list(true, item.field_nullability()),
                _ => FieldNullability::leaf(true),
            },
            Type::ListType(item) => FieldNullability::list(false, item.field_nullability()),
            Type::NamedType(_) => FieldNullability::leaf(false),
        }
    }
}