use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use std::collections::HashSet;
use std::net::SocketAddrV4;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::io::AsyncWriteExt;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use futures::stream::{self, StreamExt};

lazy_static::lazy_static! {
    static ref HTTPS_PORTS: HashSet<u16> = [443, 8443, 9443, 8883].iter().copied().collect();
    static ref PORTS: Vec<u16> = vec![80, 443, 8080, 8000, 8443, 8888, 9000, 9090, 8081];
}

#[derive(Serialize, Deserialize, Debug)]
#[serde(tag = "type")]
enum OutputMsg {
    Progress { scanned: usize, total: usize },
    Found {
        ip: String,
        port: u16,
        url: String,
        auth: String,
        status: String,
    },
    Error {
        message: String,
    },
}

/// Load credentials from combos.txt file
/// Format: user:password (one per line)
/// Returns Arc<Vec<(String, String)>> for efficient sharing across tasks
fn load_credentials(path: &str) -> Arc<Vec<(String, String)>> {
    match std::fs::read_to_string(path) {
        Ok(content) => {
            let mut combos = Vec::new();
            
            for (line_num, line) in content.lines().enumerate() {
                let trimmed = line.trim();
                
                // Skip empty lines and comments
                if trimmed.is_empty() || trimmed.starts_with('#') {
                    continue;
                }
                
                // Parse user:password format
                if let Some((user, pass)) = trimmed.split_once(':') {
                    combos.push((user.to_string(), pass.to_string()));
                } else {
                    eprintln!("Warning: Invalid format on line {}: '{}'", line_num + 1, trimmed);
                }
            }
            
            if combos.is_empty() {
                eprintln!("Warning: No valid credentials found in {}", path);
                // Fallback to default credentials
                Arc::new(vec![
                    ("admin".into(), "admin".into()),
                    ("admin".into(), "12345".into()),
                ])
            } else {
                eprintln!("Loaded {} credential combinations from {}", combos.len(), path);
                Arc::new(combos)
            }
        }
        Err(e) => {
            eprintln!("Error reading {}: {}", path, e);
            eprintln!("Using default credentials");
            Arc::new(vec![
                ("admin".into(), "admin".into()),
                ("admin".into(), "12345".into()),
                ("admin".into(), "password".into()),
                ("root".into(), "root".into()),
                ("admin".into(), "".into()),
            ])
        }
    }
}

#[tokio::main]
async fn main() {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("Usage: {} <CIDR> [combos_file]", args[0]);
        eprintln!("Example: {} 192.168.1.0/24 combos.txt", args[0]);
        std::process::exit(1);
    }

    let cidr = &args[1];
    let combos_file = args.get(2).map(|s| s.as_str()).unwrap_or("combos.txt");
    
    let net: ipnet::Ipv4Net = match cidr.parse() {
        Ok(n) => n,
        Err(e) => {
            eprintln!("Invalid CIDR '{}': {}", cidr, e);
            std::process::exit(1);
        }
    };

    // Pre-warm DNS resolver
    tokio::task::spawn(async {
        let _ = reqwest::get("https://1.1.1.1").await;
    });

    // Load credentials from file (optimized with Arc for zero-copy sharing)
    let combos = load_credentials(combos_file);

    // Optimized HTTP client with connection pooling
    let client = Arc::new(
        Client::builder()
            .pool_max_idle_per_host(100)
            .pool_idle_timeout(Duration::from_secs(90))
            .tcp_keepalive(Duration::from_secs(60))
            .http2_prior_knowledge()
            .timeout(Duration::from_secs(3))
            .danger_accept_invalid_certs(true)
            .redirect(reqwest::redirect::Policy::none())
            .build()
            .unwrap(),
    );

    let (tx, mut rx) = tokio::sync::mpsc::channel::<OutputMsg>(10000);

    // Async output writer
    let writer_task = tokio::spawn(async move {
        let mut out = tokio::io::stdout();
        while let Some(msg) = rx.recv().await {
            if let Ok(json) = serde_json::to_string(&msg) {
                let _ = out.write_all(json.as_bytes()).await;
                let _ = out.write_all(b"\n").await;
            }
        }
        let _ = out.flush().await;
    });

    // Pre-allocate targets vector
    let total_hosts = net.hosts().count();
    let estimated_size = total_hosts * PORTS.len();
    let mut targets = Vec::with_capacity(estimated_size);
    
    targets.extend(
        net.hosts()
            .flat_map(|ip| PORTS.iter().map(move |&p| (ip, p)))
    );
    
    let total_targets = targets.len();
    let scanned = Arc::new(AtomicUsize::new(0));
    let last_update = Arc::new(Mutex::new(Instant::now()));

    // Dynamic buffer sizing
    let buffer_size = std::cmp::min(10000, std::cmp::max(2000, total_targets / 10));

    stream::iter(targets)
        .map(|(ip, port)| {
            let client = Arc::clone(&client);
            let combos = Arc::clone(&combos);
            let tx = tx.clone();
            let scanned = Arc::clone(&scanned);
            let last_update = Arc::clone(&last_update);

            async move {
                let addr = std::net::SocketAddr::V4(SocketAddrV4::new(ip, port));

                // Phase 1: Lightning Raw TCP Knock (100ms)
                if tokio::time::timeout(Duration::from_millis(100), TcpStream::connect(&addr))
                    .await
                    .is_ok()
                {
                    let scheme = if HTTPS_PORTS.contains(&port) {
                        "https"
                    } else {
                        "http"
                    };
                    let base_url = format!("{}://{}:{}", scheme, ip, port);

                    // Phase 2: Rapid HEAD request
                    if let Ok(resp) = client.head(&base_url).send().await {
                        let st = resp.status();
                        if st == StatusCode::UNAUTHORIZED || st == StatusCode::OK {
                            let mut v_auth = "None".to_string();
                            let mut is_v = st == StatusCode::OK;

                            // Parallelize Auth attempts for maximum speed
                            if !is_v && !combos.is_empty() {
                                // Batch auth attempts in chunks to avoid overwhelming the target
                                // Optimal chunk size balances parallelism vs target load
                                let chunk_size = std::cmp::min(10, combos.len());
                                
                                for chunk in combos.chunks(chunk_size) {
                                    let auth_futures: Vec<_> = chunk
                                        .iter()
                                        .map(|(u, p)| {
                                            let client = Arc::clone(&client);
                                            let url = base_url.clone();
                                            let u = u.clone();
                                            let p = p.clone();
                                            async move {
                                                client
                                                    .head(&url)
                                                    .basic_auth(&u, Some(&p))
                                                    .send()
                                                    .await
                                                    .ok()
                                                    .filter(|r| r.status() == StatusCode::OK)
                                                    .map(|_| format!("{}:{}", u, p))
                                            }
                                        })
                                        .collect();

                                    let results = futures::future::join_all(auth_futures).await;
                                    if let Some(auth) = results.into_iter().find_map(|x| x) {
                                        v_auth = auth;
                                        is_v = true;
                                        break; // Found valid creds, stop checking
                                    }
                                }
                            }

                            let _ = tx
                                .send(OutputMsg::Found {
                                    ip: ip.to_string(),
                                    port,
                                    url: base_url,
                                    auth: v_auth,
                                    status: if is_v {
                                        "Verified".into()
                                    } else {
                                        "Locked".into()
                                    },
                                })
                                .await;
                        }
                    }
                }

                // Time-based progress updates (200ms throttle for smooth UI)
                let v = scanned.fetch_add(1, Ordering::Relaxed) + 1;
                let mut last = last_update.lock().await;
                if last.elapsed() > Duration::from_millis(200) || v == total_targets {
                    let _ = tx
                        .send(OutputMsg::Progress {
                            scanned: v,
                            total: total_targets,
                        })
                        .await;
                    *last = Instant::now();
                }
            }
        })
        .buffer_unordered(buffer_size)
        .for_each(|_| async {})
        .await;

    drop(tx);
    let _ = writer_task.await;
}