// <Nyanit - init program fot initrd.>
// Copyright (C) <2026>  <Kira Foundation>

// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.

// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.

// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

/*
   Functions fot system interaction
*/

use nix;
use std::fs;
use std::process::Command;

/// timeout for block deices population
const WAIT_FOR_BLK_DEVS: u64 = 5u64;

/// list of subvolumes to move from current system to snapshoot when being restored
const SBUVOL_TO_MOVE: [(&str, &str); 1] = [("var", "/tmp")];

const TMP_ROOT: &str = "/mnt/bee-root";
const TMP_ROOT_SNAPS: &str = "/mnt/bee-root/bzz-snaps";
const TMP_ROOT_USR: &str = "/mnt/bee-root/usr";
const TMP_ROOT_VAR: &str = "/mnt/bee-root/var";

/// Path where /usr subvolume got moved on snapshot restore
const TMP_ROOT_SNAPS_USR: &str = "/mnt/bee-root/bzz-snaps/tmp_usr";
/// Path where /var subvolume got moved on snapshot restore
const TMP_ROOT_SNAPS_VAR: &str = "/mnt/bee-root/bzz-snaps/tmp_var";

const TMP_ROOT_SNAPS_PREFIX: &str = "/mnt/bee-root/bzz-snaps/tmp_";

const PAC_LOCK_FILE: &str = "/var/lib/hive-up/pac-lock";

#[derive(Debug)]
pub struct CommandFailed {
    pub stderr: Vec<u8>,
    pub cmd: String,
}

impl std::fmt::Display for CommandFailed {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(
            f,
            "Command failed! Command: {}; Stderr: {}",
            self.cmd,
            String::from_utf8(self.stderr.clone())
                .unwrap_or_else(|_| { "STDEER IS NOT VALID UTF-8".to_string() })
        )
    }
}

impl CommandFailed {
    pub fn new(cmd: &str, args: &Vec<&str>, stderr: Vec<u8>) -> Self {
        Self {
            stderr: stderr,
            cmd: format!("{} {}", cmd, args.join(" ")),
        }
    }
}

#[derive(Debug)]
pub enum ExecCommandError {
    IoError(std::io::Error),
    FromUtf8Error(std::string::FromUtf8Error),
    CommandFailed(CommandFailed),
    Other(String),
}

// 2. Implement From for each wrapped error type
impl From<std::io::Error> for ExecCommandError {
    fn from(err: std::io::Error) -> Self {
        ExecCommandError::IoError(err)
    }
}

impl From<std::string::FromUtf8Error> for ExecCommandError {
    fn from(err: std::string::FromUtf8Error) -> Self {
        ExecCommandError::FromUtf8Error(err)
    }
}

impl From<CommandFailed> for ExecCommandError {
    fn from(err: CommandFailed) -> Self {
        ExecCommandError::CommandFailed(err)
    }
}

impl From<String> for ExecCommandError {
    fn from(err: String) -> Self {
        ExecCommandError::Other(err)
    }
}

impl From<&str> for ExecCommandError {
    fn from(err: &str) -> Self {
        ExecCommandError::Other(err.to_string())
    }
}

impl std::fmt::Display for ExecCommandError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        match self {
            ExecCommandError::IoError(e) => write!(f, "IO Error: {}", e),
            ExecCommandError::FromUtf8Error(e) => {
                write!(f, "IO streem to String parse Error: {}", e)
            }
            ExecCommandError::CommandFailed(ex) => {
                write!(f, "{}", ex)
            }
            ExecCommandError::Other(msg) => write!(f, "Error: {}", msg),
        }
    }
}

// Implement std::error::Error
impl std::error::Error for ExecCommandError {}

/// exec given command with arguments, capturing stderr in case or failure
pub fn exec_command(cmd: &str, args: &Vec<&str>) -> Result<(), ExecCommandError> {
    let out = Command::new(cmd).args(args).output()?;
    if out.status.success() {
        Ok(())
    } else {
        Err(CommandFailed::new(cmd, args, out.stderr).into())
    }
}

pub fn try_find_partlabel(partlabel: &str, timeout_secs: u64) -> Option<String> {
    use std::{thread, time::Duration};

    let start = std::time::Instant::now();
    while start.elapsed().as_secs() < timeout_secs {
        if let Ok(Some(label)) = find_partlabel(partlabel) {
            return Some(label);
        }
        thread::sleep(Duration::from_millis(100));
    }
    None
}

/// Gets device node aka /dev/sda1 from gpr partition label
/// blkid -s PARTLABEL --match-token PARTLABEL="bzz-primary"
pub fn find_partlabel(partlabel: &str) -> Result<Option<String>, ExecCommandError> {
    let out = Command::new("/usr/bin/blkid")
        .args([
            "-s",
            "PARTLABEL",
            "-t",
            format!("PARTLABEL=\"{}\"", partlabel).as_str(),
        ])
        .output()?;
    if out.status.success() {
        let out_text = String::from_utf8(out.stdout)?;

        let split_res = out_text.split_once(':').ok_or_else(|| {
            ExecCommandError::Other(format!("Unexpected blkid output format: \"{}\"", out_text))
        })?;

        Ok(Some(split_res.0.to_string()))
    } else {
        if let Some(code) = out.status.code()
            && code == 2
        {
            Ok(None)
        } else {
            Err(ExecCommandError::Other(format!(
                "Error while looking for partlabel={}: \"{}\"",
                partlabel,
                String::from_utf8(out.stderr).unwrap_or_else(|ex| ex.to_string())
            )))
        }
    }
}

// pub fn list_dir(dirname: &str) -> Result<String, ExecCommandError> {
//     let res: String = fs::read_dir(dirname)?
//         .filter_map(|entry| entry.ok())
//         .map(|entry| entry.file_name().to_string_lossy().to_string())
//         .collect::<Vec<String>>()
//         .join("\n");

//     Ok(res)
// }

/// exec given command with arguments, returning stdout parsed to UTF-8 string
/// capturing stderr in case or failure
// pub fn exec_command_out(cmd: &str, args: &Vec<&str>) -> Result<String, ExecCommandError> {
//     let out = Command::new(cmd).args(args).output()?;
//     if out.status.success() {
//         Ok(String::from_utf8(out.stdout)?)
//     } else {
//         Err(CommandFailed::new(cmd, args, out.stderr).into())
//     }
// }

/// bcachefs unlock --file=/tmp/bzzpsspass.txt /dev/loop0somedev
/// Ok so. It is stupid, because bcacheutils far from ideal.
/// But if we provide wrong password from file or stdin this stuff plays stupid and asks for right password.
/// Sooo.. In order to know if unlock was sucessfull we need to check if program exited with Ok status,
/// or it shits to STDERR still waiting for stdin
/// Returns true if unlock successful, and false otherwise
pub fn bcachefs_unlock(pass: String, dev: &str) -> Result<bool, ExecCommandError> {
    use std::io::Read;
    use std::io::Write;
    use std::process::Stdio;
    // "-k", "session" is important as it stores key in sessyin keyring what allows
    // system to mount filesystem after switch_root
    let mut child = Command::new("/usr/bin/bcachefs")
        .args(&vec!["unlock", "-k", "session", dev])
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()?;

    let mut c_stderr = child
        .stderr
        .take()
        .ok_or_else(|| "Unable to get child process STDERR".to_string())?;

    // Write password to STDIN in a limited scope
    {
        let mut c_stdin = child
            .stdin
            .take()
            .ok_or_else(|| "Unable to get child process STDIN".to_string())?;
        // neccessery to use thread or we may be deadlocking
        c_stdin.write_all(pass.as_bytes())?;
        let bytes_written = c_stdin.write(b"\n")?;
        println!("bytes_written: {}", bytes_written);
        // stdin is dropped here automatically when the block ends
    }

    let mut buffer = [0u8; 10];
    // This blocks.
    // if unlok fails it reads 10 bytes (piece of err message)
    // if sucess it reads 0 bytes and exit
    let read_b = c_stderr.read(&mut buffer)?;
    if read_b > 0 {
        println!("STDERR error not empty, password incorrect! Try to kill a child.");
        child.kill()?;
    }

    // Safe to wait now; child will see EOF
    let exit_status = child.wait()?;
    // if exit status is SUCCESS - unloking is successful
    Ok(exit_status.success())
}

/// Cheking if device is encripted
pub fn bzzpss_dev_encrypted(dev_name: &str) -> Result<bool, ExecCommandError> {
    match exec_command("/usr/bin/bcachefs", &vec!["unlock", "--check", dev_name]) {
        Ok(_) => Ok(true), // if status OK - dev encrypted
        Err(ExecCommandError::CommandFailed(_)) => Ok(false), // if bcachefs returned error - unencrypted
        Err(ex) => Err(ex),                                   // othervice its error
    }
}

///
/// Looking for snapshots in snaps_dir for usr and var
/// And return list of date-time part of filename if both snapshots exists
///
pub fn list_bzzpss_snaps(snaps_dir: &str) -> Result<Vec<String>, ExecCommandError> {
    use regex_lite::Regex;
    use std::path::Path;

    // usr or var snapshoot in a form of bzz-usr|var-yyyy-mm-ddTHH-MM-SS
    let re = Regex::new(r"^bzz-(usr|var)-([0-9]{4}-[0-9]{2}-[0-9]{2}T[0-9]{2}-[0-9]{2}-[0-9]{2})$")
        .unwrap();

    let mut res: Vec<String> = Vec::new();
    // if snaps dir path do not exist return empty list
    if !Path::new(snaps_dir).exists() {
        return Ok(res);
    }

    // contains only date-time part
    let mut usr_snaps: Vec<String> = Vec::new();
    let mut var_snaps: Vec<String> = Vec::new();

    // collect snapshots date-time
    for maybe_entry in fs::read_dir(snaps_dir)? {
        let entry = maybe_entry?;
        if entry.file_type()?.is_dir() {
            if let Some(dir_name) = entry.file_name().to_str() {
                if let Some(captured) = re.captures(&dir_name) {
                    let subvol_name = &captured[1];
                    let date_time = captured[2].to_string();
                    if subvol_name == "usr" {
                        usr_snaps.push(date_time);
                    } else {
                        var_snaps.push(date_time);
                    }
                }
            }
        }
    }
    // now we need to make sure snapshots of /usr and /var both exist
    // just fools safe, may reconsider in future
    usr_snaps.sort();
    for snap_time in usr_snaps {
        if var_snaps.contains(&snap_time) {
            res.push(snap_time);
        }
    }
    Ok(res)
}

///
/// creates rw snapshot of "subvolume" and places it into "target"
/// bcachefs subvolume snapshot --rw /subvolume /target
///
pub fn snapshot_rw(subvolume: &str, target: &str) -> Result<(), ExecCommandError> {
    exec_command(
        "bcachefs",
        &vec!["subvolume", "snapshot", "--rw", subvolume, target],
    )?;

    Ok(())
}

///
/// bcachefs subvolume delete shortcut
///
pub fn subvolume_delete(subvolume: &str) -> Result<(), ExecCommandError> {
    exec_command("bcachefs", &vec!["subvolume", "delete", subvolume])?;

    Ok(())
}

///
/// Helper function to get snapshot names in form of "bzz-(usr|var)-date-time"
/// vol: "usr" or "var"
/// dt: datetime
///
fn snap_name_from_dt(vol: &str, dt: &str) -> String {
    format!("bzz-{}-{}", vol, dt)
}

///
///
/// // Stable Rust example using from_fn
fn first_n(n: u64) -> impl Iterator<Item = u64> {
    let mut count = 0;
    std::iter::from_fn(move || {
        if count < n {
            let val = count;
            count += 1;
            Some(val)
        } else {
            None
        }
    })
}

enum RollbackStates {
    Start,
    RenameUsrVarAct,
    ChangeToRw,
    MoveSubvolums,
    RemoveOldVols,
    ClearSnaps,
    Finish,
    Err,
}

pub fn roll_back_snap(
    snaps_list: &Vec<String>,
    target_snap: &str,
) -> impl Iterator<Item = Result<String, ExecCommandError>> {
    let usr_ro_snap_path = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", target_snap);
    let var_ro_snap_path = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", target_snap);

    let mut roll_state = RollbackStates::Start;

    std::iter::from_fn(move || match roll_state {
        RollbackStates::Start => {
            roll_state = RollbackStates::RenameUsrVarAct;
            Some(Ok("Move usr and var to temporary location".into()))
        }
        RollbackStates::RenameUsrVarAct => {
            let move_res = fs::rename(TMP_ROOT_USR, TMP_ROOT_SNAPS_USR)
                .and_then(|_| fs::rename(TMP_ROOT_VAR, TMP_ROOT_SNAPS_VAR));
            if let Err(ex) = move_res {
                roll_state = RollbackStates::Err;
                Some(Err(ex.into()))
            } else {
                roll_state = RollbackStates::ChangeToRw;
                Some(Ok("Setting snapshot as current system ...".into()))
            }
        }
        RollbackStates::ChangeToRw => {
            let rw_res = snapshot_rw(&usr_ro_snap_path, TMP_ROOT_USR)
                .and_then(|_| snapshot_rw(&var_ro_snap_path, TMP_ROOT_VAR));
            if let Err(ex) = rw_res {
                roll_state = RollbackStates::Err;
                Some(Err(ex))
            } else {
                roll_state = RollbackStates::MoveSubvolums;
                Some(Ok(
                    "Moving subvolumes from old /usr and /var to new one".into()
                ))
            }
        }
        RollbackStates::MoveSubvolums => {
            let mut mv_res = Some(Ok("Removing old /usr and /var".into()));
            for (vol, svol_path) in SBUVOL_TO_MOVE {
                let v_from = format!("{}{}{}", TMP_ROOT_SNAPS_PREFIX, vol, svol_path);
                let v_dest = format!("{}/{}{}", TMP_ROOT, vol, svol_path);
                if let Err(ex) = fs::rename(v_from, v_dest) {
                    roll_state = RollbackStates::Err;
                    mv_res = Some(Err(ex.into()));
                    break;
                }
            }
            roll_state = RollbackStates::RemoveOldVols;
            mv_res
        }
        RollbackStates::RemoveOldVols => {
            let rem_old_res = subvolume_delete(TMP_ROOT_SNAPS_USR)
                .and_then(|_| subvolume_delete(TMP_ROOT_SNAPS_VAR));
            if let Err(ex) = rem_old_res {
                roll_state = RollbackStates::Err;
                Some(Err(ex))
            } else {
                roll_state = RollbackStates::ClearSnaps;
                Some(Ok("Clearing snapshots...".into()))
            }
        }
        RollbackStates::ClearSnaps => {
            let mut clr_snap_res = Some(Ok("All done!".into()));
            for snap_date_time in snaps_list.iter().rev() {
                let usr_snap = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", snap_date_time);
                let var_snap = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", snap_date_time);
                if let Err(ex) =
                    subvolume_delete(&usr_snap).and_then(|_| subvolume_delete(&var_snap))
                {
                    clr_snap_res = Some(Err(ex));
                    roll_state = RollbackStates::Err;
                    break;
                }
                if snap_date_time == target_snap {
                    break;
                }
            }
            clr_snap_res
        }
        RollbackStates::Finish => None,
        RollbackStates::Err => None,
    })

    // 1. Rename usr and var to TMP_ROOT_SNAPS/tmp_var|usr
    // fs::rename(TMP_ROOT_USR, TMP_ROOT_SNAPS_USR)?;
    // fs::rename(TMP_ROOT_VAR, TMP_ROOT_SNAPS_VAR)?;

    // 2. snapshot desired snaps as RW to /usr and /var
    // snapshot_rw(&usr_ro_snap_path, TMP_ROOT_USR)?;
    // snapshot_rw(&var_ro_snap_path, TMP_ROOT_VAR)?;

    // 3. move subvolumes from old /usr and /var
    // for (vol, svol_path) in SBUVOL_TO_MOVE {
    //     let v_from = format!("{}{}{}", TMP_ROOT_SNAPS_PREFIX, vol, svol_path);
    //     let v_dest = format!("{}/{}{}",TMP_ROOT,vol,svol_path);
    //     fs::rename(v_from, v_dest)?;
    // }

    // 4. removing tmp /usr /var
    // subvolume_delete(TMP_ROOT_SNAPS_USR)?;
    // subvolume_delete(TMP_ROOT_SNAPS_VAR)?;

    // 5. Removing all snapshoots made after one we restored, including inself
    // Iterate in reverse untill we get to target_snap in snaps list
    // for snap_date_time in snaps_list.iter().rev() {
    //     let usr_snap = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", snap_date_time);
    //     let var_snap = format!("{}/bzz-{}-{}", TMP_ROOT_SNAPS, "usr", snap_date_time);
    //     subvolume_delete(&usr_snap)?;
    //     subvolume_delete(&var_snap)?;
    //     if snap_date_time == target_snap  {
    //         break;
    //     }
    // }

    // Ok(())
}

///
/// Checks if pac_lock file exist and returns date-time from it if it does,
/// Otherwise return None
/// 
pub fn check_pac_lock() -> Result<Option<String>, ExecCommandError> {
    use std::path::Path;
    let pac_lock_path = Path::new(PAC_LOCK_FILE);
    if pac_lock_path.exists() {
        let snap_dt = fs::read_to_string(pac_lock_path)?;
        Ok(Some(snap_dt.trim().into()))
    }
    else {
        Ok(None)
    }
}

// pulls /sys/class/block
fn wait_for_block_devices(timeout_secs: u64) -> bool {
    use std::path::Path;
    use std::{thread, time::Duration};

    let start = std::time::Instant::now();
    while start.elapsed().as_secs() < timeout_secs {
        // Check if any block device exists (e.g., sda, nvme0n1, vda)
        // You can refine this to check for a specific device needed for root
        if Path::new("/sys/class/block")
            .read_dir()
            .ok()
            .and_then(|mut d| d.next().is_some().then_some(true))
            .unwrap_or(false)
        {
            // Optional: Verify specific device needed for root exists
            // if Path::new("/dev/nvme0n1").exists() { return true; }
            return true;
        }
        thread::sleep(Duration::from_millis(100));
    }
    false
}

// mounting process
//# Mount essential virtual filesystems
// mount -t proc none /proc
// mount -t sysfs none /sys
// modprobe bcachefs
pub fn prepare_env() -> Result<(), ExecCommandError> {
    use nix::mount::{MsFlags, mount};
    use std::path::Path;

    // Mount proc filesystem
    // mount -t proc none /proc
    let proc_mount_flags =
        MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_RELATIME;
    mount(
        Some("proc"),
        Path::new("/proc"),
        Some("proc"),
        proc_mount_flags,
        None::<&str>,
    )
    .map_err(|ex| format!("Error mounting /proc: {}", ex.to_string()))?;

    // mount -t sysfs none /sys
    let sys_mount_flags =
        MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_RELATIME;
    mount(
        Some("sysfs"),
        Path::new("/sys"),
        Some("sysfs"),
        sys_mount_flags,
        None::<&str>,
    )
    .map_err(|ex| format!("Error mounting /sys: {}", ex.to_string()))?;

    //waiting for block devs population
    if wait_for_block_devices(WAIT_FOR_BLK_DEVS) == false {
        return Err("Waiting for block devices to get populated timed out.".into());
    }

    let dev_mount_flags = MsFlags::MS_NOSUID | MsFlags::MS_RELATIME;
    mount(
        Some("devtmpfs"),
        "/dev",
        Some("devtmpfs"),
        dev_mount_flags,
        Some("mode=0755"),
    )
    .map_err(|ex| format!("Error mounting /dev: {}", ex.to_string()))?;

    //mount -t tmpfs -o nosuid,nodev,noexec,mode=0755 none /run
    let run_mount_flags =
        MsFlags::MS_NOSUID | MsFlags::MS_NODEV | MsFlags::MS_NOEXEC | MsFlags::MS_RELATIME;
    mount(
        Some("run"),
        "/run",
        Some("tmpfs"),
        run_mount_flags,
        Some("mode=0755"),
    )
    .map_err(|ex| format!("Error mounting /run: {}", ex.to_string()))?;

    // modprobe bcachefs
    // exec_command("modprobe", &vec!["bcachefs"])?;

    // Creating dir for udev so this little shit wont hang
    // fs::create_dir("/run/udev")?;

    Ok(())
}

///
///  MUST be done before spawning any threads or async runtimes
// pub fn set_env_vars_unsafe() -> () {
//     use std::env;
//     // 1. Define your custom paths
//     let custom_bin = "/usr/bin";
//     let custom_lib = "/usr/lib";

//     // 2. Retrieve existing paths to append/prepend (optional but recommended)
//     // 3. Construct new path strings
//     let new_path = match env::var("PATH") {
//         Err(_) => custom_bin.to_string(),
//         Ok(v) => format!("{}:{}", custom_bin, v),
//     };

//     let new_ld_path = match env::var("LD_LIBRARY_PATH") {
//         Err(_) => custom_lib.to_string(),
//         Ok(v) => format!("{}:{}", custom_lib, v),
//     };

//     // 4. Set environment variables (UNSAFE in Rust 2024+)
//     // MUST be done before spawning any threads or async runtimes
//     unsafe {
//         env::set_var("PATH", new_path);
//         env::set_var("LD_LIBRARY_PATH", new_ld_path);
//     }
// }

// # Mount the real root filesystem
pub fn mount_bee_root(dev_name: &str, mount_point: &str) -> Result<(), ExecCommandError> {
    use nix::mount::{MsFlags, mount};
    use std::path::Path;

    let root_mount_flags = MsFlags::MS_RDONLY | MsFlags::MS_RELATIME;
    mount(
        Some(dev_name),
        Path::new(mount_point),
        Some("bcachefs"),
        root_mount_flags,
        None::<&str>,
    )
    .map_err(|ex| {
        format!(
            "Error mounting root {} to {} :\n{}",
            dev_name,
            mount_point,
            ex.to_string()
        )
    })?;
    // exec_command("/usr/bin/mount", &vec!["-o", "ro", dev_name, mount_point])?;
    Ok(())
}

// finish boot
pub fn hand_off_control(mnt_root: &str, init_path: &str) -> Result<(), ExecCommandError> {
    //use std::env;
    use std::os::unix::process::CommandExt;

    // 2. Change directory to the new root (Mandatory for switch_root)
    // Actually this may be a problem!
    // env::set_current_dir(mnt_root)?;

    // # Clean up virtual filesystems
    // !!! do not do this, switch_root handle this stuff and fail if thay not present!
    // exec_command("umount", &vec!["/proc"])?;
    // exec_command("umount", &vec!["/sys"])?;

    // # Hand off control to the real system's init
    // exec switch_root /mnt/root /sbin/init
    // 3. Execute switch_root.
    // This call replaces the current Rust process.
    // It only returns if an error occurs.
    let err = Command::new("/usr/bin/switch_root")
        .arg(mnt_root)
        .arg(init_path)
        .exec();

    // If execution reaches here, exec() failed.
    Err(format!("switch_root failed: {}", err).into())
}