/* SPDX-License-Identifier: MIT */

use crate::config::Device;
use anyhow::{anyhow, Context, Result};
use log::{debug, log, warn, Level};
use std::cmp;
use std::collections::BTreeSet;
use std::fs;
use std::io::{self, Write};
use std::os::unix::fs::symlink;
use std::path::Path;
use std::process::Command;

fn make_parent(of: &Path) -> Result<()> {
    let parent = of
        .parent()
        .ok_or_else(|| anyhow!("Couldn't get parent of {}", of.display()))?;
    fs::create_dir_all(parent)?;
    Ok(())
}

fn make_symlink(dst: &str, src: &Path) -> Result<()> {
    make_parent(src)?;
    symlink(dst, src)
        .with_context(|| format!("Failed to create symlink {}→{}", src.display(), dst))?;
    Ok(())
}

fn virtualization_container() -> Result<bool> {
    let mut child = match Command::new("systemd-detect-virt")
        .arg("--quiet")
        .arg("--container")
        .spawn()
    {
        Ok(child) => child,
        Err(e) => {
            warn!(
                "systemd-detect-virt call failed, assuming we're not in a container: {}",
                e
            );
            return Ok(false);
        }
    };

    match child.wait() {
        Ok(status) => Ok(status.success()),
        Err(e) => Err(anyhow!("systemd-detect-virt call failed: {}", e)),
    }
}

fn modprobe(modname: &str, required: bool) {
    match Command::new("modprobe").arg(modname).status() {
        Err(e) => {
            let level = match !required && e.kind() == io::ErrorKind::NotFound {
                true => Level::Debug,
                false => Level::Warn,
            };

            log!(
                level,
                "modprobe \"{}\" cannot be spawned, ignoring: {}",
                modname,
                e
            );
        }
        Ok(status) => {
            if !status.success() {
                warn!("modprobe \"{}\" failed, ignoring: code {}", modname, status);
            }
        }
    };
}

pub fn run_generator(devices: &[Device], output_directory: &Path, fake_mode: bool) -> Result<()> {
    if devices.is_empty() {
        debug!("No devices configured, exiting.");
        return Ok(());
    }

    if virtualization_container()? && !fake_mode {
        debug!("Running in a container, exiting.");
        return Ok(());
    }

    for device in devices {
        handle_device(output_directory, device)?;
    }

    if !devices.is_empty() && !fake_mode {
        /* We created some units, let's make sure the module is loaded and the devices exist */
        if !Path::new("/sys/class/zram-control").exists() {
            modprobe("zram", true);
        }

        let max_device = devices
            .iter()
            .map(|device| {
                device.name[4..]
                    .parse()
                    .expect("already verified in read_devices()")
            })
            .fold(0, cmp::max);

        if !Path::new("/dev")
            .join(format!("zram{}", max_device))
            .exists()
        {
            while fs::read_to_string("/sys/class/zram-control/hot_add")
                .context("Adding zram device")?
                .trim_end()
                .parse::<u64>()
                .context("Fresh zram device number")?
                < max_device
            {}
        }
    }

    let compressors: BTreeSet<_> = devices
        .iter()
        .flat_map(|device| {
            device
                .compression_algorithms
                .compression_algorithms
                .iter()
                .map(|(a, _)| a.as_ref())
        })
        .collect();

    if !compressors.is_empty() {
        let proc_crypto = fs::read_to_string("/proc/crypto").unwrap_or_else(|e| {
            warn!("Failed to read /proc/crypto, proceeding as if empty: {}", e);
            String::new()
        });
        let known = parse_known_compressors(&proc_crypto);

        for comp in compressors.difference(&known) {
            modprobe(&format!("crypto-{}", comp), false);
        }
    }

    Ok(())
}

// Returns a list of names of loaded compressors
fn parse_known_compressors(proc_crypto: &str) -> BTreeSet<&str> {
    // Extract algorithm names (this includes non-compression algorithms too)
    proc_crypto
        .lines()
        .filter(|line| line.starts_with("name"))
        .map(|m| m.rsplit(':').next().unwrap().trim())
        .collect()
}

fn write_contents(output_directory: &Path, filename: &str, contents: &str) -> Result<()> {
    let path = output_directory.join(filename);
    make_parent(&path)?;

    let contents = format!(
        "\
# Automatically generated by {exe_name}

{contents}",
        exe_name = std::env::current_exe().unwrap().display(),
        contents = contents
    );

    fs::write(&path, contents).with_context(|| format!("Failed to write {}", path.display()))
}

fn handle_device(output_directory: &Path, device: &Device) -> Result<()> {
    if device.is_swap() {
        handle_zram_swap(output_directory, device)
    } else {
        handle_zram_mount_point(output_directory, device)
    }
}

fn handle_zram_bindings(output_directory: &Path, device: &Device, specific: &str) -> Result<()> {
    let wb_unit = device
        .writeback_dev
        .as_ref()
        .map(|wd| unit_name_from_path(wd, ".device"))
        .unwrap_or_default();

    /* systemd-zram-setup@.service.
     * We use the packaged unit, and only need to provide a small drop-in. */
    write_contents(
        output_directory,
        &format!("systemd-zram-setup@{}.service.d/bindings.conf", device.name),
        &format!(
            "\
[Unit]
BindsTo={}{}{}{}{}
",
            specific,
            &" "[device.writeback_dev.is_none() as usize..],
            wb_unit,
            device
                .writeback_dev
                .as_ref()
                .map(|_| "\nAfter=")
                .unwrap_or_default(),
            wb_unit,
        ),
    )
}

fn handle_zram_swap(output_directory: &Path, device: &Device) -> Result<()> {
    let swap_name = format!("dev-{}.swap", device.name);

    debug!(
        "Creating unit file {} (/dev/{} with {}MB)",
        swap_name,
        device.name,
        device.disksize / 1024 / 1024
    );

    handle_zram_bindings(output_directory, device, "dev-%i.swap")?;

    let shutdown_conflicts = if device.writeback_dev.is_some() {
        // We need to shut down the zram device to disconnect the writeback device.
        // Once https://github.com/systemd/systemd/issues/35303 is resolved, we
        // may revisit this and rely on the systemd to pull down the device stack
        // if appropriate.
        "Conflicts=shutdown.target\n"
    } else {
        ""
    };

    /* dev-zramX.swap */
    write_contents(
        output_directory,
        &swap_name,
        &format!(
            "\
[Unit]
Description=Compressed Swap on /dev/{zram_device}
Documentation=man:zram-generator(8) man:zram-generator.conf(5)

DefaultDependencies=no

Requires=systemd-zram-setup@{zram_device}.service
After=systemd-zram-setup@{zram_device}.service
Before=swap.target
{shutdown_conflicts}
[Swap]
What=/dev/{zram_device}
Priority={swap_priority}
Options={options}
",
            zram_device = device.name,
            swap_priority = device.swap_priority,
            options = device.options.replace('%', "%%"),
            shutdown_conflicts = shutdown_conflicts,
        ),
    )?;

    /* enablement symlink */
    let symlink_path = output_directory.join("swap.target.wants").join(&swap_name);
    let target_path = format!("../{}", swap_name);
    make_symlink(&target_path, &symlink_path)?;

    Ok(())
}

/// Path escaping as described in systemd.unit(5)
///
/// `/./` components stripped away when parsing `mount-point =`
fn unit_name_from_path(path: &Path, suffix: &str) -> String {
    assert!(path.is_absolute());

    let trimmed = path.to_str().unwrap().trim_matches('/');
    if trimmed.is_empty() {
        format!("-{}", suffix)
    } else {
        let mut obuf = Vec::with_capacity(path.as_os_str().len() + suffix.len());
        let mut just_slash = false;
        for (i, &b) in trimmed.as_bytes().iter().enumerate() {
            if b == b'/' && just_slash {
                continue;
            }
            just_slash = b == b'/';
            match b {
                b'/' => obuf.push(b'-'),
                b'.' if i == 0 => write!(obuf, "\\x{:02x}", b'.').unwrap(),
                b'0'..=b'9' | b'a'..=b'z' | b'A'..=b'Z' | b':' | b'_' | b'.' => obuf.push(b),
                _ => write!(obuf, "\\x{:02x}", b).unwrap(),
            }
        }
        obuf.extend_from_slice(suffix.as_bytes());
        String::from_utf8(obuf).unwrap()
    }
}

fn handle_zram_mount_point(output_directory: &Path, device: &Device) -> Result<()> {
    if device.mount_point.is_none() {
        /* In this case we don't need to generate any units. */
        return Ok(());
    }

    let mount_name = &unit_name_from_path(device.mount_point.as_ref().unwrap(), ".mount");

    debug!(
        "Creating unit file {} (/dev/{} with {}MB)",
        mount_name,
        device.name,
        device.disksize / 1024 / 1024
    );

    handle_zram_bindings(output_directory, device, mount_name)?;

    write_contents(
        output_directory,
        mount_name,
        &format!(
            "\
[Unit]
Description=Compressed Storage on /dev/{zram_device}
Documentation=man:zram-generator(8) man:zram-generator.conf(5)
Requires=systemd-zram-setup@{zram_device}.service
After=systemd-zram-setup@{zram_device}.service

[Mount]
What=/dev/{zram_device}
Where={mount_point}
Options={options}
",
            zram_device = device.name,
            mount_point = device.mount_point.as_ref().unwrap().to_str().unwrap(),
            options = device.options.replace('%', "%%"),
        ),
    )?;

    /* enablement symlink */
    let symlink_path = output_directory
        .join("local-fs.target.wants")
        .join(mount_name);
    let target_path = format!("../{}", mount_name);
    make_symlink(&target_path, &symlink_path)?;

    Ok(())
}

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

    #[test]
    fn test_parse_known_compressors() {
        let data = "\
name         : zstd
driver       : zstd-scomp
module       : zstd
priority     : 0
refcnt       : 1
selftest     : passed
internal     : no
type         : scomp

name         : zstd
driver       : zstd-generic
module       : zstd
priority     : 0
refcnt       : 1
selftest     : passed
internal     : no
type         : compression

name         : ccm(aes)
driver       : ccm_base(ctr(aes-aesni),cbcmac(aes-aesni))
module       : ccm
priority     : 300
refcnt       : 2
selftest     : passed
internal     : no
type         : aead
async        : no
geniv        : <none>

name         : ctr(aes)
driver       : ctr(aes-aesni)
module       : kernel
priority     : 300
refcnt       : 2
selftest     : passed
internal     : no
type         : skcipher
";
        let expected = ["zstd", "ccm(aes)", "ctr(aes)"];
        assert_eq!(parse_known_compressors(data), BTreeSet::from_iter(expected));
    }

    #[test]
    fn test_unit_name_from_path() {
        assert_eq!(
            unit_name_from_path(&Path::new("/waldo"), ".mount"),
            "waldo.mount"
        );
        assert_eq!(
            unit_name_from_path(&Path::new("/waldo/quuix"), ".mount"),
            "waldo-quuix.mount"
        );
        assert_eq!(
            unit_name_from_path(&Path::new("/waldo/quuix/"), ".mount"),
            "waldo-quuix.mount"
        );
        assert_eq!(
            unit_name_from_path(&Path::new("/waldo/quuix//"), ".mount"),
            "waldo-quuix.mount"
        );
        assert_eq!(unit_name_from_path(&Path::new("/"), ".mount"), "-.mount");
        assert_eq!(unit_name_from_path(&Path::new("//"), ".mount"), "-.mount");
        assert_eq!(unit_name_from_path(&Path::new("///"), ".mount"), "-.mount");
    }
}
