use crate::systemd::JobSet;
use crate::systemd::UnitManager;
mod error;
mod file_compare;
pub mod systemd;
mod unit_file;

use std::{
    collections::HashSet,
    path::{Path, PathBuf},
    rc::Rc,
    time::Duration,
};
use systemd::UnitStatus;
use unit_file::UnitFile;

use anyhow::{Context, Result};

fn pretty_unit_names<I>(unit_names: I) -> String
where
    I: IntoIterator,
    I::Item: AsRef<str>,
{
    let mut str_vec = unit_names
        .into_iter()
        .map(|s| String::from(s.as_ref()))
        .collect::<Vec<_>>();
    str_vec.sort();
    str_vec.join(", ")
}

fn is_unit_available(unit_path: &Path) -> bool {
    unit_path.exists()
        && !unit_path
            .canonicalize()
            .map(|p| p == Path::new("/dev/null"))
            .unwrap_or(true)
}

/// Given an active unit name this returns the actual unit file name. In the
/// case of a parameterized unit, e.g., `foo@bar.service` this function returns
/// `foo@.service`. For nonparameterized unit names it returns none.
fn parameterized_base_name(unit_name: &str) -> Option<String> {
    let res = unit_name.splitn(2, '@').collect::<Vec<_>>();
    match res[..] {
        [base_name, arg_and_suffix] => {
            let res = arg_and_suffix.rsplitn(2, '.').collect::<Vec<_>>();
            match res[..] {
                [suffix, arg] if !arg.is_empty() => Some(format!("{}@.{}", base_name, suffix)),
                _ => None,
            }
        }
        _ => None,
    }
}

/// Returns the file path of the given unit name within the given directory.
///
/// If no matching file is found then `None` is returned.
///
/// If the given unit name is a parameterized named then an exactly matching
/// file is returned, if it exists, otherwise the template file path is
/// returned.
fn find_unit_file_path(unit_directory: &Path, unit_name: &str) -> Option<PathBuf> {
    Some(unit_directory.join(unit_name))
        .filter(|e| is_unit_available(e))
        .or_else(|| {
            parameterized_base_name(unit_name)
                .map(|n| unit_directory.join(n))
                .filter(|e| is_unit_available(e))
        })
}

/// A plan of unit actions needed to accomplish the switch.
#[derive(Debug)]
struct SwitchPlan {
    stop_units: HashSet<Rc<str>>,
    start_units: HashSet<Rc<str>>,
    reload_units: HashSet<Rc<str>>,
    restart_units: HashSet<Rc<str>>,
    keep_old_units: HashSet<Rc<str>>,
    unchanged_units: HashSet<Rc<str>>,
}

struct UnitWithTarget {
    unit_name: Rc<str>,
    target_name: Rc<str>,
}

fn build_switch_plan(
    old_dir: Option<&Path>,
    new_dir: &Path,
    service_manager: &impl systemd::ServiceManager,
) -> Result<SwitchPlan> {
    let mut stop_units = HashSet::new();
    let mut start_units = HashSet::new();
    let mut reload_units = HashSet::new();
    let mut restart_units = HashSet::new();
    let mut keep_old_units = HashSet::new();
    let mut unchanged_units = HashSet::new();

    let mut active_unit_names = HashSet::new();

    let active_units = service_manager
        .list_units_by_states(&["active", "activating"])
        .context("Failed to list active and activating units")?;

    // Handle units that are currently active, typically this implies restarting
    // the unit in some way.
    for active_unit in active_units {
        let new_unit_path_opt = find_unit_file_path(new_dir, active_unit.name());
        let old_unit_path_opt = old_dir
            .as_ref()
            .and_then(|d| find_unit_file_path(d, active_unit.name()));

        let active_unit_name: Rc<str> = active_unit.name().into();
        active_unit_names.insert(active_unit_name.clone());

        if let Some(new_unit_path) = new_unit_path_opt {
            let new_unit_file = UnitFile::load(&new_unit_path).with_context(|| {
                format!("Failed load of new unit file {}", new_unit_path.display())
            })?;

            if let Some(old_unit_path) = old_unit_path_opt {
                let old_unit_file = UnitFile::load(&old_unit_path).with_context(|| {
                    format!("Failed load of old unit file {}", old_unit_path.display())
                })?;

                if new_unit_file.unit_type() == unit_file::UnitType::Target {
                    if !new_unit_file.refuse_manual_start()? {
                        start_units.insert(active_unit_name);
                    }
                } else if unit_file::unit_eq(&old_unit_file, &new_unit_file)? {
                    unchanged_units.insert(active_unit_name);
                } else {
                    match new_unit_file.switch_method()? {
                        unit_file::UnitSwitchMethod::Reload => {
                            reload_units.insert(active_unit_name);
                        }
                        unit_file::UnitSwitchMethod::Restart => {
                            restart_units.insert(active_unit_name);
                        }
                        unit_file::UnitSwitchMethod::StopStart => {
                            if service_manager
                                .unit_manager(&active_unit)?
                                .refuse_manual_stop()?
                            {
                                keep_old_units.insert(active_unit_name);
                            } else if new_unit_file.refuse_manual_start()? {
                                stop_units.insert(active_unit_name);
                            } else {
                                stop_units.insert(active_unit_name.clone());
                                start_units.insert(active_unit_name);
                            }
                        }
                        unit_file::UnitSwitchMethod::KeepOld => {
                            keep_old_units.insert(active_unit_name);
                        }
                    };
                }
            } else {
                stop_units.insert(active_unit_name.clone());
                start_units.insert(active_unit_name);
            }
        } else if old_unit_path_opt.is_some() {
            stop_units.insert(active_unit_name);
        }
    }

    // Handle units that are not currently active but are wanted by an active
    // target. Typically this is simply a matter of starting the new unit.
    for wanted_unit in find_wanted_units(new_dir)? {
        // Skip if the unit is not wanted by an active target.
        if !active_unit_names.contains(&wanted_unit.target_name) {
            continue;
        }

        // Skip if the unit is actually active.
        if active_unit_names.contains(&wanted_unit.unit_name) {
            continue;
        }

        let new_unit_path_opt = find_unit_file_path(new_dir, &wanted_unit.unit_name);
        let old_unit_path_opt = old_dir
            .as_ref()
            .and_then(|d| find_unit_file_path(d, &wanted_unit.unit_name));

        if let Some(new_unit_path) = new_unit_path_opt {
            let new_unit_file = UnitFile::load(&new_unit_path).with_context(|| {
                format!("Failed load of new unit file {}", new_unit_path.display())
            })?;

            if let Some(old_unit_path) = old_unit_path_opt {
                let old_unit_file = UnitFile::load(&old_unit_path).with_context(|| {
                    format!("Failed load of old unit file {}", old_unit_path.display())
                })?;

                if unit_file::unit_eq(&old_unit_file, &new_unit_file)? {
                    unchanged_units.insert(wanted_unit.unit_name);
                } else if !new_unit_file.refuse_manual_start()? {
                    start_units.insert(wanted_unit.unit_name);
                }
            } else {
                start_units.insert(wanted_unit.unit_name);
            }
        }
    }

    Ok(SwitchPlan {
        stop_units,
        start_units,
        reload_units,
        restart_units,
        keep_old_units,
        unchanged_units,
    })
}

fn find_wanted_units(new_dir: &Path) -> Result<Vec<UnitWithTarget>> {
    let mut result = Vec::new();

    for entry in std::fs::read_dir(new_dir)? {
        let entry = entry?;
        let entry_file_name = entry.file_name().into_string().unwrap();
        if entry.metadata()?.is_dir() && entry_file_name.ends_with(".target.wants") {
            let dir_name = entry_file_name;
            let target_name: Rc<str> = dir_name.strip_suffix(".wants").unwrap().into();
            for entry in std::fs::read_dir(entry.path())? {
                let unit_name = entry?.file_name().into_string().unwrap().into();
                result.push(UnitWithTarget {
                    unit_name,
                    target_name: target_name.clone(),
                });
            }
        }
    }

    Ok(result)
}

fn exec_pre_reload<F>(
    plan: &SwitchPlan,
    service_manager: &impl systemd::ServiceManager,
    job_handler: F,
    dry_run: bool,
    timeout: Duration,
) -> Result<()>
where
    F: Fn(&str, &str) + Send + 'static,
{
    if !plan.stop_units.is_empty() {
        println!("Stopping units: {}", pretty_unit_names(&plan.stop_units));
        if !dry_run {
            let mut job_set = service_manager.new_job_set()?;

            for uf in plan.stop_units.iter() {
                job_set
                    .stop_unit(uf)
                    .with_context(|| format!("Failed to stop unit {}", uf))?;
            }

            job_set.wait_for_all(job_handler, timeout)?
        }
    }

    Ok(())
}

fn exec_reload(
    service_manager: &impl systemd::ServiceManager,
    dry_run: bool,
    verbose: bool,
) -> Result<()> {
    if !dry_run {
        if verbose {
            println!("Resetting failed units");
        }
        service_manager
            .reset_failed()
            .context("Failed to reset failed systemd units")?;

        if verbose {
            println!("Reloading systemd");
        }
        service_manager
            .daemon_reload()
            .context("Failed to reload systemd")?;
    }

    Ok(())
}

fn exec_post_reload<F>(
    plan: &SwitchPlan,
    service_manager: &impl systemd::ServiceManager,
    job_handler: F,
    dry_run: bool,
    verbose: bool,
    timeout: Duration,
) -> Result<()>
where
    F: Fn(&str, &str) + Send + 'static,
{
    let mut job_set = service_manager.new_job_set()?;

    if !plan.reload_units.is_empty() {
        println!("Reloading units: {}", pretty_unit_names(&plan.reload_units));
        if !dry_run {
            for uf in plan.reload_units.iter() {
                job_set
                    .reload_unit(uf)
                    .with_context(|| format!("Failed to reload unit {}", uf))?;
            }
        }
    }

    if !plan.restart_units.is_empty() {
        println!(
            "Restarting units: {}",
            pretty_unit_names(&plan.restart_units)
        );
        if !dry_run {
            for uf in plan.restart_units.iter() {
                job_set
                    .restart_unit(uf)
                    .with_context(|| format!("Failed to restart unit {}", uf))?;
            }
        }
    }

    if !plan.keep_old_units.is_empty() {
        println!(
            "Keeping old units: {}",
            pretty_unit_names(&plan.keep_old_units)
        );
    }

    if !plan.unchanged_units.is_empty() && verbose {
        println!(
            "Keeping units: {}",
            pretty_unit_names(&plan.unchanged_units)
        );
    }

    if !plan.start_units.is_empty() {
        println!("Starting units: {}", pretty_unit_names(&plan.start_units));
        if !dry_run {
            for uf in plan.start_units.iter() {
                job_set
                    .start_unit(uf)
                    .with_context(|| format!("Failed to start unit {}", uf))?;
            }
        }
    }

    job_set.wait_for_all(job_handler, timeout)?;

    Ok(())
}

/// Performs a systemd unit "switch".
pub fn switch(
    service_manager: impl systemd::ServiceManager,
    //connection: &zbus::blocking::Connection,
    old_dir: Option<&Path>,
    new_dir: &Path,
    dry_run: bool,
    verbose: bool,
    timeout: Duration,
) -> Result<()> {
    let system_status = service_manager.system_status()?;
    let do_switch = match system_status {
        systemd::SystemStatus::Initializing => true,
        systemd::SystemStatus::Starting => true,
        systemd::SystemStatus::Running => true,
        systemd::SystemStatus::Degraded => {
            let units_by_states = service_manager.list_units_by_states(&["failed"])?;
            let failed: Vec<&str> = units_by_states.iter().map(|status| status.name()).collect();
            let failed = failed.join(", ");
            eprintln!(
                "The service manager is degraded.\n\
                 Failed services: {failed}\n\
                 Attempting to continue anyway..."
            );
            true
        }
        systemd::SystemStatus::Maintenance => false,
        systemd::SystemStatus::Stopping => false,
    };

    if !do_switch {
        if verbose {
            println!("Skipping switch since systemd has {system_status} status");
        }
        return Ok(());
    }

    let plan = build_switch_plan(old_dir, new_dir, &service_manager)
        .context("Failed to build switch plan")?;

    let job_handler = move |name: &str, state: &str| {
        if verbose || state != "done" {
            println!("{} {}", name, state)
        }
    };

    exec_pre_reload(&plan, &service_manager, job_handler, dry_run, timeout)
        .context("Failed to perform pre-reload tasks")?;

    exec_reload(&service_manager, dry_run, verbose)?;

    exec_post_reload(
        &plan,
        &service_manager,
        job_handler,
        dry_run,
        verbose,
        timeout,
    )
    .context("Failed to perform post-reload tasks")?;

    Ok(())
}

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

    #[test]
    fn can_get_base_name_for_parameterized_unit() {
        assert_eq!(
            parameterized_base_name("foo@bar.service"),
            Some(String::from("foo@.service"))
        );
        assert_eq!(
            parameterized_base_name("foo@bar.baz.service"),
            Some(String::from("foo@.service"))
        );
    }

    #[test]
    fn no_base_name_for_nonparameterized_units() {
        assert_eq!(parameterized_base_name("foo@.service"), None);
        assert_eq!(parameterized_base_name("foo.service"), None);
        assert_eq!(parameterized_base_name("foo@barservice"), None);
    }
}