#![allow(clippy::arithmetic_side_effects)]
use {
    clap::{crate_description, crate_name, crate_version, Arg, ArgMatches, Command},
    rand::{thread_rng, Rng},
    serde::{Deserialize, Serialize},
    std::{fs, io, path::PathBuf},
};

#[derive(Deserialize, Serialize, Debug)]
struct NetworkInterconnect {
    pub a: u8,
    pub b: u8,
    pub config: String,
}

#[derive(Deserialize, Serialize, Debug)]
struct NetworkTopology {
    pub partitions: Vec<u8>,
    pub interconnects: Vec<NetworkInterconnect>,
}

impl Default for NetworkTopology {
    fn default() -> Self {
        Self {
            partitions: vec![100],
            interconnects: vec![],
        }
    }
}

impl NetworkTopology {
    pub fn verify(&self) -> bool {
        let sum: u8 = self.partitions.iter().sum();
        if sum != 100 {
            return false;
        }

        for x in self.interconnects.iter() {
            if x.a as usize > self.partitions.len() || x.b as usize > self.partitions.len() {
                return false;
            }
        }

        true
    }

    pub fn new_from_stdin() -> Self {
        let mut input = String::new();
        println!("Configure partition map (must add up to 100, e.g. [70, 20, 10]):");
        let partitions_str = match io::stdin().read_line(&mut input) {
            Ok(_) => input,
            Err(error) => panic!("error: {error}"),
        };

        let partitions: Vec<u8> = serde_json::from_str(&partitions_str)
            .expect("Failed to parse input. It must be a JSON string");

        let mut interconnects: Vec<NetworkInterconnect> = vec![];

        for i in 0..partitions.len() - 1 {
            for j in i + 1..partitions.len() {
                println!("Configure interconnect ({i} <-> {j}):");
                let mut input = String::new();
                let mut interconnect_config = match io::stdin().read_line(&mut input) {
                    Ok(_) => input,
                    Err(error) => panic!("error: {error}"),
                };

                if interconnect_config.ends_with('\n') {
                    interconnect_config.pop();
                    if interconnect_config.ends_with('\r') {
                        interconnect_config.pop();
                    }
                }

                if !interconnect_config.is_empty() {
                    let interconnect = NetworkInterconnect {
                        a: i as u8,
                        b: j as u8,
                        config: interconnect_config.clone(),
                    };
                    interconnects.push(interconnect);
                    let interconnect = NetworkInterconnect {
                        a: j as u8,
                        b: i as u8,
                        config: interconnect_config,
                    };
                    interconnects.push(interconnect);
                }
            }
        }

        Self {
            partitions,
            interconnects,
        }
    }

    fn new_random(max_partitions: usize, max_packet_drop: u8, max_packet_delay: u32) -> Self {
        let mut rng = thread_rng();
        let num_partitions = rng.gen_range(0..max_partitions + 1);

        if num_partitions == 0 {
            return NetworkTopology::default();
        }

        let mut partitions = vec![];
        let mut used_partition = 0;
        for i in 0..num_partitions {
            let partition = if i == num_partitions - 1 {
                100 - used_partition
            } else {
                rng.gen_range(0..100 - used_partition - num_partitions + i)
            };
            used_partition += partition;
            partitions.push(partition as u8);
        }

        let mut interconnects: Vec<NetworkInterconnect> = vec![];
        for i in 0..partitions.len() - 1 {
            for j in i + 1..partitions.len() {
                let drop_config = if max_packet_drop > 0 {
                    let packet_drop = rng.gen_range(0..max_packet_drop + 1);
                    format!("loss {packet_drop}% 25% ")
                } else {
                    String::default()
                };

                let config = if max_packet_delay > 0 {
                    let packet_delay = rng.gen_range(0..max_packet_delay + 1);
                    format!("{drop_config}delay {packet_delay}ms 10ms")
                } else {
                    drop_config
                };

                let interconnect = NetworkInterconnect {
                    a: i as u8,
                    b: j as u8,
                    config: config.clone(),
                };
                interconnects.push(interconnect);
                let interconnect = NetworkInterconnect {
                    a: j as u8,
                    b: i as u8,
                    config,
                };
                interconnects.push(interconnect);
            }
        }
        Self {
            partitions,
            interconnects,
        }
    }
}

fn run(
    cmd: &str,
    args: &[&str],
    launch_err_msg: &str,
    status_err_msg: &str,
    ignore_err: bool,
) -> bool {
    println!("Running {:?}", std::process::Command::new(cmd).args(args));
    let output = std::process::Command::new(cmd)
        .args(args)
        .output()
        .expect(launch_err_msg);

    if ignore_err {
        return true;
    }

    if !output.status.success() {
        eprintln!(
            "{} command failed with exit code: {}",
            status_err_msg, output.status
        );
        use std::str::from_utf8;
        println!("stdout: {}", from_utf8(&output.stdout).unwrap_or("?"));
        println!("stderr: {}", from_utf8(&output.stderr).unwrap_or("?"));
        false
    } else {
        true
    }
}

fn insert_iptables_rule(tos: u8) -> bool {
    let my_tos = tos.to_string();

    // iptables -t mangle -A PREROUTING -p udp -j TOS --set-tos <my_parition_index>
    run(
        "iptables",
        &[
            "-t",
            "mangle",
            "-A",
            "OUTPUT",
            "-p",
            "udp",
            "-j",
            "TOS",
            "--set-tos",
            my_tos.as_str(),
        ],
        "Failed to add iptables rule",
        "iptables",
        false,
    )
}

fn flush_iptables_rule() {
    run(
        "iptables",
        &["-F", "-t", "mangle"],
        "Failed to flush iptables",
        "iptables flush",
        true,
    );
}

fn setup_ifb(interface: &str) -> bool {
    // modprobe ifb numifbs=1
    run(
        "modprobe",
        &[
            "ifb", "numifbs=1",
        ],
        "Failed to load ifb module",
        "modprobe ifb numifbs=1",
        false
    ) &&
    // ip link set dev ifb0 up
    run(
        "ip",
        &[
            "link", "set", "dev", "ifb0", "up"
        ],
        "Failed to bring ifb0 online",
        "ip link set dev ifb0 up",
        false
    ) &&
    // tc qdisc add dev <if> handle ffff: ingress
    run(
        "tc",
        &[
            "qdisc", "add", "dev", interface, "handle", "ffff:", "ingress"
        ],
        "Failed to setup ingress qdisc",
        "tc qdisc add dev <if> handle ffff: ingress",
        false
    )
    &&
    // tc filter add dev <if> parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev ifb0
    run(
        "tc",
        &[
            "filter", "add", "dev", interface, "parent", "ffff:", "protocol", "ip", "u32", "match", "u32", "0", "0", "action", "mirred", "egress", "redirect", "dev", "ifb0"
        ],
        "Failed to redirect ingress traffc",
        "tc filter add dev <if> parent ffff: protocol ip u32 match u32 0 0 action mirred egress redirect dev ifb0",
        false
    )
}

fn delete_ifb(interface: &str) -> bool {
    run(
        "tc",
        &[
            "qdisc", "delete", "dev", interface, "handle", "ffff:", "ingress",
        ],
        "Failed to setup ingress qdisc",
        "tc qdisc delete dev <if> handle ffff: ingress",
        true,
    ) && run(
        "modprobe",
        &["ifb", "--remove"],
        "Failed to delete ifb module",
        "modprobe ifb --remove",
        true,
    )
}

fn insert_tc_ifb_root(num_bands: &str) -> bool {
    // tc qdisc add dev ifb0 root handle 1: prio bands <num_bands>
    run(
        "tc",
        &[
            "qdisc", "add", "dev", "ifb0", "root", "handle", "1:", "prio", "bands", num_bands,
        ],
        "Failed to add root ifb qdisc",
        "tc qdisc add dev ifb0 root handle 1: prio bands <num_bands>",
        false,
    )
}

fn insert_tc_ifb_netem(class: &str, handle: &str, filter: &str) -> bool {
    let mut filters: Vec<&str> = filter.split(' ').collect();
    let mut args = vec![
        "qdisc", "add", "dev", "ifb0", "parent", class, "handle", handle, "netem",
    ];
    args.append(&mut filters);
    // tc qdisc add dev ifb0 parent <class> handle <handle> netem <filters>
    run("tc", &args, "Failed to add tc child", "tc add child", false)
}

fn insert_tos_ifb_filter(class: &str, tos: &str) -> bool {
    // tc filter add dev ifb0 protocol ip parent 1: prio 1 u32 match ip tos <tos> 0xff flowid <class>
    run(
        "tc",
        &[
            "filter", "add", "dev", "ifb0", "protocol", "ip", "parent", "1:", "prio", "1",
            "u32", "match", "ip", "tos", tos, "0xff", "flowid", class,
        ],
        "Failed to add tos filter",
        "tc filter add dev ifb0 protocol ip parent 1: prio 1 u32 match ip tos <tos> 0xff flowid <class>",
        false,
    )
}

fn insert_default_ifb_filter(class: &str) -> bool {
    // tc filter add dev ifb0 parent 1: protocol all prio 2 u32 match u32 0 0 flowid 1:<class>
    run(
        "tc",
        &[
            "filter", "add", "dev", "ifb0", "parent", "1:", "protocol", "all", "prio", "2", "u32",
            "match", "u32", "0", "0", "flowid", class,
        ],
        "Failed to add catch-all filter",
        "tc filter add dev ifb0 parent 1: protocol all prio 2 u32 match u32 0 0 flowid 1:<class>",
        false,
    )
}

fn delete_all_filters(interface: &str) {
    // tc filter delete dev <if>
    run(
        "tc",
        &["filter", "delete", "dev", interface],
        "Failed to delete all filters",
        "tc delete all filters",
        true,
    );
}

fn identify_my_partition(partitions: &[u8], index: u64, size: u64) -> usize {
    let mut my_partition = 0;
    let mut watermark = 0;
    for (i, p) in partitions.iter().enumerate() {
        watermark += *p;
        if u64::from(watermark) >= index * 100 / size {
            my_partition = i;
            break;
        }
    }

    my_partition
}

fn partition_id_to_tos(partition: usize) -> u8 {
    if partition < 4 {
        2u8.pow(partition as u32 + 1)
    } else {
        0
    }
}

fn shape_network(matches: &ArgMatches) {
    let config_path = PathBuf::from(matches.value_of_t_or_exit::<String>("file"));
    let config = fs::read_to_string(config_path).expect("Unable to read config file");
    let topology: NetworkTopology =
        serde_json::from_str(&config).expect("Failed to parse log as JSON");
    let interface: String = matches.value_of_t_or_exit("iface");
    let network_size: u64 = matches.value_of_t_or_exit("size");
    let my_index: u64 = matches.value_of_t_or_exit("position");
    if !shape_network_steps(&topology, &interface, network_size, my_index) {
        delete_ifb(interface.as_str());
        flush_iptables_rule();
    }
}

fn shape_network_steps(
    topology: &NetworkTopology,
    interface: &str,
    network_size: u64,
    my_index: u64,
) -> bool {
    // Integrity checks
    assert!(topology.verify(), "Failed to verify the configuration file");
    assert!(my_index < network_size);

    // Figure out partition we belong in
    let my_partition = identify_my_partition(&topology.partitions, my_index + 1, network_size);

    // Clear any lingering state
    println!(
        "my_index: {}, network_size: {}, partitions: {:?}",
        my_index, network_size, topology.partitions
    );
    println!("My partition is {my_partition}");

    cleanup_network(interface);

    // Mark egress packets with our partition id
    if !insert_iptables_rule(partition_id_to_tos(my_partition)) {
        return false;
    }

    let num_bands = topology.partitions.len() + 1;
    let default_filter_class = format!("1:{num_bands}");
    if !topology.interconnects.is_empty() {
        let num_bands_str = num_bands.to_string();
        // Redirect ingress traffic to the virtual interface ifb0 so we can
        // apply egress rules
        if !setup_ifb(interface)
            // Setup root qdisc on ifb0
            || !insert_tc_ifb_root(num_bands_str.as_str())
            // Catch all so regular traffic/traffic within the same partition
            // is not filtered out
            || !insert_default_ifb_filter(default_filter_class.as_str())
        {
            return false;
        }
    }

    println!("Setting up interconnects");
    for i in &topology.interconnects {
        if i.b as usize == my_partition {
            println!("interconnects: {i:#?}");
            let tos = partition_id_to_tos(i.a as usize);
            if tos == 0 {
                println!("Incorrect value of TOS/Partition in config {}", i.a);
                return false;
            }
            let tos_string = tos.to_string();
            // First valid class is 1:1
            let class = format!("1:{}", i.a + 1);
            if !insert_tc_ifb_netem(class.as_str(), tos_string.as_str(), i.config.as_str()) {
                return false;
            }

            if !insert_tos_ifb_filter(class.as_str(), tos_string.as_str()) {
                return false;
            }
        }
    }

    true
}

fn parse_interface(interfaces: &str) -> &str {
    for line in interfaces.lines() {
        if line != "ifb0" {
            return line;
        }
    }

    panic!("No valid interfaces");
}

fn cleanup_network(interface: &str) {
    delete_all_filters("ifb0");
    delete_ifb(interface);
    flush_iptables_rule();
}

fn configure(matches: &ArgMatches) {
    let config = if !matches.is_present("random") {
        NetworkTopology::new_from_stdin()
    } else {
        let max_partitions: usize = matches.value_of_t("max-partitions").unwrap_or(4);
        let max_drop: u8 = matches.value_of_t("max-drop").unwrap_or(100);
        let max_delay: u32 = matches.value_of_t("max-delay").unwrap_or(50);
        NetworkTopology::new_random(max_partitions, max_drop, max_delay)
    };

    assert!(config.verify(), "Failed to verify the configuration");

    let topology = serde_json::to_string(&config).expect("Failed to write as JSON");

    println!("{topology}");
}

fn main() {
    solana_logger::setup();

    let matches = Command::new(crate_name!())
        .about(crate_description!())
        .version(crate_version!())
        .subcommand(
            Command::new("shape")
                .about("Shape the network using config file")
                .arg(
                    Arg::new("file")
                        .short('f')
                        .long("file")
                        .value_name("config file")
                        .takes_value(true)
                        .required(true)
                        .help("Location of the network config file"),
                )
                .arg(
                    Arg::new("size")
                        .short('s')
                        .long("size")
                        .value_name("network size")
                        .takes_value(true)
                        .required(true)
                        .help("Number of nodes in the network"),
                )
                .arg(
                    Arg::new("iface")
                        .short('i')
                        .long("iface")
                        .value_name("network interface name")
                        .takes_value(true)
                        .required(true)
                        .help("Name of network interface"),
                )
                .arg(
                    Arg::new("position")
                        .short('p')
                        .long("position")
                        .value_name("position of node")
                        .takes_value(true)
                        .required(true)
                        .help("Position of current node in the network"),
                ),
        )
        .subcommand(
            Command::new("cleanup")
                .about("Remove the network filters using config file")
                .arg(
                    Arg::new("file")
                        .short('f')
                        .long("file")
                        .value_name("config file")
                        .takes_value(true)
                        .required(true)
                        .help("Location of the network config file"),
                )
                .arg(
                    Arg::new("size")
                        .short('s')
                        .long("size")
                        .value_name("network size")
                        .takes_value(true)
                        .required(true)
                        .help("Number of nodes in the network"),
                )
                .arg(
                    Arg::new("iface")
                        .short('i')
                        .long("iface")
                        .value_name("network interface name")
                        .takes_value(true)
                        .required(true)
                        .help("Name of network interface"),
                )
                .arg(
                    Arg::new("position")
                        .short('p')
                        .long("position")
                        .value_name("position of node")
                        .takes_value(true)
                        .required(true)
                        .help("Position of current node in the network"),
                ),
        )
        .subcommand(
            Command::new("configure")
                .about("Generate a config file")
                .arg(
                    Arg::new("random")
                        .short('r')
                        .long("random")
                        .required(false)
                        .help("Generate a random config file"),
                )
                .arg(
                    Arg::new("max-partitions")
                        .short('p')
                        .long("max-partitions")
                        .value_name("count")
                        .takes_value(true)
                        .required(false)
                        .help("Maximum number of partitions. Used only with random configuration generation"),
                )
                .arg(
                    Arg::new("max-drop")
                        .short('d')
                        .long("max-drop")
                        .value_name("percentage")
                        .takes_value(true)
                        .required(false)
                        .help("Maximum amount of packet drop. Used only with random configuration generation"),
                )
                .arg(
                    Arg::new("max-delay")
                        .short('y')
                        .long("max-delay")
                        .value_name("ms")
                        .takes_value(true)
                        .required(false)
                        .help("Maximum amount of packet delay. Used only with random configuration generation"),
                ),
        )
        .get_matches();

    match matches.subcommand() {
        Some(("shape", args_matches)) => shape_network(args_matches),
        Some(("cleanup", args_matches)) => {
            let interfaces: String = args_matches.value_of_t_or_exit("iface");
            let iface = parse_interface(&interfaces);
            cleanup_network(iface)
        }
        Some(("configure", args_matches)) => configure(args_matches),
        _ => {}
    };
}