use std::ffi::OsString;
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Command;

use anyhow::{bail, Context, Result};
use sha2::{Digest, Sha256};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct KernelVersion(String);

impl KernelVersion {
    pub fn discover_rt(rootfs: &Path) -> Result<Self> {
        let boot = rootfs.join("boot");
        let entries = fs::read_dir(&boot)
            .with_context(|| format!("read boot directory {}", boot.display()))?;
        let mut kernels = entries
            .filter_map(|entry| entry.ok())
            .filter_map(|entry| {
                let name = entry.file_name();
                let name = name.to_str()?;
                let version = name.strip_prefix("vmlinuz-")?;
                (version.contains("rt")).then(|| version.to_owned())
            })
            .collect::<Vec<_>>();
        kernels.sort();
        match kernels.as_slice() {
            [] => bail!("no RT kernel image found in {}", boot.display()),
            [version] => Ok(Self(version.clone())),
            _ => bail!("multiple RT kernel images found in {}: {}", boot.display(), kernels.join(", ")),
        }
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsRecipe {
    features: Vec<String>,
    modules: Vec<String>,
}

impl InitramfsRecipe {
    pub fn new(
        features: impl IntoIterator<Item = impl Into<String>>,
        modules: impl IntoIterator<Item = impl Into<String>>,
    ) -> Result<Self> {
        let mut features = normalize("initrd feature", features)?;
        let mut modules = normalize("initrd module", modules)?;
        features.sort();
        features.dedup();
        modules.sort();
        modules.dedup();
        Ok(Self { features, modules })
    }

    pub fn render(&self) -> String {
        format!(
            "FEATURES += {}\nMODULES += {}\n",
            self.features.join(" "),
            self.modules.join(" ")
        )
    }
}

fn normalize(kind: &str, values: impl IntoIterator<Item = impl Into<String>>) -> Result<Vec<String>> {
    values
        .into_iter()
        .map(Into::into)
        .map(|value: String| {
            if value.trim().is_empty() || value.chars().any(char::is_whitespace) {
                bail!("{kind} cannot be empty or contain whitespace");
            }
            Ok(value)
        })
        .collect()
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsRequest {
    rootfs: PathBuf,
    kernel: KernelVersion,
}

impl InitramfsRequest {
    pub fn discover(rootfs: impl AsRef<Path>) -> Result<Self> {
        let rootfs = rootfs.as_ref();
        if !rootfs.is_dir() {
            bail!("rootfs is not a directory: {}", rootfs.display());
        }
        Ok(Self {
            rootfs: rootfs.to_path_buf(),
            kernel: KernelVersion::discover_rt(rootfs)?,
        })
    }

    pub fn rootfs(&self) -> &Path {
        &self.rootfs
    }

    pub fn kernel(&self) -> &KernelVersion {
        &self.kernel
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Invocation {
    program: OsString,
    arguments: Vec<OsString>,
}

impl Invocation {
    pub fn new(program: impl Into<OsString>, arguments: impl IntoIterator<Item = impl Into<OsString>>) -> Self {
        Self { program: program.into(), arguments: arguments.into_iter().map(Into::into).collect() }
    }
}

pub trait CommandRunner {
    fn run(&mut self, invocation: Invocation) -> Result<()>;
}

#[derive(Debug, Default)]
pub struct ProcessRunner;

impl CommandRunner for ProcessRunner {
    fn run(&mut self, invocation: Invocation) -> Result<()> {
        let status = Command::new(&invocation.program)
            .args(&invocation.arguments)
            .status()
            .with_context(|| format!("run {}", invocation.program.to_string_lossy()))?;
        if !status.success() {
            bail!("{} exited with {status}", invocation.program.to_string_lossy());
        }
        Ok(())
    }
}

pub trait InitramfsBuilder {
    fn build(&mut self, request: &InitramfsRequest) -> Result<InitramfsResult>;
}

#[derive(Debug)]
pub struct MakeInitrdBuilder<R> {
    runner: R,
}

impl<R> MakeInitrdBuilder<R> {
    pub fn new(runner: R) -> Self {
        Self { runner }
    }

    pub fn runner(&self) -> &R {
        &self.runner
    }
}

impl<R: CommandRunner> InitramfsBuilder for MakeInitrdBuilder<R> {
    fn build(&mut self, request: &InitramfsRequest) -> Result<InitramfsResult> {
        self.runner.run(Invocation::new(
            "chroot",
            [
                request.rootfs().as_os_str().to_owned(),
                "make-initrd".into(),
                "-k".into(),
                request.kernel().as_str().into(),
            ],
        ))?;
        InitramfsResult::from_rootfs(request.rootfs(), request.kernel())
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct InitramfsResult {
    initrd_path: PathBuf,
    sha256: String,
}

impl InitramfsResult {
    fn from_rootfs(rootfs: &Path, kernel: &KernelVersion) -> Result<Self> {
        let initrd_path = PathBuf::from(format!("boot/initrd-{}.img", kernel.as_str()));
        let contents = fs::read(rootfs.join(&initrd_path))
            .with_context(|| format!("read generated initrd {}", initrd_path.display()))?;
        Ok(Self { initrd_path, sha256: format!("{:x}", Sha256::digest(contents)) })
    }

    pub fn initrd_path(&self) -> &Path {
        &self.initrd_path
    }

    pub fn sha256(&self) -> &str {
        &self.sha256
    }
}