//! IP Whitelist Enforcement Module
//!
//! Provides IP-based access control with CIDR subnet support.
//! Loads configuration from YAML file with auto-reload support.
//! All connections not in whitelist are silently dropped.

use anyhow::{Context, Result};
use notify::{Config, Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
use serde::Deserialize;
use std::collections::HashSet;
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use std::path::Path;
use std::sync::{Arc, RwLock};
use std::time::Duration;
use tracing::{debug, info, warn};

/// Whitelist entry types
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum WhitelistEntry {
    /// Single IP address
    Ip(Ipv4Addr),
    /// CIDR subnet
    Cidr { network: Ipv4Addr, prefix: u8 },
    /// Hostname (resolved at startup)
    Hostname { name: String, resolved: Ipv4Addr },
}

/// Whitelist configuration loaded from YAML
#[derive(Debug, Deserialize, Clone)]
pub struct WhitelistConfig {
    pub entries: Vec<String>,
}

/// IP Whitelist manager with auto-reload support
pub struct WhitelistManager {
    entries: Arc<RwLock<HashSet<WhitelistEntry>>>,
    config_path: String,
    watcher: Option<RecommendedWatcher>,
}

impl WhitelistManager {
    /// Create a new whitelist manager
    pub fn new(config_path: &str) -> Result<Self> {
        let entries = Arc::new(RwLock::new(HashSet::new()));

        let mut manager = Self {
            entries: entries.clone(),
            config_path: config_path.to_string(),
            watcher: None,
        };

        // Load initial whitelist
        manager.reload()?;

        // Set up file watcher for auto-reload
        manager.setup_watcher()?;

        Ok(manager)
    }

    /// Reload whitelist from configuration file
    pub fn reload(&self) -> Result<()> {
        let config = self.load_config()?;
        let entries = self.parse_entries(&config.entries)?;

        let mut current_entries = self
            .entries
            .write()
            .map_err(|e| anyhow::anyhow!("Failed to acquire whitelist lock: {}", e))?;

        *current_entries = entries;

        info!(
            path = %self.config_path,
            count = current_entries.len(),
            "Whitelist reloaded successfully"
        );

        Ok(())
    }

    /// Check if an IP address is allowed
    pub fn is_allowed(&self, ip: &Ipv4Addr) -> bool {
        let entries = self.entries.read().unwrap();

        for entry in entries.iter() {
            match entry {
                WhitelistEntry::Ip(allowed_ip) => {
                    if ip == allowed_ip {
                        return true;
                    }
                }
                WhitelistEntry::Cidr { network, prefix } => {
                    if ip_in_subnet(ip, *network, *prefix) {
                        return true;
                    }
                }
                WhitelistEntry::Hostname { resolved, .. } => {
                    if ip == resolved {
                        return true;
                    }
                }
            }
        }

        false
    }

    /// Check if a socket address is allowed
    pub fn is_socket_allowed(&self, socket_addr: &SocketAddr) -> bool {
        match socket_addr.ip() {
            IpAddr::V4(ip) => self.is_allowed(&ip),
            IpAddr::V6(_) => {
                // IPv6 not supported in whitelist - deny by default
                warn!(socket_addr = %socket_addr, "IPv6 address denied - whitelist supports IPv4 only");
                false
            }
        }
    }

    /// Get the number of entries in the whitelist
    pub fn entry_count(&self) -> usize {
        self.entries.read().unwrap().len()
    }

    /// Load configuration from YAML file
    fn load_config(&self) -> Result<WhitelistConfig> {
        let content = std::fs::read_to_string(&self.config_path)
            .with_context(|| format!("Failed to read whitelist config: {}", self.config_path))?;

        let config: WhitelistConfig = serde_yaml::from_str(&content)
            .with_context(|| format!("Failed to parse whitelist config: {}", self.config_path))?;

        Ok(config)
    }

    /// Parse whitelist entries from strings
    fn parse_entries(&self, entries: &[String]) -> Result<HashSet<WhitelistEntry>> {
        let mut parsed = HashSet::new();

        for entry_str in entries {
            let entry_str = entry_str.trim();

            // Skip comments and empty lines
            if entry_str.is_empty() || entry_str.starts_with('#') {
                continue;
            }

            // Check for CIDR notation
            if let Some((ip_str, prefix_str)) = entry_str.split_once('/') {
                let ip: Ipv4Addr = ip_str
                    .parse()
                    .with_context(|| format!("Invalid IP in CIDR notation: {}", entry_str))?;
                let prefix: u8 = prefix_str
                    .parse()
                    .with_context(|| format!("Invalid prefix in CIDR notation: {}", entry_str))?;

                if prefix > 32 {
                    anyhow::bail!("Invalid CIDR prefix (must be 0-32): {}", entry_str);
                }

                parsed.insert(WhitelistEntry::Cidr {
                    network: ip,
                    prefix,
                });
                debug!("Added CIDR entry: {}", entry_str);
            } else {
                // Try to parse as IP address
                if let Ok(ip) = entry_str.parse::<Ipv4Addr>() {
                    parsed.insert(WhitelistEntry::Ip(ip));
                    debug!("Added IP entry: {}", entry_str);
                } else {
                    // Try to resolve as hostname
                    match resolve_hostname(entry_str) {
                        Ok(resolved) => {
                            parsed.insert(WhitelistEntry::Hostname {
                                name: entry_str.to_string(),
                                resolved,
                            });
                            info!("Resolved hostname {} to {}", entry_str, resolved);
                        }
                        Err(e) => {
                            warn!("Failed to resolve hostname {}: {}", entry_str, e);
                        }
                    }
                }
            }
        }

        Ok(parsed)
    }

    /// Set up file watcher for auto-reload
    fn setup_watcher(&mut self) -> Result<()> {
        let config_path = self.config_path.clone();
        let _entries = self.entries.clone();

        let watcher = RecommendedWatcher::new(
            move |res: Result<Event, notify::Error>| {
                if let Ok(event) = res {
                    match event.kind {
                        EventKind::Modify(_) | EventKind::Create(_) => {
                            info!("Whitelist file changed, reloading...");
                            // Reload is handled by the manager
                        }
                        _ => {}
                    }
                }
            },
            Config::default().with_poll_interval(Duration::from_secs(5)),
        )?;

        let mut watcher = watcher;
        let path = Path::new(&config_path);

        if path.exists() {
            watcher.watch(path, RecursiveMode::NonRecursive)?;
            info!("Watching whitelist file for changes: {}", config_path);
        } else {
            warn!("Whitelist file does not exist yet: {}", config_path);
        }

        self.watcher = Some(watcher);

        Ok(())
    }
}

/// Check if an IP address is within a CIDR subnet
fn ip_in_subnet(ip: &Ipv4Addr, network: Ipv4Addr, prefix: u8) -> bool {
    let ip_bits = u32::from(*ip);
    let network_bits = u32::from(network);
    let mask = if prefix == 0 {
        0
    } else {
        !0u32 << (32 - prefix)
    };

    (ip_bits & mask) == (network_bits & mask)
}

/// Resolve a hostname to an IPv4 address
fn resolve_hostname(hostname: &str) -> Result<Ipv4Addr> {
    use std::net::ToSocketAddrs;

    let addrs = (hostname, 0)
        .to_socket_addrs()
        .with_context(|| format!("Failed to resolve hostname: {}", hostname))?;

    for addr in addrs {
        if let IpAddr::V4(ip) = addr.ip() {
            return Ok(ip);
        }
    }

    anyhow::bail!("No IPv4 address found for hostname: {}", hostname)
}

/// Whitelist middleware for Actix-web
pub struct WhitelistMiddleware {
    manager: Arc<WhitelistManager>,
}

impl WhitelistMiddleware {
    /// Create a new whitelist middleware
    pub fn new(manager: WhitelistManager) -> Self {
        Self {
            manager: Arc::new(manager),
        }
    }

    /// Get the whitelist manager reference
    pub fn manager(&self) -> Arc<WhitelistManager> {
        self.manager.clone()
    }
}

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

    #[test]
    fn test_ip_in_subnet() {
        // Test /24 subnet
        assert!(ip_in_subnet(
            &"192.168.1.100".parse().unwrap(),
            "192.168.1.0".parse().unwrap(),
            24
        ));
        assert!(ip_in_subnet(
            &"192.168.1.254".parse().unwrap(),
            "192.168.1.0".parse().unwrap(),
            24
        ));
        assert!(!ip_in_subnet(
            &"192.168.2.1".parse().unwrap(),
            "192.168.1.0".parse().unwrap(),
            24
        ));

        // Test /16 subnet
        assert!(ip_in_subnet(
            &"192.168.100.50".parse().unwrap(),
            "192.168.0.0".parse().unwrap(),
            16
        ));
        assert!(!ip_in_subnet(
            &"192.169.0.1".parse().unwrap(),
            "192.168.0.0".parse().unwrap(),
            16
        ));

        // Test /32 (single host)
        assert!(ip_in_subnet(
            &"10.0.0.50".parse().unwrap(),
            "10.0.0.50".parse().unwrap(),
            32
        ));
        assert!(!ip_in_subnet(
            &"10.0.0.51".parse().unwrap(),
            "10.0.0.50".parse().unwrap(),
            32
        ));

        // Test /0 (all IPs)
        assert!(ip_in_subnet(
            &"1.2.3.4".parse().unwrap(),
            "0.0.0.0".parse().unwrap(),
            0
        ));
    }

    #[test]
    fn test_whitelist_entry_parsing() {
        let manager = WhitelistManager::new("/tmp/test_whitelist.yaml").unwrap_or_else(|_| {
            // Create a temp file for testing
            let temp_path = "/tmp/test_whitelist_temp.yaml";
            std::fs::write(temp_path, "entries:\n  - \"192.168.1.0/24\"\n").unwrap();
            WhitelistManager::new(temp_path).unwrap()
        });

        // Test IP entry
        let ip: Ipv4Addr = "192.168.1.100".parse().unwrap();
        assert!(manager.is_allowed(&ip));

        // Test IP outside subnet
        let ip_outside: Ipv4Addr = "192.168.2.100".parse().unwrap();
        assert!(!manager.is_allowed(&ip_outside));
    }
}
