use spvirit_codec::spvirit_state::PvaStateStats;
use dns_lookup::lookup_addr;
use std::collections::HashMap;

use gethostname::gethostname;

//serde_json::json;

#[derive(Clone)]
pub struct PVEvent {
    pub system_time: std::time::SystemTime, // Real-world time
    pub timestamp: std::time::Instant,      // Monotonic clock
    pub source_ip: String,
    pub destination_ip: String,
    pub request_type: String,
    pub context: String,
    pub resolved_source_ip: String,
    pub resolved_destination_ip: String,
    pub host_name: String,
    // Structured fields for normalized Loki output
    pub pv_name: Option<String>,
    pub ioid: Option<u32>,
    pub sid: Option<u32>,
    pub cid: Option<u32>,
    pub value_text: Option<String>,
    pub mid_stream: bool,
}
impl PVEvent {
    pub fn new(source_ip: String, destination_ip: String, request_type: String) -> Self {
        Self {
            system_time: std::time::SystemTime::now(),
            timestamp: std::time::Instant::now(),
            source_ip,
            destination_ip,
            request_type,
            context: String::new(),
            resolved_source_ip: "Unknown".to_string(),
            resolved_destination_ip: "Unknown".to_string(),
            host_name: gethostname().to_string_lossy().to_string(),
            pv_name: None,
            ioid: None,
            sid: None,
            cid: None,
            value_text: None,
            mid_stream: false,
        }
    }

    pub fn new_resolved(
        source_ip: String,
        destination_ip: String,
        request_type: String,
        resolved_source_ip: String,
        resolved_destination_ip: String,
        context: String,
    ) -> Self {
        Self {
            system_time: std::time::SystemTime::now(),
            timestamp: std::time::Instant::now(),
            source_ip,
            destination_ip,
            request_type,
            context,
            resolved_source_ip,
            resolved_destination_ip,
            host_name: gethostname().to_string_lossy().to_string(),
            pv_name: None,
            ioid: None,
            sid: None,
            cid: None,
            value_text: None,
            mid_stream: false,
        }
    }

    pub fn print(&self) {
        println!(
            "Source IP: {}, Destination IP: {}, Request type: {}",
            self.source_ip, self.destination_ip, self.request_type
        );
    }

    pub fn to_string_with_host(&self) -> String {
        self.internal_to_string(true)
    }

    pub fn to_string(&self) -> String {
        self.internal_to_string(false)
    }
    fn internal_to_string(&self, show_host: bool) -> String {
        // let show_host = show_host;

        let req_type = self.request_type.clone();

        // we need to pad all ips to be 15 characters long
        //  right align
        let source_ip_fn = || {
            if self.resolved_source_ip == "Unknown" {
                self.source_ip.clone().to_string()
            } else {
                self.resolved_source_ip.clone().to_string()
            }
        };
        // left align
        // let destination_ip = format!("{: >15}", self.destination_ip);
        let destination_ip_fn = || {
            if self.resolved_destination_ip == "Unknown" {
                self.destination_ip.clone().to_string()
            } else {
                // println!("Resolved IP: {} -> {}", self.destination_ip, self.resolved_destination_ip);
                self.resolved_destination_ip.clone().to_string()
            }
        };
        // let source_ip = format!("{: <15}", source_ip_fn());
        // let destination_ip = format!("{: <15}", destination_ip_fn());

        // need limit the names to 15 characters
        let source_ip = || {
            let mut source_ip = source_ip_fn();
            if source_ip.len() > 15 {
                source_ip.truncate(15);
            }
            source_ip
        };
        let destination_ip = || {
            let mut destination_ip = destination_ip_fn();
            if destination_ip.len() > 15 {
                destination_ip.truncate(15);
            }
            destination_ip
        };

        let source_ip = format!("{: <15}", source_ip());
        let destination_ip = format!("{: <15}", destination_ip());

        let human_readable_time = || {
            // padded to 15 characters
            let duration = self.timestamp.elapsed();
            let secs = duration.as_secs();
            let millis = duration.subsec_millis();
            let micros = duration.subsec_micros();
            let nanos = duration.subsec_nanos();
            if secs > 60 {
                format!("{}m", secs / 60)
            } else if secs > 0 {
                format!("{}s", secs)
            } else if millis > 0 {
                format!("{}ms", millis)
            } else if micros > 0 {
                format!("{}μs", micros)
            } else {
                format!("{}ns", nanos)
            }
        };
        let human_readable_time = format!("{: <5}", human_readable_time());
        let mut hostname = "".to_string();
        if show_host {
            hostname = self.host_name.clone();
        }

        match req_type.as_str() {
            "SEARCH" => format!(
                "{}:{} -> {}:{} {} 🔍",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "GET" => format!(
                "{}:{} -> {}:{} {} 📥",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "POST" => format!(
                "{}:{} -> {}:{} {} 📤",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "PUT" => format!(
                "{}:{} -> {}:{} {} 📤",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "PUT_GET" => format!(
                "{}:{} -> {}:{} {} 📥📤",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "MONITOR" => format!(
                "{}:{} -> {}:{} {} 🖥️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "ARRAY" => format!(
                "{}:{} -> {}:{} {} 📦",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "DESTROY_REQUEST" => format!(
                "{}:{} -> {}:{} {} 🗑️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "PROCESS" => format!(
                "{}:{} -> {}:{} {} 🔄",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "GET_FIELD" => format!(
                "{}:{} -> {}:{} {} 📄",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "MESSAGE" => format!(
                "{}:{} -> {}:{} {} 📧",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "MULTIPLE_DATA" => format!(
                "{}:{} -> {}:{} {} 📦📦",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "RPC" => format!(
                "{}:{} -> {}:{} {} 📞",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "CANCEL_REQUEST" => format!(
                "{}:{} -> {}:{} {} 🚫",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "ORIGIN_TAG" => format!(
                "{}:{} -> {}:{} {} 🏷️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "BEACON" => format!(
                "{}:{} -> {}:{} {} 🚨",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "CONNECTION_VALIDATION" => format!(
                "{}:{} -> {}:{} {} 🛡️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "ECHO" => format!(
                "{}:{} -> {}:{} {} 🗣️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "SEARCH_RESPONSE" => format!(
                "{}:{} -> {}:{} {} 🔍📤",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "AUTHNZ" => format!(
                "{}:{} -> {}:{} {} 🛡️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "ACL_CHANGE" => format!(
                "{}:{} -> {}:{} {} 🛡️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "CREATE_CHANNEL" => format!(
                "{}:{} -> {}:{} {} 📡",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "DESTROY_CHANNEL" => format!(
                "{}:{} -> {}:{} {} 💣📡",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "CONNECTION_VALIDATED" => format!(
                "{}:{} -> {}:{} {} ✔️",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            "UNKNOWN" => format!(
                "{}:{} -> {}:{} {} ❓",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
            _ => format!(
                "{}:{} -> {}:{} {}",
                human_readable_time, source_ip, destination_ip, req_type, hostname
            ),
        }
    }
}
#[derive(Clone)]
pub struct IfStats {
    pub data_rate: f64,
    pub packet_rate: f64,
    pub timestamp: std::time::Instant,
    pub source_ips: HashMap<String, u64>, // we will show ip and number of packets
    pub destination_ips: HashMap<String, u64>,
    pub request_types: HashMap<String, u64>,
    pub resolved_ips: HashMap<String, String>,
    pub events: Vec<PVEvent>,
    pub tracked_connections_summary: Vec<TrackedConnectionEntry>,
    pub tracked_connections_detail: Vec<TrackedConnectionEntry>,
    pub resolve: bool,
    pub buffer_size: usize,
    // Self-monitoring gauge fields
    pub pcap_received: u64,
    pub pcap_dropped: u64,
    pub pcap_if_dropped: u64,
    pub pva_stats: PvaStateStats,
    pub pva_stats_prev: PvaStateStats,
    pub pva_total_channels: usize,
    pub pva_total_connections: usize,
    pub pva_seg_started: u64,
    pub pva_seg_completed: u64,
    pub pva_seg_dropped: u64,
    pub pva_seg_drop_short: u64,
    pub pva_seg_drop_invalid: u64,
    pub pva_seg_drop_partial: u64,
    pub pva_seg_drop_missing_first: u64,
    pub pva_stream_dropped: u64,
    pub pva_seg_log: Vec<String>,
    // pub client_ids: HashMap<String, String>,
}

#[derive(Debug, Clone)]
pub struct TrackedConnectionEntry {
    pub id: String,
    pub endpoint_a: String,
    pub endpoint_b: String,
    pub cid: Option<u32>,
    pub sid: Option<u32>,
    pub pv_name: Option<String>,
    pub channel_count: usize,
    pub operation_count: usize,
    pub updates_per_sec: f64,
    pub last_seen_secs: u64,
    pub mid_stream: bool,
    pub is_beacon: bool,
    pub is_broadcast: bool,
    pub recent_messages: Vec<String>,
}

impl Default for IfStats {
    fn default() -> Self {
        IfStats {
            data_rate: 0.0,
            packet_rate: 0.0,
            timestamp: std::time::Instant::now(),
            source_ips: HashMap::new(),
            destination_ips: HashMap::new(),
            request_types: HashMap::new(),
            events: Vec::new(),
            resolved_ips: HashMap::new(),
            tracked_connections_summary: Vec::new(),
            tracked_connections_detail: Vec::new(),
            resolve: false,
            buffer_size: 5000,
            pcap_received: 0,
            pcap_dropped: 0,
            pcap_if_dropped: 0,
            pva_stats: PvaStateStats::default(),
            pva_stats_prev: PvaStateStats::default(),
            pva_total_channels: 0,
            pva_total_connections: 0,
            pva_seg_started: 0,
            pva_seg_completed: 0,
            pva_seg_dropped: 0,
            pva_seg_drop_short: 0,
            pva_seg_drop_invalid: 0,
            pva_seg_drop_partial: 0,
            pva_seg_drop_missing_first: 0,
            pva_stream_dropped: 0,
            pva_seg_log: Vec::new(),
        }
    }
}

impl IfStats {
    pub fn new() -> Self {
        Self {
            data_rate: 0.0,
            packet_rate: 0.0,
            timestamp: std::time::Instant::now(),
            source_ips: HashMap::new(),
            destination_ips: HashMap::new(),
            request_types: HashMap::new(),
            events: Vec::new(),
            resolved_ips: HashMap::new(),
            tracked_connections_summary: Vec::new(),
            tracked_connections_detail: Vec::new(),
            resolve: false,
            buffer_size: 5000,
            pcap_received: 0,
            pcap_dropped: 0,
            pcap_if_dropped: 0,
            pva_stats: PvaStateStats::default(),
            pva_stats_prev: PvaStateStats::default(),
            pva_total_channels: 0,
            pva_total_connections: 0,
            pva_seg_started: 0,
            pva_seg_completed: 0,
            pva_seg_dropped: 0,
            pva_seg_drop_short: 0,
            pva_seg_drop_invalid: 0,
            pva_seg_drop_partial: 0,
            pva_seg_drop_missing_first: 0,
            pva_stream_dropped: 0,
            pva_seg_log: Vec::new(),
        }
    }

    pub fn print(&self) {
        println!(
            "Data rate: {:.2} Mbit/s, Packet rate: {:.2} pps",
            self.data_rate, self.packet_rate
        );
    }

    pub fn to_line_protocol(&self) -> Vec<String> {
        let mut data_line_protocol = vec![];
        let host_name = gethostname().to_string_lossy().to_string();
        data_line_protocol.push(format!(
            "rates,host={} pps={},data_rate={}\n",
            host_name, self.packet_rate, self.data_rate
        ));

        for (ip, count) in self.source_ips.iter() {
            data_line_protocol.push(format!(
                "source_ips,host={},ip={} count={}\n",
                host_name, ip, count
            ));
        }
        for (ip, count) in self.destination_ips.iter() {
            data_line_protocol.push(format!(
                "destination_ips,host={},ip={} count={}\n",
                host_name, ip, count
            ));
        }
        for (req, count) in self.request_types.iter() {
            data_line_protocol.push(format!(
                "request_types,host={},request_type={} count={}\n",
                host_name, req, count
            ));
        }

        for event in self.events.iter() {
            // println!("{}",event.timestamp.elapsed().as_secs());
            data_line_protocol.push(format!("events,host={},source_ip={},destination_ip={},request_type={} payload_size={},count=1\n",event.host_name,event.source_ip,event.destination_ip,event.request_type,event.context.len()));
        }

        // for node graphs on grafana we need a nodes table
        // id	title
        // lets merge sources and destinations
        let combined_set = self.source_ips.keys().chain(self.destination_ips.keys());
        // strip of repeated ips
        let combined_set: std::collections::HashSet<&String> = combined_set.collect();
        let mut combined_set_counts_per_ip = HashMap::new();
        for ip in combined_set.clone() {
            let count =
                self.source_ips.get(ip).unwrap_or(&0) + self.destination_ips.get(ip).unwrap_or(&0);
            combined_set_counts_per_ip.insert(ip.clone(), count);
        }

        for ip in combined_set {
            let resolved_ip = self.get_resolved_ip(ip.clone());
            if resolved_ip == "Unknown" {
                data_line_protocol.push(format!(
                    "node,host={},id={},title={} count={}\n",
                    host_name,
                    ip,
                    ip,
                    combined_set_counts_per_ip.get(ip).unwrap_or(&0)
                ));
            } else {
                data_line_protocol.push(format!(
                    "node,host={},id={},title={} count={}\n",
                    host_name,
                    ip,
                    resolved_ip,
                    combined_set_counts_per_ip.get(ip).unwrap_or(&0)
                ));
            }
        }
        // id	source	target	mainstat	detail

        //event dict that will summarise all unique events and count them
        let mut event_dict = HashMap::new();
        for event in self.events.iter() {
            let id = format!("{}-{}", event.source_ip, event.destination_ip);
            let count = event_dict.entry(id).or_insert(0);
            *count += 1;
        }

        for (id, count) in event_dict.iter() {
            let parts: Vec<&str> = id.split("-").collect();
            let source_ip = parts[0];
            let destination_ip = parts[1];
            data_line_protocol.push(format!(
                "edge,host={},id={},source={},target={} mainStat={}\n",
                host_name, id, source_ip, destination_ip, count
            ));
        }

        //format!("edges,host={},id={},source={},target={},mainstat={}, count={}\n",host_name,id,event.source_ip,event.destination_ip,event.request_type,event.context.len())

        // --- Self-monitoring gauge measurements ---

        // Capture engine gauges
        data_line_protocol.push(format!(
            "capture_gauge,host={} received={}i,dropped={}i,if_dropped={}i,data_rate_bps={},packet_rate_pps={},unique_sources={}i,unique_destinations={}i\n",
            host_name,
            self.pcap_received,
            self.pcap_dropped,
            self.pcap_if_dropped,
            self.data_rate,
            self.packet_rate,
            self.source_ips.len(),
            self.destination_ips.len(),
        ));

        // PVA state gauges (cumulative counters — use non_negative_derivative in Grafana)
        data_line_protocol.push(format!(
            "pva_gauge,host={} active_channels={}i,active_connections={}i,channels_created={}i,channels_destroyed={}i,channels_expired={}i,channels_evicted={}i,operations_created={}i,operations_completed={}i,create_channel_requests={}i,create_channel_responses={}i,seg_started={}i,seg_completed={}i,seg_dropped={}i,seg_drop_short={}i,seg_drop_invalid={}i,seg_drop_partial={}i,seg_drop_missing_first={}i,stream_dropped={}i\n",
            host_name,
            self.pva_total_channels,
            self.pva_total_connections,
            self.pva_stats.channels_created,
            self.pva_stats.channels_destroyed,
            self.pva_stats.channels_expired,
            self.pva_stats.channels_evicted,
            self.pva_stats.operations_created,
            self.pva_stats.operations_completed,
            self.pva_stats.create_channel_requests,
            self.pva_stats.create_channel_responses,
            self.pva_seg_started,
            self.pva_seg_completed,
            self.pva_seg_dropped,
            self.pva_seg_drop_short,
            self.pva_seg_drop_invalid,
            self.pva_seg_drop_partial,
            self.pva_seg_drop_missing_first,
            self.pva_stream_dropped,
        ));

        // PVA state deltas (per-interval rates — instant rate in Grafana)
        let delta_channels_created = self.pva_stats.channels_created.saturating_sub(self.pva_stats_prev.channels_created);
        let delta_channels_destroyed = self.pva_stats.channels_destroyed.saturating_sub(self.pva_stats_prev.channels_destroyed);
        let delta_channels_expired = self.pva_stats.channels_expired.saturating_sub(self.pva_stats_prev.channels_expired);
        let delta_channels_evicted = self.pva_stats.channels_evicted.saturating_sub(self.pva_stats_prev.channels_evicted);
        let delta_operations_created = self.pva_stats.operations_created.saturating_sub(self.pva_stats_prev.operations_created);
        let delta_operations_completed = self.pva_stats.operations_completed.saturating_sub(self.pva_stats_prev.operations_completed);
        data_line_protocol.push(format!(
            "pva_gauge_delta,host={} channels_created={}i,channels_destroyed={}i,channels_expired={}i,channels_evicted={}i,operations_created={}i,operations_completed={}i\n",
            host_name,
            delta_channels_created,
            delta_channels_destroyed,
            delta_channels_expired,
            delta_channels_evicted,
            delta_operations_created,
            delta_operations_completed,
        ));

        // Collector process gauges
        let rss_bytes = std::fs::read_to_string("/proc/self/statm")
            .ok()
            .and_then(|s| s.split_whitespace().nth(1)?.parse::<u64>().ok())
            .map(|pages| pages * 4096)
            .unwrap_or(0);
        let event_buffer_pct = if self.buffer_size > 0 {
            (self.events.len() as f64 / self.buffer_size as f64) * 100.0
        } else {
            0.0
        };
        data_line_protocol.push(format!(
            "collector_gauge,host={} rss_bytes={}i,event_buffer_pct={:.1},event_count={}i,buffer_size={}i\n",
            host_name,
            rss_bytes,
            event_buffer_pct,
            self.events.len(),
            self.buffer_size,
        ));

        data_line_protocol
    }

    pub fn to_loki_stream(&self) -> Vec<(String, String, HashMap<String, String>)> {
        let database = "pv_events";
        self.events
            .iter()
            .map(|event| {
                let timestamp_ns = match event
                    .system_time
                    .checked_add(event.timestamp.elapsed())
                    .and_then(|time| time.duration_since(std::time::UNIX_EPOCH).ok())
                {
                    Some(duration) => duration.as_nanos().to_string(),
                    None => "0".to_string(),
                };

                // Per-category normalized log line
                let log_line = Self::format_loki_log_line(event);

                let mut metadata = HashMap::new();
                metadata.insert("source_ip".to_string(), event.source_ip.clone());
                metadata.insert("destination_ip".to_string(), event.destination_ip.clone());
                metadata.insert("request_type".to_string(), event.request_type.clone());
                metadata.insert("context".to_string(), event.context.clone());
                metadata.insert(
                    "resolved_source_ip".to_string(),
                    event.resolved_source_ip.clone(),
                );
                metadata.insert(
                    "resolved_destination_ip".to_string(),
                    event.resolved_destination_ip.clone(),
                );
                metadata.insert("host_name".to_string(), event.host_name.clone());
                metadata.insert("database".to_string(), database.to_string());

                // Add structured metadata for filtering
                if let Some(ref pv) = event.pv_name {
                    metadata.insert("pv_name".to_string(), pv.clone());
                }
                if let Some(ioid) = event.ioid {
                    metadata.insert("ioid".to_string(), ioid.to_string());
                }
                if let Some(sid) = event.sid {
                    metadata.insert("sid".to_string(), sid.to_string());
                }
                if let Some(cid) = event.cid {
                    metadata.insert("cid".to_string(), cid.to_string());
                }
                if event.mid_stream {
                    metadata.insert("mid_stream".to_string(), "true".to_string());
                }

                (timestamp_ns, log_line, metadata)
            })
            .collect()
    }

    /// Format a normalized Loki log line based on request type category
    fn format_loki_log_line(event: &PVEvent) -> String {
        let header = format!(
            "{} from {} to {}:",
            event.request_type, event.source_ip, event.destination_ip
        );

        match event.request_type.as_str() {
            // Op commands: MONITOR, GET, PUT, PUT_GET, ARRAY, PROCESS, RPC
            "MONITOR" | "GET" | "PUT" | "PUT_GET" | "ARRAY" | "PROCESS" | "RPC" => {
                let pv_display = match &event.pv_name {
                    Some(name) => format!(" pv={}", name),
                    None => {
                        // Extract data=[...] hint from context to help identify unresolved PVs
                        let data_hint = event.context.find("data=[").and_then(|start| {
                            event.context[start..].find("])").map(|end| {
                                format!(" {}", &event.context[start..start + end + 1])
                            })
                        }).unwrap_or_default();
                        match event.ioid {
                            Some(id) => format!(" pv=<unresolved:ioid={}>{}", id, data_hint),
                            None => data_hint,
                        }
                    }
                };
                let value_display = event
                    .value_text
                    .as_ref()
                    .map(|v| format!(" value={}", v))
                    .unwrap_or_default();
                let ioid_display = event
                    .ioid
                    .map(|id| format!(" ioid={}", id))
                    .unwrap_or_default();
                format!("{}{}{}{}", header, pv_display, value_display, ioid_display)
            }

            // Channel lifecycle
            "CREATE_CHANNEL" => {
                let pv_display = event
                    .pv_name
                    .as_ref()
                    .map(|n| format!(" pv={}", n))
                    .unwrap_or_default();
                let cid_display = event
                    .cid
                    .map(|c| format!(" cid={}", c))
                    .unwrap_or_default();
                let sid_display = event
                    .sid
                    .map(|s| format!(" sid={}", s))
                    .unwrap_or_default();
                format!("{}{}{}{}", header, pv_display, cid_display, sid_display)
            }
            "DESTROY_CHANNEL" => {
                let cid_display = event
                    .cid
                    .map(|c| format!(" cid={}", c))
                    .unwrap_or_default();
                let sid_display = event
                    .sid
                    .map(|s| format!(" sid={}", s))
                    .unwrap_or_default();
                format!("{}{}{}", header, cid_display, sid_display)
            }

            // All other types: use the raw context but without redundant prefix
            _ => {
                format!("{} {}", header, event.context)
            }
        }
    }

    pub fn clear_events(&mut self) {
        self.events.clear();
    }

    /// Rotate PVA stats snapshot for delta computation.
    /// Call after writing to InfluxDB each cycle.
    pub fn rotate_stats(&mut self) {
        self.pva_stats_prev = self.pva_stats.clone();
    }

    pub fn diable_resolve(&mut self) {
        self.resolve = false;
    }
    pub fn enable_resolve(&mut self) {
        self.resolve = true;
    }
    pub fn add_source_ip(&mut self, ip: String) {
        if !self.resolved_ips.contains_key(&ip) && self.resolve {
            let host = self.dns_addr_lookup(ip.clone());
            self.resolved_ips.insert(ip.clone(), host);
        }
        let count = self.source_ips.entry(ip).or_insert(0);
        *count += 1;
    }
    pub fn add_destination_ip(&mut self, ip: String) {
        // println!("{:?}", self.resolved_ips);
        if !self.resolved_ips.contains_key(&ip) && self.resolve {
            let host = self.dns_addr_lookup(ip.clone());
            self.resolved_ips.insert(ip.clone(), host);
        }
        let count = self.destination_ips.entry(ip).or_insert(0);
        *count += 1;
    }
    pub fn add_request_type(&mut self, request_type: String) {
        let count = self.request_types.entry(request_type).or_insert(0);
        *count += 1;
    }
    pub fn add_event(&mut self, event: PVEvent) {
        self.events.insert(0, event);
        if self.events.len() > self.buffer_size {
            self.events.pop();
        }
    }
    fn dns_addr_lookup(&mut self, ip: String) -> String {
        let ip_addr: std::net::IpAddr = ip.parse().unwrap();

        let host: String = match lookup_addr(&ip_addr) {
            Ok(host) => host,
            Err(_) => "Unknown".to_string(),
        };
        // println!("Resolved IP: {}:{} -> {}", ip, host);
        host
    }
    pub fn get_resolved_ip(&self, ip: String) -> String {
        match self.resolved_ips.get(&ip) {
            Some(host) => host.clone(),
            None => "Unknown".to_string(),
        }
    }
    pub fn set_buffer_size(&mut self, size: usize) {
        // 100k worked fine in testing but moving ifstats between threads caused a lot of CPU usage, realisticaly we should never need a buffer size of 100k events
        if size > 50000 {
            self.buffer_size = 50000;
        } else if size < 100 {
            self.buffer_size = 100;
        } else {
            self.buffer_size = size;
        }
    }
    fn get_memory_size(&self) -> usize {
        self.events.len() * std::mem::size_of::<PVEvent>()
    }
    pub fn get_memory_usage(&self) -> String {
        let mem_size = self.get_memory_size();
        if mem_size > 1_000_000 {
            format!("{:.2} MB", mem_size as f64 / 1_000_000.0)
        } else if mem_size > 1_000 {
            format!("{:.2} KB", mem_size as f64 / 1_000.0)
        } else {
            format!("{:.2} B", mem_size as f64)
        }
    }

    pub fn aggregate(&mut self, other: IfStats) {
        self.timestamp = std::time::Instant::now();
        for (ip, count) in other.source_ips.iter() {
            let entry = self.source_ips.entry(ip.clone()).or_insert(0);
            *entry += count;
        }
        for (ip, count) in other.destination_ips.iter() {
            let entry = self.destination_ips.entry(ip.clone()).or_insert(0);
            *entry += count;
        }
        for (req, count) in other.request_types.iter() {
            let entry = self.request_types.entry(req.clone()).or_insert(0);
            *entry += count;
        }
        for event in other.events.iter() {
            self.add_event(event.clone());
        }

        for (ip, host) in other.resolved_ips.iter() {
            self.resolved_ips.insert(ip.clone(), host.clone());
        }
        self.tracked_connections_summary = other.tracked_connections_summary.clone();
        self.tracked_connections_detail = other.tracked_connections_detail.clone();

        // Merge gauge fields (take latest values from other)
        self.pcap_received = other.pcap_received;
        self.pcap_dropped = other.pcap_dropped;
        self.pcap_if_dropped = other.pcap_if_dropped;
        self.pva_stats = other.pva_stats.clone();
        self.pva_stats_prev = other.pva_stats_prev.clone();
        self.pva_total_channels = other.pva_total_channels;
        self.pva_total_connections = other.pva_total_connections;

        // iterate through events and count how many occured in the last second
        let mut last_second = 0;
        let agg_data_rate = 0.0;
        for event in self.events.iter() {
            if event.timestamp.elapsed().as_secs() < 1 {
                last_second += 1;
            }
        }
        self.packet_rate = last_second as f64;
        self.data_rate = agg_data_rate;
    }
    pub fn clear(&mut self) {
        self.data_rate = 0.0;
        self.packet_rate = 0.0;
        self.source_ips.clear();
        self.destination_ips.clear();
        self.request_types.clear();
        self.events.clear();
        self.resolved_ips.clear();
        self.tracked_connections_summary.clear();
        self.tracked_connections_detail.clear();
        self.pcap_received = 0;
        self.pcap_dropped = 0;
        self.pcap_if_dropped = 0;
        self.pva_stats = PvaStateStats::default();
        self.pva_stats_prev = PvaStateStats::default();
        self.pva_total_channels = 0;
        self.pva_total_connections = 0;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_pv_event_new() {
        let event = PVEvent::new(
            "192.168.1.1".to_string(),
            "192.168.1.2".to_string(),
            "GET".to_string(),
        );
        assert_eq!(event.source_ip, "192.168.1.1");
        assert_eq!(event.destination_ip, "192.168.1.2");
        assert_eq!(event.request_type, "GET");
        assert_eq!(event.resolved_source_ip, "Unknown");
        assert_eq!(event.resolved_destination_ip, "Unknown");
    }

    #[test]
    fn test_pv_event_to_string() {
        let event = PVEvent::new(
            "192.168.1.1".to_string(),
            "192.168.1.2".to_string(),
            "GET".to_string(),
        );
        let event_string = event.to_string();
        assert!(event_string.contains("192.168.1.1"));
        assert!(event_string.contains("192.168.1.2"));
        assert!(event_string.contains("GET"));
    }

    #[test]
    fn test_IfStats_new() {
        let stats = IfStats::new();
        assert_eq!(stats.data_rate, 0.0);
        assert_eq!(stats.packet_rate, 0.0);
        assert!(stats.source_ips.is_empty());
        assert!(stats.destination_ips.is_empty());
        assert!(stats.request_types.is_empty());
        assert!(stats.events.is_empty());
        assert!(stats.resolved_ips.is_empty());
        assert!(!stats.resolve);
        assert_eq!(stats.buffer_size, 5000);
    }

    #[test]
    fn test_IfStats_add_source_ip() {
        let mut stats = IfStats::new();
        stats.add_source_ip("192.168.1.1".to_string());
        assert_eq!(*stats.source_ips.get("192.168.1.1").unwrap(), 1);
    }

    #[test]
    fn test_IfStats_add_destination_ip() {
        let mut stats = IfStats::new();
        stats.add_destination_ip("192.168.1.2".to_string());
        assert_eq!(*stats.destination_ips.get("192.168.1.2").unwrap(), 1);
    }

    #[test]
    fn test_IfStats_add_request_type() {
        let mut stats = IfStats::new();
        stats.add_request_type("GET".to_string());
        assert_eq!(*stats.request_types.get("GET").unwrap(), 1);
    }

    #[test]
    fn test_IfStats_add_event() {
        let mut stats = IfStats::new();
        let event = PVEvent::new(
            "192.168.1.1".to_string(),
            "192.168.1.2".to_string(),
            "GET".to_string(),
        );
        stats.add_event(event.clone());
        assert_eq!(stats.events.len(), 1);
        assert_eq!(stats.events[0].source_ip, event.source_ip);
    }

    #[test]
    fn test_IfStats_memory_usage() {
        let mut stats = IfStats::new();
        let event = PVEvent::new(
            "192.168.1.1".to_string(),
            "192.168.1.2".to_string(),
            "GET".to_string(),
        );
        stats.add_event(event);
        let memory_usage = stats.get_memory_usage();
        assert!(memory_usage.ends_with("B"));
    }

    #[test]
    fn test_IfStats_set_buffer_size() {
        let mut stats = IfStats::new();
        stats.set_buffer_size(100000);
        assert_eq!(stats.buffer_size, 50000);
        stats.set_buffer_size(50);
        assert_eq!(stats.buffer_size, 100);
        stats.set_buffer_size(1000);
        assert_eq!(stats.buffer_size, 1000);
    }

    #[test]
    fn test_IfStats_resolve_ip() {
        let mut stats = IfStats::new();
        stats.enable_resolve();
        stats.add_source_ip("8.8.8.8".to_string());
        let resolved_ip = stats.get_resolved_ip("8.8.8.8".to_string());
        assert!(stats.resolved_ips.contains_key("8.8.8.8"));
        assert!(!resolved_ip.is_empty());
    }
}
