use std::collections::BTreeSet;
use std::fs;
use std::os::unix::fs::symlink;
use std::path::{Path, PathBuf};

use anyhow::{Context, Result, bail};

use crate::files::{RootfsPath, relative_symlink_is_safe};

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CopyTree {
    source: PathBuf,
    destination: RootfsPath,
}

impl CopyTree {
    pub fn new(source: impl AsRef<Path>, destination: impl AsRef<Path>) -> Result<Self> {
        Ok(Self {
            source: source.as_ref().to_path_buf(),
            destination: RootfsPath::new(destination)?,
        })
    }
}

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

impl InitrdOem {
    pub fn new(
        features: impl IntoIterator<Item = impl Into<String>>,
        modules: impl IntoIterator<Item = impl Into<String>>,
    ) -> Self {
        Self {
            features: features.into_iter().map(Into::into).collect(),
            modules: modules.into_iter().map(Into::into).collect(),
        }
    }

    fn render(&self) -> String {
        format!(
            "FEATURES += {}\nMODULES += {}\n",
            self.features.iter().cloned().collect::<Vec<_>>().join(" "),
            self.modules.iter().cloned().collect::<Vec<_>>().join(" "),
        )
    }
}

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

impl ServiceName {
    pub fn new(name: impl Into<String>) -> Result<Self> {
        let name = name.into();
        if !name.ends_with(".service")
            || name.contains('/')
            || name.chars().any(char::is_whitespace)
            || name == ".service"
        {
            bail!("invalid systemd service name: {name}");
        }
        Ok(Self(name))
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RootfsFinalization {
    copies: Vec<CopyTree>,
    initrd_oem: InitrdOem,
    enabled_services: Vec<ServiceName>,
}

impl RootfsFinalization {
    pub fn new(
        copies: Vec<CopyTree>,
        initrd_oem: InitrdOem,
        enabled_services: Vec<ServiceName>,
    ) -> Self {
        Self {
            copies,
            initrd_oem,
            enabled_services,
        }
    }

    pub fn apply(&self, rootfs: &Path) -> Result<MutationManifest> {
        if !rootfs.is_dir() {
            bail!("rootfs is not a directory: {}", rootfs.display());
        }

        let mut manifest = MutationManifest::default();
        for copy in &self.copies {
            copy_tree(rootfs, copy, &mut manifest)?;
        }
        write_file(
            rootfs,
            &RootfsPath::new("etc/initrd.mk.oem")?,
            self.initrd_oem.render().as_bytes(),
            &mut manifest,
        )?;
        for service in &self.enabled_services {
            enable_service(rootfs, service, &mut manifest)?;
        }
        manifest.paths.sort();
        Ok(manifest)
    }
}

#[derive(Debug, Default, Clone, PartialEq, Eq)]
pub struct MutationManifest {
    paths: Vec<String>,
}

impl MutationManifest {
    pub fn created_paths(&self) -> &[String] {
        &self.paths
    }

    fn record(&mut self, path: &Path) {
        self.paths.push(path.display().to_string());
    }
}

fn copy_tree(rootfs: &Path, copy: &CopyTree, manifest: &mut MutationManifest) -> Result<()> {
    if !copy.source.is_dir() {
        bail!("copy source is not a directory: {}", copy.source.display());
    }
    copy_entry(
        rootfs,
        &copy.source,
        copy.destination.as_path(),
        Path::new(""),
        manifest,
    )
}

fn copy_entry(
    rootfs: &Path,
    source: &Path,
    destination: &Path,
    source_relative: &Path,
    manifest: &mut MutationManifest,
) -> Result<()> {
    let metadata = fs::symlink_metadata(source)
        .with_context(|| format!("inspect copy source {}", source.display()))?;
    let output = rootfs.join(destination);

    if metadata.file_type().is_dir() {
        fs::create_dir_all(&output).with_context(|| format!("create {}", output.display()))?;
        manifest.record(destination);
        let mut entries = fs::read_dir(source)
            .with_context(|| format!("read copy source directory {}", source.display()))?
            .collect::<std::result::Result<Vec<_>, _>>()?;
        entries.sort_by_key(|entry| entry.file_name());
        for entry in entries {
            let name = entry.file_name();
            copy_entry(
                rootfs,
                &entry.path(),
                &destination.join(&name),
                &source_relative.join(&name),
                manifest,
            )?;
        }
    } else if metadata.file_type().is_symlink() {
        let target = fs::read_link(source)
            .with_context(|| format!("read source symlink {}", source.display()))?;
        if !relative_symlink_is_safe(source_relative.parent().unwrap_or(Path::new("")), &target) {
            bail!("source symlink escapes copied tree: {}", source.display());
        }
        if let Some(parent) = output.parent() {
            fs::create_dir_all(parent)?;
        }
        replace_with_symlink(&target, &output)?;
        manifest.record(destination);
    } else if metadata.file_type().is_file() {
        if let Some(parent) = output.parent() {
            fs::create_dir_all(parent)?;
        }
        replace_with_file(source, &output)?;
        fs::set_permissions(&output, metadata.permissions())?;
        manifest.record(destination);
    } else {
        bail!("unsupported copy source type: {}", source.display());
    }
    Ok(())
}

fn write_file(
    rootfs: &Path,
    destination: &RootfsPath,
    contents: &[u8],
    manifest: &mut MutationManifest,
) -> Result<()> {
    let output = rootfs.join(destination.as_path());
    let parent = output.parent().expect("rootfs-relative path has a parent");
    fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
    fs::write(&output, contents).with_context(|| format!("write {}", output.display()))?;
    manifest.record(destination.as_path());
    Ok(())
}

fn enable_service(
    rootfs: &Path,
    service: &ServiceName,
    manifest: &mut MutationManifest,
) -> Result<()> {
    let destination = RootfsPath::new(format!(
        "etc/systemd/system/multi-user.target.wants/{}",
        service.0
    ))?;
    let output = rootfs.join(destination.as_path());
    fs::create_dir_all(output.parent().expect("service path has a parent"))?;
    replace_with_symlink(
        Path::new(&format!("/usr/lib/systemd/system/{}", service.0)),
        &output,
    )
    .with_context(|| format!("enable service {}", service.0))?;
    manifest.record(destination.as_path());
    Ok(())
}

fn replace_with_file(source: &Path, output: &Path) -> Result<()> {
    match fs::symlink_metadata(output) {
        Ok(metadata) if metadata.file_type().is_dir() => {
            bail!("cannot replace directory with file: {}", output.display());
        }
        Ok(_) => {
            fs::remove_file(output).with_context(|| format!("replace {}", output.display()))?
        }
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
        Err(error) => return Err(error).with_context(|| format!("inspect {}", output.display())),
    }
    fs::copy(source, output)
        .with_context(|| format!("copy {} to {}", source.display(), output.display()))?;
    Ok(())
}

fn replace_with_symlink(target: &Path, output: &Path) -> Result<()> {
    match fs::symlink_metadata(output) {
        Ok(metadata) if metadata.file_type().is_dir() => {
            bail!(
                "cannot replace directory with symlink: {}",
                output.display()
            );
        }
        Ok(_) => {
            fs::remove_file(output).with_context(|| format!("replace {}", output.display()))?
        }
        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
        Err(error) => return Err(error).with_context(|| format!("inspect {}", output.display())),
    }
    symlink(target, output).with_context(|| format!("create symlink {}", output.display()))
}