//! Worker-binary CLI parsing helpers.
//!
//! The `supermachine-worker` binary's argument loop inlined all of its
//! parsing + validation and called `std::process::exit` on every bad
//! value, which made none of it unit-testable. These pure functions
//! return `Result<_, String>` so the binary stays a thin
//! "parse → on Err, print + exit(2)" wrapper while the parsing and
//! validation — including the security-sensitive vsock auth token — are
//! exercised by the tests below.

use crate::vmm::resources::{MountSpec, SymlinkPolicy, VolumeSpec};

/// The vsock TSI control-channel auth token.
///
/// `bytes` is installed on `VmResources` (the muxer enforces it on every
/// control dispatch); `hex` is the canonical lowercase form appended to
/// the kernel cmdline as `supermachine.tsi_token=<hex>`. The two are kept
/// together so a caller can't accidentally feed the cmdline a different
/// rendering than the bytes it enforces.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct TsiToken {
    pub hex: String,
    pub bytes: [u8; 32],
}

/// Parse a `--tsi-token` value: exactly 64 hex characters (= 32 bytes),
/// case-insensitive on input. Returns the canonical lowercase hex plus
/// the decoded bytes. Anything else — wrong length, a non-hex digit — is
/// an error (the binary turns it into exit code 2).
pub fn parse_tsi_token(s: &str) -> Result<TsiToken, String> {
    if s.len() != 64 || !s.bytes().all(|b| b.is_ascii_hexdigit()) {
        return Err(format!(
            "--tsi-token: expected 64 hex chars (32 bytes), got {} chars",
            s.len()
        ));
    }
    let hex = s.to_ascii_lowercase();
    let mut bytes = [0u8; 32];
    for (i, b) in bytes.iter_mut().enumerate() {
        *b = u8::from_str_radix(&hex[i * 2..i * 2 + 2], 16)
            .map_err(|e| format!("--tsi-token: malformed hex at byte {i}: {e}"))?;
    }
    Ok(TsiToken { hex, bytes })
}

impl TsiToken {
    /// Generate a fresh random control-channel token from the OS CSPRNG
    /// (`/dev/urandom`). Used by the KVM cold-boot path, which mints a
    /// per-VM token rather than receiving one over the worker CLI.
    pub fn generate() -> std::io::Result<Self> {
        use std::io::Read;
        let mut bytes = [0u8; 32];
        std::fs::File::open("/dev/urandom")?.read_exact(&mut bytes)?;
        Ok(Self {
            hex: hex_lower(&bytes),
            bytes,
        })
    }
}

/// Lowercase-hex render of a byte slice (no separators).
pub fn hex_lower(bytes: &[u8]) -> String {
    let mut s = String::with_capacity(bytes.len() * 2);
    for b in bytes {
        s.push(char::from_digit((b >> 4) as u32, 16).unwrap());
        s.push(char::from_digit((b & 0xf) as u32, 16).unwrap());
    }
    s
}

/// Append `supermachine.tsi_token=<hex>` to a kernel cmdline, inserting a
/// single space separator only when the existing cmdline needs one.
pub fn append_tsi_token_cmdline(cmdline: &mut String, hex: &str) {
    if !cmdline.is_empty() && !cmdline.ends_with(' ') {
        cmdline.push(' ');
    }
    cmdline.push_str("supermachine.tsi_token=");
    cmdline.push_str(hex);
}

/// Parse a `--volume HOST:GUEST[:SIZE_BYTES]` spec. Size is optional
/// (`VolumeSpec` applies the 1 GiB default). Splits on at most two `:`
/// so the host path keeps any later colons (host paths on macOS never
/// contain `:`, but the guest path / size never do either).
pub fn parse_volume_spec(raw: &str) -> Result<VolumeSpec, String> {
    let parts: Vec<&str> = raw.splitn(3, ':').collect();
    if parts.len() < 2 {
        return Err(format!(
            "--volume expects HOST:GUEST[:SIZE_BYTES], got {raw:?}"
        ));
    }
    let mut spec = VolumeSpec::new(parts[0], parts[1]);
    if let Some(s) = parts.get(2) {
        let sz = s
            .parse::<u64>()
            .map_err(|_| format!("--volume SIZE_BYTES not a u64: {s:?}"))?;
        spec = spec.with_size_bytes(sz);
    }
    Ok(spec)
}

/// Parse a `--mount HOST:TAG:GUEST_PATH[:POLICY]` spec. POLICY is one of
/// `deny` / `opaque` / `follow` (default `Opaque`). Validates the tag
/// (1..=35 bytes) and an absolute guest path.
pub fn parse_mount_spec(raw: &str) -> Result<MountSpec, String> {
    let parts: Vec<&str> = raw.splitn(4, ':').collect();
    let (host, tag, guest_path, policy_str) = match parts.len() {
        3 => (parts[0], parts[1], parts[2], None),
        4 => (parts[0], parts[1], parts[2], Some(parts[3])),
        _ => {
            return Err(format!(
                "--mount expects HOST:TAG:GUEST_PATH[:POLICY], got {raw:?}"
            ))
        }
    };
    if tag.is_empty() {
        return Err(format!("--mount tag is empty: {raw:?}"));
    }
    if tag.len() > 35 {
        return Err(format!(
            "--mount tag too long (max 35 bytes, got {}): {raw:?}",
            tag.len()
        ));
    }
    if guest_path.is_empty() {
        return Err(format!("--mount guest_path is empty: {raw:?}"));
    }
    if !guest_path.starts_with('/') {
        return Err(format!(
            "--mount guest_path must be absolute (start with `/`), got {guest_path:?}"
        ));
    }
    let policy = match policy_str {
        None => SymlinkPolicy::default(),
        Some("deny") => SymlinkPolicy::Deny,
        Some("opaque") => SymlinkPolicy::Opaque,
        Some("follow") => SymlinkPolicy::Follow,
        Some(other) => {
            return Err(format!(
                "--mount policy must be one of: deny, opaque, follow (got {other:?})"
            ))
        }
    };
    Ok(MountSpec::new(host, tag, guest_path).with_symlinks(policy))
}

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

    // ── --tsi-token (vsock auth secret) ────────────────────────────────

    #[test]
    fn tsi_token_decodes_64_hex_to_32_bytes() {
        let hex = "00112233445566778899aabbccddeeff00112233445566778899aabbccddeeff";
        let t = parse_tsi_token(hex).unwrap();
        assert_eq!(t.hex, hex);
        assert_eq!(t.bytes[0], 0x00);
        assert_eq!(t.bytes[1], 0x11);
        assert_eq!(t.bytes[15], 0xff);
        assert_eq!(t.bytes[31], 0xff);
    }

    #[test]
    fn tsi_token_uppercase_is_canonicalised_to_lowercase() {
        let up = "AABBCCDDEEFF00112233445566778899AABBCCDDEEFF00112233445566778899";
        let t = parse_tsi_token(up).unwrap();
        assert_eq!(t.hex, up.to_ascii_lowercase());
        assert_eq!(t.bytes[0], 0xaa);
        // The cmdline rendering must round-trip back to the same bytes.
        assert_eq!(parse_tsi_token(&t.hex).unwrap().bytes, t.bytes);
    }

    #[test]
    fn tsi_token_all_zero_is_valid() {
        let z = "0".repeat(64);
        assert_eq!(parse_tsi_token(&z).unwrap().bytes, [0u8; 32]);
    }

    #[test]
    fn hex_lower_renders_each_byte_as_two_lowercase_digits() {
        assert_eq!(hex_lower(&[]), "");
        assert_eq!(hex_lower(&[0x00, 0x0f, 0xa5, 0xff]), "000fa5ff");
        // hex_lower is the exact inverse of parse_tsi_token's decode.
        let bytes: [u8; 32] = std::array::from_fn(|i| (i * 7 + 1) as u8);
        let hex = hex_lower(&bytes);
        assert_eq!(hex.len(), 64);
        assert_eq!(parse_tsi_token(&hex).unwrap().bytes, bytes);
    }

    #[test]
    fn generate_yields_canonical_distinct_tokens() {
        let a = TsiToken::generate().expect("urandom");
        let b = TsiToken::generate().expect("urandom");
        // 64 lowercase hex chars that re-parse to the same 32 bytes.
        assert_eq!(a.hex.len(), 64);
        assert_eq!(a.hex, hex_lower(&a.bytes));
        assert_eq!(parse_tsi_token(&a.hex).unwrap().bytes, a.bytes);
        // Two CSPRNG draws must not collide (1 / 2^256 chance otherwise).
        assert_ne!(a.bytes, b.bytes);
    }

    #[test]
    fn tsi_token_rejects_bad_length() {
        assert!(parse_tsi_token("").is_err());
        assert!(parse_tsi_token("abcd").is_err()); // too short
        assert!(parse_tsi_token(&"a".repeat(63)).is_err()); // off by one
        assert!(parse_tsi_token(&"a".repeat(65)).is_err());
    }

    #[test]
    fn tsi_token_rejects_non_hex() {
        // Right length, but contains non-hex chars.
        let mut s = "a".repeat(63);
        s.push('z');
        assert!(parse_tsi_token(&s).is_err());
        // Spaces / punctuation rejected too.
        let s2 = format!("{} ", "a".repeat(63));
        assert!(parse_tsi_token(&s2).is_err());
    }

    // ── cmdline append ────────────────────────────────────────────────

    #[test]
    fn append_token_inserts_space_only_when_needed() {
        let mut c = String::from("console=ttyAMA0 root=/dev/vda");
        append_tsi_token_cmdline(&mut c, "deadbeef");
        assert_eq!(
            c,
            "console=ttyAMA0 root=/dev/vda supermachine.tsi_token=deadbeef"
        );

        let mut c2 = String::from("trailing ");
        append_tsi_token_cmdline(&mut c2, "ab");
        assert_eq!(c2, "trailing supermachine.tsi_token=ab"); // no double space

        let mut c3 = String::new();
        append_tsi_token_cmdline(&mut c3, "ff");
        assert_eq!(c3, "supermachine.tsi_token=ff"); // no leading space
    }

    // ── --volume ──────────────────────────────────────────────────────

    #[test]
    fn volume_two_and_three_field_forms() {
        let v = parse_volume_spec("/host/db:/var/lib/db").unwrap();
        assert_eq!(v.host_path, "/host/db");
        assert_eq!(v.guest_path, "/var/lib/db");

        let v3 = parse_volume_spec("/h:/g:4096").unwrap();
        assert_eq!(v3.size_bytes, 4096);
    }

    #[test]
    fn volume_rejects_missing_guest_and_bad_size() {
        assert!(parse_volume_spec("/only-host").is_err());
        assert!(parse_volume_spec("").is_err());
        assert!(parse_volume_spec("/h:/g:notanum").is_err());
    }

    // ── --mount ───────────────────────────────────────────────────────

    #[test]
    fn mount_three_field_defaults_to_opaque() {
        let m = parse_mount_spec("/host/share:work:/mnt/work").unwrap();
        assert_eq!(m.host_path, "/host/share");
        assert_eq!(m.guest_tag, "work");
        assert_eq!(m.guest_path, "/mnt/work");
        assert_eq!(m.symlinks, SymlinkPolicy::Opaque);
    }

    #[test]
    fn mount_policy_each_value_parses() {
        assert_eq!(
            parse_mount_spec("/h:t:/g:deny").unwrap().symlinks,
            SymlinkPolicy::Deny
        );
        assert_eq!(
            parse_mount_spec("/h:t:/g:opaque").unwrap().symlinks,
            SymlinkPolicy::Opaque
        );
        assert_eq!(
            parse_mount_spec("/h:t:/g:follow").unwrap().symlinks,
            SymlinkPolicy::Follow
        );
    }

    #[test]
    fn mount_rejects_malformed() {
        assert!(parse_mount_spec("/h:t").is_err(), "too few fields");
        assert!(parse_mount_spec("/h::/g").is_err(), "empty tag");
        assert!(
            parse_mount_spec(&format!("/h:{}:/g", "x".repeat(36))).is_err(),
            "tag > 35 bytes"
        );
        assert!(
            parse_mount_spec("/h:t:relative").is_err(),
            "non-absolute guest path"
        );
        assert!(parse_mount_spec("/h:t:/g:bogus").is_err(), "unknown policy");
    }
}
