mod sensors;

use self::sensors::Sensors;
use serde::Deserialize;
use std::cmp::Ordering;
use std::fs;
use std::io;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::str::FromStr;
use std::thread;
use std::time::Duration;
use std::sync::Mutex;
use std::collections::HashSet;
use lazy_static::lazy_static;
use zbus::{blocking::Connection, proxy, Result, proxy::CacheProperties};

#[proxy(
    interface = "com.steampowered.SteamOSManager1.BatteryChargeLimit1",
    default_service = "com.steampowered.SteamOSManager1",
    default_path = "/com/steampowered/SteamOSManager1"
)]
trait ChargeLevel {
    #[zbus(property)]
    fn max_charge_level(&self) -> Result<i32>;
}

fn get_maxchargelevel_dbus(proxy: &ChargeLevelProxyBlocking) -> Result<i32> {
    let reply = proxy.max_charge_level().unwrap_or(-1);
    Ok(reply)
}

#[derive(Deserialize)]
struct Config {
    request_shutdown_battery_percent: Option<f64>,
    force_shutdown_timeout_secs: Option<f64>,
}

lazy_static! {
    static ref failed: Mutex<HashSet<String>> = Default::default();
}

fn read_battery_string(path_bat: &PathBuf, var_name: &str) -> Option<String> {
    let path = format!("{}/{var_name}", path_bat.display());
    match fs::read_to_string(&path) {
        Err(err) => {
            if !failed.lock().unwrap().contains(&path) {
                eprintln!("read {path}: {err}");
                failed.lock().unwrap().insert(path);
            }
            None
        }
        Ok(string) => Some(string.trim().to_owned()),
    }
}

fn read_battery_f64(path_bat: &PathBuf, var_name: &str) -> Option<f64> {
    let path = format!("{}/{var_name}", path_bat.display());
    match fs::read_to_string(&path) {
        Err(err) => {
            if !failed.lock().unwrap().contains(&path) {
                eprintln!("read {path}: {err}");
                failed.lock().unwrap().insert(path);
            }
            None
        }
        Ok(string) => match f64::from_str(string.trim()) {
            Err(err) => {
                eprintln!("read {path}: {err}");
                None
            }
            Ok(val) => {
                if !val.is_finite() {
                    eprintln!("read {path}: {val} is not finite");
                    None
                } else {
                    Some(val)
                }
            }
        },
    }
}

fn read_battery_maxchargelevel(path: &str) -> Option<f64> {
    // retry 3 times, as there seems to be a strange bug in which some
    // /sys files sometimes disappear, so not adding to the problem by
    // also failing and adding noise to the logs
    for _i in 1..3 {
	let bat_maxchargelevel_from_file = fs::read_to_string(path).unwrap_or("-1.0".to_string());
	let bat_maxchargelevel = i32::from_str(&bat_maxchargelevel_from_file.trim()).unwrap_or(-1);

	if bat_maxchargelevel == 0 {
	    // limit is disabled, returning 100% instead
	    return Some(100.0);
	}
	else if bat_maxchargelevel > 0 {
	    // success, returning supposedly good value
	    return Some(bat_maxchargelevel as f64);
	}
	else {
	    // problem, sleep and retry
	    thread::sleep(Duration::from_millis(333));
	}
    }

    // default
    if !failed.lock().unwrap().contains(path) {
	eprintln!("read '{path}': could not read from file 3 times in a row");
        failed.lock().unwrap().insert(path.to_string());
    }
    None
}

fn write_str(dir_path: &str, var_name: &str, val: Option<&str>) {
    let val = match val {
        Some(val) => val,
        None => return,
    };

    if let Err(err) = fs::create_dir(dir_path) {
        if err.kind() != io::ErrorKind::AlreadyExists {
            eprintln!("mkdir {dir_path}: {err}");
            return;
        }
    }

    // Write to a temporary path first.
    let dot_path = format!("{dir_path}/.{var_name}");
    if let Err(err) = fs::write(&dot_path, format!("{val}\n")) {
        eprintln!("write {dot_path}: {err}");
        return;
    }

    // Then move into place for atomicity.
    let final_path = format!("{dir_path}/{var_name}");
    if let Err(err) = fs::rename(&dot_path, &final_path) {
        eprintln!("rename {dot_path} -> {final_path}: {err}");
    }
}

fn write_f64(dir_path: &str, var_name: &str, val: Option<f64>) {
    if let Some(val) = val {
        write_str(dir_path, var_name, Some(&val.to_string()))
    }
}

fn main() {
    // Mains/AC
    let mut path_ac = PathBuf::from("");
    let power_supply_paths = fs::read_dir("/sys/class/power_supply/").unwrap();
    for ps in power_supply_paths {
	let path_string_test_base = PathBuf::from(ps.unwrap().path());
	let path_string_test = format!("{}/type", path_string_test_base.display());
	let path_test = Path::new(&path_string_test);
	if ! path_test.exists() {
	    continue;
	}
	let path_test_type: String = fs::read_to_string(path_test).expect("Cannot read path");
	if path_test_type.contains("Mains") {
	    path_ac = PathBuf::from(path_string_test_base);
	    println!("Found AC power supply: '{}'", path_ac.display());
	    break;
	}
    }
    if ! path_ac.exists() {
	println!("Warning: Could not find device for AC/Mains, some functionality might be missing or not accurate.");
    }

    // Try to find reasonable BATn to use (stop at the first),
    // otherwise it's a system without battery -- bail-out
    let mut path_bat = PathBuf::from("");
    for i in 0..9 {
	let path_string_test_base = format!("/sys/class/power_supply/BAT{i}");
	let path_string_test = format!("{path_string_test_base}/type");
	let path_bat_test = Path::new(&path_string_test);
	if ! path_bat_test.exists() {
	    continue;
	}

	let path_bat_test_type: String = fs::read_to_string(path_bat_test).expect("Cannot read path");
	if path_bat_test_type.contains("Battery") {
	    path_bat = PathBuf::from(path_string_test_base);
	    println!("Found battery: {}", path_bat.display());
	    break;
	}
    }
    if ! path_bat.exists() {
	println!("This system does not use batteries, stopping.");
	return;
    }

    // Some files that the code further below will attempt to read
    // every second (not all devices might provide them, probably
    // better to keep running for partial functionality than stopping
    // completely)
    let bat_values_filenames = vec!["status", "voltage_min_design", "voltage_now"];
    for expected_file in bat_values_filenames.into_iter() {
	let path_expected_file = PathBuf::from(format!("{}/{expected_file}", path_bat.display()));
	if ! path_expected_file.exists() {
	    println!("Warning: missing expected file: {}", path_expected_file.display());
	}
    }
    // for the following files, names vary between charge_full/now
    // (SteamDeck for example) and energy_full/now
    let mut files_named_charge = true;
    let bat_values_filenames_charge = vec!["charge_full", "charge_now"];
    for expected_file in bat_values_filenames_charge.into_iter() {
	let path_expected_file = PathBuf::from(format!("{}/{expected_file}", path_bat.display()));
	if ! path_expected_file.exists() {
	    // assume files are named energy_*
	    files_named_charge = false;
	    let expected_file_subst = expected_file.replace("charge_", "energy_");
	    let path_expected_file_subst = PathBuf::from(format!("{}/{expected_file_subst}", path_bat.display()));
	    if ! path_expected_file_subst.exists() {
		println!("Warning: missing expected files: {} or {}", path_expected_file.display(), path_expected_file_subst.display());
	    }
	    else {
		println!("Info: using {} (instead of '{}')", path_expected_file_subst.display(), expected_file);
	    }
	}
    }
    // the following name varies between current_now and power_now
    let mut files_named_current = true;
    let bat_values_filenames_current = vec!["current_now"];
    for expected_file in bat_values_filenames_current.into_iter() {
	let path_expected_file = PathBuf::from(format!("{}/{expected_file}", path_bat.display()));
	if ! path_expected_file.exists() {
	    // assume files are named power_*
	    files_named_current = false;
	    let expected_file_subst = expected_file.replace("current_", "power_");
	    let path_expected_file_subst = PathBuf::from(format!("{}/{expected_file_subst}", path_bat.display()));
	    if ! path_expected_file_subst.exists() {
		println!("Warning: missing expected files: {} or {}", path_expected_file.display(), path_expected_file_subst.display());
	    }
	    else {
		println!("Info: using {} (instead of '{}')", path_expected_file_subst.display(), expected_file);
	    }
	}
    }

    // MaxChargeLevel files
    let maxchargelevel_path_std = path_bat.display().to_string() + "/charge_control_end_threshold";
    let maxchargelevel_filenames = vec![
	// SteamDeck, LCD and OLED models
	"/sys/devices/pci0000:00/0000:00:14.3/PNP0C09:00/VLV0100:00/steamdeck-hwmon/hwmon/hwmon3/max_battery_charge_level",
	// generic value supported by e.g. many consumer laptops
	&maxchargelevel_path_std,
    ];
    let mut path_maxchargelevel_file = PathBuf::from("");
    for maxchargelevel_file in maxchargelevel_filenames.into_iter() {
	path_maxchargelevel_file = PathBuf::from(maxchargelevel_file);
	if path_maxchargelevel_file.exists() {
	    println!("Info: using {} file for reading battery's MaxChargeLevel feature", path_maxchargelevel_file.display());
	    break;
	}
	else {
	    // reset to default for later use (empty means that file was not found)
	    path_maxchargelevel_file = PathBuf::from("");
	}
    }
    let path_maxchargelevel_file_found = if path_maxchargelevel_file.display().to_string().is_empty() {
	println!("Warning: cound not find suitable file for reading battery's MaxChargeLevel feature, assuming MaxChargeLevel=100%");
	false
    }
    else {
	true
    };

    // Read /etc/vpower.toml
    let config_path = "/etc/vpower.toml";
    let mut request_shutdown_battery_percent = 0.49999998;
    let mut force_shutdown_timeout_secs = 10.0;

    match fs::read(config_path) {
        Err(err) => eprintln!("read {config_path}: {err}"),

        Ok(bytes) => match toml::from_slice::<Config>(&bytes) {
            Err(err) => eprintln!("read {config_path}: {err}"),

            Ok(config) => {
                if let Some(value) = config.request_shutdown_battery_percent {
                    request_shutdown_battery_percent = value;
                }
                if let Some(value) = config.force_shutdown_timeout_secs {
                    force_shutdown_timeout_secs = value;
                }
            }
        },
    }

    println!("request_shutdown_battery_percent: {request_shutdown_battery_percent}");
    println!("force_shutdown_timeout_secs: {force_shutdown_timeout_secs}");

    // Initialize libsensors.
    let sensors = Sensors::new();

    // Keep for heuristics.
    let mut prev_ac_status: Option<&str> = None;
    let mut prev_battery_percent: Option<f64> = None;

    let mut last_bat_maxchargelevel = -999.9;

    let mut loop_counter : u64 = 0;
    let mut attempts_to_read_current_now : u64 = 0;
    let mut attempts_to_read_power_now : u64 = 0;
    let mut failed_to_read_current_now : u64 = 0;
    let mut failed_to_read_power_now : u64 = 0;

    // Start.
    println!("Running.");

    // Every second:
    loop {
	loop_counter += 1;

	// Get max charge battery level, if set
	let mut bat_maxchargelevel = match path_maxchargelevel_file_found {
	    false => 100.0,
	    true  => match read_battery_maxchargelevel(&path_maxchargelevel_file.display().to_string()) {
		None       => -999.9,
		Some(val)  => val
	    },
	};

	// sanity check, if out of bounds either take from previous
	// value (if looks ok-ish) or otherwise clamp to sane default
	if bat_maxchargelevel < 0.0 || bat_maxchargelevel > 100.0 {
	    if last_bat_maxchargelevel >= 0.0 && last_bat_maxchargelevel <= 100.0 {
		bat_maxchargelevel = last_bat_maxchargelevel;
	    }
	    else {
		bat_maxchargelevel = 100.0;
	    }
	}

	// update value for next iteration
	if bat_maxchargelevel != last_bat_maxchargelevel {
	    last_bat_maxchargelevel = bat_maxchargelevel;

	    // print new detected value, skipping first time (uninitialized)
	    if last_bat_maxchargelevel >= 0.0 {
		println!("New MaxChargeLevel value detected for battery = '{}'", last_bat_maxchargelevel);
	    }
	}

        // Read battery variables.
	let (charge_full, charge_now) = if files_named_charge {
	    // SteamDeck (and others)
            ( read_battery_f64(&path_bat, "charge_full"), read_battery_f64(&path_bat, "charge_now") )
	} else {
	    // Units compared to charge_* files are different, but
	    // these are used in values as ratios =now/full or
	    // percentages, so should be fine as long as it's not
	    // mixed or used in other ways
            ( read_battery_f64(&path_bat, "energy_full"), read_battery_f64(&path_bat, "energy_now") )
	};
        let (current_now, power_now_from_file) = if files_named_current {
	    // SteamDeck (and others)
	    attempts_to_read_current_now += 1;
	    let current_now_value = match read_battery_f64(&path_bat, "current_now") {
		Some(v) => v.abs(), // use absolute value
		None => {
		    failed_to_read_current_now += 1;
		    if (failed_to_read_current_now % 60) == 1 {
			eprintln!("read {}/current_now: {} failed out of {} attempts, loop counter={loop_counter} (throttled message)",
				  path_bat.display(),
				  failed_to_read_current_now,
				  attempts_to_read_current_now);
		    }
		    0.0
		}
	    };
	    ( Some(current_now_value), None )
	}
	else {
	    attempts_to_read_power_now += 1;
	    let power_now_value = match read_battery_f64(&path_bat, "power_now") {
		Some(v) => v,
		None => {
		    failed_to_read_power_now += 1;
		    if (failed_to_read_power_now % 60) == 1 {
			eprintln!("read {}/power_now: {} failed out of {} attempts, loop counter={loop_counter} (throttled message)",
				  path_bat.display(),
				  failed_to_read_power_now,
				  attempts_to_read_power_now);
		    }
		    0.0
		}
	    };
	    ( None, Some(power_now_value) )
	};
        let pdam = sensors.pdam();
        let pdcs = sensors.pdcs();
        let pdvl = sensors.pdvl();
        let status = read_battery_string(&path_bat, "status");
        let voltage_min_design = read_battery_f64(&path_bat, "voltage_min_design");
        let voltage_now = read_battery_f64(&path_bat, "voltage_now");

        // Derive battery variables.
        let charge_shutdown = charge_full.map(|charge_full| {
            let rsbp = request_shutdown_battery_percent;
            charge_full * (rsbp / 100.0)
        });

        let power_now = match (voltage_now, current_now, power_now_from_file) {
            (Some(voltage_now), Some(current_now), _) => Some(voltage_now * current_now),
            (Some(voltage_now), None, Some(power_now_from_file)) => Some(voltage_now * power_now_from_file),
            (Some(voltage_now), None, None) => Some(voltage_now * 0.0),
            _ => None,
        };

        // Calculate ac_status.
        let ac_status = if let Some(pdcs) = pdcs {
            let connected = (pdcs & (1 << 0)) != 0;
            let sink = (pdcs & (1 << 4)) == 0;
            if connected && sink {
                let was_disconnected = prev_ac_status == Some("Disconnected");
                let pd_power = match (pdvl, pdam) {
                    (Some(pdvl), Some(pdam)) => pdvl * pdam, // Watts.
                    _ => 0.0,
                };

                // Basically all power supplies get reported as low power for ~0.5 seconds
                // after connecting, so ignore it on the first iteration after connecting.
                if !was_disconnected && pd_power > 0.0 && pd_power < 30.0 {
                    Some("Connected slow")
                } else {
                    Some("Connected")
                }
            } else {
                Some("Disconnected")
            }
        } else {
            let ac = read_battery_string(&path_ac, "online");
            match ac.as_deref() {
                Some("0") => Some("Disconnected"),
                Some("1") => Some("Connected"),
                None => {
                    match status.as_deref() {
                        Some("Full" | "Charging") => Some("Connected"),
                        Some("Discharging") => Some("Disconnected"),
                        _ => None,
                    }
                },
		_ => Some("Disconnected"),
            }
        };

        // Calculate battery_percent.
        let battery_percent = match (charge_now, charge_full) {
            (Some(charge_now), Some(charge_full)) => Some(charge_now / charge_full * 100.0),
            _ => None,
        };
	let battery_reached_maxchargelevel : bool = battery_percent > Some(f64::from(bat_maxchargelevel) - 0.51);

        // Calculate battery_status.
        let battery_status = match (ac_status, status.as_deref()) {
            (_, Some("Full")) => Some("Full"),
            (_, Some("Discharging")) => Some("Discharging"),
	    // Connected to AC/Mains but Battery 'Not charging', whether "Max Charge Level" reached (represented as "Full") or not
            (Some("Connected"), Some("Not charging")) =>
		if battery_reached_maxchargelevel { Some("Full") } else { Some("Not charging") },
	    // Connected to AC/Mains and Battery 'Charging', whether "Max Charge Level" reached (="Full"), otherwise "Charging"
            (Some("Connected"), Some("Charging")) =>
		if battery_reached_maxchargelevel { Some("Full") } else { Some("Charging") },
            _ => {
                // Probably "Unknown" or "Not charging". Use heuristics as a fallback.
                let ordering = match (battery_percent, prev_battery_percent) {
                    (Some(lhs), Some(rhs)) => lhs.partial_cmp(&rhs),
                    _ => None,
                };
                match ordering {
                    Some(Ordering::Less) => Some("Discharging"),
                    Some(Ordering::Greater) => Some("Charging"),
                    _ => {
                        if battery_percent.unwrap_or(0.0) >= 89.5 {
                            // Some batteries won't charge when plugged in above ~90%.
                            // We call this "Full".
                            Some("Full")
                        } else {
                            None
                        }
                    }
                }
            }
        };

        // Calculate secs_until_battery_full.
        let vars = (charge_full, charge_now, voltage_min_design, power_now);
        let secs_until_battery_full = match vars {
            (Some(charge_full), Some(charge_now), Some(voltage_min_design), Some(power_now)) => {
		let charge_maxlevel = charge_full * (bat_maxchargelevel / 100.0);
                let charge_delta = if charge_now < charge_maxlevel { charge_maxlevel - charge_now } else { 0.0 };
                let hours = if charge_delta == 0.0 { 0.0 } else { charge_delta * voltage_min_design / power_now };
                Some(hours * 3600.0)
            }
            _ => None,
        };

        // Calcuate secs_until_shutdown_request.
        let vars = (charge_now, charge_shutdown, voltage_min_design, power_now);
        let secs_until_shutdown_request = match vars {
            (
                Some(charge_now),
                Some(charge_shutdown),
                Some(voltage_min_design),
                Some(power_now),
            ) => {
                if charge_now > charge_shutdown {
                    let charge_delta = charge_now - charge_shutdown;
                    let hours = charge_delta * voltage_min_design / power_now;
                    Some(hours * 3600.0)
                } else {
                    match ac_status {
                        // Avoid shutdown request while connected.
                        Some("Connected") => Some(1.0),
                        _ => Some(0.0),
                    }
                }
            }
            _ => None,
        };

        // Write to /run/vpower/*
        let dir_path = "/run/vpower";
        write_str(dir_path, "ac_status", ac_status);
        write_f64(dir_path, "battery_percent", battery_percent);
        write_str(dir_path, "battery_status", battery_status);

        let val = secs_until_battery_full;
        write_f64(dir_path, "secs_until_battery_full", val);

        let val = secs_until_shutdown_request;
        write_f64(dir_path, "secs_until_shutdown_request", val);

        // Force shutdown after timeout.
        if secs_until_shutdown_request.map_or(false, |x| x == 0.0) {
            println!("Reached {request_shutdown_battery_percent}% battery.");
            println!("Forcing shutdown in {force_shutdown_timeout_secs} seconds.");
            thread::sleep(Duration::from_secs_f64(force_shutdown_timeout_secs));

            println!("Shutting down now.");
            match Command::new("poweroff").status() {
                Err(err) => panic!("poweroff: {err}"),
                Ok(status) => match status.success() {
                    false => panic!("poweroff: {status}"),
                    true => return,
                },
            }
        }

        // Update prev_*.
        prev_ac_status = ac_status;
        prev_battery_percent = battery_percent;

        // Sleep until next iteration.
        thread::sleep(Duration::from_secs(1));
    }
}
