//! Query Whitelisting for production security
//!
//! This module provides query whitelisting (also known as "Stored Operations" or "Persisted Queries")
//! to restrict which GraphQL queries can be executed. This is a critical security feature for
//! public-facing GraphQL APIs.
//!
//! # Features
//!
//! - **Hash-based validation**: Queries identified by SHA-256 hash
//! - **ID-based validation**: Queries identified by custom operation IDs
//! - **Multiple modes**: Enforce, Warn, or Disabled
//! - **Flexible loading**: From files, inline maps, or runtime registration
//! - **APQ compatible**: Works alongside Automatic Persisted Queries
//!
//! # Example
//!
//! ```rust,no_run
//! use grpc_graphql_gateway::{Gateway, QueryWhitelistConfig, WhitelistMode};
//! use std::collections::HashMap;
//!
//! # async fn example() -> Result<(), Box<dyn std::error::Error>> {
//! let mut allowed_queries = HashMap::new();
//! allowed_queries.insert(
//!     "getUserById".to_string(),
//!     "query getUserById($id: ID!) { user(id: $id) { id name email } }".to_string()
//! );
//!
//! let gateway = Gateway::builder()
//!     .with_descriptor_set_bytes(&[])
//!     .with_query_whitelist(QueryWhitelistConfig {
//!         mode: WhitelistMode::Enforce,
//!         allowed_queries,
//!         allow_introspection: true,
//!     })
//!     .build()?;
//! # Ok(())
//! # }
//! ```

use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::{Arc, RwLock};
use thiserror::Error;

/// Mode for query whitelist enforcement
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum WhitelistMode {
    /// Enforce whitelist - reject non-whitelisted queries with error
    Enforce,
    /// Warn mode - log warnings but allow non-whitelisted queries
    Warn,
    /// Disabled - whitelist is not checked (default behavior)
    Disabled,
}

/// Configuration for query whitelisting
#[derive(Debug, Clone)]
pub struct QueryWhitelistConfig {
    /// Enforcement mode
    pub mode: WhitelistMode,
    /// Map of operation ID/name -> query string
    /// The key can be a custom operation ID or the query hash
    pub allowed_queries: HashMap<String, String>,
    /// Whether to allow introspection queries even in Enforce mode
    pub allow_introspection: bool,
}

impl Default for QueryWhitelistConfig {
    fn default() -> Self {
        Self {
            mode: WhitelistMode::Disabled,
            allowed_queries: HashMap::new(),
            allow_introspection: true,
        }
    }
}

impl QueryWhitelistConfig {
    /// Create a new disabled whitelist config
    pub fn disabled() -> Self {
        Self {
            mode: WhitelistMode::Disabled,
            ..Default::default()
        }
    }

    /// Create a new enforcing whitelist config
    pub fn enforce() -> Self {
        Self {
            mode: WhitelistMode::Enforce,
            ..Default::default()
        }
    }

    /// Create a new warning whitelist config
    pub fn warn() -> Self {
        Self {
            mode: WhitelistMode::Warn,
            ..Default::default()
        }
    }

    /// Load allowed queries from a JSON file
    ///
    /// Expected format:
    /// ```json
    /// {
    ///   "getUserById": "query getUserById($id: ID!) { user(id: $id) { id name } }",
    ///   "listProducts": "query { products { id name price } }"
    /// }
    /// ```
    pub fn from_json_file<P: AsRef<Path>>(path: P, mode: WhitelistMode) -> Result<Self, QueryWhitelistError> {
        let content = std::fs::read_to_string(path.as_ref())
            .map_err(|e| QueryWhitelistError::FileLoadError(e.to_string()))?;
        
        let allowed_queries: HashMap<String, String> = serde_json::from_str(&content)
            .map_err(|e| QueryWhitelistError::ParseError(e.to_string()))?;

        Ok(Self {
            mode,
            allowed_queries,
            allow_introspection: true,
        })
    }

    /// Load allowed queries from a YAML file
    #[cfg(feature = "yaml")]
    pub fn from_yaml_file<P: AsRef<Path>>(path: P, mode: WhitelistMode) -> Result<Self, QueryWhitelistError> {
        let content = std::fs::read_to_string(path.as_ref())
            .map_err(|e| QueryWhitelistError::FileLoadError(e.to_string()))?;
        
        let allowed_queries: HashMap<String, String> = serde_yaml::from_str(&content)
            .map_err(|e| QueryWhitelistError::ParseError(e.to_string()))?;

        Ok(Self {
            mode,
            allowed_queries,
            allow_introspection: true,
        })
    }

    /// Set enforcement mode
    pub fn with_mode(mut self, mode: WhitelistMode) -> Self {
        self.mode = mode;
        self
    }

    /// Set whether to allow introspection queries
    pub fn with_introspection(mut self, allow: bool) -> Self {
        self.allow_introspection = allow;
        self
    }

    /// Add a single allowed query
    pub fn add_query(mut self, id: String, query: String) -> Self {
        self.allowed_queries.insert(id, query);
        self
    }

    /// Add multiple queries
    pub fn add_queries(mut self, queries: HashMap<String, String>) -> Self {
        self.allowed_queries.extend(queries);
        self
    }
}

/// Errors related to query whitelisting
#[derive(Debug, Error)]
pub enum QueryWhitelistError {
    #[error("Query not in whitelist: {0}")]
    QueryNotWhitelisted(String),
    
    #[error("Failed to load whitelist file: {0}")]
    FileLoadError(String),
    
    #[error("Failed to parse whitelist file: {0}")]
    ParseError(String),
    
    #[error("Query validation failed: {0}")]
    ValidationError(String),
}

/// Query whitelist validator
///
/// This struct is responsible for validating incoming GraphQL queries against
/// a whitelist of allowed operations.
#[derive(Clone)]
pub struct QueryWhitelist {
    config: QueryWhitelistConfig,
    /// Hash lookup: SHA-256 hash -> query string
    query_by_hash: Arc<RwLock<HashMap<String, String>>>,
    /// ID lookup: operation ID/name -> query string
    query_by_id: Arc<RwLock<HashMap<String, String>>>,
    /// Reverse lookup: query hash -> operation IDs
    ids_by_hash: Arc<RwLock<HashMap<String, HashSet<String>>>>,
}

impl QueryWhitelist {
    /// Create a new query whitelist from config
    pub fn new(config: QueryWhitelistConfig) -> Self {
        let mut query_by_hash = HashMap::new();
        let mut query_by_id = HashMap::new();
        let mut ids_by_hash: HashMap<String, HashSet<String>> = HashMap::new();

        // Build indexes
        for (id, query) in &config.allowed_queries {
            let hash = Self::hash_query(query);
            
            // Store by hash
            query_by_hash.insert(hash.clone(), query.clone());
            
            // Store by ID
            query_by_id.insert(id.clone(), query.clone());
            
            // Track which IDs map to this hash
            ids_by_hash.entry(hash).or_insert_with(HashSet::new).insert(id.clone());
        }

        Self {
            config,
            query_by_hash: Arc::new(RwLock::new(query_by_hash)),
            query_by_id: Arc::new(RwLock::new(query_by_id)),
            ids_by_hash: Arc::new(RwLock::new(ids_by_hash)),
        }
    }

    /// Hash a query string using SHA-256
    pub fn hash_query(query: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(query.trim().as_bytes());
        hex::encode(hasher.finalize())
    }

    /// Check if the whitelist is enabled
    pub fn is_enabled(&self) -> bool {
        self.config.mode != WhitelistMode::Disabled
    }

    /// Get the enforcement mode
    pub fn mode(&self) -> WhitelistMode {
        self.config.mode
    }

    /// Validate a query against the whitelist
    ///
    /// Returns Ok(()) if the query is allowed, or an error if it's not whitelisted.
    pub fn validate_query(&self, query: &str, operation_id: Option<&str>) -> Result<(), QueryWhitelistError> {
        // Disabled mode - always allow
        if self.config.mode == WhitelistMode::Disabled {
            return Ok(());
        }

        // Check if this is an introspection query
        if self.config.allow_introspection && self.is_introspection_query(query) {
            tracing::debug!("Allowing introspection query");
            return Ok(());
        }

        // Try to validate by operation ID first if provided
        if let Some(id) = operation_id {
            if self.is_allowed_by_id(id) {
                tracing::debug!("Query allowed by operation ID: {}", id);
                return Ok(());
            }
        }

        // Validate by query hash
        let hash = Self::hash_query(query);
        if self.is_allowed_by_hash(&hash) {
            tracing::debug!("Query allowed by hash: {}", hash);
            return Ok(());
        }

        // Query not whitelisted
        let error_msg = operation_id
            .map(|id| format!("Operation '{}' (hash: {})", id, &hash[..16]))
            .unwrap_or_else(|| format!("Query hash: {}", &hash[..16]));

        match self.config.mode {
            WhitelistMode::Enforce => {
                tracing::warn!("Query rejected by whitelist: {}", error_msg);
                Err(QueryWhitelistError::QueryNotWhitelisted(error_msg))
            }
            WhitelistMode::Warn => {
                tracing::warn!("Query not in whitelist (allowed in Warn mode): {}", error_msg);
                Ok(())
            }
            WhitelistMode::Disabled => Ok(()),
        }
    }

    /// Check if a query is allowed by its hash
    fn is_allowed_by_hash(&self, hash: &str) -> bool {
        self.query_by_hash.read().unwrap().contains_key(hash)
    }

    /// Check if a query is allowed by its operation ID
    fn is_allowed_by_id(&self, id: &str) -> bool {
        self.query_by_id.read().unwrap().contains_key(id)
    }

    /// Check if a query is an introspection query
    fn is_introspection_query(&self, query: &str) -> bool {
        let normalized = query.trim().to_lowercase();
        normalized.contains("__schema") 
            || normalized.contains("__type")
            || normalized.contains("introspectionquery")
    }

    /// Get a whitelisted query by operation ID
    pub fn get_query_by_id(&self, id: &str) -> Option<String> {
        self.query_by_id.read().unwrap().get(id).cloned()
    }

    /// Get a whitelisted query by hash
    pub fn get_query_by_hash(&self, hash: &str) -> Option<String> {
        self.query_by_hash.read().unwrap().get(hash).cloned()
    }

    /// Register a new allowed query at runtime (useful for dynamic whitelists)
    pub fn register_query(&self, id: String, query: String) {
        let hash = Self::hash_query(&query);
        
        // Add to hash lookup
        self.query_by_hash.write().unwrap().insert(hash.clone(), query.clone());
        
        // Add to ID lookup
        self.query_by_id.write().unwrap().insert(id.clone(), query);
        
        // Add to reverse lookup
        self.ids_by_hash.write().unwrap()
            .entry(hash)
            .or_insert_with(HashSet::new)
            .insert(id);
    }

    /// Remove a query from the whitelist
    pub fn remove_query(&self, id: &str) -> bool {
        let mut query_by_id = self.query_by_id.write().unwrap();
        
        if let Some(query) = query_by_id.remove(id) {
            let hash = Self::hash_query(&query);
            
            // Update reverse lookup
            let mut ids_by_hash = self.ids_by_hash.write().unwrap();
            if let Some(ids) = ids_by_hash.get_mut(&hash) {
                ids.remove(id);
                
                // If no more IDs reference this hash, remove the hash entry
                if ids.is_empty() {
                    ids_by_hash.remove(&hash);
                    self.query_by_hash.write().unwrap().remove(&hash);
                }
            }
            
            true
        } else {
            false
        }
    }

    /// Get statistics about the whitelist
    pub fn stats(&self) -> WhitelistStats {
        WhitelistStats {
            mode: self.config.mode,
            total_queries: self.query_by_id.read().unwrap().len(),
            total_hashes: self.query_by_hash.read().unwrap().len(),
            allow_introspection: self.config.allow_introspection,
        }
    }
}

/// Statistics about the query whitelist
#[derive(Debug, Clone, Serialize)]
pub struct WhitelistStats {
    pub mode: WhitelistMode,
    pub total_queries: usize,
    pub total_hashes: usize,
    pub allow_introspection: bool,
}

/// Shared query whitelist (thread-safe)
pub type SharedQueryWhitelist = Arc<QueryWhitelist>;

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

    #[test]
    fn test_query_hashing() {
        let query1 = "query { user(id: 1) { name } }";
        let query2 = "query { user(id: 1) { name } }"; // Same
        let query3 = "query { user(id: 2) { name } }"; // Different

        let hash1 = QueryWhitelist::hash_query(query1);
        let hash2 = QueryWhitelist::hash_query(query2);
        let hash3 = QueryWhitelist::hash_query(query3);

        assert_eq!(hash1, hash2);
        assert_ne!(hash1, hash3);
        assert_eq!(hash1.len(), 64); // SHA-256 produces 64 hex chars
    }

    #[test]
    fn test_whitelist_enforce_mode() {
        let config = QueryWhitelistConfig {
            mode: WhitelistMode::Enforce,
            allowed_queries: vec![
                ("getUser".to_string(), "query { user { id name } }".to_string()),
            ].into_iter().collect(),
            allow_introspection: false,
        };

        let whitelist = QueryWhitelist::new(config);

        // Allowed query
        assert!(whitelist.validate_query("query { user { id name } }", None).is_ok());

        // Allowed by operation ID
        assert!(whitelist.validate_query("query { user { id name } }", Some("getUser")).is_ok());

        // Not allowed
        assert!(whitelist.validate_query("query { posts { title } }", None).is_err());
    }

    #[test]
    fn test_whitelist_warn_mode() {
        let config = QueryWhitelistConfig {
            mode: WhitelistMode::Warn,
            allowed_queries: HashMap::new(),
            allow_introspection: false,
        };

        let whitelist = QueryWhitelist::new(config);

        // Any query should pass in Warn mode
        assert!(whitelist.validate_query("query { unknown { field } }", None).is_ok());
    }

    #[test]
    fn test_whitelist_disabled() {
        let config = QueryWhitelistConfig::disabled();
        let whitelist = QueryWhitelist::new(config);

        // Any query should pass when disabled
        assert!(whitelist.validate_query("query { anything }", None).is_ok());
    }

    #[test]
    fn test_introspection_allowed() {
        let config = QueryWhitelistConfig {
            mode: WhitelistMode::Enforce,
            allowed_queries: HashMap::new(),
            allow_introspection: true,
        };

        let whitelist = QueryWhitelist::new(config);

        // Introspection queries should be allowed
        assert!(whitelist.validate_query("query { __schema { types { name } } }", None).is_ok());
        assert!(whitelist.validate_query("query { __type(name: \"User\") { fields { name } } }", None).is_ok());

        // Regular queries should still be blocked
        assert!(whitelist.validate_query("query { user { id } }", None).is_err());
    }

    #[test]
    fn test_runtime_registration() {
        let config = QueryWhitelistConfig::enforce();
        let whitelist = QueryWhitelist::new(config);

        // Initially not allowed
        let query = "query { newQuery { field } }";
        assert!(whitelist.validate_query(query, None).is_err());

        // Register the query
        whitelist.register_query("newQuery".to_string(), query.to_string());

        // Now it should be allowed
        assert!(whitelist.validate_query(query, None).is_ok());
        assert!(whitelist.validate_query(query, Some("newQuery")).is_ok());
    }

    #[test]
    fn test_remove_query() {
        let mut queries = HashMap::new();
        queries.insert("getUser".to_string(), "query { user { id } }".to_string());

        let config = QueryWhitelistConfig {
            mode: WhitelistMode::Enforce,
            allowed_queries: queries,
            allow_introspection: false,
        };

        let whitelist = QueryWhitelist::new(config);

        // Initially allowed
        assert!(whitelist.validate_query("query { user { id } }", Some("getUser")).is_ok());

        // Remove it
        assert!(whitelist.remove_query("getUser"));

        // Now should be blocked
        assert!(whitelist.validate_query("query { user { id } }", Some("getUser")).is_err());

        // Removing again should return false
        assert!(!whitelist.remove_query("getUser"));
    }

    #[test]
    fn test_stats() {
        let mut queries = HashMap::new();
        queries.insert("q1".to_string(), "query { a }".to_string());
        queries.insert("q2".to_string(), "query { b }".to_string());

        let config = QueryWhitelistConfig {
            mode: WhitelistMode::Enforce,
            allowed_queries: queries,
            allow_introspection: true,
        };

        let whitelist = QueryWhitelist::new(config);
        let stats = whitelist.stats();

        assert_eq!(stats.mode, WhitelistMode::Enforce);
        assert_eq!(stats.total_queries, 2);
        assert_eq!(stats.total_hashes, 2);
        assert!(stats.allow_introspection);
    }
}
