use crate::{bridge::Bridge, config::Config};

use curve25519_dalek::scalar::Scalar;
use lox_cli::{get_lox_pub, networking::Networking};
use lox_library::{cred::Lox, scalar_u32};
use rand::Rng;
use serde::Deserialize;
use std::{
    cmp::min,
    collections::{HashMap, HashSet},
};
use troll_patrol::{get_date, positive_report::PositiveReport};

pub struct Censor {
    // If we have a bootstrapping period, the censor does not begin
    // until this date.
    pub start_date: u32,

    pub known_bridges: HashSet<[u8; 20]>,

    // We don't actually implement the technical restriction to prevent
    // one Lox credential from being used to submit many reports, so we
    // just implement this as a map of bridge fingerprint to (most
    // recent Lox credential for this bridge, count of unique level 3+
    // credentials we have for this bridge).
    pub lox_credentials: HashMap<[u8; 20], (Lox, u32)>,

    // Map of buckets to count of censor users with that bucket. Note
    // that this is the count of users that have *ever* had that bucket,
    // so this variable should NOT be used to count the overall number
    // of censor agents.
    pub agents: HashMap<Scalar, u32>,

    // If censor implements random blocking, this is the date when it
    // will start blocking all the bridges it knows.
    pub delay_date: u32,

    // If censor implements partial blocking, what percent of
    // connections are blocked?
    pub partial_blocking_percent: f64,
}

impl Censor {
    pub fn new(config: &Config) -> Self {
        let start_date = get_date() + config.bootstrapping_period_duration;
        let mut rng = rand::thread_rng();
        let delay_date = if config.censor_speed == Speed::Random {
            let num: u32 = rng.gen_range(1..365);
            start_date + num
        } else {
            0
        };
        let partial_blocking_percent = if config.censor_totality == Totality::Partial {
            config.censor_partial_blocking_percent
        } else {
            1.0
        };
        Censor {
            start_date,
            known_bridges: HashSet::<[u8; 20]>::new(),
            lox_credentials: HashMap::<[u8; 20], (Lox, u32)>::new(),
            agents: HashMap::<Scalar, u32>::new(),
            delay_date: delay_date,
            partial_blocking_percent: partial_blocking_percent,
        }
    }

    pub fn is_active(&self) -> bool {
        get_date() >= self.start_date
    }

    pub fn knows_bridge(&self, fingerprint: &[u8; 20]) -> bool {
        self.known_bridges.contains(fingerprint)
    }

    pub fn blocks_bridge(&self, config: &Config, fingerprint: &[u8; 20]) -> bool {
        self.knows_bridge(fingerprint)
            && (config.censor_speed == Speed::Fast
                || config.censor_speed == Speed::Lox && self.has_lox_cred(fingerprint)
                || config.censor_speed == Speed::Random && self.delay_date <= get_date())
    }

    pub fn learn_bridge(&mut self, fingerprint: &[u8; 20]) {
        self.known_bridges.insert(*fingerprint);
    }

    pub fn has_lox_cred(&self, fingerprint: &[u8; 20]) -> bool {
        self.lox_credentials.contains_key(fingerprint)
            && self.lox_credentials.get(fingerprint).unwrap().1 > 0
    }

    pub fn give_lox_cred(&mut self, fingerprint: &[u8; 20], cred: &Lox, new_cred: bool) {
        // We only need one level 3+ credential per bridge. (This will
        // change if we restrict positive reports to one per bridge per
        // credential.)
        if scalar_u32(&cred.trust_level).unwrap() >= 3 {
            // We want to clone the credential, but that's not allowed,
            // so we're going to serialize it and then deserialize it.
            let cloned_cred = bincode::deserialize(&bincode::serialize(&cred).unwrap()).unwrap();

            // Insert the new credential and add to the count of unique
            // credentials we have. We assume that a duplicate
            // credential will never be given. If we don't want to make
            // this assumption, we could change the count from a u32 to
            // a set of credential IDs and get the count as its length.
            let count = match self.lox_credentials.get(fingerprint) {
                Some((_cred, count)) => *count,
                None => 0,
            };
            // Sometimes we want to add a fresh credential but assume it
            // has the same gamma and does not give the ability to make
            // more reports.
            let new_count = if new_cred {
                // If we're adding a new gamma, increase count
                count + 1
            } else {
                // If we're adding a fresh cred with old gamma, don't
                count
            };
            // Insert the new credential. We always want the newest
            // credential to ensure freshness.
            self.lox_credentials
                .insert(*fingerprint, (cloned_cred, new_count));
        }
    }

    // Get the number of agents the censor has with a given bucket
    pub fn num_agents(&self, bucket: &Scalar) -> u32 {
        match self.agents.get(bucket) {
            Some(v) => *v,
            None => 0,
        }
    }

    // Add to the number of agents a censor has for a given bucket
    pub fn add_agent(&mut self, bucket: &Scalar) {
        self.agents.insert(
            *bucket,
            match self.agents.get(bucket) {
                Some(v) => *v + 1,
                None => 1,
            },
        );
    }

    // Censor sends a positive report for the given bridge. Returns true
    // if successful, false otherwise.
    pub async fn send_positive_report(&self, config: &Config, fingerprint: &[u8; 20]) -> bool {
        // If we don't have an appropriate Lox credential, we can't send
        // a report. Return false.
        if !self.has_lox_cred(fingerprint) {
            return false;
        }

        let (cred, _) = &self.lox_credentials.get(fingerprint).unwrap();
        if let Ok(pr) = PositiveReport::from_lox_credential(
            *fingerprint,
            None,
            cred,
            get_lox_pub(&config.la_pubkeys),
            config.country.clone(),
        ) {
            if config
                .tp_net
                .request("/positivereport".to_string(), pr.to_json().into_bytes())
                .await
                .is_err()
            {
                // failed to send positive report
                return false;
            }
        }
        true
    }

    // Make a bunch of connections and submit positive reports if possible
    async fn flood(&self, config: &Config, bridges: &mut HashMap<[u8; 20], Bridge>) {
        // Only do this if Flooding censor
        if config.censor_secrecy == Secrecy::Flooding {
            for fingerprint in &self.known_bridges {
                // Only do this if we're blocking the bridge
                if config.censor_speed == Speed::Fast
                    || config.censor_speed == Speed::Lox && self.has_lox_cred(fingerprint)
                    || config.censor_speed == Speed::Random && self.delay_date <= get_date()
                {
                    let bridge = bridges.get_mut(fingerprint).unwrap();

                    // Make a bunch of connections to the bridge
                    bridge.censor_flood(config.censor_max_connections);

                    // If we have a lv3+ credential, submit a bunch of
                    // positive reports
                    if self.has_lox_cred(fingerprint) {
                        let (_cred, cred_count) =
                            &self.lox_credentials.get(&bridge.fingerprint).unwrap();

                        // Reduce the number of these for efficiency. It
                        // does not benefit the censor to submit more
                        // than 25 positive reports under the current
                        // detection algorithm. In practice, the censor
                        // should submit as many reports as possible.
                        let num_prs = if config.one_positive_report_per_cred {
                            min(*cred_count, config.censor_max_pr)
                        } else {
                            config.censor_max_pr
                        };
                        for _ in 0..num_prs {
                            self.send_positive_report(config, &bridge.fingerprint).await;
                        }
                    }
                }
            }
        }
    }

    // Send one positive report per connection we blocked
    async fn send_positive_reports(
        &self,
        config: &Config,
        bridges: &mut HashMap<[u8; 20], Bridge>,
    ) {
        // Only do this if Hiding censor. Flooding censors should use
        // flood() instead.
        if config.censor_secrecy == Secrecy::Hiding {
            for fingerprint in &self.known_bridges {
                // Only do this if we're blocking the bridge
                if self.blocks_bridge(config, fingerprint) && self.has_lox_cred(fingerprint) {
                    let bridge = bridges.get_mut(fingerprint).unwrap();

                    // We may be restricted to one positive report per
                    // credential
                    let num_reports_to_send = if config.one_positive_report_per_cred {
                        min(
                            bridge.total_connections - bridge.real_connections,
                            self.lox_credentials.get(fingerprint).unwrap().1,
                        )
                    } else {
                        bridge.total_connections - bridge.real_connections
                    };
                    for _ in 0..num_reports_to_send {
                        self.send_positive_report(config, fingerprint).await;
                    }
                }
            }
        }
    }

    fn recompute_delay(&mut self, config: &Config) {
        // Only do this if Random censor
        if config.censor_speed == Speed::Random
            && self.delay_date + config.censor_event_duration <= get_date()
        {
            // Compute new delay date
            self.delay_date = {
                let mut rng = rand::thread_rng();
                let num: u32 = rng.gen_range(1..365);
                get_date() + num
            }
        }
    }

    pub async fn end_of_day_tasks(
        &mut self,
        config: &Config,
        bridges: &mut HashMap<[u8; 20], Bridge>,
    ) {
        // This check is redundant because the simulation driver only
        // calls this function if the censor is active.
        if self.is_active() {
            if config.censor_secrecy == Secrecy::Flooding
                && !(config.censor_speed == Speed::Random && self.delay_date <= get_date())
            {
                self.flood(config, bridges).await;
            } else if config.censor_secrecy == Secrecy::Hiding
                && !(config.censor_speed == Speed::Random && self.delay_date <= get_date())
            {
                self.send_positive_reports(config, bridges).await;
            }

            self.recompute_delay(config);
        }
    }
}

#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub enum Speed {
    Fast,
    Lox,
    Random,
}

#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub enum Secrecy {
    Overt,
    Hiding,
    Flooding,
}

#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
pub enum Totality {
    Full,
    Partial,
    Throttling,
}