use std::cell::RefCell;
use std::fs::File;
use std::io::prelude::*;
use std::io::BufWriter;
use std::path::Path;
use std::rc::Rc;

use clap::Parser;
use rand::prelude::*;
use rand_distr::Exp;
use rand_pcg::Pcg64;
use simcore::simulation::Simulation;

mod balancer;
mod config;
mod events;
mod host;
mod log;
mod metrics;

use crate::balancer::{create_selector_by_name, LoadBalancer};
use crate::config::Config;
use crate::events::{RequestArrivalEvent, SyncEvent};
use crate::host::Host;
use crate::log::EventLog;
use crate::metrics::{windowed_p99, WINDOW};

#[derive(Parser, Debug)]
#[clap(about, long_about = None)]
struct Args {
    /// Results dump path.
    #[clap(long)]
    dump: String,

    /// Configuration JSON file.
    #[clap(long)]
    config: String,

    /// Number of servers.
    #[clap(long, default_value = "100")]
    servers: u32,

    /// Last arrival time shown on the graph (seconds).
    #[clap(long, default_value = "3600")]
    time_limit: f64,

    /// Number of senders generating requests.
    #[clap(long, default_value = "200")]
    senders: u32,

    /// Load distribution preset (balanced/skewed).
    #[clap(long, default_value = "balanced")]
    preset: String,

    /// Adds slower servers to simulate heterogenous cluster.
    #[clap(long)]
    different_servers: bool,

    /// Random seed.
    #[clap(long, default_value = "123")]
    seed: u64,

    /// Reproduce correlated selector streams for the teaching experiment.
    #[clap(long)]
    correlated_selectors: bool,
}

fn run_simulation(args: &Args, config: &Config) -> Vec<(f64, f64)> {
    assert!(
        args.servers > 0 && args.senders > 0,
        "servers and senders must be positive"
    );
    assert!(
        args.time_limit.is_finite() && args.time_limit > 0.0,
        "time-limit must be finite and positive"
    );
    assert!(config.n_balancers > 0, "n_balancers must be positive");
    assert!(
        config.sync_interval.is_finite() && config.sync_interval > 0.0,
        "sync_interval must be finite and positive"
    );
    // Generate all arrivals needed by the last complete measurement window.
    let arrivals_end = args.time_limit + WINDOW;
    let mut sim = Simulation::new(args.seed);
    let ctx = sim.create_context("entry point");
    let log = Rc::new(RefCell::new(EventLog::new()));
    let mut hosts = Vec::with_capacity(args.servers as usize);
    for i in 0..args.servers {
        let slowdown = if args.different_servers && i < args.servers / 20 {
            5.0
        } else {
            1.0
        };
        hosts.push(Rc::new(RefCell::new(Host::new(
            i as usize,
            slowdown,
            log.clone(),
            Rc::new(RefCell::new(sim.create_context(format!("host_{i}")))),
        ))));
        sim.add_handler(format!("host_{i}"), hosts[i as usize].clone());
    }
    let mut balancers = Vec::with_capacity(config.n_balancers as usize);
    // Keep selector randomness separate from the workload generator.
    let mut selector_seeds = Pcg64::seed_from_u64(args.seed);
    for i in 0..config.n_balancers {
        let seed = if args.correlated_selectors {
            111
        } else {
            selector_seeds.gen()
        };
        balancers.push(Rc::new(RefCell::new(LoadBalancer::new(
            hosts.clone(),
            create_selector_by_name(&config.balancer, seed, hosts.len()),
            config.sync_interval,
            arrivals_end,
            Rc::new(RefCell::new(sim.create_context(format!("balancer_{i}")))),
        ))));
        sim.add_handler(format!("balancer_{i}"), balancers[i as usize].clone());
        ctx.emit(SyncEvent {}, sim.lookup_id(&format!("balancer_{i}")), 0.0);
    }
    let mut rng = Pcg64::seed_from_u64(args.seed);
    let mut routing_rng = Pcg64::seed_from_u64(args.seed.wrapping_add(1));
    for sender in 0..(args.senders as usize) {
        let distr = Exp::new(if args.preset == "skewed" && sender == 0 {
            50.0
        } else {
            1.0
        })
        .unwrap();
        let mut t = 0f64;
        loop {
            t += distr.sample(&mut rng) * 5.;
            if t >= arrivals_end {
                break;
            }
            let processing_time = rng.gen_range(0.5..2.5);
            ctx.emit(
                RequestArrivalEvent {
                    processing_time,
                    sender,
                },
                sim.lookup_id(&format!(
                    "balancer_{}",
                    routing_rng.gen_range(0..balancers.len())
                )),
                t,
            );
        }
    }
    // Periodic updates stop at arrivals_end. Drain every queued request so that
    // slow responses are measured in full, not truncated at a time cutoff.
    sim.step_until_no_events();
    let requests = log
        .borrow()
        .iter()
        .map(|&(end, start)| (start, end - start))
        .collect();
    windowed_p99(requests, args.time_limit)
}

fn main() {
    let args = Args::parse();
    let configs: Vec<Config> =
        serde_json::from_reader(File::open(Path::new(&args.config)).unwrap()).unwrap();
    let mut out = BufWriter::new(File::create(&args.dump).unwrap());
    for config in configs.iter() {
        let result = run_simulation(&args, config);
        out.write_all(config.balancer.as_bytes()).unwrap();
        out.write_all(b"\n").unwrap();
        out.write_all(
            result
                .iter()
                .map(|x| format!("{:.3}", x.0))
                .collect::<Vec<_>>()
                .join(",")
                .as_bytes(),
        )
        .unwrap();
        out.write_all(b"\n").unwrap();
        out.write_all(
            result
                .iter()
                .map(|x| format!("{}", x.1))
                .collect::<Vec<_>>()
                .join(",")
                .as_bytes(),
        )
        .unwrap();
        out.write_all(b"\n").unwrap();
    }
    out.flush().unwrap();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn overloaded_server_finishes_requests_after_arrival_horizon() {
        let args = Args::parse_from([
            "lab",
            "--dump",
            "unused",
            "--config",
            "unused",
            "--servers",
            "1",
            "--senders",
            "100",
            "--time-limit",
            "1",
        ]);
        let config = Config {
            balancer: "Hash".into(),
            n_balancers: 1,
            sync_interval: 0.01,
        };
        let result = run_simulation(&args, &config);
        assert!(!result.is_empty());
        assert!(result.iter().all(|(t, _)| *t <= args.time_limit));
        assert!(result[0].1 > args.time_limit + WINDOW);
        assert_eq!(result, run_simulation(&args, &config));
    }
}