use eyre::{Result, format_err};
use log::{error, info, warn};
use std::collections::BTreeSet as Set;
use std::convert::Infallible;
use std::os::unix::fs::symlink;
use tokio::sync::Mutex;
use tokio::{fs, io::AsyncReadExt, process::Command};

use crate::{cmd::version::version_string, dklog, input, utils};
use dkl::bootstrap::Config;

mod bootstrap;
mod dmcrypt;
mod lvm;
mod networks;
mod sshd;

const INIT_LOG: &str = "/var/log/init.log";

// devices used by any block target (DM, LVM...)
static USED_DEVS: Mutex<Set<String>> = Mutex::const_new(Set::new());

pub async fn run() {
    if std::process::id() != 1 {
        error!("init must run as PID 1, not {}", std::process::id());
        std::process::exit(1);
    }

    unsafe {
        use std::env;
        env::set_var("PATH", "/bin:/sbin:/usr/bin:/usr/sbin");
        env::set_var("HOME", "/root");
    }

    // open input channels
    tokio::spawn(crate::input::answer_requests_from_stdin());
    tokio::spawn(crate::input::answer_requests_from_socket());

    // also log to file
    dklog::LOG.spawn(async {
        let Ok(log_file) = (tokio::fs::File::create(INIT_LOG).await)
            .inspect_err(|e| warn!("failed to open init.log: {e}"))
        else {
            return;
        };
        dklog::LOG.copy_to(log_file).await;
    });

    // welcome
    info!("Welcome to {}", version_string());

    let uname = nix::sys::utsname::uname().expect("uname should work");
    let kernel_version = uname.release().to_string_lossy();
    info!("Linux version {kernel_version}");

    // mount basic filesystems
    mount(None::<&str>, "/proc", "proc", None).await;
    mount(None::<&str>, "/sys", "sysfs", None).await;
    mount(None::<&str>, "/dev", "devtmpfs", None).await;
    mount(None::<&str>, "/dev/pts", "devpts", Some("gid=5,mode=620")).await;

    if utils::bool_param("debug") {
        log::set_max_level(log::LevelFilter::Debug);
    }

    // load config
    let cfg: Config = retry(async || {
        let cfg = (fs::read("config.yaml").await)
            .map_err(|e| format_err!("failed to read config: {e}"))?;
        serde_yaml::from_slice(cfg.as_slice())
            .map_err(|e| format_err!("failed to parse config: {e}"))
    })
    .await;

    info!("config loaded");
    info!("anti-phishing-code: {}", cfg.anti_phishing_code);

    // seem to occasionaly race with Child::wait
    // tokio::spawn(child_reaper());

    // mount modules
    if let Some(modules) = cfg.modules.as_deref() {
        retry_or_ignore(async || mount_modules(modules, &kernel_version).await).await;
    } else {
        warn!("modules NOT mounted (not configured)");
    }

    // init devices
    info!("initializing devices");
    start_daemon("udevd", &[]).await;

    exec("udevadm", &["trigger", "-c", "add", "-t", "devices"]).await;
    exec("udevadm", &["trigger", "-c", "add", "-t", "subsystems"]).await;
    exec("udevadm", &["settle"]).await;

    // resolv.conf
    if let Some(ref resolv_conf) = cfg.resolv_conf {
        retry_or_ignore(async || {
            let path = "/etc/resolv.conf";
            info!("writing /etc/resolv.conf");
            fs::write(path, resolv_conf.as_bytes()).await?;
            Ok(())
        })
        .await;
    }

    // Wireguard VPNs
    for (name, conf) in &cfg.vpns {
        retry_or_ignore(async || {
            info!("starting VPN {name}");
            let dir = "/etc/wireguard";
            fs::create_dir_all(dir).await?;

            let path = &format!("{dir}/{name}.conf");
            fs::write(path, conf.as_bytes()).await?;
            chmod(path, 0o600).await?;

            try_exec("wg-quick", &["up", name]).await?;
            Ok(())
        })
        .await
    }

    // SSH service
    sshd::start(&cfg).await;

    // networks
    networks::setup(&cfg).await;

    // pre-lvm dmcrypt devs
    dmcrypt::setup(&cfg.pre_lvm_crypt).await;

    // LVM
    lvm::setup(&cfg).await;

    // post-lvm dmcrypt devs
    dmcrypt::setup(&cfg.crypt).await;

    // bootstrap the system
    bootstrap::bootstrap(cfg).await;

    // finalize
    if let Err(e) = tokio::fs::copy(INIT_LOG, format!("/system{INIT_LOG}")).await {
        warn!("failed to copy {INIT_LOG} to system: {e}");
    }

    if let Err(e) = nix::mount::umount2("/modules", nix::mount::MntFlags::MNT_DETACH) {
        warn!("failed to umount /modules: {e}");
    }
    retry(async || switch_root("/system").await).await;
}

use std::path::Path;

async fn mount_modules(modules: &str, kernel_version: &str) -> Result<()> {
    info!("mounting modules");

    mount_ro_fs(modules, "/modules").await?;

    fs::create_dir_all("/lib/modules").await?;
    let modules_path = &format!("/modules/lib/modules/{kernel_version}");

    if !std::fs::exists(modules_path)? {
        return Err(format_err!(
            "invalid modules package: {modules_path} should exist"
        ));
    }

    symlink(modules_path, format!("/lib/modules/{kernel_version}"))?;

    let firmware_path = &format!("/modules/lib/firmware/{kernel_version}");
    if std::fs::exists(firmware_path)? {
        fs::create_dir_all("/lib/firmware").await?;
        symlink(firmware_path, format!("/lib/firmware/{kernel_version}"))?;
    }

    Ok(())
}

async fn chmod(path: impl AsRef<Path>, mode: u32) -> std::io::Result<()> {
    use std::fs::Permissions;
    use std::os::unix::fs::PermissionsExt;

    let perms = Permissions::from_mode(mode);
    fs::set_permissions(path, perms).await
}

async fn mount<S: AsRef<Path>>(
    src: Option<S>,
    dst: impl AsRef<Path>,
    fstype: &str,
    opts: Option<&str>,
) {
    let src = src.as_ref().map(|s| s.as_ref());
    let src_str = src.map(|s| s.display().to_string());
    let src_str = src_str.as_deref();

    let dst = dst.as_ref();
    let dst_str = &dst.display().to_string();

    if let Err(e) = fs::create_dir_all(dst).await {
        error!("failed to create dir {dst_str}: {e}");
    }

    retry_or_ignore(async || {
        let mut is_file = false;

        if let Some(src) = src_str {
            is_file = (fs::metadata(src).await)
                .map_err(|e| format_err!("stat {src} failed: {e}"))?
                .is_file();

            #[allow(clippy::single_match)] // may be extended later
            match fstype {
                "ext4" => {
                    exec("fsck.ext4", &["-p", src]).await;
                }
                _ => {}
            }
        }

        let mut args = vec![src_str.unwrap_or("none"), dst_str, "-t", fstype];
        if let Some(opts) = opts {
            args.extend(["-o", opts]);
        }

        // if it's a file, we need to use a loopdev
        if is_file {
            // loopdev crate has annoying dependencies, just use the normal mount program
            return try_exec("mount", &args).await;
        }

        let (cmd_str, _) = cmd_str("mount", &args);

        info!("# {cmd_str}",);

        let mount = |flags| nix::mount::mount(src, dst, Some(fstype), flags, opts);

        use nix::{errno::Errno, mount::MsFlags};
        match mount(MsFlags::empty()) {
            Err(Errno::EACCES) => mount(MsFlags::MS_RDONLY),
            r => r,
        }
        .map_err(|e| format_err!("mount {dst_str} failed: {e}"))
    })
    .await
}

async fn mount_ro_fs(src: impl AsRef<Path>, dst: impl AsRef<Path>) -> Result<()> {
    let src = src.as_ref();
    let dst = dst.as_ref();

    // identify RO fs type
    let mut buf = [0u8; 1028];

    fs::File::open(src)
        .await
        .map_err(|e| format_err!("open {}: {e}", src.display()))?
        .read_exact(&mut buf)
        .await
        .map_err(|e| format_err!("read {}: {e}", src.display()))?;

    let fstype = if buf[1024..1028] == 0xE0F5E1E2u32.to_le_bytes() {
        "erofs"
    } else {
        "squashfs"
    };

    if let Err(e) = fs::create_dir_all(dst).await {
        error!("failed to create dir {dst}: {e}", dst = dst.display());
    }

    let mut cmd = Command::new("mount");
    cmd.args(["-t", fstype]).arg(src).arg(dst);
    try_exec_cmd(cmd).await
}

async fn start_daemon(prog: &str, args: &[&str]) {
    let (cmd_str, mut cmd) = cmd_str(prog, args);
    retry_or_ignore(async || {
        info!("starting as daemon: {cmd_str}");
        cmd.spawn()?;
        Ok(())
    })
    .await;
}

async fn try_exec_cmd(mut cmd: Command) -> Result<()> {
    info!(
        "# {} {}",
        cmd.as_std().get_program().to_string_lossy(),
        cmd.as_std()
            .get_args()
            .map(|a| a.to_string_lossy())
            .collect::<Vec<_>>()
            .join(" ")
    );

    let s = cmd.status().await?;
    if s.success() {
        Ok(())
    } else {
        Err(format_err!("command failed: {s}"))
    }
}

async fn try_exec(prog: &str, args: &[&str]) -> Result<()> {
    let (cmd_str, mut cmd) = cmd_str(prog, args);
    info!("# {cmd_str}");

    let s = cmd.status().await?;
    if s.success() {
        Ok(())
    } else {
        Err(format_err!("command failed: {s}"))
    }
}

async fn exec(prog: &str, args: &[&str]) {
    retry_or_ignore(async || try_exec(prog, args).await).await;
}

async fn retry_or_ignore(mut action: impl AsyncFnMut() -> Result<()>) {
    loop {
        match action().await {
            Ok(_) => return,
            Err(e) => {
                error!("{e}");

                match input::read_choice(["[r]etry", "[i]gnore", "[s]hell"]).await {
                    'r' => {}
                    'i' => return,
                    's' => exec_shell().await,
                    _ => unreachable!(),
                }
            }
        }
    }
}

async fn retry<T>(mut action: impl AsyncFnMut() -> Result<T>) -> T {
    loop {
        match action().await {
            Ok(v) => return v,
            Err(e) => {
                error!("{e}");

                match input::read_choice(["[r]etry", "[s]hell"]).await {
                    'r' => {}
                    's' => exec_shell().await,
                    _ => unreachable!(),
                }
            }
        }
    }
}

async fn exec_shell() {
    let mut child = match Command::new("ash").spawn() {
        Ok(c) => c,
        Err(e) => {
            error!("failed to exec shell: {e}");
            return;
        }
    };

    let _ = child.wait().await;
}

fn cmd_str(prog: &str, args: &[&str]) -> (String, Command) {
    use std::borrow::Cow;

    let mut buf = String::new();

    buf.push_str(&shell_escape::escape(Cow::Borrowed(prog)));

    for &arg in args {
        buf.push(' ');
        buf.push_str(&shell_escape::escape(Cow::Borrowed(arg)));
    }

    let mut cmd = Command::new(prog);
    cmd.args(args);

    (buf, cmd)
}

#[allow(unused)]
async fn child_reaper() {
    use nix::sys::wait::{WaitPidFlag, waitpid};
    use nix::unistd::Pid;
    use tokio::signal::unix::{SignalKind, signal};

    let Ok(mut sigs) =
        signal(SignalKind::child()).inspect_err(|e| warn!("failed to setup SIGCHLD handler: {e}"))
    else {
        return;
    };

    loop {
        sigs.recv().await;
        let _ = waitpid(Some(Pid::from_raw(-1)), Some(WaitPidFlag::WNOHANG));
    }
}

async fn switch_root(root: &str) -> Result<Infallible> {
    info!("killing all processes and switching root");
    dklog::LOG.close().await;

    use nix::sys::signal::{SIGKILL, kill};
    use nix::unistd::Pid;

    if let Err(e) = kill(Pid::from_raw(-1), SIGKILL) {
        eprintln!("failed to kill processes: {e}");
    }

    macro_rules! cstr {
        ($s:expr) => {
            std::ffi::CString::new($s)?.as_c_str()
        };
    }

    Ok(nix::unistd::execv(
        cstr!("/sbin/switch_root"),
        &[
            cstr!("switch_root"),
            cstr!("-c"),
            cstr!("/dev/console"),
            cstr!(root),
            cstr!("/sbin/init"),
        ],
    )?)
}
