// Before running this, run:
// 1. rdsys
// 2. lox-distributor
// 3. troll-patrol with the feature "simulation"

use troll_patrol::{
    extra_info::ExtraInfo,
    get_date, increment_simulated_date, set_simulated_date,
    simulation::{
        bridge::Bridge,
        censor::{self, Censor},
        config::Config as SConfig,
        extra_infos_server,
        user::User,
    },
};

use clap::Parser;
use lox_cli::{networking::*, *};
use lox_library::proto::{level_up::LEVEL_INTERVAL, trust_promotion::UNTRUSTED_INTERVAL};
use memory_stats::memory_stats;
use rand::{prelude::SliceRandom, Rng};
use serde::Deserialize;
use std::{
    collections::{HashMap, HashSet},
    fs::File,
    io::BufReader,
    path::PathBuf,
    time::Duration,
};
use tokio::{spawn, time::sleep};

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// Name/path of the configuration file
    #[arg(short, long, default_value = "simulation_config.json")]
    config: PathBuf,
}

#[derive(Debug, Deserialize)]
pub struct Config {
    pub la_port: u16,
    pub la_test_port: u16,
    pub tp_port: u16,
    pub tp_test_port: u16,
    pub censor_hides: censor::Hides,
    pub censor_speed: censor::Speed,
    pub censor_event_duration: u32,
    pub censor_totality: censor::Totality,
    pub censor_partial_blocking_percent: f64,
    pub country: String,
    pub min_new_users_per_day: u32,
    pub max_new_users_per_day: u32,
    // How many days to simulate
    pub num_days: u32,
    // We start with this many level 4 users
    pub num_initial_trusted_users: u32,
    pub one_positive_report_per_cred: bool,
    pub prob_connection_fails: f64,
    pub prob_user_invites_friend: f64,
    pub prob_user_is_censor: f64,
    pub prob_user_submits_reports: f64,
    pub prob_user_treats_throttling_as_blocking: f64,
}

#[tokio::main]
pub async fn main() {
    let args: Args = Args::parse();

    let config: Config = serde_json::from_reader(BufReader::new(
        File::open(&args.config).expect("Could not read config file"),
    ))
    .expect("Reading config file from JSON failed");

    let la_net = HyperNet {
        hostname: format!("http://localhost:{}", config.la_port),
    };
    let la_net_test = HyperNet {
        hostname: format!("http://localhost:{}", config.la_test_port),
    };
    let tp_net = HyperNet {
        hostname: format!("http://localhost:{}", config.tp_port),
    };
    let tp_net_test = HyperNet {
        hostname: format!("http://localhost:{}", config.tp_test_port),
    };
    let extra_infos_net = HyperNet {
        hostname: "http://localhost:8004".to_string(),
    };

    let la_pubkeys = get_lox_auth_keys(&la_net).await.unwrap();

    let sconfig = SConfig {
        la_pubkeys,
        la_net,
        tp_net,
        censor_hides: config.censor_hides,
        censor_speed: config.censor_speed,
        censor_event_duration: config.censor_event_duration,
        censor_totality: config.censor_totality,
        censor_partial_blocking_percent: config.censor_partial_blocking_percent,
        country: config.country,
        one_positive_report_per_cred: config.one_positive_report_per_cred,
        prob_connection_fails: config.prob_connection_fails,
        prob_user_invites_friend: config.prob_user_invites_friend,
        prob_user_is_censor: config.prob_user_is_censor,
        prob_user_submits_reports: config.prob_user_submits_reports,
        prob_user_treats_throttling_as_blocking: config.prob_user_treats_throttling_as_blocking,
    };

    let mut rng = rand::thread_rng();

    // Set up censor
    let mut censor = Censor::new(&sconfig);

    // Set up bridges (no bridges yet)
    let mut bridges = HashMap::<[u8; 20], Bridge>::new();

    // Set up users
    let mut users = Vec::<User>::new();

    if config.num_initial_trusted_users > 0 {
        // Add some number of trusted users initially
        for _ in 0..config.num_initial_trusted_users {
            let new_user = User::trusted_user(&sconfig).await;
            if new_user.is_ok() {
                users.push(new_user.unwrap());
            }
        }

        // Level trusted users up to level 4

        // Advance LA's time
        la_net_test
            .request(
                "/advancedays".to_string(),
                serde_json::to_string(&(UNTRUSTED_INTERVAL as u16))
                    .unwrap()
                    .into(),
            )
            .await;
        // Advance simulated time
        set_simulated_date(get_date() + UNTRUSTED_INTERVAL);

        for user in &mut users {
            let migcred = trust_promotion(
                &sconfig.la_net,
                &user.primary_cred,
                get_lox_pub(&sconfig.la_pubkeys),
            )
            .await;
            if migcred.is_ok() {
                let new_cred = trust_migration(
                    &sconfig.la_net,
                    &user.primary_cred,
                    &migcred.unwrap(),
                    get_lox_pub(&sconfig.la_pubkeys),
                    get_migration_pub(&sconfig.la_pubkeys),
                )
                .await;
                if new_cred.is_ok() {
                    user.primary_cred = new_cred.unwrap();
                }
            }
        }

        for i in 1..LEVEL_INTERVAL.len() - 2 {
            // Advance LA's time
            la_net_test
                .request(
                    "/advancedays".to_string(),
                    serde_json::to_string(&(LEVEL_INTERVAL[i] as u16))
                        .unwrap()
                        .into(),
                )
                .await;
            // Advance simulated time
            set_simulated_date(get_date() + LEVEL_INTERVAL[i]);

            for user in &mut users {
                let reachcred_res = get_bucket(&sconfig.la_net, &user.primary_cred).await;
                if reachcred_res.is_ok() {
                    let reachcred = reachcred_res.unwrap().1;
                    let new_cred = level_up(
                        &sconfig.la_net,
                        &user.primary_cred,
                        &reachcred.unwrap(),
                        get_lox_pub(&sconfig.la_pubkeys),
                        get_reachability_pub(&sconfig.la_pubkeys),
                    )
                    .await;
                    if new_cred.is_ok() {
                        user.primary_cred = new_cred.unwrap();
                    } else {
                        eprintln!("Failed to level up trusted user's credential at start");
                    }
                }
            }
        }

        // Advance LA's time to tomorrow
        la_net_test
            .request(
                "/advancedays".to_string(),
                serde_json::to_string(&(1 as u16)).unwrap().into(),
            )
            .await;

        // Advance simulated time to tomorrow
        increment_simulated_date();
    }

    // Set up extra-infos server
    spawn(async move {
        extra_infos_server::server().await;
    });
    sleep(Duration::from_millis(1)).await;

    let mut false_neg = 0;
    let mut false_pos = 0;
    let mut true_neg = 0;
    let mut true_pos = 0;

    // Track memory use during simulation
    let mut max_physical_mem = 0;
    let mut max_virtual_mem = 0;

    // Main loop
    for day in 1..=config.num_days {
        println!("Starting day {} of the simulation", day);
        println!(
            "    We have {} users and {} bridges",
            users.len(),
            bridges.len()
        );
        println!(
            "    The censor has learned {} bridges",
            censor.known_bridges.len()
        );

        if let Some(usage) = memory_stats() {
            if usage.physical_mem > max_physical_mem {
                max_physical_mem = usage.physical_mem;
            }
            if usage.virtual_mem > max_virtual_mem {
                max_virtual_mem = usage.virtual_mem;
            }
        } else {
            println!("Failed to get the current memory usage");
        }

        // USER TASKS

        // Number of users who want to join today
        let mut num_users_requesting_invites: u32 =
            rng.gen_range(config.min_new_users_per_day..=config.max_new_users_per_day);

        let mut new_users = Vec::<User>::new();

        // Shuffle users so they act in a random order
        users.shuffle(&mut rng);

        // Users do daily actions
        for user in &mut users {
            let invited_friends = user
                .daily_tasks(
                    &sconfig,
                    num_users_requesting_invites,
                    &mut bridges,
                    &mut censor,
                )
                .await;

            if invited_friends.is_ok() {
                let mut invited_friends = invited_friends.unwrap();
                if invited_friends.len() > 0 {
                    if !user.is_censor {
                        // Users should never invite more friends than
                        // need invitations, so this should never become
                        // negative
                        num_users_requesting_invites -= invited_friends.len() as u32;
                    }
                    // If this user invited any friends, add them to the
                    // list of users
                    new_users.append(&mut invited_friends);
                }
            }
        }

        // Add new users
        users.append(&mut new_users);

        // If any users couldn't get invites, they join with open-entry
        // invitations
        for _ in 0..num_users_requesting_invites {
            let user = User::new(&sconfig).await;
            if user.is_ok() {
                users.push(user.unwrap());
            } else {
                eprintln!("Failed to create new user.");
            }
        }

        // CENSOR TASKS
        censor.end_of_day_tasks(&sconfig, &mut bridges).await;

        // BRIDGE TASKS
        let mut new_extra_infos = HashSet::<ExtraInfo>::new();
        for (_, bridge) in bridges.iter_mut() {
            // Bridge reports its connections for the day
            new_extra_infos.insert(bridge.gen_extra_info(&sconfig.country));

            // Bridge resets for tomorrow
            bridge.reset_for_tomorrow();
        }

        // Publish all the bridges' extra-infos for today
        extra_infos_net
            .request(
                "/add".to_string(),
                serde_json::to_string(&new_extra_infos).unwrap().into(),
            )
            .await;

        // TROLL PATROL TASKS
        let new_blockages_resp = tp_net_test.request("/update".to_string(), vec![]).await;
        let new_blockages: HashMap<String, HashSet<String>> =
            serde_json::from_slice(&new_blockages_resp).unwrap();

        // Since we have only one censor, just convert to a set of bridges
        let mut blocked_bridges = HashSet::<[u8; 20]>::new();
        for (bridge, ccs) in new_blockages {
            let fingerprint = array_bytes::hex2array(bridge).unwrap();
            if ccs.contains(&sconfig.country) {
                blocked_bridges.insert(fingerprint);
            }
        }

        for (fingerprint, bridge) in &mut bridges {
            let detected_blocked = blocked_bridges.contains(fingerprint);

            // If this is the first day Troll Patrol has determined this
            // bridge is blocked, note that for stats
            if detected_blocked && bridge.first_detected_blocked == 0 {
                bridge.first_detected_blocked = get_date();
            }

            // Check if censor actually blocks this bridge
            let really_blocked = censor.blocks_bridge(&sconfig, fingerprint);
            if really_blocked && bridge.first_blocked == 0 {
                bridge.first_blocked = get_date();
            }
            if detected_blocked && really_blocked {
                true_pos += 1;
            } else if detected_blocked {
                false_pos += 1;
            } else if really_blocked {
                false_neg += 1;
            } else {
                true_neg += 1;
            }
        }

        // LOX AUTHORITY TASKS

        // Advance LA's time to tomorrow
        la_net_test
            .request(
                "/advancedays".to_string(),
                serde_json::to_string(&(1 as u16)).unwrap().into(),
            )
            .await;

        // SIMULATION TASKS

        // Advance simulated time to tomorrow
        increment_simulated_date();
    }

    // Print various information about the simulation run
    println!(
        "\nSimulation ended with {} users and {} bridges",
        users.len(),
        bridges.len()
    );
    println!("The censor learned {} bridges", censor.known_bridges.len());

    println!(
        "\nMaximum physical memory usage during simulation: {}",
        max_physical_mem
    );
    println!(
        "Maximum virtual memory usage during simulation: {}\n",
        max_virtual_mem
    );

    println!("True Positives: {}", true_pos);
    println!("True Negatives: {}", true_neg);
    println!("False Positives: {}", false_pos);
    println!("False Negatives: {}", false_neg);

    println!("\nFull stats per bridge:");

    println!(
        "Fingerprint,first_distributed,first_blocked,first_detected_blocked,first_positive_report"
    );
    for (fingerprint, bridge) in bridges {
        println!(
            "{},{},{},{},{}",
            array_bytes::bytes2hex("", fingerprint),
            bridge.first_distributed,
            bridge.first_blocked,
            bridge.first_detected_blocked,
            bridge.first_positive_report
        );
    }
}