//! Real initramfs builder producing valid cpio archives (newc format).
//!
//! The cpio "newc" (SVR4 with no CRC) format is what the Linux kernel expects
//! for initramfs archives. Each entry consists of a 110-byte ASCII header
//! followed by the filename (NUL-terminated, padded to 4-byte boundary),
//! followed by the file data (padded to 4-byte boundary).
//!
//! The archive is terminated by a trailer entry with the name "TRAILER!!!".

use flate2::write::GzEncoder;
use flate2::Compression;
use std::io::Write;

use crate::error::KernelError;

/// Magic number for cpio newc format headers.
const CPIO_NEWC_MAGIC: &str = "070701";

/// Default /init script that mounts essential filesystems, configures
/// networking, and starts services.
pub fn default_init_script(services: &[&str]) -> String {
    let mut script = String::from(
        r#"#!/bin/sh
# RVF initramfs init script
# Generated by rvf-kernel

set -e

echo "RVF initramfs booting..."

# Mount essential filesystems
mount -t proc proc /proc
mount -t sysfs sysfs /sys
mount -t devtmpfs devtmpfs /dev
mkdir -p /dev/pts /dev/shm
mount -t devpts devpts /dev/pts
mount -t tmpfs tmpfs /dev/shm
mount -t tmpfs tmpfs /tmp
mount -t tmpfs tmpfs /run

# Set hostname
hostname rvf

# Configure loopback
ip link set lo up

# Configure primary network interface
for iface in eth0 ens3 enp0s3; do
    if [ -d "/sys/class/net/$iface" ]; then
        ip link set "$iface" up
        # Try DHCP if udhcpc is available
        if command -v udhcpc >/dev/null 2>&1; then
            udhcpc -i "$iface" -s /etc/udhcpc/simple.script -q &
        fi
        break
    fi
done

# Set up minimal /etc
echo "root:x:0:0:root:/root:/bin/sh" > /etc/passwd
echo "root:x:0:" > /etc/group
echo "nameserver 8.8.8.8" > /etc/resolv.conf
echo "rvf" > /etc/hostname

"#,
    );

    // Start requested services
    for service in services {
        script.push_str(&format!(
            "# Start service: {service}\n\
             echo \"Starting {service}...\"\n"
        ));
        match *service {
            "sshd" | "dropbear" => {
                script.push_str(
                    "mkdir -p /etc/dropbear\n\
                     if command -v dropbear >/dev/null 2>&1; then\n\
                         dropbear -R -F -E -p 2222 &\n\
                     elif command -v sshd >/dev/null 2>&1; then\n\
                         mkdir -p /etc/ssh\n\
                         ssh-keygen -A 2>/dev/null || true\n\
                         /usr/sbin/sshd -p 2222\n\
                     fi\n",
                );
            }
            "rvf-server" => {
                script.push_str(
                    "if command -v rvf-server >/dev/null 2>&1; then\n\
                         rvf-server --listen 0.0.0.0:8080 &\n\
                     fi\n",
                );
            }
            other => {
                script.push_str(&format!(
                    "if command -v {other} >/dev/null 2>&1; then\n\
                         {other} &\n\
                     fi\n"
                ));
            }
        }
        script.push('\n');
    }

    script.push_str(
        "echo \"RVF initramfs ready.\"\n\
         \n\
         # Drop to shell or wait\n\
         exec /bin/sh\n",
    );

    script
}

/// A builder for creating cpio newc archives suitable for Linux initramfs.
pub struct CpioBuilder {
    /// Accumulated cpio archive bytes (uncompressed).
    data: Vec<u8>,
    /// Monotonically increasing inode counter.
    next_ino: u32,
}

impl CpioBuilder {
    /// Create a new, empty cpio archive builder.
    pub fn new() -> Self {
        Self {
            data: Vec::with_capacity(64 * 1024),
            next_ino: 1,
        }
    }

    /// Add a directory entry to the archive.
    ///
    /// `path` must not include a leading `/` in the archive (e.g., "bin", "etc").
    /// Mode 0o755 is used for directories.
    pub fn add_dir(&mut self, path: &str) {
        self.add_entry(path, 0o040755, &[]);
    }

    /// Add a regular file to the archive.
    ///
    /// `path` is the archive-internal path (e.g., "init", "bin/busybox").
    /// `mode` is the Unix file mode (e.g., 0o100755 for executable).
    pub fn add_file(&mut self, path: &str, mode: u32, content: &[u8]) {
        self.add_entry(path, mode, content);
    }

    /// Add a symlink to the archive.
    ///
    /// The symlink `path` will point to `target`.
    pub fn add_symlink(&mut self, path: &str, target: &str) {
        self.add_entry(path, 0o120777, target.as_bytes());
    }

    /// Add a device node (character or block).
    ///
    /// `mode` should include the device type bits (e.g., 0o020666 for char device).
    /// `devmajor` and `devminor` are the device major/minor numbers.
    pub fn add_device(&mut self, path: &str, mode: u32, devmajor: u32, devminor: u32) {
        let ino = self.next_ino;
        self.next_ino += 1;

        let name_with_nul = format!("{path}\0");
        let name_len = name_with_nul.len() as u32;
        let filesize = 0u32;

        let header = format!(
            "{CPIO_NEWC_MAGIC}\
             {ino:08X}\
             {mode:08X}\
             {uid:08X}\
             {gid:08X}\
             {nlink:08X}\
             {mtime:08X}\
             {filesize:08X}\
             {devmajor:08X}\
             {devminor:08X}\
             {rdevmajor:08X}\
             {rdevminor:08X}\
             {namesize:08X}\
             {check:08X}",
            uid = 0u32,
            gid = 0u32,
            nlink = 1u32,
            mtime = 0u32,
            rdevmajor = devmajor,
            rdevminor = devminor,
            namesize = name_len,
            check = 0u32,
        );

        self.data.extend_from_slice(header.as_bytes());
        self.data.extend_from_slice(name_with_nul.as_bytes());
        pad4(&mut self.data);
        // No file data for device nodes
    }

    /// Finalize the archive by appending the TRAILER!!! entry.
    ///
    /// Returns the raw (uncompressed) cpio archive bytes.
    pub fn finish(mut self) -> Vec<u8> {
        self.add_entry("TRAILER!!!", 0, &[]);
        self.data
    }

    /// Finalize and gzip-compress the archive.
    ///
    /// Returns the gzipped cpio archive suitable for use as a Linux initramfs.
    pub fn finish_gzipped(self) -> Result<Vec<u8>, KernelError> {
        let raw = self.finish();
        let mut encoder = GzEncoder::new(Vec::new(), Compression::best());
        encoder
            .write_all(&raw)
            .map_err(|e| KernelError::CompressionFailed(e.to_string()))?;
        encoder
            .finish()
            .map_err(|e| KernelError::CompressionFailed(e.to_string()))
    }

    fn add_entry(&mut self, path: &str, mode: u32, content: &[u8]) {
        let ino = self.next_ino;
        self.next_ino += 1;

        let name_with_nul = format!("{path}\0");
        let name_len = name_with_nul.len() as u32;
        let filesize = content.len() as u32;

        // cpio newc header is 110 bytes of ASCII hex fields
        let header = format!(
            "{CPIO_NEWC_MAGIC}\
             {ino:08X}\
             {mode:08X}\
             {uid:08X}\
             {gid:08X}\
             {nlink:08X}\
             {mtime:08X}\
             {filesize:08X}\
             {devmajor:08X}\
             {devminor:08X}\
             {rdevmajor:08X}\
             {rdevminor:08X}\
             {namesize:08X}\
             {check:08X}",
            uid = 0u32,
            gid = 0u32,
            nlink = if mode & 0o040000 != 0 { 2u32 } else { 1u32 },
            mtime = 0u32,
            devmajor = 0u32,
            devminor = 0u32,
            rdevmajor = 0u32,
            rdevminor = 0u32,
            namesize = name_len,
            check = 0u32,
        );

        debug_assert_eq!(header.len(), 110);

        self.data.extend_from_slice(header.as_bytes());
        self.data.extend_from_slice(name_with_nul.as_bytes());
        pad4(&mut self.data);

        if !content.is_empty() {
            self.data.extend_from_slice(content);
            pad4(&mut self.data);
        }
    }
}

impl Default for CpioBuilder {
    fn default() -> Self {
        Self::new()
    }
}

/// Pad `data` to the next 4-byte boundary with NUL bytes.
fn pad4(data: &mut Vec<u8>) {
    let rem = data.len() % 4;
    if rem != 0 {
        let padding = 4 - rem;
        data.extend(std::iter::repeat_n(0u8, padding));
    }
}

/// Build a complete initramfs with standard directory structure and an /init script.
///
/// `services` are the names of services to start in the init script
/// (e.g., "sshd", "rvf-server").
///
/// `extra_binaries` are (archive_path, content) pairs for additional binaries
/// to include (e.g., ("bin/busybox", &busybox_bytes)).
///
/// Returns a gzipped cpio archive.
pub fn build_initramfs(
    services: &[&str],
    extra_binaries: &[(&str, &[u8])],
) -> Result<Vec<u8>, KernelError> {
    let mut cpio = CpioBuilder::new();

    // Create directory structure
    let dirs = [
        ".",
        "bin",
        "sbin",
        "etc",
        "etc/udhcpc",
        "dev",
        "proc",
        "sys",
        "tmp",
        "var",
        "var/log",
        "var/run",
        "run",
        "root",
        "lib",
        "usr",
        "usr/bin",
        "usr/sbin",
        "usr/lib",
        "mnt",
        "opt",
    ];
    for dir in &dirs {
        cpio.add_dir(dir);
    }

    // Create essential device nodes
    cpio.add_device("dev/console", 0o020600, 5, 1);
    cpio.add_device("dev/ttyS0", 0o020660, 4, 64);
    cpio.add_device("dev/null", 0o020666, 1, 3);
    cpio.add_device("dev/zero", 0o020666, 1, 5);
    cpio.add_device("dev/urandom", 0o020444, 1, 9);

    // Create /init script
    let init_script = default_init_script(services);
    cpio.add_file("init", 0o100755, init_script.as_bytes());

    // Create a minimal udhcpc script
    let udhcpc_script = r#"#!/bin/sh
case "$1" in
    bound|renew)
        ip addr add "$ip/$mask" dev "$interface"
        if [ -n "$router" ]; then
            ip route add default via "$router"
        fi
        if [ -n "$dns" ]; then
            echo "nameserver $dns" > /etc/resolv.conf
        fi
        ;;
esac
"#;
    cpio.add_file(
        "etc/udhcpc/simple.script",
        0o100755,
        udhcpc_script.as_bytes(),
    );

    // Add extra binaries
    for (path, content) in extra_binaries {
        cpio.add_file(path, 0o100755, content);
    }

    cpio.finish_gzipped()
}

/// Build an ultra-fast boot initramfs optimized for minimal startup time.
///
/// Compared to `build_initramfs`, this:
/// - Skips network interface enumeration/DHCP
/// - Mounts only /proc, /sys, /dev (no /dev/pts, /dev/shm, /tmp, /run)
/// - No /etc setup (no passwd, resolv.conf, hostname)
/// - Starts services immediately without probing
/// - Uses minimal directory structure
///
/// Target: kernel-to-service in under 50ms of userspace init time.
pub fn build_fast_initramfs(
    services: &[&str],
    extra_binaries: &[(&str, &[u8])],
) -> Result<Vec<u8>, KernelError> {
    let mut cpio = CpioBuilder::new();

    // Minimal directory structure
    let dirs = [".", "bin", "sbin", "dev", "proc", "sys", "tmp", "run"];
    for dir in &dirs {
        cpio.add_dir(dir);
    }

    // Essential device nodes only
    cpio.add_device("dev/console", 0o020600, 5, 1);
    cpio.add_device("dev/ttyS0", 0o020660, 4, 64);
    cpio.add_device("dev/null", 0o020666, 1, 3);
    cpio.add_device("dev/urandom", 0o020444, 1, 9);

    // Ultra-fast /init script
    let mut script = String::from(
        "#!/bin/sh\n\
         mount -t proc proc /proc\n\
         mount -t sysfs sysfs /sys\n\
         mount -t devtmpfs devtmpfs /dev\n",
    );

    for service in services {
        match *service {
            "sshd" | "dropbear" => {
                script.push_str(
                    "mkdir -p /etc/dropbear\n\
                     dropbear -R -F -E -p 2222 &\n",
                );
            }
            "rvf-server" => {
                script.push_str("rvf-server --listen 0.0.0.0:8080 &\n");
            }
            other => {
                script.push_str(&format!("{other} &\n"));
            }
        }
    }

    script.push_str("exec /bin/sh\n");
    cpio.add_file("init", 0o100755, script.as_bytes());

    // Add extra binaries
    for (path, content) in extra_binaries {
        cpio.add_file(path, 0o100755, content);
    }

    cpio.finish_gzipped()
}

/// Parse a cpio newc archive and return the list of entries.
///
/// Each entry is returned as (path, mode, filesize, data_offset_in_archive).
/// This is primarily used for testing/verification.
pub fn parse_cpio_entries(data: &[u8]) -> Result<Vec<(String, u32, u32)>, KernelError> {
    let mut entries = Vec::new();
    let mut offset = 0;

    loop {
        if offset + 110 > data.len() {
            break;
        }

        let header = &data[offset..offset + 110];
        let header_str = std::str::from_utf8(header).map_err(|_| {
            KernelError::InitramfsBuildFailed("invalid cpio header encoding".into())
        })?;

        // Verify magic
        if &header_str[..6] != CPIO_NEWC_MAGIC {
            return Err(KernelError::InitramfsBuildFailed(format!(
                "invalid cpio magic at offset {offset}: {:?}",
                &header_str[..6]
            )));
        }

        let mode = u32::from_str_radix(&header_str[14..22], 16)
            .map_err(|_| KernelError::InitramfsBuildFailed("invalid mode".into()))?;
        let filesize = u32::from_str_radix(&header_str[54..62], 16)
            .map_err(|_| KernelError::InitramfsBuildFailed("invalid filesize".into()))?;
        let namesize = u32::from_str_radix(&header_str[94..102], 16)
            .map_err(|_| KernelError::InitramfsBuildFailed("invalid namesize".into()))?;

        // Name starts right after the 110-byte header
        let name_start = offset + 110;
        let name_end = name_start + namesize as usize;
        if name_end > data.len() {
            break;
        }

        let name_bytes = &data[name_start..name_end];
        // Strip trailing NUL
        let name = std::str::from_utf8(name_bytes)
            .map_err(|_| KernelError::InitramfsBuildFailed("invalid name encoding".into()))?
            .trim_end_matches('\0')
            .to_string();

        if name == "TRAILER!!!" {
            break;
        }

        // Advance past header + name (padded to 4 bytes)
        let after_name = align4(name_end);
        // Advance past data (padded to 4 bytes)
        let data_end = after_name + filesize as usize;
        let next_entry = if filesize > 0 {
            align4(data_end)
        } else {
            after_name
        };

        entries.push((name, mode, filesize));
        offset = next_entry;
    }

    Ok(entries)
}

fn align4(val: usize) -> usize {
    (val + 3) & !3
}

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

    #[test]
    fn cpio_builder_produces_valid_archive() {
        let mut cpio = CpioBuilder::new();
        cpio.add_dir("bin");
        cpio.add_file("init", 0o100755, b"#!/bin/sh\necho hello\n");
        cpio.add_symlink("bin/sh", "/bin/busybox");
        let archive = cpio.finish();

        // Verify it starts with the cpio magic
        assert!(archive.starts_with(b"070701"));

        // Parse it back
        let entries = parse_cpio_entries(&archive).expect("parse should succeed");
        assert_eq!(entries.len(), 3);

        // Check directory
        assert_eq!(entries[0].0, "bin");
        assert_eq!(entries[0].1 & 0o040000, 0o040000); // is a directory

        // Check file
        assert_eq!(entries[1].0, "init");
        assert_eq!(entries[1].1 & 0o100000, 0o100000); // is a regular file
        assert_eq!(entries[1].1 & 0o755, 0o755); // executable
        assert_eq!(entries[1].2, 21); // file size: "#!/bin/sh\necho hello\n"

        // Check symlink
        assert_eq!(entries[2].0, "bin/sh");
        assert_eq!(entries[2].1 & 0o120000, 0o120000); // is a symlink
    }

    #[test]
    fn cpio_archive_ends_with_trailer() {
        let cpio = CpioBuilder::new();
        let archive = cpio.finish();
        let as_str = String::from_utf8_lossy(&archive);
        assert!(as_str.contains("TRAILER!!!"));
    }

    #[test]
    fn build_initramfs_produces_gzipped_cpio() {
        let result = build_initramfs(&["sshd"], &[]).expect("build should succeed");

        // gzip magic bytes
        assert_eq!(result[0], 0x1F);
        assert_eq!(result[1], 0x8B);

        // Decompress and verify
        use flate2::read::GzDecoder;
        use std::io::Read;
        let mut decoder = GzDecoder::new(&result[..]);
        let mut decompressed = Vec::new();
        decoder
            .read_to_end(&mut decompressed)
            .expect("decompress should succeed");

        // Verify it's a valid cpio archive
        let entries = parse_cpio_entries(&decompressed).expect("parse should succeed");

        // Should have directories + devices + init + udhcpc script
        assert!(
            entries.len() >= 20,
            "expected at least 20 entries, got {}",
            entries.len()
        );

        // Check that /init exists
        let init_entry = entries.iter().find(|(name, _, _)| name == "init");
        assert!(init_entry.is_some(), "must have /init");

        // Check that directories exist
        let dir_names: Vec<&str> = entries
            .iter()
            .filter(|(_, mode, _)| mode & 0o040000 != 0)
            .map(|(name, _, _)| name.as_str())
            .collect();
        assert!(dir_names.contains(&"bin"));
        assert!(dir_names.contains(&"etc"));
        assert!(dir_names.contains(&"proc"));
        assert!(dir_names.contains(&"sys"));
        assert!(dir_names.contains(&"dev"));
    }

    #[test]
    fn build_initramfs_with_extra_binaries() {
        let fake_binary = b"\x7FELF fake binary content";
        let result = build_initramfs(&["rvf-server"], &[("bin/rvf-server", fake_binary)])
            .expect("build should succeed");

        // Decompress
        use flate2::read::GzDecoder;
        use std::io::Read;
        let mut decoder = GzDecoder::new(&result[..]);
        let mut decompressed = Vec::new();
        decoder.read_to_end(&mut decompressed).unwrap();

        let entries = parse_cpio_entries(&decompressed).unwrap();
        let binary_entry = entries.iter().find(|(name, _, _)| name == "bin/rvf-server");
        assert!(binary_entry.is_some(), "must have bin/rvf-server");
        assert_eq!(binary_entry.unwrap().2, fake_binary.len() as u32);
    }

    #[test]
    fn default_init_script_mounts_filesystems() {
        let script = default_init_script(&[]);
        assert!(script.contains("mount -t proc proc /proc"));
        assert!(script.contains("mount -t sysfs sysfs /sys"));
        assert!(script.contains("mount -t devtmpfs devtmpfs /dev"));
    }

    #[test]
    fn default_init_script_includes_services() {
        let script = default_init_script(&["sshd", "rvf-server"]);
        assert!(script.contains("dropbear") || script.contains("sshd"));
        assert!(script.contains("rvf-server"));
    }

    #[test]
    fn cpio_header_is_110_bytes() {
        let mut cpio = CpioBuilder::new();
        cpio.add_file("x", 0o100644, b"");
        let archive = cpio.finish();
        // First entry header should be exactly 110 ASCII chars
        let header_str = std::str::from_utf8(&archive[..110]).unwrap();
        assert!(header_str.starts_with(CPIO_NEWC_MAGIC));
    }

    #[test]
    fn build_fast_initramfs_is_smaller() {
        let normal = build_initramfs(&["sshd", "rvf-server"], &[]).unwrap();
        let fast = build_fast_initramfs(&["sshd", "rvf-server"], &[]).unwrap();

        // Fast initramfs should be smaller (fewer dirs, shorter init script)
        assert!(
            fast.len() < normal.len(),
            "fast ({}) should be smaller than normal ({})",
            fast.len(),
            normal.len()
        );

        // Both should be valid gzip
        assert_eq!(fast[0], 0x1F);
        assert_eq!(fast[1], 0x8B);

        // Decompress and verify it has /init
        use flate2::read::GzDecoder;
        use std::io::Read;
        let mut decoder = GzDecoder::new(&fast[..]);
        let mut decompressed = Vec::new();
        decoder.read_to_end(&mut decompressed).unwrap();
        let entries = parse_cpio_entries(&decompressed).unwrap();
        let has_init = entries.iter().any(|(name, _, _)| name == "init");
        assert!(has_init, "fast initramfs must have /init");
    }

    #[test]
    fn device_nodes_are_parseable() {
        let mut cpio = CpioBuilder::new();
        cpio.add_device("dev/null", 0o020666, 1, 3);
        let archive = cpio.finish();
        let entries = parse_cpio_entries(&archive).unwrap();
        assert_eq!(entries.len(), 1);
        assert_eq!(entries[0].0, "dev/null");
        // Character device bit
        assert_eq!(entries[0].1 & 0o020000, 0o020000);
    }
}