mod bpf_skel;
use benchmark::{BenchmarkScheduler, Mode};
pub use bpf_skel::*;
pub mod bpf_intf;

mod benchmark;
mod e_core_selector;
mod energy;
mod freq;
mod model;
mod scheduler;
mod socket;

#[rustfmt::skip]
mod bpf;

use anyhow::Result;
use burn::tensor::Tensor;
use clap::{Arg, ArgAction, Command};
use model::load_model;
use scheduler::Scheduler;
use std::mem::MaybeUninit;

type Pid = i32;

fn main() -> Result<()> {
    let matches = Command::new("Energy User Space Scheduler")
        .arg(
            Arg::new("mock")
                .short('m')
                .long("mock")
                .help("Use this flag if to activate kernel module mocking")
                .action(ArgAction::SetTrue)
                .required(false),
        )
        .arg(
            Arg::new("benchmark")
                .short('b')
                .long("benchmark")
                .help("Enable benchmarking mode. Use \"e\" for e cores and \"p\" for p cores")
                .required(false)
                .value_name("mode"),
        )
        .arg(
            Arg::new("power_cap")
                .long("energy_cap")
                .help("Set a power cap for the processor")
                .required(false)
                .value_name("power in watts"),
        )
        .get_matches();

    let device = Default::default();
    let model = load_model("perf.pt");
    let tensor = Tensor::from_floats(
        [
            800., 90678., 54734., 153646., 20354478., 40948418.,
            89103105.,
            //5200., 148947., 322426., 498965., 62340773., 144451046., 41976480.,
        ],
        &device,
    );
    let result = model.forward(tensor);
    let energy: f32 = result.into_scalar();
    println!("energy: {energy}");
    // panic!();

    let power_cap = *matches.get_one::<u64>("power_cap").unwrap_or(&u64::MAX);
    let use_mocking = matches.get_flag("mock");
    let benchmark = matches.get_one::<String>("benchmark");

    // Initialize and load the scheduler.
    let mut open_object = MaybeUninit::uninit();
    let log_path = "/tmp/logs.csv";
    if let Some(mode) = benchmark {
        let mode = mode.trim().chars().next().unwrap();
        let mut sched = BenchmarkScheduler::init(&mut open_object, log_path, Mode::from(mode))?;
        sched.run()?;
        return Ok(());
    }
    loop {
        let mut sched = Scheduler::init(&mut open_object, use_mocking, power_cap)?;
        if !sched.run()?.should_restart() {
            break;
        }
    }

    Ok(())
}