use config::{Config as ConfigBuilder, File};
use serde::Deserialize;
use signal_hook::{consts, flag};
use std::env;
use std::fs;
use std::io::Write;
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::thread;
use std::time::{Duration, Instant};

#[derive(Deserialize, Debug)]
struct Config {
    #[serde(default)]
    global: GlobalConfig,
    device: Vec<DeviceConfig>,
}

#[derive(Deserialize, Debug, Clone)]
struct GlobalConfig {
    #[serde(default = "default_hysteresis")]
    spindown_hysteresis_seconds: u32,
    #[serde(default = "default_pwm_steps")]
    pwm_steps: Vec<u8>,
    #[serde(default = "default_thresholds_t_edge")]
    thresholds_t_edge: Vec<f32>,
    #[serde(default = "default_thresholds_t_junction")]
    thresholds_t_junction: Vec<f32>,
    #[serde(default = "default_thresholds_t_memory")]
    thresholds_t_memory: Vec<f32>,
}

impl Default for GlobalConfig {
    fn default() -> Self {
        Self {
            spindown_hysteresis_seconds: 30,
            pwm_steps: vec![51, 102, 153, 180, 204, 225, 255],
            thresholds_t_edge: vec![70.0, 80.0, 90.0, 92.0, 94.0, 96.0],
            thresholds_t_junction: vec![75.0, 92.0, 94.0, 97.0, 98.0, 100.0],
            thresholds_t_memory: vec![75.0, 85.0, 90.0, 91.0, 92.0, 94.0],
        }
    }
}
// somewhat boilerplaty, but want to support partial specification of global options while keeping defaults in one place.
fn default_hysteresis() -> u32 {
    GlobalConfig::default().spindown_hysteresis_seconds
}

fn default_pwm_steps() -> Vec<u8> {
    GlobalConfig::default().pwm_steps
}

fn default_thresholds_t_edge() -> Vec<f32> {
    GlobalConfig::default().thresholds_t_edge
}

fn default_thresholds_t_junction() -> Vec<f32> {
    GlobalConfig::default().thresholds_t_junction
}

fn default_thresholds_t_memory() -> Vec<f32> {
    GlobalConfig::default().thresholds_t_memory
}

fn check_global_config(gc: &GlobalConfig) -> bool {
    if !gc.pwm_steps.iter().is_sorted() {
        eprintln!("pwm steps not monotonic");
        return false;
    }
    if !gc.thresholds_t_edge.iter().is_sorted() {
        eprintln!("t_edge thresholds not sorted");
        return false;
    }
    if !gc.thresholds_t_junction.iter().is_sorted() {
        eprintln!("t_junction thresholds not sorted");
        return false;
    }
    if !gc.thresholds_t_memory.iter().is_sorted() {
        eprintln!("t_memory thresholds not sorted");
        return false;
    }
    if gc.thresholds_t_edge.len() != gc.thresholds_t_junction.len()
        || gc.thresholds_t_junction.len() != gc.thresholds_t_memory.len()
        || gc.thresholds_t_edge.len() + 1 != gc.pwm_steps.len()
    {
        eprintln!("Temperature threshold vectors need to be all of same length and number of pwm steps needs to be one more.");
        return false;
    }
    true
}

#[derive(Deserialize, Debug)]
struct DeviceConfig {
    name: String,
    pci_id: String,
    fan_control: FanControlType,
}

#[derive(Deserialize, Debug)]
#[serde(untagged)]
pub enum FanControlType {
    Ipmi { ipmi_command: String },
    Pwm { device: String, pwm_number: u32 },
}

#[derive(Debug)]
struct ExpandedDeviceConfig {
    device: String,
    power: PowerPath,
    pwm: PwmPath,
    temperature: TemperatureFiles,
}

#[derive(Default, Debug)]
struct TemperatureFiles {
    junction: String,
    overall: String,
    memory: String,
}

#[derive(Deserialize, Debug)]
struct PwmPath {
    file: String,
}

#[derive(Deserialize, Debug)]
struct PowerPath {
    file: String,
}

struct PwmControlState {
    value: u8, // 1 manual control
}

const MANUAL_CONTROL_STATE: PwmControlState = PwmControlState { value: 1 };

fn find_hwmon_for_device(pci_addr_or_device_name: &str) -> Result<String, String> {
    let hwmon_dir = Path::new("/sys/class/hwmon");

    for entry in fs::read_dir(hwmon_dir).map_err(|e| format!("HWMON read error: {}", e))? {
        let entry = entry.map_err(|e| format!("HWMON entry error: {}", e))?;
        let path = entry.path();

        if !path.is_dir() {
            continue;
        }

        let device_path = path.join("device");
        if !device_path.exists() {
            continue;
        }

        let target = fs::read_link(&device_path).map_err(|e| format!("Link read error: {}", e))?;

        if target.to_string_lossy().contains(pci_addr_or_device_name) {
            return Ok(path
                .to_str()
                .map(|s| s.to_string())
                .ok_or("Invalid hwmon name")?);
        }
    }

    Err(format!(
        "HWMON not found for device: {}",
        pci_addr_or_device_name
    ))
}

fn determine_min_pwm_setpoint(
    global_options: &GlobalConfig,
    t_edge: f32,
    t_junction: f32,
    t_memory: f32,
    previous_setpoint: u8,
) -> u8 {
    fn lookup_pwm_level(value: f32, levels: &Vec<u8>, thresholds: &Vec<f32>) -> u8 {
        let mut level = levels[0];
        for (i, &th) in thresholds.iter().enumerate() {
            if value >= th {
                level = levels[i + 1];
            } else {
                break;
            }
        }
        level
    }

    let l_edge = lookup_pwm_level(
        t_edge,
        &global_options.pwm_steps,
        &global_options.thresholds_t_edge,
    );
    let l_junction = lookup_pwm_level(
        t_junction,
        &global_options.pwm_steps,
        &global_options.thresholds_t_junction,
    );
    let l_memory = lookup_pwm_level(
        t_memory,
        &global_options.pwm_steps,
        &global_options.thresholds_t_memory,
    );
    let l_previous_minus_one = *global_options
        .pwm_steps
        .iter()
        .filter(|x| **x < previous_setpoint)
        .max()
        .unwrap_or(&0);
    [l_edge, l_junction, l_memory, l_previous_minus_one]
        .into_iter()
        .max()
        .unwrap()
}

fn monitor_temperatures(
    config: ExpandedDeviceConfig,
    global_options: GlobalConfig,
    shutdown_flag: Arc<AtomicBool>,
) {
    let initial_power_limit = read_power_limit(&config.power.file).unwrap_or(0.0);
    let mut current_power_limit = initial_power_limit;

    let junction_temp = read_temperature(&config.temperature.junction).unwrap_or(105.0);
    let overall_temp = read_temperature(&config.temperature.overall).unwrap_or(105.0);
    let memory_temp = read_temperature(&config.temperature.memory).unwrap_or(94.0);

    let mut fan_level =
        determine_min_pwm_setpoint(&global_options, overall_temp, junction_temp, memory_temp, 0);
    let mut time_fan_level_set: Instant = Instant::now();

    // Set initial fan speed
    let original_pwm_mode = enable_pwm_control_and_set_init(&config.pwm.file, fan_level).unwrap();

    loop {
        if shutdown_flag.load(Ordering::SeqCst) {
            // Set fan to 100% on exit
            restore_pwm_mode(&config.pwm.file, original_pwm_mode).unwrap();
            break;
        }

        let junction_temp = read_temperature(&config.temperature.junction).unwrap_or(105.0);
        let overall_temp = read_temperature(&config.temperature.overall).unwrap_or(105.0);
        let memory_temp = read_temperature(&config.temperature.memory).unwrap_or(94.0);

        // Check if any temperature is above critical thresholds
        if junction_temp >= 102.0 || memory_temp >= 94.0 {
            current_power_limit = current_power_limit - 3.0;
            println!(
                "High temperature detected, reducing power limit by 3 watt to {}",
                current_power_limit
            );
            set_power_limit(&config.power.file, current_power_limit).unwrap();
        }
        let new_fan_level = determine_min_pwm_setpoint(
            &global_options,
            overall_temp,
            junction_temp,
            memory_temp,
            fan_level,
        );
        if new_fan_level < fan_level
            && time_fan_level_set.elapsed()
                < Duration::from_secs(global_options.spindown_hysteresis_seconds.into())
        {
            //println!("Temperature lower but waiting for hysteresis time to expire before applying");
        } else if new_fan_level == fan_level {
            // reset hysteris
            time_fan_level_set = Instant::now();
        } else {
            println!("Adjusting PWM from {} to {}", fan_level, new_fan_level);
            fan_level = new_fan_level;
            time_fan_level_set = Instant::now();
            set_pwm(&config.pwm.file, fan_level).unwrap();
        }
        thread::sleep(Duration::from_millis(1000));
    }
}

fn read_temperature(path: &str) -> Result<f32, std::io::Error> {
    let content = fs::read_to_string(path)?;
    let value: f32 = content
        .trim()
        .parse()
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
    Ok(value as f32 * 1e-3)
}

fn enable_pwm_control_and_set_init(
    path: &str,
    value: u8,
) -> Result<PwmControlState, std::io::Error> {
    println!("Reading old status from {}", format!("{}_enable", path));
    let old_status_text = fs::read_to_string(format!("{}_enable", path))?;
    let old_status_val: u8 = old_status_text.trim().parse().map_err(|_e| {
        std::io::Error::new(std::io::ErrorKind::Other, "Failed to parse old pwm status.")
    })?;
    println!(
        "Read old status {} from {}",
        old_status_val,
        format!("{}_enable", path)
    );
    let old_status = PwmControlState {
        value: old_status_val,
    };

    std::fs::write(
        format!("{}_enable", path),
        format!("{}", MANUAL_CONTROL_STATE.value),
    )?;
    if let Err(e) = std::fs::write(path, format!("{}", value)) {
        eprintln!("Failed to write initial PWM value, restoring previous control state.");
        std::fs::write(format!("{}_enable", path), format!("{}", old_status.value))?;
        return Err(e.into());
    }
    println!("Took over fan control and set PWM to {}", value);
    Ok(old_status)
}

fn restore_pwm_mode(path: &str, previous_mode: PwmControlState) -> Result<(), std::io::Error> {
    println!("Restoring PWM control mode to {}", previous_mode.value);
    std::fs::write(
        format!("{}_enable", path),
        format!("{}", previous_mode.value),
    )
}

fn set_pwm(path: &str, value: u8) -> Result<(), std::io::Error> {
    //let mut file = fs::File::create(path)?;
    std::fs::write(path, format!("{}", value))
}

fn read_power_limit(path: &str) -> Result<f32, std::io::Error> {
    let content = fs::read_to_string(path)?;
    let value: f32 = content
        .trim()
        .parse()
        .map(|x: f32| x * 1e-6)
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
    Ok(value)
}

fn set_power_limit(path: &str, value: f32) -> Result<(), std::io::Error> {
    let mut file = fs::File::create(path)?;
    writeln!(file, "{}", (value * 1e6) as u32)
}

fn find_temperature_sensors(card_hwmon: &str) -> Result<TemperatureFiles, std::io::Error> {
    let mut sensor_idx: [u8; 3] = Default::default(); // index 0 unused by sysfs
    (0..20).for_each(|i| {
        let label_path = format!("{:}/temp{}_label", card_hwmon, i);
        if let Ok(label) = fs::read_to_string(&label_path) {
            let label = label.trim();
            println!("Read label {:} for temp {:}", &label, i);
            match label {
                "edge" => sensor_idx[0] = i,
                "junction" => sensor_idx[1] = i,
                "mem" => sensor_idx[2] = i,
                _ => (),
            }
        }
    });
    println!("Determined sensor indices {:?}", sensor_idx);
    let overall = format!("{:}/temp{:}_input", card_hwmon, sensor_idx[0]);
    let junction = format!("{:}/temp{:}_input", card_hwmon, sensor_idx[1]);
    let memory = format!("{:}/temp{:}_input", card_hwmon, sensor_idx[2]);
    if !Path::new(&overall).exists() {
        Err(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!("Overall temp file {:} not found", overall),
        ))
    } else if !Path::new(&junction).exists() {
        Err(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!("Overall temp file {:} not found", junction),
        ))
    } else if !Path::new(&memory).exists() {
        Err(std::io::Error::new(
            std::io::ErrorKind::NotFound,
            format!("Overall temp file {:} not found", memory),
        ))
    } else {
        Ok(TemperatureFiles {
            junction,
            overall,
            memory,
        })
    }
}

fn check_and_expand_config(config: &Config) -> Vec<ExpandedDeviceConfig> {
    config
        .device
        .iter()
        .map(|dev| -> ExpandedDeviceConfig {
            println!("Checking config for {:}...", &dev.name);
            let card_hwmon =
                find_hwmon_for_device(&dev.pci_id).expect("Hwmon not found, aborting...");
            let (fan_device, fan_number) = match &dev.fan_control {
                FanControlType::Pwm { device, pwm_number } => (device, pwm_number),
                _ => unimplemented!("IPMI fan control not yet fully implementd."),
            };

            let pwm_hwmon = find_hwmon_for_device(&fan_device)
                .expect("Hwmon not found for fan control device, aborting...");
            let power_file = format!("{:}/power1_cap", card_hwmon);
            let pwm_file = format!("{:}/pwm{:}", pwm_hwmon, fan_number);
            assert!(
                Path::new(&power_file).exists(),
                "Power limit control file {:} not found.",
                power_file
            );
            assert!(
                Path::new(&pwm_file).exists(),
                "Pwm control file {:} not found.",
                pwm_file
            );
            let temperature =
                find_temperature_sensors(&card_hwmon).expect("Temperature sensor not found.");
            ExpandedDeviceConfig {
                device: dev.name.clone(),
                power: PowerPath { file: power_file },
                pwm: PwmPath { file: pwm_file },
                temperature,
            }
        })
        .collect()
}

fn main() {
    let shutdown_flag = Arc::new(AtomicBool::new(false));

    // Do clean shutdown on all catchable termination signals
    let signals = [
        consts::SIGABRT,
        consts::SIGALRM,
        consts::SIGBUS,
        consts::SIGCONT,
        consts::SIGINT,
        consts::SIGQUIT,
        consts::SIGTERM,
    ];
    for sig in signals.iter() {
        flag::register(*sig, shutdown_flag.clone()).unwrap();
    }

    let args: Vec<String> = env::args().collect();
    let mut config_path = "config.toml".to_string();

    for i in 0..args.len() {
        if args[i] == "--config" && i + 1 < args.len() {
            config_path = args[i + 1].clone();
        }
    }

    let settings = ConfigBuilder::builder()
        .add_source(File::with_name(&config_path))
        .build()
        .unwrap();

    let config: Config = settings.try_deserialize().unwrap();
    println!("Config {:?}", config);
    assert!(check_global_config(&config.global), "Global options invalid. Make sure thresholds are monotonic and pwm steps one more than thresholds.");
    let expanded_configs = check_and_expand_config(&config);
    let join_handles = expanded_configs
        .into_iter()
        .map(|ec| {
            println!("Spawning fan control thread for device {:}", &ec.device);
            let gc = config.global.clone();
            let shutdown_flag = shutdown_flag.clone();
            std::thread::spawn(|| monitor_temperatures(ec, gc, shutdown_flag))
        })
        .collect::<Vec<_>>();
    join_handles.into_iter().for_each(|h| h.join().unwrap());
}
