use std::net::{IpAddr, SocketAddr};
use std::time::Duration;

use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpStream;
use tokio::task::JoinSet;
use tokio::time::timeout;

use crate::models::host::{PortResult, PortState, Protocol};

fn service_name(port: u16) -> Option<String> {
    let name = match port {
        21 => "ftp",
        22 => "ssh",
        23 => "telnet",
        25 => "smtp",
        53 => "dns",
        80 => "http",
        110 => "pop3",
        111 => "rpcbind",
        135 => "msrpc",
        139 => "netbios-ssn",
        143 => "imap",
        443 => "https",
        445 => "microsoft-ds",
        993 => "imaps",
        995 => "pop3s",
        1433 => "mssql",
        1521 => "oracle",
        3306 => "mysql",
        3389 => "rdp",
        5432 => "postgresql",
        5900 => "vnc",
        5985 => "winrm",
        5986 => "winrm-s",
        8080 => "http-alt",
        8088 => "http-alt",
        8443 => "https-alt",
        _ => return None,
    };
    Some(name.to_string())
}

pub const DEFAULT_PORTS: &[u16] = &[
    21, 22, 23, 25, 53, 80, 110, 111, 135, 139, 143,
    443, 445, 993, 995, 1433, 1521, 3306, 3389,
    5432, 5900, 5985, 5986, 8080, 8088, 8443,
];

async fn scan_port(ip: IpAddr, port: u16, timeout_ms: u64) -> PortResult {
    let addr = SocketAddr::new(ip, port);
    let dur = Duration::from_millis(timeout_ms);

    let stream = match timeout(dur, TcpStream::connect(addr)).await {
        Ok(Ok(stream)) => stream,
        Ok(Err(_)) => {
            return PortResult {
                port,
                protocol: Protocol::Tcp,
                state: PortState::Closed,
                service: service_name(port),
                banner: None,
            };
        }
        Err(_) => {
            return PortResult {
                port,
                protocol: Protocol::Tcp,
                state: PortState::Filtered,
                service: service_name(port),
                banner: None,
            };
        }
    };

    // Port is open, try to grab a banner
    let banner = grab_banner(stream, port).await;

    PortResult {
        port,
        protocol: Protocol::Tcp,
        state: PortState::Open,
        service: service_name(port),
        banner,
    }
}

async fn grab_banner(mut stream: TcpStream, port: u16) -> Option<String> {
    let banner_timeout = Duration::from_millis(800);

    // Some services need us to send something first
    let probe = match port {
        80 | 8080 | 8088 | 8443 => Some(b"HEAD / HTTP/1.0\r\nHost: localhost\r\n\r\n".as_slice()),
        443 => None, // TLS, can't grab without handshake
        _ => None,
    };

    if let Some(data) = probe {
        if timeout(banner_timeout, stream.write_all(data)).await.is_err() {
            return None;
        }
    }

    // Try to read whatever the service sends back
    let mut buf = vec![0u8; 1024];
    match timeout(banner_timeout, stream.read(&mut buf)).await {
        Ok(Ok(n)) if n > 0 => {
            buf.truncate(n);
            // Strip to printable ASCII, replace control chars
            let banner: String = buf.iter()
                .map(|&b| if b >= 0x20 && b < 0x7f { b as char } else { ' ' })
                .collect();
            let banner = banner.trim().to_string();
            if banner.is_empty() {
                None
            } else {
                // Cap length so it doesn't blow up the UI
                Some(banner.chars().take(200).collect())
            }
        }
        _ => None,
    }
}

const MAX_PORT_CONCURRENT: usize = 64;

pub async fn scan_host_ports(
    ip: IpAddr,
    ports: &[u16],
    timeout_ms: u64,
) -> Vec<PortResult> {
    let mut results = Vec::with_capacity(ports.len());

    for chunk in ports.chunks(MAX_PORT_CONCURRENT) {
        let mut join_set = JoinSet::new();

        for &port in chunk {
            join_set.spawn(scan_port(ip, port, timeout_ms));
        }

        while let Some(result) = join_set.join_next().await {
            match result {
                Ok(port_result) => results.push(port_result),
                Err(e) => tracing::warn!("Port scan task panicked: {}", e),
            }
        }
    }

    results.sort_by_key(|p| p.port);
    results
}