use crate::oracle::Oracle;
use crate::server::SimulatedServer;
use fuzzytags::{RootSecret, TaggingKey};
use rand::distributions::Distribution;
use rand::Rng;
use rand_distr::num_traits::ToPrimitive;
use tracing::event;
use tracing::span;
use tracing::Level;

pub struct SimulatedParties {
    parties: Vec<RootSecret<24>>,
}

impl SimulatedParties {
    pub fn new_simulation(num_parties: usize) -> SimulatedParties {
        let mut parties = vec![];
        for _p in 0..num_parties {
            let key = RootSecret::<24>::generate();
            parties.push(key);
        }
        SimulatedParties { parties }
    }

    pub fn register_with_server<R>(&self, server: &mut SimulatedServer, rng: &mut R, min_p: usize, max_p: usize, oracle: &mut Oracle)
    where
        R: Rng,
    {
        for party in self.parties.iter() {
            let n = rng.gen_range(min_p..max_p);
            let span = span!(Level::INFO, "register", party = party.tagging_key().id().as_str());
            let _enter = span.enter();
            let detection_key = party.extract_detection_key(n);
            event!(Level::TRACE, "create detection key {detection_key}", detection_key = detection_key.id().as_str());
            event!(Level::TRACE, "register with server");
            server.register_key(&detection_key, &party.tagging_key());
            oracle.register_party(party.tagging_key().id());
        }
    }

    pub fn sample_traffic<R, D>(&self, server: &mut SimulatedServer, rng: &mut R, distribution: D, probs_entangled: f64, oracle: &mut Oracle)
    where
        D: Distribution<f64>,
        R: Rng,
    {
        let span = span!(Level::INFO, "sample_traffic");
        let _enter = span.enter();
        let v = distribution.sample(rng).to_u16().unwrap();

        let sender = rng.gen_range(0..self.parties.len());
        let sender_public_key = self.parties.get(sender).unwrap().tagging_key();

        let receiver = rng.gen_range(0..self.parties.len());
        let receiver_public_key = self.parties.get(receiver).unwrap().tagging_key();

        if sender != receiver {
            let entangle = rng.gen_bool(probs_entangled);
            if entangle {
                let receiver_2 = rng.gen_range(0..self.parties.len());
                let receiver_public_key_2 = self.parties.get(receiver_2).unwrap().tagging_key();
                event!(
                    Level::INFO,
                    "entangled send {party_1} {party_2}",
                    party_1 = receiver_public_key.id().as_str(),
                    party_2 = receiver_public_key_2.id().as_str()
                );

                for _i in 0..v {
                    let tag = TaggingKey::<24>::generate_entangled_tag(vec![receiver_public_key.clone(), receiver_public_key_2.clone()], 8);
                    event!(Level::TRACE, "message sent to server {tag}", tag = tag.to_string());
                    server.add_message(tag, &sender_public_key);
                }

                oracle.add_event(sender_public_key.id().clone(), receiver_public_key.id(), Some(receiver_public_key_2.id()), 1.0);
            } else {
                event!(Level::INFO, "regular send {party}", party = receiver_public_key.id().as_str());
                for _i in 0..v {
                    let tag = receiver_public_key.generate_tag();
                    event!(Level::INFO, "message sent server {tag}", tag = tag.to_string());
                    server.add_message(tag, &sender_public_key);
                }
                oracle.add_event(sender_public_key.id().clone(), receiver_public_key.id(), None, 1.0);
            }
        }
    }
}