use anyhow::{anyhow, Context, Result};
use futures::future::join_all;
use std::{collections::HashSet, path::{Path, PathBuf}, sync::{Arc, RwLock}, time::Duration};
use tokio::{process::Command, time::{sleep}};
use walkdir::WalkDir;

pub struct DockerVpnProxyPool {
    container_names: Vec<String>,
    proxy_ports: Vec<u16>,  // e.g., [10801, 10802, ...]
    dead_proxies: Arc<RwLock<HashSet<usize>>>,
}

impl DockerVpnProxyPool {
    pub async fn new(
        ovpn_dir: &Path, 
        username: String, 
        password: String,
        instances_per_ovpn: usize,
    ) -> Result<Self> {
        if instances_per_ovpn == 0 {
            return Err(anyhow!("instances_per_ovpn must be at least 1"));
        }

        // Count hostnames (subdirs in ovpn_dir)
        let hostnames: Vec<_> = std::fs::read_dir(ovpn_dir)?
            .filter_map(Result::ok)
            .filter(|e| e.path().is_dir())
            .map(|e| e.file_name().into_string().unwrap())
            .collect();

        let num_servers = hostnames.len();
        if num_servers == 0 {
            return Err(anyhow!("No VPN hostnames found in {:?}", ovpn_dir));
        }

        // Calculate total containers: hostnames × instances_per_ovpn
        let total_containers = num_servers * instances_per_ovpn;

        crate::util::logger::log_info(&format!(
            "Found {} VPN hostnames × {} instances = {} total containers",
            num_servers, instances_per_ovpn, total_containers
        )).await;

        let mut container_names = Vec::with_capacity(total_containers);
        let mut proxy_ports = Vec::with_capacity(total_containers);
        let base_port: u16 = 10800;
        let mut port_counter = 0u16;

        // === STEP 1: Start ALL containers first ===
        for hostname in hostnames.iter() {
            // Pick tcp443.ovpn if exists, else first .ovpn
            let hostname_dir = ovpn_dir.join(hostname);
            let mut ovpn_path: Option<PathBuf> = None;
            for entry in WalkDir::new(&hostname_dir).max_depth(1) {
                let entry = entry?;
                if entry.path().extension().map_or(false, |ext| ext == "ovpn") {
                    if entry.file_name().to_str().unwrap_or("").contains("tcp443") {
                        ovpn_path = Some(entry.path().to_path_buf());
                        break;
                    } else if ovpn_path.is_none() {
                        ovpn_path = Some(entry.path().to_path_buf());
                    }
                }
            }

            let ovpn_path = ovpn_path.ok_or_else(|| anyhow!("No .ovpn found for {}", hostname))?;

            // Spawn multiple instances for this .ovpn file
            for instance_num in 0..instances_per_ovpn {
                let name = format!("vpn-proxy-{}-{}", hostname, instance_num);
                let port = base_port + port_counter + 1;
                port_counter += 1;

                // Clean up any existing container with the same name
                let _ = Command::new("docker")
                    .args(["rm", "-f", &name])
                    .status()
                    .await;

                // Run Docker container
                let status = Command::new("docker")
                    .args([
                        "run", "-d",
                        "--name", &name,
                        "--cap-add=NET_ADMIN",
                        "--device", "/dev/net/tun",
                        "--sysctl", "net.ipv4.ip_forward=1",
                        "-v", &format!("{}:/vpn/config.ovpn", ovpn_path.display()),
                        "-e", &format!("VPN_USERNAME={}", username),
                        "-e", &format!("VPN_PASSWORD={}", password),
                        "-p", &format!("{}:1080", port),
                        "rust-vpn-proxy",
                    ])
                    .status()
                    .await
                    .context("Failed to run Docker")?;

                if !status.success() {
                    return Err(anyhow!("Docker run failed for {}", name));
                }

                crate::util::logger::log_info(&format!(
                    "Started container {} on port {} (using {})", 
                    name, port, ovpn_path.file_name().unwrap().to_string_lossy()
                )).await;

                container_names.push(name);
                proxy_ports.push(port);
            }
        }

        // Brief pause to let containers start
        sleep(Duration::from_secs(8)).await;
        crate::util::logger::log_info(&format!(
            "All {} containers started, beginning health checks...", 
            container_names.len()
        )).await;

        // === STEP 2: Test ALL proxies in parallel ===
        let results = Self::test_all_proxies_parallel(&container_names, &proxy_ports).await;
        
        // Filter out failed containers
        let mut working_containers = Vec::new();
        let mut working_ports = Vec::new();
        let mut failed_count = 0;

        for (i, (container_name, port)) in container_names.into_iter().zip(proxy_ports.into_iter()).enumerate() {
            match &results[i] {
                Ok(Some(ip)) => {
                    crate::util::logger::log_info(&format!(
                        "✓ Container {} on port {} ready with IP: {}", 
                        container_name, port, ip
                    )).await;
                    working_containers.push(container_name);
                    working_ports.push(port);
                }
                Ok(None) => {
                    let logs = Command::new("docker")
                        .args(["logs", "--tail", "20", &container_name])
                        .output()
                        .await
                        .ok()
                        .and_then(|output| String::from_utf8_lossy(&output.stdout).to_string().into());
                    
                    crate::util::logger::log_error(&format!(
                        "✗ Container {} on port {} ready but IP detection failed. Logs: {:?}", 
                        container_name, port, logs
                    )).await;
                    failed_count += 1;
                    // Clean up failed container
                    let _ = Self::cleanup_container(&container_name).await;
                }
                Err(e) => {
                    let logs = Command::new("docker")
                        .args(["logs", "--tail", "20", &container_name])
                        .output()
                        .await
                        .ok()
                        .and_then(|output| String::from_utf8_lossy(&output.stdout).to_string().into());
                    
                    crate::util::logger::log_error(&format!(
                        "✗ Container {} on port {} failed: {}. Logs: {:?}", 
                        container_name, port, e, logs
                    )).await;
                    failed_count += 1;
                    // Clean up failed container
                    let _ = Self::cleanup_container(&container_name).await;
                }
            }
        }

        if working_containers.is_empty() {
            return Err(anyhow!("All {} VPN proxy containers failed to start", total_containers));
        }

        crate::util::logger::log_info(&format!(
            "Started {}/{} VPN proxy containers successfully ({} hostnames × {} instances)", 
            working_containers.len(), total_containers, num_servers, instances_per_ovpn
        )).await;
        
        if failed_count > 0 {
            crate::util::logger::log_warn(&format!(
                "{} containers failed and were cleaned up", 
                failed_count
            )).await;
        }

        Ok(Self {
            container_names: working_containers,
            proxy_ports: working_ports,
            dead_proxies: Arc::new(RwLock::new(HashSet::new())),
        })
    }

    /// Test all proxies in parallel with 10-second intervals between tests
    async fn test_all_proxies_parallel(container_names: &[String], proxy_ports: &[u16]) -> Vec<Result<Option<String>>> {
        let mut tasks = Vec::new();
        
        for (_i, (container_name, port)) in container_names.iter().zip(proxy_ports.iter()).enumerate() {
            let name = container_name.clone();
            let port = *port;
            
            tasks.push(tokio::spawn(async move {
                // Try up to 6 times with 10-second intervals (total 60 seconds)
                for attempt in 1..=6 {
                    crate::util::logger::log_info(&format!("Testing proxy {} (port {}) - Attempt {}/6", 
                        name, port, attempt)).await;
                    
                    match Self::test_single_proxy(port).await {
                        Ok(Some(ip)) => {
                            return Ok(Some(ip));
                        }
                        Ok(None) => {
                            // Connection works but IP detection failed
                            return Ok(None);
                        }
                        Err(e) if attempt < 6 => {
                            crate::util::logger::log_info(&format!("Attempt {}/6 for {}: {} - retrying in 10s", 
                                attempt, name, e)).await;
                            sleep(Duration::from_secs(10)).await;
                        }
                        Err(e) => {
                            return Err(anyhow!("Failed after 6 attempts: {}", e));
                        }
                    }
                }
                Err(anyhow!("Unexpected exit from retry loop"))
            }));
        }
        
        // Wait for all tasks to complete
        join_all(tasks)
            .await
            .into_iter()
            .map(|result| match result {
                Ok(inner) => inner,
                Err(e) => Err(anyhow!("Task panicked: {}", e)),
            })
            .collect()
    }

    /// Test a single proxy connection
    async fn test_single_proxy(port: u16) -> Result<Option<String>> {
        use std::io::{Read, Write};
        use std::net::TcpStream;
        use std::time::Duration as StdDuration;
        
        // First, test SOCKS5 handshake directly
        crate::util::logger::log_info(&format!("Testing SOCKS5 handshake on port {}...", port)).await;
        
        // Use spawn_blocking for synchronous I/O
        let test_result = tokio::task::spawn_blocking(move || {
            // Connect to SOCKS5 proxy
            let mut stream = match TcpStream::connect_timeout(
                &format!("127.0.0.1:{}", port).parse().unwrap(),
                StdDuration::from_secs(5)
            ) {
                Ok(stream) => stream,
                Err(e) => return Err(anyhow!("Failed to connect: {}", e)),
            };
            
            // Send SOCKS5 greeting: version 5, 1 method (no auth)
            let greeting: [u8; 3] = [0x05, 0x01, 0x00]; // SOCKS5, 1 method, no auth
            if let Err(e) = stream.write_all(&greeting) {
                return Err(anyhow!("Failed to send greeting: {}", e));
            }
            
            // Read response
            let mut response = [0u8; 2];
            if let Err(e) = stream.read_exact(&mut response) {
                return Err(anyhow!("Failed to read response: {}", e));
            }
            
            // Check response: should be [0x05, 0x00] for no auth required
            if response[0] != 0x05 || response[1] != 0x00 {
                return Err(anyhow!("Unexpected SOCKS5 response: {:?}", response));
            }
            
            Ok(())
        }).await;
        
        match test_result {
            Ok(Ok(())) => {
                crate::util::logger::log_info(&format!("✓ SOCKS5 proxy on port {} accepts connections", port)).await;
                
                // Try to get IP through proxy using curl (fallback method)
                let curl_result = tokio::process::Command::new("curl")
                    .args([
                        "-s",
                        "--socks5", &format!("localhost:{}", port),
                        "--max-time", "10",
                        "https://checkip.amazonaws.com"
                    ])
                    .output()
                    .await;
                
                match curl_result {
                    Ok(output) if output.status.success() => {
                        let ip = String::from_utf8_lossy(&output.stdout).trim().to_string();
                        if Self::is_valid_ip(&ip) {
                            crate::util::logger::log_info(&format!("✓ Got IP via proxy: {}", ip)).await;
                            return Ok(Some(ip));
                        } else {
                            crate::util::logger::log_info(&format!("✓ Proxy works, invalid IP format: {}", ip)).await;
                            return Ok(None);
                        }
                    }
                    _ => {
                        // Proxy accepts connections but curl failed - still acceptable
                        crate::util::logger::log_info(&format!("✓ Proxy accepts connections (curl test failed)")).await;
                        return Ok(None);
                    }
                }
            }
            Ok(Err(e)) => {
                return Err(anyhow!("SOCKS5 test failed: {}", e));
            }
            Err(e) => {
                return Err(anyhow!("Task failed: {}", e));
            }
        }
    }

    /// Clean up a failed container
    async fn cleanup_container(container_name: &str) -> Result<()> {
        let _ = Command::new("docker")
            .args(["stop", container_name])
            .status()
            .await;
        
        let _ = Command::new("docker")
            .args(["rm", container_name])
            .status()
            .await;
        
        Ok(())
    }

    fn is_valid_ip(ip: &str) -> bool {
        let parts: Vec<&str> = ip.split('.').collect();
        if parts.len() != 4 {
            return false;
        }
        
        for part in parts {
            if let Ok(num) = part.parse::<u8>() {
                if part != num.to_string() {
                    return false;
                }
            } else {
                return false;
            }
        }
        
        true
    }

    pub fn get_proxy_url(&self, index: usize) -> String {
        let port = self.proxy_ports[index % self.proxy_ports.len()];
        format!("socks5h://localhost:{}", port)
    }

    pub fn num_proxies(&self) -> usize {
        self.proxy_ports.len()
    }

    pub async fn shutdown(&self) -> Result<()> {
        crate::util::logger::log_info(&format!("Shutting down {} Docker proxy containers...", 
            self.container_names.len())).await;
            
        for name in &self.container_names {
            let _ = Command::new("docker")
                .args(["stop", name])
                .status()
                .await;
            let _ = Command::new("docker")
                .args(["rm", name])
                .status()
                .await;
        }
        Ok(())
    }

    /// Get ProxyInfo for monitoring dashboard
    pub fn get_proxy_info(&self, index: usize) -> Option<crate::monitoring::ProxyInfo> {
        if index >= self.container_names.len() {
            return None;
        }
        
        Some(crate::monitoring::ProxyInfo {
            container_name: self.container_names[index].clone(),
            ip_address: "127.0.0.1".to_string(), // SOCKS5 proxy on localhost
            port: self.proxy_ports[index],
            status: crate::monitoring::ProxyStatus::Connected,
        })
    }

    /// Get container name by index
    pub fn get_container_name(&self, index: usize) -> Option<String> {
        self.container_names.get(index).cloned()
    }

    // Get a healthy proxy URL (skips dead proxies)
    pub async fn get_healthy_proxy_url(&self, start_index: usize) -> Option<(usize, String)> {
        let dead = match self.dead_proxies.read() {
            Ok(value) => value,
            Err(_) => return None,
        };
        let total = self.proxy_ports.len();
        
        // Try up to 'total' proxies starting from start_index
        for attempt in 0..total {
            let index = (start_index + attempt) % total;
            
            // Skip if dead
            if dead.contains(&index) {
                continue;
            }
            
            let port = self.proxy_ports[index];
            return Some((index, format!("socks5h://localhost:{}", port)));
        }
        
        None
    }
    
    // Mark a proxy as dead
    pub async fn mark_proxy_dead(&self, index: usize) -> Option<bool> {
        // Acquire lock, perform mutation, and get values for logging
        let (port, remaining, total) = {
            let mut dead = match self.dead_proxies.write() {
                Ok(value) => value,
                Err(_) => return None,
            };
            dead.insert(index);
            
            let port = self.proxy_ports.get(index).copied().unwrap_or(0);
            let remaining = self.proxy_ports.len() - dead.len();
            let total = self.proxy_ports.len();
            
            // Lock is automatically dropped here when the scope ends
            (port, remaining, total)
        };
        
        // Now we can await without holding the lock
        crate::util::logger::log_warn(&format!(
            "⚠ Marked proxy {} (port {}) as DEAD ({}/{} proxies remaining)",
            index,
            port,
            remaining,
            total
        )).await;

        Some(true)
    }
    
    // Get count of healthy proxies
    pub async fn num_healthy_proxies(&self) -> Option<usize> {
        let dead = match self.dead_proxies.read() {
            Ok(value) => value,
            Err(_) => return None,
        };
        Some(self.proxy_ports.len() - dead.len())
    }
}

pub async fn cleanup_all_proxy_containers() -> Result<()> {
    // Step 1: List all container IDs that match our pattern
    let output = Command::new("docker")
        .args(["ps", "-a", "--format", "{{.ID}} {{.Names}} {{.Image}}"])
        .output()
        .await?;

    let stdout = String::from_utf8_lossy(&output.stdout);

    let mut containers_to_kill = Vec::new();

    for line in stdout.lines() {
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() >= 2 {
            let name_or_id = parts[0];
            let name = parts[1];
            let image = if parts.len() >= 3 { parts[2] } else { "" };

            // Match by name prefix OR by image name
            if name.starts_with("vpn-proxy-") || image.contains("rust-vpn-proxy") {
                containers_to_kill.push(name_or_id.to_string());
            }
        }
    }

    if containers_to_kill.is_empty() {
        crate::util::logger::log_info("No old rust-vpn-proxy containers found").await;
        return Ok(());
    }

    // Step 2: Kill and remove them all at once
    let status = Command::new("docker")
        .arg("rm")
        .arg("-f")
        .args(&containers_to_kill)
        .status()
        .await?;

    if status.success() {
        crate::util::logger::log_info(&format!(
            "Successfully removed {} old rust-vpn-proxy container(s)",
            containers_to_kill.len()
        ))
        .await;
    } else {
        crate::util::logger::log_warn("Some containers may still remain (non-critical)").await;
    }

    Ok(())
}