#[cfg(unix)]
use pnet::packet::icmp::{IcmpPacket, IcmpTypes};
#[cfg(unix)]
use pnet::packet::icmpv6::{Icmpv6Packet, Icmpv6Types};
#[cfg(unix)]
use pnet::packet::ip::IpNextHeaderProtocols;
#[cfg(unix)]
use pnet::transport::{
    TransportChannelType, TransportProtocol, icmp_packet_iter, icmpv6_packet_iter,
    transport_channel,
};
use serde::{Deserialize, Serialize};
use socket2::{Domain, Protocol, Socket, Type};
use std::net::{IpAddr, SocketAddr};
use std::time::{Duration, Instant};
use thiserror::Error;
use tokio::net::{TcpStream, lookup_host};
use tokio::time::timeout;
use wtfnet_geoip::GeoIpRecord;

#[derive(Debug, Error)]
pub enum ProbeError {
    #[error("resolution failed: {0}")]
    Resolve(String),
    #[error("io error: {0}")]
    Io(String),
    #[error("timeout")]
    Timeout,
    #[error("ping error: {0}")]
    Ping(String),
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PingResult {
    pub seq: u16,
    pub rtt_ms: Option<u128>,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PingSummary {
    pub sent: u32,
    pub received: u32,
    pub loss_pct: f64,
    pub min_ms: Option<u128>,
    pub avg_ms: Option<f64>,
    pub max_ms: Option<u128>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct PingReport {
    pub target: String,
    pub ip: Option<String>,
    pub geoip: Option<GeoIpRecord>,
    pub timeout_ms: u64,
    pub count: u32,
    pub results: Vec<PingResult>,
    pub summary: PingSummary,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TcpPingResult {
    pub seq: u16,
    pub rtt_ms: Option<u128>,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TcpPingReport {
    pub target: String,
    pub ip: Option<String>,
    pub geoip: Option<GeoIpRecord>,
    pub port: u16,
    pub timeout_ms: u64,
    pub count: u32,
    pub results: Vec<TcpPingResult>,
    pub summary: PingSummary,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceHop {
    pub ttl: u8,
    pub addr: Option<String>,
    pub rtt_ms: Option<u128>,
    pub note: Option<String>,
    pub geoip: Option<GeoIpRecord>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TraceReport {
    pub target: String,
    pub ip: Option<String>,
    pub geoip: Option<GeoIpRecord>,
    pub port: u16,
    pub max_hops: u8,
    pub timeout_ms: u64,
    pub protocol: String,
    pub hops: Vec<TraceHop>,
}

pub async fn ping(
    target: &str,
    count: u32,
    timeout_ms: u64,
    interval_ms: u64,
) -> Result<PingReport, ProbeError> {
    let addr = resolve_one(target).await?;
    let mut results = Vec::new();
    let mut received = 0u32;
    let mut min = None;
    let mut max = None;
    let mut sum = 0u128;

    let config = match addr {
        IpAddr::V4(_) => surge_ping::Config::default(),
        IpAddr::V6(_) => surge_ping::Config::builder()
            .kind(surge_ping::ICMP::V6)
            .build(),
    };
    let client =
        surge_ping::Client::new(&config).map_err(|err| ProbeError::Ping(err.to_string()))?;
    let mut pinger = client.pinger(addr, surge_ping::PingIdentifier(0)).await;
    let timeout_dur = Duration::from_millis(timeout_ms);

    for seq in 0..count {
        let seq = seq as u16;
        let start = Instant::now();
        let response = timeout(
            timeout_dur,
            pinger.ping(surge_ping::PingSequence(seq), &[0; 8]),
        )
        .await;
        match response {
            Ok(Ok((_packet, _))) => {
                let rtt = start.elapsed().as_millis();
                received += 1;
                min = Some(min.map_or(rtt, |value: u128| value.min(rtt)));
                max = Some(max.map_or(rtt, |value: u128| value.max(rtt)));
                sum += rtt;
                results.push(PingResult {
                    seq,
                    rtt_ms: Some(rtt),
                    error: None,
                });
            }
            Ok(Err(err)) => {
                results.push(PingResult {
                    seq,
                    rtt_ms: None,
                    error: Some(err.to_string()),
                });
            }
            Err(_) => {
                results.push(PingResult {
                    seq,
                    rtt_ms: None,
                    error: Some("timeout".to_string()),
                });
            }
        }
        if interval_ms > 0 {
            tokio::time::sleep(Duration::from_millis(interval_ms)).await;
        }
    }

    let summary = build_summary(count, received, min, max, sum);
    Ok(PingReport {
        target: target.to_string(),
        ip: Some(addr.to_string()),
        geoip: None,
        timeout_ms,
        count,
        results,
        summary,
    })
}

pub async fn tcp_ping(
    target: &str,
    port: u16,
    count: u32,
    timeout_ms: u64,
) -> Result<TcpPingReport, ProbeError> {
    let addr = resolve_one(target).await?;
    let socket_addr = SocketAddr::new(addr, port);
    let timeout_dur = Duration::from_millis(timeout_ms);
    let mut results = Vec::new();
    let mut received = 0u32;
    let mut min = None;
    let mut max = None;
    let mut sum = 0u128;

    for seq in 0..count {
        let seq = seq as u16;
        let start = Instant::now();
        let attempt = timeout(timeout_dur, TcpStream::connect(socket_addr)).await;
        match attempt {
            Ok(Ok(_stream)) => {
                let rtt = start.elapsed().as_millis();
                received += 1;
                min = Some(min.map_or(rtt, |value: u128| value.min(rtt)));
                max = Some(max.map_or(rtt, |value: u128| value.max(rtt)));
                sum += rtt;
                results.push(TcpPingResult {
                    seq,
                    rtt_ms: Some(rtt),
                    error: None,
                });
            }
            Ok(Err(err)) => {
                results.push(TcpPingResult {
                    seq,
                    rtt_ms: None,
                    error: Some(err.to_string()),
                });
            }
            Err(_) => {
                results.push(TcpPingResult {
                    seq,
                    rtt_ms: None,
                    error: Some("timeout".to_string()),
                });
            }
        }
    }

    let summary = build_summary(count, received, min, max, sum);
    Ok(TcpPingReport {
        target: target.to_string(),
        ip: Some(addr.to_string()),
        geoip: None,
        port,
        timeout_ms,
        count,
        results,
        summary,
    })
}

pub async fn tcp_trace(
    target: &str,
    port: u16,
    max_hops: u8,
    timeout_ms: u64,
) -> Result<TraceReport, ProbeError> {
    let addr = resolve_one(target).await?;
    let socket_addr = SocketAddr::new(addr, port);
    let timeout_dur = Duration::from_millis(timeout_ms);
    let mut hops = Vec::new();

    for ttl in 1..=max_hops {
        let addr = socket_addr;
        let start = Instant::now();
        let result =
            tokio::task::spawn_blocking(move || tcp_connect_with_ttl(addr, ttl, timeout_dur))
                .await
                .map_err(|err| ProbeError::Io(err.to_string()))?;

        match result {
            Ok(()) => {
                let rtt = start.elapsed().as_millis();
                hops.push(TraceHop {
                    ttl,
                    addr: Some(socket_addr.ip().to_string()),
                    rtt_ms: Some(rtt),
                    note: None,
                    geoip: None,
                });
                break;
            }
            Err(err) => {
                let rtt = start.elapsed().as_millis();
                hops.push(TraceHop {
                    ttl,
                    addr: None,
                    rtt_ms: Some(rtt),
                    note: Some(err.to_string()),
                    geoip: None,
                });
            }
        }
    }

    Ok(TraceReport {
        target: target.to_string(),
        ip: Some(addr.to_string()),
        geoip: None,
        port,
        max_hops,
        timeout_ms,
        protocol: "tcp".to_string(),
        hops,
    })
}

pub async fn udp_trace(
    target: &str,
    port: u16,
    max_hops: u8,
    timeout_ms: u64,
) -> Result<TraceReport, ProbeError> {
    let addr = resolve_one(target).await?;

    let timeout_dur = Duration::from_millis(timeout_ms);
    let mut hops = Vec::new();

    for ttl in 1..=max_hops {
        let addr = SocketAddr::new(addr, port);
        let start = Instant::now();
        let result = tokio::task::spawn_blocking(move || udp_trace_hop(addr, ttl, timeout_dur))
            .await
            .map_err(|err| ProbeError::Io(err.to_string()))?;

        match result {
            Ok((hop_addr, reached)) => {
                let rtt = start.elapsed().as_millis();
                hops.push(TraceHop {
                    ttl,
                    addr: hop_addr.map(|ip| ip.to_string()),
                    rtt_ms: Some(rtt),
                    note: None,
                    geoip: None,
                });
                if reached {
                    break;
                }
            }
            Err(err) => {
                hops.push(TraceHop {
                    ttl,
                    addr: None,
                    rtt_ms: None,
                    note: Some(err.to_string()),
                    geoip: None,
                });
            }
        }
    }

    Ok(TraceReport {
        target: target.to_string(),
        ip: Some(addr.to_string()),
        geoip: None,
        port,
        max_hops,
        timeout_ms,
        protocol: "udp".to_string(),
        hops,
    })
}

fn build_summary(
    sent: u32,
    received: u32,
    min: Option<u128>,
    max: Option<u128>,
    sum: u128,
) -> PingSummary {
    let loss_pct = if sent == 0 {
        0.0
    } else {
        ((sent - received) as f64 / sent as f64) * 100.0
    };
    let avg_ms = if received == 0 {
        None
    } else {
        Some(sum as f64 / received as f64)
    };
    PingSummary {
        sent,
        received,
        loss_pct,
        min_ms: min,
        avg_ms,
        max_ms: max,
    }
}

async fn resolve_one(target: &str) -> Result<IpAddr, ProbeError> {
    let mut iter = lookup_host((target, 0))
        .await
        .map_err(|err| ProbeError::Resolve(err.to_string()))?;
    iter.next()
        .map(|addr| addr.ip())
        .ok_or_else(|| ProbeError::Resolve("no address found".to_string()))
}

fn tcp_connect_with_ttl(addr: SocketAddr, ttl: u8, timeout: Duration) -> Result<(), ProbeError> {
    let domain = match addr.ip() {
        IpAddr::V4(_) => Domain::IPV4,
        IpAddr::V6(_) => Domain::IPV6,
    };
    let socket = Socket::new(domain, Type::STREAM, Some(Protocol::TCP))
        .map_err(|err| ProbeError::Io(err.to_string()))?;
    match addr.ip() {
        IpAddr::V4(_) => socket
            .set_ttl_v4(u32::from(ttl))
            .map_err(|err| ProbeError::Io(err.to_string()))?,
        IpAddr::V6(_) => socket
            .set_unicast_hops_v6(u32::from(ttl))
            .map_err(|err| ProbeError::Io(err.to_string()))?,
    }
    socket
        .connect_timeout(&addr.into(), timeout)
        .map_err(|err| ProbeError::Io(err.to_string()))?;
    Ok(())
}

#[cfg(unix)]
fn udp_trace_hop(
    addr: SocketAddr,
    ttl: u8,
    timeout: Duration,
) -> Result<(Option<IpAddr>, bool), ProbeError> {
    match addr.ip() {
        IpAddr::V4(_) => udp_trace_hop_v4(addr, ttl, timeout),
        IpAddr::V6(_) => udp_trace_hop_v6(addr, ttl, timeout),
    }
}

#[cfg(not(unix))]
fn udp_trace_hop(
    _addr: SocketAddr,
    _ttl: u8,
    _timeout: Duration,
) -> Result<(Option<IpAddr>, bool), ProbeError> {
    Err(ProbeError::Io(
        "udp trace not supported on this platform".to_string(),
    ))
}

#[cfg(unix)]
fn udp_trace_hop_v4(
    addr: SocketAddr,
    ttl: u8,
    timeout: Duration,
) -> Result<(Option<IpAddr>, bool), ProbeError> {
    let protocol =
        TransportChannelType::Layer4(TransportProtocol::Ipv4(IpNextHeaderProtocols::Icmp));
    let (_tx, mut rx) =
        transport_channel(4096, protocol).map_err(|err| ProbeError::Io(err.to_string()))?;

    let socket =
        std::net::UdpSocket::bind("0.0.0.0:0").map_err(|err| ProbeError::Io(err.to_string()))?;
    socket
        .set_ttl(u32::from(ttl))
        .map_err(|err| ProbeError::Io(err.to_string()))?;
    let _ = socket.send_to(&[0u8; 4], addr);

    let mut iter = icmp_packet_iter(&mut rx);
    match iter.next_with_timeout(timeout) {
        Ok(Some((packet, addr))) => {
            if let Some(result) = interpret_icmp_v4(&packet) {
                return Ok((Some(addr), result));
            }
            Ok((Some(addr), false))
        }
        Ok(None) => Err(ProbeError::Timeout),
        Err(err) => Err(ProbeError::Io(err.to_string())),
    }
}

#[cfg(unix)]
fn udp_trace_hop_v6(
    addr: SocketAddr,
    ttl: u8,
    timeout: Duration,
) -> Result<(Option<IpAddr>, bool), ProbeError> {
    let protocol =
        TransportChannelType::Layer4(TransportProtocol::Ipv6(IpNextHeaderProtocols::Icmpv6));
    let (_tx, mut rx) =
        transport_channel(4096, protocol).map_err(|err| ProbeError::Io(err.to_string()))?;

    let socket =
        std::net::UdpSocket::bind("[::]:0").map_err(|err| ProbeError::Io(err.to_string()))?;
    socket
        .set_multicast_hops_v6(u32::from(ttl))
        .map_err(|err| ProbeError::Io(err.to_string()))?;
    let _ = socket.send_to(&[0u8; 4], addr);

    let mut iter = icmpv6_packet_iter(&mut rx);
    match iter.next_with_timeout(timeout) {
        Ok(Some((packet, addr))) => {
            if let Some(result) = interpret_icmp_v6(&packet) {
                return Ok((Some(addr), result));
            }
            Ok((Some(addr), false))
        }
        Ok(None) => Err(ProbeError::Timeout),
        Err(err) => Err(ProbeError::Io(err.to_string())),
    }
}

#[cfg(unix)]
fn interpret_icmp_v4(packet: &IcmpPacket) -> Option<bool> {
    let icmp_type = packet.get_icmp_type();
    if icmp_type == IcmpTypes::TimeExceeded {
        return Some(false);
    }
    if icmp_type == IcmpTypes::DestinationUnreachable {
        return Some(true);
    }
    None
}

#[cfg(unix)]
fn interpret_icmp_v6(packet: &Icmpv6Packet) -> Option<bool> {
    let icmp_type = packet.get_icmpv6_type();
    if icmp_type == Icmpv6Types::TimeExceeded {
        return Some(false);
    }
    if icmp_type == Icmpv6Types::DestinationUnreachable {
        return Some(true);
    }
    None
}