auto_discovery/

discovery.rs

//! Main service discovery implementation

use crate::{
    config::DiscoveryConfig,
    error::{DiscoveryError, Result},
    protocols::ProtocolManager,
    service::ServiceInfo,
    types::ProtocolType,
};
use std::{
    collections::HashMap,
    sync::Arc,
};
use tokio::sync::Mutex;
use tracing::{debug, info};

/// Main service discovery interface
pub struct ServiceDiscovery {
    config: DiscoveryConfig,
    protocol_manager: ProtocolManager,
    discovered_services: Arc<Mutex<HashMap<String, ServiceInfo>>>,
    registered_services: Arc<Mutex<HashMap<String, ServiceInfo>>>,
}

impl ServiceDiscovery {
    /// Create a new service discovery instance with the given configuration
    /// 
    /// # Arguments
    /// 
    /// * `config` - The discovery configuration to use
    /// 
    /// # Errors
    /// 
    /// Returns an error if the configuration is invalid or if protocol initialization fails
    pub async fn new(config: DiscoveryConfig) -> Result<Self> {
        // Validate configuration before proceeding
        config.validate()?;
        
        let protocol_manager = ProtocolManager::new(config.clone()).await?;

        Ok(Self {
            config,
            protocol_manager,
            discovered_services: Arc::new(Mutex::new(HashMap::new())),
            registered_services: Arc::new(Mutex::new(HashMap::new())),
        })
    }

    /// Discover services with optional protocol type filter
    pub async fn discover_services(&self, protocol_type: Option<ProtocolType>) -> Result<Vec<ServiceInfo>> {
        debug!("Starting service discovery");
        
        let service_types = self.config.service_types().to_vec();
        if service_types.is_empty() {
            return Err(DiscoveryError::configuration("No service types configured for discovery"));
        }

        let timeout = Some(self.config.protocol_timeout());
        let mut services = match protocol_type {
            Some(protocol) => {
                if !self.config.is_protocol_enabled(protocol) {
                    return Err(DiscoveryError::protocol(format!("Protocol {protocol:?} is not enabled")));
                }
                self.protocol_manager.discover_services_with_protocol(protocol, service_types, timeout).await?
            }
            None => self.protocol_manager.discover_services(service_types, timeout).await?,
        };

        // Apply service filtering
        if let Some(filter) = self.config.filter() {
            services.retain(|service| filter.matches(service));
        }

        // Limit number of services if configured
        let max_services = self.config.max_services();
        if max_services > 0 && services.len() > max_services {
            services.truncate(max_services);
        }

        // Update discovered services cache
        let mut discovered = self.discovered_services.lock().await;
        for service in &services {
            discovered.insert(service.name().to_string(), service.clone());
        }

        info!("Discovered {} services", services.len());
        info!("Discovered {} services", services.len());
        Ok(services)
    }

    /// Discover services with filtering by service types
    /// 
    /// This provides more granular control over service discovery than the basic
    /// `discover_services` method.
    /// 
    /// # Arguments
    /// 
    /// * `service_types` - Optional list of specific service types to discover. 
    ///   If None, uses all configured service types.
    /// * `protocol_type` - Optional protocol filter. If None, uses all enabled protocols.
    /// 
    /// # Example
    /// 
    /// ```rust
    /// use auto_discovery::{ServiceDiscovery, types::{ServiceType, ProtocolType}};
    /// 
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// # let discovery = ServiceDiscovery::new(auto_discovery::config::DiscoveryConfig::new()).await?;
    /// // Discover only HTTP services using mDNS
    /// let http_services = discovery.discover_services_filtered(
    ///     Some(vec![ServiceType::new("_http._tcp")?]),
    ///     Some(ProtocolType::Mdns)
    /// ).await?;
    /// 
    /// // Discover HTTP services using any protocol
    /// let all_services = discovery.discover_services_filtered(
    ///     Some(vec![ServiceType::new("_http._tcp")?]), 
    ///     None
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn discover_services_filtered(&self, 
        service_types: Option<Vec<crate::types::ServiceType>>, 
        protocol_type: Option<ProtocolType>
    ) -> Result<Vec<ServiceInfo>> {
        debug!("Starting filtered service discovery");
        
        let target_service_types = match service_types {
            Some(types) => types,
            None => self.config.service_types().to_vec()
        };

        if target_service_types.is_empty() {
            return Err(DiscoveryError::configuration("No service types specified for discovery"));
        }

        let timeout = Some(self.config.protocol_timeout());
        let mut services = match protocol_type {
            Some(protocol) => {
                if !self.config.is_protocol_enabled(protocol) {
                    return Err(DiscoveryError::protocol(format!("Protocol {protocol:?} is not enabled")));
                }
                self.protocol_manager.discover_services_with_protocol(protocol, target_service_types, timeout).await?
            }
            None => self.protocol_manager.discover_services(target_service_types, timeout).await?,
        };

        // Apply service filtering
        if let Some(filter) = self.config.filter() {
            services.retain(|service| filter.matches(service));
        }

        // Update discovered services cache
        {
            let mut discovered = self.discovered_services.lock().await;
            for service in &services {
                discovered.insert(service.id.to_string(), service.clone());
            }
        }

        info!("Discovered {} filtered services", services.len());
        Ok(services)
    }

    /// Register a service
    pub async fn register_service(&self, service: ServiceInfo) -> Result<()> {
        let service_name = service.name().to_string();
        debug!("Registering service: {}", service_name);

        self.protocol_manager.register_service(service.clone()).await?;

        let mut registered = self.registered_services.lock().await;
        registered.insert(service_name.clone(), service);

        info!("Successfully registered service: {}", service_name);
        Ok(())
    }

    /// Unregister a service
    pub async fn unregister_service(&self, service: &ServiceInfo) -> Result<()> {
        let service_name = service.name().to_string();
        debug!("Unregistering service: {}", service_name);

        self.protocol_manager.unregister_service(service).await?;

        let mut registered = self.registered_services.lock().await;
        registered.remove(&service_name);

        info!("Successfully unregistered service: {}", service_name);
        Ok(())
    }

    /// Verify a service is still available
    pub async fn verify_service(&self, service: &ServiceInfo) -> Result<bool> {
        debug!("Verifying service: {}", service.name());

        self.protocol_manager.verify_service(service).await
    }

    /// Get all discovered services
    pub async fn get_discovered_services(&self) -> Vec<ServiceInfo> {
        self.discovered_services.lock().await
            .values()
            .cloned()
            .collect()
    }

    /// Get all registered services
    pub async fn get_registered_services(&self) -> Vec<ServiceInfo> {
        self.registered_services.lock().await
            .values()
            .cloned()
            .collect()
    }

    /// Check if a service exists
    pub async fn service_exists(&self, service_name: &str) -> bool {
        self.discovered_services.lock().await.contains_key(service_name) ||
        self.registered_services.lock().await.contains_key(service_name)
    }

    /// Update discovery configuration
    pub async fn update_config(&mut self, config: DiscoveryConfig) -> Result<()> {
        self.config = config.clone();
        self.protocol_manager = ProtocolManager::new(config).await?;
        Ok(())
    }
}



#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::ServiceType;
    use std::time::Duration;

    #[tokio::test]
    async fn test_service_discovery_creation() {
        let config = DiscoveryConfig::new()
            .with_service_type(ServiceType::new("_test._tcp").unwrap())
            .with_timeout(Duration::from_secs(1));

        let discovery = ServiceDiscovery::new(config).await;
        assert!(discovery.is_ok());
    }

    #[tokio::test]
    async fn test_service_registration() {
        let config = DiscoveryConfig::new();
        let discovery = ServiceDiscovery::new(config).await.unwrap();

        let service = ServiceInfo::new("Test Service", "_test._tcp", 8080, None).unwrap();
        let result = discovery.register_service(service).await;
        
        // Registration might fail due to missing protocol implementation
        // This is expected in unit tests
        match result {
            Ok(_) => {
                assert_eq!(discovery.get_registered_services().await.len(), 1);
            }
            Err(_) => {
                // Expected in unit test environment
            }
        }
    }

    #[tokio::test]
    async fn test_event_subscription() {
        let config = DiscoveryConfig::new();
        let discovery = ServiceDiscovery::new(config).await.unwrap();

        // Event subscription is not explicitly tested here as it depends on the underlying protocol implementation
        // Just ensure it can be called without error
        let _ = discovery.discover_services(None).await;
    }

    #[tokio::test]
    async fn test_config_validation() {
        let invalid_config = DiscoveryConfig::new().with_timeout(Duration::ZERO);
        let discovery = ServiceDiscovery::new(invalid_config).await;
        assert!(discovery.is_err());
    }
}