// SPDX-License-Identifier: GPL-2.0
//
// Copyright (c) 2024-2025 Vladislav Nepogodin <vnepogodin@cachyos.org>

// This software may be used and distributed according to the terms of the
// GNU General Public License version 2.

use std::collections::HashMap;
use std::fs;
use std::path::Path;

use anyhow::Result;
use serde::Deserialize;
use serde::Serialize;

use crate::SchedMode;
use crate::SupportedSched;

#[derive(Debug, PartialEq, Default, Serialize, Deserialize)]
#[serde(default)]
pub struct Config {
    pub default_sched: Option<SupportedSched>,
    pub default_mode: Option<SchedMode>,
    pub scheds: HashMap<String, Sched>,
}

#[derive(Debug, PartialEq, Default, Serialize, Deserialize)]
pub struct Sched {
    pub auto_mode: Option<Vec<String>>,
    pub gaming_mode: Option<Vec<String>>,
    pub lowlatency_mode: Option<Vec<String>>,
    pub powersave_mode: Option<Vec<String>>,
    pub server_mode: Option<Vec<String>>,
}

/// Initialize config from first found config path, overwise fallback to default config
///
/// # Errors
///
/// This function will return an error if a config file is found but fails to be parsed.
pub fn init_config() -> Result<Config> {
    if let Ok(config_path) = get_config_path() {
        parse_config_file(&config_path)
    } else {
        Ok(get_default_config())
    }
}

/// Parses the config file at the given path.
///
/// # Errors
///
/// This function will return an error if:
/// - The file cannot be read (e.g., permissions, not found).
/// - The file content is empty.
/// - The file content is not valid TOML.
pub fn parse_config_file(filepath: &str) -> Result<Config> {
    let file_content = fs::read_to_string(filepath)?;
    parse_config_content(&file_content)
}

/// Searches for and returns the path to the configuration file.
///
/// # Errors
///
/// This function will return an error if no config file is found in any of the
/// predefined locations.
pub fn get_config_path() -> Result<String> {
    let vendordir = option_env!("VENDORDIR").unwrap_or("/usr/share");
    // Search in system directories
    let check_paths = [
        // locations for user config
        "/etc/scx_loader/config.toml".to_owned(),
        "/etc/scx_loader.toml".to_owned(),
        // locations for distributions to ship default configuration
        format!("{vendordir}/scx_loader/config.toml").to_owned(),
        format!("{vendordir}/scx_loader.toml").to_owned(),
    ];
    for check_path in check_paths {
        if !Path::new(&check_path).exists() {
            continue;
        }
        // we found config path
        return Ok(check_path);
    }

    anyhow::bail!("Failed to find config!");
}

fn parse_config_content(file_content: &str) -> Result<Config> {
    if file_content.is_empty() {
        anyhow::bail!("The config file is empty!")
    }
    let config: Config = toml::from_str(file_content)?;
    Ok(config)
}

pub fn get_default_config() -> Config {
    let supported_scheds = [
        SupportedSched::Bpfland,
        SupportedSched::Rusty,
        SupportedSched::Lavd,
        SupportedSched::Flash,
        SupportedSched::P2DQ,
        SupportedSched::Tickless,
        SupportedSched::Rustland,
        SupportedSched::Cosmos,
        SupportedSched::Beerland,
        SupportedSched::Cake,
        SupportedSched::Pandemonium,
        SupportedSched::Flow,
        SupportedSched::Forge,
    ];
    let scheds_map = HashMap::from(supported_scheds.map(init_default_config_entry));
    Config {
        default_sched: None,
        default_mode: Some(SchedMode::Auto),
        scheds: scheds_map,
    }
}

/// Get the scx flags for the given sched mode
pub fn get_scx_flags_for_mode(
    config: &Config,
    scx_sched: &SupportedSched,
    sched_mode: SchedMode,
) -> Vec<String> {
    let scx_name: &str = scx_sched.clone().into();
    if let Some(sched_config) = config.scheds.get(scx_name) {
        let scx_flags = extract_scx_flags_from_config(sched_config, sched_mode);

        // try to exact flags from config, otherwise fallback to hardcoded default
        scx_flags.unwrap_or({
            get_default_scx_flags_for_mode(scx_sched, sched_mode)
                .into_iter()
                .map(String::from)
                .collect()
        })
    } else {
        get_default_scx_flags_for_mode(scx_sched, sched_mode)
            .into_iter()
            .map(String::from)
            .collect()
    }
}

/// Extract the scx flags from config
fn extract_scx_flags_from_config(
    sched_config: &Sched,
    sched_mode: SchedMode,
) -> Option<Vec<String>> {
    match &sched_mode {
        SchedMode::Gaming => sched_config.gaming_mode.clone(),
        SchedMode::LowLatency => sched_config.lowlatency_mode.clone(),
        SchedMode::PowerSave => sched_config.powersave_mode.clone(),
        SchedMode::Server => sched_config.server_mode.clone(),
        SchedMode::Auto => sched_config.auto_mode.clone(),
    }
}

/// Get Sched object for configuration object
fn get_default_sched_for_config(scx_sched: &SupportedSched) -> Sched {
    Sched {
        auto_mode: Some(
            get_default_scx_flags_for_mode(scx_sched, SchedMode::Auto)
                .into_iter()
                .map(String::from)
                .collect(),
        ),
        gaming_mode: Some(
            get_default_scx_flags_for_mode(scx_sched, SchedMode::Gaming)
                .into_iter()
                .map(String::from)
                .collect(),
        ),
        lowlatency_mode: Some(
            get_default_scx_flags_for_mode(scx_sched, SchedMode::LowLatency)
                .into_iter()
                .map(String::from)
                .collect(),
        ),
        powersave_mode: Some(
            get_default_scx_flags_for_mode(scx_sched, SchedMode::PowerSave)
                .into_iter()
                .map(String::from)
                .collect(),
        ),
        server_mode: Some(
            get_default_scx_flags_for_mode(scx_sched, SchedMode::Server)
                .into_iter()
                .map(String::from)
                .collect(),
        ),
    }
}

/// Get the default scx flags for the given sched mode
fn get_default_scx_flags_for_mode(
    scx_sched: &SupportedSched,
    sched_mode: SchedMode,
) -> Vec<&'static str> {
    match &scx_sched {
        SupportedSched::Bpfland => match sched_mode {
            SchedMode::LowLatency => {
                vec!["-m", "performance", "-w"]
            }
            SchedMode::PowerSave => {
                vec!["-s", "20000", "-m", "powersave", "-I", "100", "-t", "100"]
            }
            SchedMode::Server => vec!["-s", "20000", "-S"],
            SchedMode::Gaming => vec!["-m", "all"],
            SchedMode::Auto => vec!["-m", "auto"],
        },
        SupportedSched::Lavd => match sched_mode {
            SchedMode::Gaming | SchedMode::LowLatency => {
                vec!["--performance", "--pinned-slice-us", "500"]
            }
            SchedMode::PowerSave => vec!["--powersave", "--pinned-slice-us", "500"],
            SchedMode::Server => vec![
                "--performance",
                "--slice-min-us",
                "3000",
                "--slice-max-us",
                "10000",
                "--pinned-slice-us",
                "3000",
            ],
            SchedMode::Auto => vec!["--autopilot", "--pinned-slice-us", "500"],
        },
        SupportedSched::P2DQ => match sched_mode {
            SchedMode::Gaming => vec!["--task-slice", "true", "-f", "--sched-mode", "performance"],
            SchedMode::LowLatency => vec!["-y", "-f", "--task-slice", "true"],
            SchedMode::PowerSave => vec!["--sched-mode", "efficiency"],
            SchedMode::Server => vec!["--keep-running"],
            SchedMode::Auto => vec!["--sched-mode", "default"],
        },
        SupportedSched::Tickless => match sched_mode {
            SchedMode::Gaming => vec!["-f", "5000", "-s", "5000"],
            SchedMode::LowLatency => vec!["-f", "5000", "-s", "1000"],
            SchedMode::PowerSave => vec!["-f", "50"],
            SchedMode::Server => vec!["-f", "100"],
            SchedMode::Auto => vec![],
        },
        SupportedSched::Cosmos => match sched_mode {
            SchedMode::Gaming => vec!["-s", "700"],
            SchedMode::LowLatency => vec!["-s", "700", "-m", "performance", "-w"],
            SchedMode::PowerSave => vec!["-m", "powersave"],
            SchedMode::Server | SchedMode::Auto => vec![],
        },
        SupportedSched::Cake => match sched_mode {
            SchedMode::Gaming | SchedMode::Server => vec!["--profile", "gaming"],
            SchedMode::LowLatency => vec!["--profile", "esports"],
            SchedMode::PowerSave => vec!["--profile", "battery"],
            SchedMode::Auto => vec!["--profile", "default"],
        },
        // The below Schedulers haven't defined any modes
        SupportedSched::Rusty
        | SupportedSched::Rustland
        | SupportedSched::Beerland
        | SupportedSched::Pandemonium
        | SupportedSched::Flash
        | SupportedSched::Flow
        | SupportedSched::Forge => vec![],
    }
}

/// Initializes entry for config sched map
fn init_default_config_entry(scx_sched: SupportedSched) -> (String, Sched) {
    let default_modes = get_default_sched_for_config(&scx_sched);
    (
        <SupportedSched as Into<&str>>::into(scx_sched).to_owned(),
        default_modes,
    )
}

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

    #[test]
    fn test_default_config() {
        let config_str = r#"
default_mode = "Auto"

[scheds.scx_bpfland]
auto_mode = ["-m", "auto"]
gaming_mode = ["-m", "all"]
lowlatency_mode = ["-m", "performance", "-w"]
powersave_mode = ["-s", "20000", "-m", "powersave", "-I", "100", "-t", "100"]
server_mode = ["-s", "20000", "-S"]

[scheds.scx_rusty]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_lavd]
auto_mode = ["--autopilot", "--pinned-slice-us", "500"]
gaming_mode = ["--performance", "--pinned-slice-us", "500"]
lowlatency_mode = ["--performance", "--pinned-slice-us", "500"]
powersave_mode = ["--powersave", "--pinned-slice-us", "500"]
server_mode = ["--performance", "--slice-min-us", "3000", "--slice-max-us", "10000", "--pinned-slice-us", "3000"]

[scheds.scx_flash]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_p2dq]
auto_mode = ["--sched-mode", "default"]
gaming_mode = ["--task-slice", "true", "-f", "--sched-mode", "performance"]
lowlatency_mode = ["-y", "-f", "--task-slice", "true"]
powersave_mode = ["--sched-mode", "efficiency"]
server_mode = ["--keep-running"]

[scheds.scx_tickless]
auto_mode = []
gaming_mode = ["-f", "5000", "-s", "5000"]
lowlatency_mode = ["-f", "5000", "-s", "1000"]
powersave_mode = ["-f", "50"]
server_mode = ["-f", "100"]

[scheds.scx_rustland]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_cosmos]
auto_mode = []
gaming_mode = ["-s", "700"]
lowlatency_mode = ["-s", "700", "-m", "performance", "-w"]
powersave_mode = ["-m", "powersave"]
server_mode = []

[scheds.scx_beerland]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_cake]
auto_mode = ["--profile", "default"]
gaming_mode = ["--profile", "gaming"]
lowlatency_mode = ["--profile", "esports"]
powersave_mode = ["--profile", "battery"]
server_mode = ["--profile", "gaming"]

[scheds.scx_pandemonium]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_flow]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []

[scheds.scx_forge]
auto_mode = []
gaming_mode = []
lowlatency_mode = []
powersave_mode = []
server_mode = []
"#;

        let parsed_config = parse_config_content(config_str).expect("Failed to parse config");
        let expected_config = get_default_config();

        assert_eq!(parsed_config, expected_config);
    }

    #[test]
    fn test_simple_fallback_config_flags() {
        let config_str = r#"
default_mode = "Auto"
"#;

        let parsed_config = parse_config_content(config_str).expect("Failed to parse config");

        let bpfland_flags =
            get_scx_flags_for_mode(&parsed_config, &SupportedSched::Bpfland, SchedMode::Gaming);
        let expected_flags =
            get_default_scx_flags_for_mode(&SupportedSched::Bpfland, SchedMode::Gaming);
        assert_eq!(
            bpfland_flags
                .iter()
                .map(std::string::String::as_str)
                .collect::<Vec<&str>>(),
            expected_flags
        );
    }

    #[test]
    fn test_sched_fallback_config_flags() {
        let config_str = r#"
default_mode = "Auto"

[scheds.scx_lavd]
auto_mode = ["--help"]
"#;

        let parsed_config = parse_config_content(config_str).expect("Failed to parse config");

        let lavd_flags =
            get_scx_flags_for_mode(&parsed_config, &SupportedSched::Lavd, SchedMode::Gaming);
        let expected_flags =
            get_default_scx_flags_for_mode(&SupportedSched::Lavd, SchedMode::Gaming);
        assert_eq!(
            lavd_flags
                .iter()
                .map(std::string::String::as_str)
                .collect::<Vec<&str>>(),
            expected_flags
        );

        let lavd_flags =
            get_scx_flags_for_mode(&parsed_config, &SupportedSched::Lavd, SchedMode::Auto);
        assert_eq!(
            lavd_flags
                .iter()
                .map(std::string::String::as_str)
                .collect::<Vec<&str>>(),
            vec!["--help"]
        );
    }

    #[test]
    fn test_empty_config() {
        let config_str = "";
        let result = parse_config_content(config_str);
        assert!(result.is_err());
    }
}
