ftr 0.7.0

//! Traceroute engine
//!
//! This implementation maximizes parallelism by:
//! 1. Starting ISP detection (STUN) immediately on launch
//! 2. Sending all probes in parallel
//! 3. Starting ASN/rDNS enrichment immediately as each response arrives
//! 4. Using caches to avoid duplicate lookups

use crate::enrichment::EnrichmentService;
use crate::probe::{ProbeInfo, ProbeResponse};
use crate::services::Services;
use crate::socket::traits::ProbeSocket;
use crate::socket::{ProbeProtocol, SocketMode};
use crate::traceroute::TracerouteError;
use crate::traceroute::{AsnInfo, ClassifiedHopInfo, SegmentType, TracerouteResult};
use std::collections::HashMap;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
use tokio::task::JoinSet;

#[cfg(test)]
#[path = "engine_test.rs"]
mod engine_test;
#[cfg(test)]
#[path = "sandwich_test.rs"]
mod sandwich_test;
use tokio::time::{sleep, timeout};

/// Enrichment result for an address
#[derive(Clone)]
struct EnrichmentResult {
    hostname: Option<String>,
    asn_info: Option<AsnInfo>,
}

/// Fully parallel async traceroute engine
pub struct TracerouteEngine {
    socket: Arc<Box<dyn ProbeSocket>>,
    config: crate::TracerouteConfig,
    target: IpAddr,
    enrichment_service: Arc<EnrichmentService>,
    enrichment_cache: Arc<Mutex<HashMap<IpAddr, EnrichmentResult>>>,
    services: Option<Arc<Services>>,
}

impl TracerouteEngine {
    /// Create a new fully parallel async traceroute engine with injected services
    pub async fn new_with_services(
        socket: Box<dyn ProbeSocket>,
        config: crate::TracerouteConfig,
        target: IpAddr,
        services: Arc<Services>,
    ) -> Result<Self, TracerouteError> {
        // Create enrichment service upfront
        let enrichment_service = Arc::new(
            EnrichmentService::new()
                .await
                .map_err(|e| TracerouteError::Other(e.to_string()))?,
        );

        Ok(Self {
            socket: Arc::new(socket),
            config,
            target,
            enrichment_service,
            enrichment_cache: Arc::new(Mutex::new(HashMap::new())),
            services: Some(services),
        })
    }

    /// Create a new fully parallel async traceroute engine (uses global caches)
    pub async fn new(
        socket: Box<dyn ProbeSocket>,
        config: crate::TracerouteConfig,
        target: IpAddr,
    ) -> Result<Self, TracerouteError> {
        // Create enrichment service upfront
        let enrichment_service = Arc::new(
            EnrichmentService::new()
                .await
                .map_err(|e| TracerouteError::Other(e.to_string()))?,
        );

        Ok(Self {
            socket: Arc::new(socket),
            config,
            target,
            enrichment_service,
            enrichment_cache: Arc::new(Mutex::new(HashMap::new())),
            services: None,
        })
    }

    /// Run the fully parallel async traceroute
    pub async fn run(&self) -> Result<TracerouteResult, TracerouteError> {
        let start_time = Instant::now();
        trace_time!(
            self.config.verbose,
            "Starting fully parallel async traceroute to {}",
            self.target
        );

        // 1. Start ISP detection immediately (if enabled)
        let isp_future = if self.config.enable_asn_lookup {
            if let Some(public_ip) = self.config.public_ip {
                // Use provided public IP
                trace_time!(
                    self.config.verbose,
                    "Using provided public IP: {}",
                    public_ip
                );
                let verbose = self.config.verbose;
                if let Some(ref services) = self.services {
                    // Use injected services
                    let services = services.clone();
                    Some(tokio::spawn(async move {
                        let isp_start = Instant::now();
                        let result = crate::public_ip::detect_isp_from_ip_with_services(
                            public_ip, &services,
                        )
                        .await;
                        trace_time!(
                            verbose,
                            "ISP detection from provided IP completed in {:?}",
                            isp_start.elapsed()
                        );
                        result
                    }))
                } else {
                    // Cannot detect ISP without services
                    None
                }
            } else {
                // Use STUN detection
                trace_time!(
                    self.config.verbose,
                    "Starting ISP detection (STUN) in parallel"
                );
                let verbose = self.config.verbose;
                if let Some(ref services) = self.services {
                    // Use injected services
                    let services = services.clone();
                    Some(tokio::spawn(async move {
                        let isp_start = Instant::now();
                        let result =
                            crate::public_ip::detect_isp_stun_with_services(&services).await;
                        trace_time!(
                            verbose,
                            "ISP detection completed in {:?}",
                            isp_start.elapsed()
                        );
                        result
                    }))
                } else {
                    // Cannot detect ISP without services
                    None
                }
            }
        } else {
            None
        };

        // Start destination ASN lookup early in parallel (if services and enrichment enabled)
        let dest_asn_future = if self.config.enable_asn_lookup {
            if let Some(ref services) = self.services {
                let services = services.clone();
                let target = self.target;
                let verbose = self.config.verbose;
                Some(tokio::spawn(async move {
                    trace_time!(verbose, "Starting destination ASN lookup for {}", target);
                    match services.asn.lookup(target).await {
                        Ok(info) if info.asn != 0 => Some(info),
                        _ => None,
                    }
                }))
            } else {
                None
            }
        } else {
            None
        };

        // 2. Create futures for all probes with immediate enrichment
        let mut probe_futures = JoinSet::new();
        let mut sequence = 1u16;

        trace_time!(
            self.config.verbose,
            "Creating probe futures with immediate enrichment"
        );
        for ttl in self.config.start_ttl..=self.config.max_hops {
            for _query in 0..self.config.queries_per_hop {
                let probe = ProbeInfo {
                    sequence,
                    ttl,
                    sent_at: Instant::now(),
                };
                sequence += 1;

                let socket = Arc::clone(&self.socket);
                let target = self.target;
                let verbose = self.config.verbose;
                let enrichment_service = Arc::clone(&self.enrichment_service);
                let enrichment_cache = Arc::clone(&self.enrichment_cache);
                let enable_asn = self.config.enable_asn_lookup;
                let enable_rdns = self.config.enable_rdns;

                // Create future that sends probe AND enriches response immediately
                let probe_future = async move {
                    trace_time!(
                        verbose,
                        "Sending probe seq={} ttl={}",
                        probe.sequence,
                        probe.ttl
                    );

                    match socket.send_probe_and_recv(target, probe).await {
                        Ok(response) => {
                            trace_time!(
                                verbose,
                                "Received response seq={} ttl={} from={} rtt={:?}",
                                response.sequence,
                                response.ttl,
                                response.from_addr,
                                response.rtt
                            );

                            // Start enrichment immediately if not cached
                            if enable_asn || enable_rdns {
                                let cache = enrichment_cache.lock().await;
                                if !cache.contains_key(&response.from_addr) {
                                    drop(cache); // Release lock before enrichment

                                    trace_time!(
                                        verbose,
                                        "Starting immediate enrichment for {}",
                                        response.from_addr
                                    );
                                    let enrich_start = Instant::now();

                                    // Enrich this single address
                                    let enrichment_results = enrichment_service
                                        .enrich_addresses(vec![response.from_addr])
                                        .await;

                                    if let Some(enrichment) =
                                        enrichment_results.get(&response.from_addr)
                                    {
                                        trace_time!(
                                            verbose,
                                            "Enrichment for {} completed in {:?}",
                                            response.from_addr,
                                            enrich_start.elapsed()
                                        );

                                        // Cache the result
                                        let mut cache = enrichment_cache.lock().await;
                                        cache.insert(
                                            response.from_addr,
                                            EnrichmentResult {
                                                hostname: enrichment.hostname.clone(),
                                                asn_info: enrichment.asn_info.clone(),
                                            },
                                        );
                                    }
                                }
                            }

                            Some(response)
                        }
                        Err(_) => None,
                    }
                };

                probe_futures.spawn(probe_future);
            }
        }

        trace_time!(
            self.config.verbose,
            "Started {} probe futures",
            sequence - 1
        );

        // 3. Collect responses as they arrive
        let mut responses: Vec<ProbeResponse> = Vec::new();
        let mut ttl_responses: HashMap<u8, usize> = HashMap::new();
        let mut destination_ttl: Option<u8> = None;

        let collection_start = Instant::now();
        trace_time!(self.config.verbose, "Starting response collection");

        let collection_future = async {
            while let Some(Ok(response)) = probe_futures.join_next().await {
                if let Some(resp) = response {
                    let ttl = resp.ttl;
                    let is_destination = resp.is_destination;

                    // Track responses per TTL
                    *ttl_responses.entry(ttl).or_insert(0) += 1;

                    // Update destination TTL if we found it
                    if is_destination && destination_ttl.is_none() {
                        destination_ttl = Some(ttl);
                    }

                    responses.push(resp);

                    // Check if we can exit early
                    if let Some(dest_ttl) = destination_ttl {
                        let mut can_exit = true;
                        for check_ttl in self.config.start_ttl..=dest_ttl {
                            if ttl_responses.get(&check_ttl).copied().unwrap_or(0) == 0 {
                                can_exit = false;
                                break;
                            }
                        }
                        if can_exit {
                            trace_time!(
                                self.config.verbose,
                                "Early exit - all TTLs up to destination responded"
                            );
                            // Wait briefly for late responses
                            sleep(Duration::from_millis(25)).await;

                            // Collect any remaining
                            while let Ok(Some(Ok(response))) =
                                timeout(Duration::from_millis(10), probe_futures.join_next()).await
                            {
                                if let Some(resp) = response {
                                    responses.push(resp);
                                }
                            }
                            break;
                        }
                    }
                }
            }
        };

        // Execute with overall timeout
        let _ = timeout(self.config.overall_timeout, collection_future).await;
        trace_time!(
            self.config.verbose,
            "Response collection completed in {:?}",
            collection_start.elapsed()
        );

        // 4. Build result with cached enrichment data
        let elapsed = start_time.elapsed();
        self.build_result(responses, elapsed, isp_future, dest_asn_future)
            .await
    }

    /// Build the final traceroute result with enriched data
    async fn build_result(
        &self,
        responses: Vec<ProbeResponse>,
        elapsed: Duration,
        isp_future: Option<
            tokio::task::JoinHandle<
                Result<crate::traceroute::IspInfo, crate::public_ip::PublicIpError>,
            >,
        >,
        dest_asn_future: Option<tokio::task::JoinHandle<Option<AsnInfo>>>,
    ) -> Result<TracerouteResult, TracerouteError> {
        let mut hops: HashMap<u8, Vec<ProbeResponse>> = HashMap::new();

        // Group responses by TTL
        for response in responses {
            hops.entry(response.ttl).or_default().push(response);
        }

        // Check if destination was reached
        let destination_reached = hops
            .values()
            .any(|ttl_responses| ttl_responses.iter().any(|r| r.is_destination));

        let destination_ttl = hops
            .values()
            .flat_map(|ttl_responses| ttl_responses.iter())
            .filter(|r| r.is_destination)
            .map(|r| r.ttl)
            .min()
            .unwrap_or(self.config.max_hops);

        // Wait for ISP detection to complete
        let isp_info = if let Some(future) = isp_future {
            trace_time!(self.config.verbose, "Waiting for ISP detection to complete");
            match future.await {
                Ok(Ok(isp)) => Some(isp),
                Ok(Err(e)) => {
                    trace_time!(self.config.verbose, "ISP detection failed: {}", e);
                    None
                }
                Err(e) => {
                    trace_time!(self.config.verbose, "ISP detection task failed: {}", e);
                    None
                }
            }
        } else {
            None
        };

        // Get ISP ASN for segment classification
        let isp_asn = isp_info.as_ref().map(|isp| isp.asn);

        // Wait for destination ASN
        let dest_asn = if let Some(fut) = dest_asn_future {
            fut.await.unwrap_or_default()
        } else {
            None
        };

        // Build classified hops with enrichment data
        let enrichment_cache = self.enrichment_cache.lock().await;
        let mut hop_infos: Vec<ClassifiedHopInfo> = Vec::new();
        let mut in_isp_segment = false;

        let display_max_ttl = if destination_reached {
            destination_ttl
        } else {
            self.config.max_hops
        };

        for ttl in self.config.start_ttl..=display_max_ttl {
            if let Some(ttl_responses) = hops.get(&ttl) {
                if !ttl_responses.is_empty() {
                    // Get unique addresses
                    let mut unique_addrs: Vec<IpAddr> = ttl_responses
                        .iter()
                        .filter(|r| !r.is_timeout)
                        .map(|r| r.from_addr)
                        .collect();
                    unique_addrs.sort();
                    unique_addrs.dedup();

                    for addr in unique_addrs {
                        let addr_responses: Vec<&ProbeResponse> = ttl_responses
                            .iter()
                            .filter(|r| r.from_addr == addr)
                            .collect();

                        // Calculate average RTT
                        let rtt = if !addr_responses.is_empty() {
                            let total_rtt: Duration = addr_responses.iter().map(|r| r.rtt).sum();
                            Some(total_rtt / addr_responses.len() as u32)
                        } else {
                            None
                        };

                        // Get cached enrichment data
                        let enrichment = enrichment_cache.get(&addr);
                        let hostname = enrichment.and_then(|e| e.hostname.clone());
                        let asn_info = enrichment.and_then(|e| e.asn_info.clone());

                        // Classify segment
                        let segment = if let IpAddr::V4(ipv4) = addr {
                            if crate::traceroute::is_internal_ip(&ipv4) {
                                SegmentType::Lan
                            } else if crate::traceroute::is_cgnat(&ipv4) {
                                in_isp_segment = true;
                                SegmentType::Isp
                            } else {
                                // Public IP classification requires ISP and/or destination ASN
                                if let Some(ref asn) = asn_info {
                                    // ISP boundary check if known
                                    if let Some(isp) = isp_asn {
                                        if asn.asn == isp {
                                            in_isp_segment = true;
                                            SegmentType::Isp
                                        } else if let Some(ref dest) = dest_asn {
                                            if asn.asn == dest.asn {
                                                SegmentType::Destination
                                            } else {
                                                SegmentType::Transit
                                            }
                                        } else {
                                            // ISP known but not this AS; without dest ASN, mark as TRANSIT
                                            SegmentType::Transit
                                        }
                                    } else if let Some(ref dest) = dest_asn {
                                        if asn.asn == dest.asn {
                                            SegmentType::Destination
                                        } else if in_isp_segment {
                                            SegmentType::Transit
                                        } else {
                                            SegmentType::Unknown
                                        }
                                    } else if in_isp_segment {
                                        SegmentType::Transit
                                    } else {
                                        SegmentType::Unknown
                                    }
                                } else if in_isp_segment {
                                    // Public IP after ISP without ASN = likely IXP/peering point
                                    SegmentType::Transit
                                } else {
                                    // Public IP without ASN before ISP boundary = assume transit
                                    SegmentType::Transit
                                }
                            }
                        } else {
                            SegmentType::Unknown
                        };

                        hop_infos.push(ClassifiedHopInfo {
                            ttl,
                            segment,
                            hostname,
                            addr: Some(addr),
                            asn_info,
                            rtt,
                        });
                    }

                    // Add timeout responses
                    let timeout_count = ttl_responses.iter().filter(|r| r.is_timeout).count();
                    for _ in 0..timeout_count {
                        hop_infos.push(ClassifiedHopInfo {
                            ttl,
                            segment: SegmentType::Unknown,
                            hostname: None,
                            addr: None,
                            asn_info: None,
                            rtt: None,
                        });
                    }
                }
            } else {
                // Add blank hop
                hop_infos.push(ClassifiedHopInfo {
                    ttl,
                    segment: SegmentType::Unknown,
                    hostname: None,
                    addr: None,
                    asn_info: None,
                    rtt: None,
                });
            }
        }

        // Sort by TTL
        hop_infos.sort_by_key(|h| h.ttl);

        // Apply sandwich logic: fill Unknown/Transit hops between same-type segments
        Self::apply_sandwich_logic(&mut hop_infos);

        // Determine protocol and socket mode
        let (protocol_used, socket_mode_used) = match self.socket.mode() {
            crate::socket::traits::ProbeMode::DgramIcmp => (ProbeProtocol::Icmp, SocketMode::Dgram),
            crate::socket::traits::ProbeMode::WindowsIcmp => (ProbeProtocol::Icmp, SocketMode::Raw),
            crate::socket::traits::ProbeMode::UdpWithRecverr => {
                (ProbeProtocol::Udp, SocketMode::Dgram)
            }
            crate::socket::traits::ProbeMode::RawIcmp => (ProbeProtocol::Icmp, SocketMode::Raw),
        };

        trace_time!(
            self.config.verbose,
            "Total fully parallel async traceroute completed in {:?}",
            elapsed
        );

        Ok(TracerouteResult {
            target: self.config.target.clone(),
            target_ip: self.target,
            hops: hop_infos,
            isp_info,
            destination_asn: dest_asn,
            protocol_used,
            socket_mode_used,
            destination_reached,
            total_duration: elapsed,
        })
    }

    /// Apply sandwich logic: if Unknown/Transit hops are between two hops of the same segment type,
    /// convert them to that segment type. This handles cases where probes fail or lack ASN data
    /// but are clearly part of the same network segment.
    pub(crate) fn apply_sandwich_logic(hops: &mut [ClassifiedHopInfo]) {
        if hops.len() < 3 {
            return; // Need at least 3 hops for sandwiching
        }

        for i in 1..hops.len() - 1 {
            // Only modify Unknown or Transit segments that have addresses
            if hops[i].addr.is_some()
                && (hops[i].segment == SegmentType::Unknown
                    || hops[i].segment == SegmentType::Transit)
            {
                // Look for surrounding hops of the same type
                let mut check_segment = |segment_type: SegmentType| {
                    // Find previous hop with this segment type
                    let has_before = (0..i).rev().any(|j| hops[j].segment == segment_type);

                    // Find next hop with this segment type
                    let has_after = (i + 1..hops.len()).any(|j| hops[j].segment == segment_type);

                    if has_before && has_after {
                        hops[i].segment = segment_type;
                        true
                    } else {
                        false
                    }
                };

                // Check ISP first, then Destination
                // ISP takes precedence as it's typically earlier in the path
                if check_segment(SegmentType::Isp) {
                    continue;
                }
                check_segment(SegmentType::Destination);
            }
        }
    }
}