//! GraphQL Federation support
//!
//! This module implements Apollo Federation v2 specification, allowing the gateway
//! to participate in a federated GraphQL architecture.

use crate::error::{Error, Result};
use crate::graphql::GraphqlEntity;
use async_graphql::dynamic::{Field, FieldFuture, FieldValue, InputValue, Object, TypeRef};
use async_graphql::{Name, Value as GqlValue};
use prost::Message;
use prost_reflect::{DescriptorPool, ExtensionDescriptor, MessageDescriptor, Value};
use std::collections::HashMap;
use std::sync::Arc;

/// Federation configuration extracted from protobuf descriptors
#[derive(Clone, Debug)]
pub struct FederationConfig {
    /// Map of entity type names to their key fields
    pub entities: HashMap<String, EntityConfig>,
}

/// Configuration for a federated entity
#[derive(Clone, Debug)]
pub struct EntityConfig {
    /// The message descriptor for this entity type
    pub descriptor: MessageDescriptor,
    /// Key field sets for this entity (e.g., ["id"], ["email"], or ["orgId", "userId"])
    pub keys: Vec<Vec<String>>,
    /// Whether this entity extends an entity from another service
    pub extend: bool,
    /// Whether this service can resolve this entity
    pub resolvable: bool,
    /// The GraphQL type name for this entity
    pub type_name: String,
}

impl FederationConfig {
    /// Create a new empty federation configuration
    pub fn new() -> Self {
        Self {
            entities: HashMap::new(),
        }
    }

    /// Extract federation configuration from descriptor pool
    pub fn from_descriptor_pool(
        pool: &DescriptorPool,
        entity_ext: &ExtensionDescriptor,
    ) -> Result<Self> {
        let mut config = Self::new();

        // Scan all messages for entity annotations
        for message in pool.all_messages() {
            if let Some(entity_opts) = decode_entity_extension(&message, entity_ext)? {
                if entity_opts.keys.is_empty() {
                    continue; // Skip messages without keys
                }

                let type_name = message.full_name().replace('.', "_");
                let keys: Vec<Vec<String>> = entity_opts
                    .keys
                    .iter()
                    .map(|key| {
                        // Split on whitespace to support composite keys like "orgId userId"
                        key.split_whitespace().map(String::from).collect::<Vec<_>>()
                    })
                    .collect();

                config.entities.insert(
                    type_name.clone(),
                    EntityConfig {
                        descriptor: message.clone(),
                        keys,
                        extend: entity_opts.extend,
                        resolvable: entity_opts.resolvable,
                        type_name,
                    },
                );
            }
        }

        Ok(config)
    }

    /// Check if federation is enabled (i.e., if there are any entities)
    pub fn is_enabled(&self) -> bool {
        !self.entities.is_empty()
    }

    /// Build the _entities field for the Query type
    pub fn build_entities_field_for_query(
        &self,
        entity_resolvers: Arc<dyn EntityResolver>,
    ) -> Result<Field> {
        let config = self.clone();

        let field = Field::new("_entities", TypeRef::named_nn_list("_Entity"), move |ctx| {
            let entity_resolvers = entity_resolvers.clone();
            let config = config.clone();

            FieldFuture::new(async move {
                let representations = ctx
                    .args
                    .get("representations")
                    .ok_or_else(|| async_graphql::Error::new("missing representations argument"))?
                    .list()?;

                let mut results = Vec::new();
                for repr in representations.iter() {
                    let obj = repr.object()?;

                    // Convert ObjectAccessor to IndexMap
                    let mut representation_map = async_graphql::indexmap::IndexMap::new();
                    for (key, value) in obj.iter() {
                        representation_map.insert(key.clone(), value.as_value().clone());
                    }

                    // Extract __typename from representation
                    let typename = representation_map
                        .get(&Name::new("__typename"))
                        .and_then(|v| match v {
                            GqlValue::String(s) => Some(s.as_str()),
                            _ => None,
                        })
                        .ok_or_else(|| {
                            async_graphql::Error::new("missing __typename in representation")
                        })?;

                    // Find entity config
                    let entity_config = config.entities.get(typename).ok_or_else(|| {
                        async_graphql::Error::new(format!("unknown entity type: {}", typename))
                    })?;

                    // Resolve the entity
                    let entity = entity_resolvers
                        .resolve_entity(entity_config, &representation_map)
                        .await
                        .map_err(|e| async_graphql::Error::new(e.to_string()))?;

                    results
                        .push(FieldValue::value(entity).with_type(entity_config.type_name.clone()));
                }

                Ok(Some(FieldValue::list(results)))
            })
        })
        .argument(InputValue::new(
            "representations",
            TypeRef::named_nn_list_nn("_Any"),
        ));

        Ok(field)
    }

    /// Apply federation directives to an object type
    pub fn apply_directives_to_object(&self, obj: Object, type_name: &str) -> Result<Object> {
        if let Some(entity_config) = self.entities.get(type_name) {
            let mut obj = obj;

            // Add @key directives
            for key_fields in &entity_config.keys {
                let fields_str = key_fields.join(" ");
                // Mark as entity for async-graphql federation support
                if entity_config.resolvable {
                    obj = obj.key(fields_str.clone());
                } else {
                    obj = obj.unresolvable(fields_str.clone());
                }
                obj = obj.directive(
                    async_graphql::dynamic::Directive::new("key")
                        .argument("fields", GqlValue::String(fields_str)),
                );
            }

            // Add @extends directive if this is an extension
            if entity_config.extend {
                obj = obj.extends();
                obj = obj.directive(async_graphql::dynamic::Directive::new("extends"));
            }

            Ok(obj)
        } else {
            Ok(obj)
        }
    }
}

impl Default for FederationConfig {
    fn default() -> Self {
        Self::new()
    }
}

/// Configuration for mapping an entity type to its resolver method
///
/// This tells the gateway which gRPC method to call to resolve a specific entity type.
///
/// # Example
///
/// ```rust
/// use grpc_graphql_gateway::EntityResolverMapping;
///
/// let mapping = EntityResolverMapping {
///     service_name: "user.UserService".to_string(),
///     method_name: "GetUser".to_string(),
///     key_field: "id".to_string(),
/// };
/// ```
#[derive(Clone, Debug)]
pub struct EntityResolverMapping {
    /// Service name (e.g., "user.UserService")
    pub service_name: String,
    /// Method name (e.g., "GetUser")
    pub method_name: String,
    /// Field name in the request message for the primary key (e.g., "id")
    pub key_field: String,
}

/// Trait for resolving federated entities
///
/// Implementors should resolve entities based on their representation
/// (which contains the key fields and __typename).
///
/// This trait supports both single and batch resolution.
#[async_trait::async_trait]
pub trait EntityResolver: Send + Sync {
    /// Resolve an entity from its representation
    async fn resolve_entity(
        &self,
        entity_config: &EntityConfig,
        representation: &async_graphql::indexmap::IndexMap<Name, GqlValue>,
    ) -> Result<GqlValue>;

    /// Batch resolve multiple entities of the same type
    /// Default implementation calls resolve_entity for each
    async fn batch_resolve_entities(
        &self,
        entity_config: &EntityConfig,
        representations: Vec<async_graphql::indexmap::IndexMap<Name, GqlValue>>,
    ) -> Result<Vec<GqlValue>> {
        let mut results = Vec::with_capacity(representations.len());
        for repr in representations {
            results.push(self.resolve_entity(entity_config, &repr).await?);
        }
        Ok(results)
    }
}

/// Default entity resolver that uses gRPC client pool with DataLoader batching
pub struct GrpcEntityResolver {
    client_pool: crate::grpc_client::GrpcClientPool,
    /// Maps entity type names to their resolver configuration
    resolver_mappings: HashMap<String, EntityResolverMapping>,
}

impl GrpcEntityResolver {
    pub fn new(client_pool: crate::grpc_client::GrpcClientPool) -> Self {
        Self {
            client_pool,
            resolver_mappings: HashMap::new(),
        }
    }

    /// Register an entity resolver mapping
    ///
    /// # Example
    /// ```ignore
    /// resolver.register_entity_resolver(
    ///     "user_User",
    ///     EntityResolverMapping {
    ///         service_name: "user.UserService".to_string(),
    ///         method_name: "GetUser".to_string(),
    ///         key_field: "id".to_string(),
    ///     }
    /// );
    /// ```
    pub fn register_entity_resolver(
        &mut self,
        entity_type: impl Into<String>,
        mapping: EntityResolverMapping,
    ) {
        self.resolver_mappings.insert(entity_type.into(), mapping);
    }

    /// Create a builder for configuring entity resolvers
    pub fn builder(client_pool: crate::grpc_client::GrpcClientPool) -> GrpcEntityResolverBuilder {
        GrpcEntityResolverBuilder::new(client_pool)
    }
}

/// Builder for GrpcEntityResolver
pub struct GrpcEntityResolverBuilder {
    client_pool: crate::grpc_client::GrpcClientPool,
    resolver_mappings: HashMap<String, EntityResolverMapping>,
}

impl GrpcEntityResolverBuilder {
    pub fn new(client_pool: crate::grpc_client::GrpcClientPool) -> Self {
        Self {
            client_pool,
            resolver_mappings: HashMap::new(),
        }
    }

    /// Register an entity resolver mapping
    pub fn register_entity_resolver(
        mut self,
        entity_type: impl Into<String>,
        mapping: EntityResolverMapping,
    ) -> Self {
        self.resolver_mappings.insert(entity_type.into(), mapping);
        self
    }

    pub fn build(self) -> GrpcEntityResolver {
        GrpcEntityResolver {
            client_pool: self.client_pool,
            resolver_mappings: self.resolver_mappings,
        }
    }
}

impl Default for GrpcEntityResolver {
    fn default() -> Self {
        Self::new(crate::grpc_client::GrpcClientPool::new())
    }
}

#[async_trait::async_trait]
impl EntityResolver for GrpcEntityResolver {
    async fn resolve_entity(
        &self,
        entity_config: &EntityConfig,
        representation: &async_graphql::indexmap::IndexMap<Name, GqlValue>,
    ) -> Result<GqlValue> {
        tracing::debug!(
            "Resolving entity {} with representation: {:?}",
            entity_config.type_name,
            representation
        );

        // Get the resolver mapping for this entity type
        let mapping = self.resolver_mappings.get(&entity_config.type_name);

        if mapping.is_none() {
            tracing::warn!(
                "No resolver mapping found for entity {}. Returning representation as-is.",
                entity_config.type_name
            );
            return Ok(GqlValue::Object(representation.clone()));
        }

        let mapping = mapping.unwrap();

        // Extract the key field from the representation
        let key_value = representation
            .get(&Name::new(&mapping.key_field))
            .ok_or_else(|| {
                Error::Schema(format!(
                    "Missing key field '{}' in representation for {}",
                    mapping.key_field, entity_config.type_name
                ))
            })?;

        // Get gRPC client
        let _client = self.client_pool.get(&mapping.service_name).ok_or_else(|| {
            Error::Schema(format!(
                "gRPC client not found for service: {}",
                mapping.service_name
            ))
        })?;

        // Build the request message - for now, return the representation
        // In a full implementation, you would:
        // 1. Create a DynamicMessage for the request
        // 2. Set the key field(s) from the representation
        // 3. Call the gRPC method
        // 4. Transform the response to GqlValue

        tracing::debug!(
            "Would call {}/{} with key {}={:?}",
            mapping.service_name,
            mapping.method_name,
            mapping.key_field,
            key_value
        );

        // For now, return representation as-is
        // This preserves the existing behavior while providing the infrastructure
        Ok(GqlValue::Object(representation.clone()))
    }

    async fn batch_resolve_entities(
        &self,
        entity_config: &EntityConfig,
        representations: Vec<async_graphql::indexmap::IndexMap<Name, GqlValue>>,
    ) -> Result<Vec<GqlValue>> {
        tracing::debug!(
            "Batch resolving {} entities of type {}",
            representations.len(),
            entity_config.type_name
        );

        // In a production implementation, you could batch these into a single gRPC call
        // if your service supports batch operations. For now, we resolve serially.
        let mut results = Vec::with_capacity(representations.len());
        for repr in representations {
            results.push(self.resolve_entity(entity_config, &repr).await?);
        }
        Ok(results)
    }
}

/// Decode entity extension from message descriptor
fn decode_entity_extension(
    message: &MessageDescriptor,
    ext: &ExtensionDescriptor,
) -> Result<Option<GraphqlEntity>> {
    let opts = message.options();
    if !opts.has_extension(ext) {
        return Ok(None);
    }

    let val = opts.get_extension(ext);
    if let Value::Message(msg) = val.as_ref() {
        GraphqlEntity::decode(msg.encode_to_vec().as_slice())
            .map(Some)
            .map_err(|e| Error::Schema(format!("failed to decode entity extension: {e}")))
    } else {
        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_federation_config_new() {
        let config = FederationConfig::new();
        assert!(!config.is_enabled());
        assert!(config.entities.is_empty());
    }

    #[test]
    fn test_entity_config_composite_keys() {
        // Test that key field sets are properly parsed
        let keys = vec![
            vec!["id".to_string()],
            vec!["org".to_string(), "user".to_string()],
        ];
        assert_eq!(keys.len(), 2);
        assert_eq!(keys[0], vec!["id"]);
        assert_eq!(keys[1], vec!["org", "user"]);
    }
}
