use crate::{extra_info::ExtraInfo, get_date};
use lox_library::bridge_table::BridgeLine;
use std::collections::{BTreeMap, HashMap};

// The Bridge struct only tracks data for today
pub struct Bridge {
    pub fingerprint: [u8; 20],
    real_connections: HashMap<String, u32>,
    total_connections: BTreeMap<String, u32>,
}

impl Bridge {
    pub fn new(fingerprint: &[u8; 20]) -> Self {
        Self {
            fingerprint: *fingerprint,
            real_connections: HashMap::<String, u32>::new(),
            total_connections: BTreeMap::<String, u32>::new(),
        }
    }

    pub fn from_bridge_line(bridgeline: &BridgeLine) -> Self {
        Self::new(&bridgeline.fingerprint)
    }

    pub fn connect_real(&mut self, country: &str) {
        if self.real_connections.contains_key(country) {
            let prev = self.real_connections.get(country).unwrap();
            self.real_connections
                .insert(country.to_string(), prev + 1)
                .unwrap();
        } else {
            self.real_connections
                .insert(country.to_string(), 1)
                .unwrap();
        }
        self.connect_total(country);
    }

    pub fn connect_total(&mut self, country: &str) {
        if self.total_connections.contains_key(country) {
            let prev = self.total_connections.get(country).unwrap();
            self.total_connections
                .insert(country.to_string(), prev + 1)
                .unwrap();
        } else {
            self.total_connections
                .insert(country.to_string(), 1)
                .unwrap();
        }
    }

    // Let the censor simulate a bunch of connections at once
    pub fn censor_flood(&mut self, country: &str, num_connections: u32) {
        if self.total_connections.contains_key(country) {
            let prev = self.total_connections.get(country).unwrap();
            self.total_connections
                .insert(country.to_string(), prev + num_connections)
                .unwrap();
        } else {
            self.total_connections
                .insert(country.to_string(), num_connections)
                .unwrap();
        }
    }

    // Generate an extra-info report for today
    pub fn gen_extra_info(&self) -> ExtraInfo {
        ExtraInfo {
            nickname: String::default(),
            fingerprint: self.fingerprint,
            date: get_date(),
            bridge_ips: self.total_connections.clone(),
        }
    }

    fn reset_for_tomorrow(&mut self) {
        self.real_connections = HashMap::<String, u32>::new();
        self.total_connections = BTreeMap::<String, u32>::new();
    }
}