use crate::lease::PeerId;
use std::collections::HashMap;
use std::net::IpAddr;
use std::time::{Duration, Instant};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TransportProtocol {
    Tcp,
    Udp,
    Icmp,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct FlowKey {
    pub protocol: TransportProtocol,
    pub source: IpAddr,
    pub source_port: u16,
    pub destination: IpAddr,
    pub destination_port: u16,
}

#[derive(Debug, Clone)]
pub struct FlowEntry {
    pub key: FlowKey,
    pub peer_id: PeerId,
    pub created_at: Instant,
    pub last_seen_at: Instant,
    pub bytes_from_peer: u64,
    pub bytes_to_peer: u64,
}

#[derive(Debug, Default)]
pub struct FlowTable {
    entries: HashMap<FlowKey, FlowEntry>,
}

impl FlowTable {
    pub fn upsert(&mut self, key: FlowKey, peer_id: PeerId, now: Instant, bytes_from_peer: u64) {
        self.entries
            .entry(key.clone())
            .and_modify(|entry| {
                entry.last_seen_at = now;
                entry.bytes_from_peer = entry.bytes_from_peer.saturating_add(bytes_from_peer);
            })
            .or_insert(FlowEntry {
                key,
                peer_id,
                created_at: now,
                last_seen_at: now,
                bytes_from_peer,
                bytes_to_peer: 0,
            });
    }

    pub fn record_return_bytes(&mut self, key: &FlowKey, bytes_to_peer: u64) {
        if let Some(entry) = self.entries.get_mut(key) {
            entry.bytes_to_peer = entry.bytes_to_peer.saturating_add(bytes_to_peer);
        }
    }

    pub fn expire_idle(&mut self, now: Instant, idle_timeout: Duration) -> usize {
        let before = self.entries.len();
        self.entries
            .retain(|_, entry| now.duration_since(entry.last_seen_at) <= idle_timeout);
        before - self.entries.len()
    }

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

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::net::{IpAddr, Ipv4Addr};

    #[test]
    fn expires_idle_flows() {
        let now = Instant::now();
        let mut table = FlowTable::default();
        let key = FlowKey {
            protocol: TransportProtocol::Tcp,
            source: IpAddr::V4(Ipv4Addr::new(10, 241, 0, 2)),
            source_port: 1234,
            destination: IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)),
            destination_port: 443,
        };
        table.upsert(key, PeerId::from_u128(1), now, 10);

        assert_eq!(table.len(), 1);
        let expired = table.expire_idle(now + Duration::from_secs(61), Duration::from_secs(60));
        assert_eq!(expired, 1);
        assert!(table.is_empty());
    }
}