use base64::engine::general_purpose::STANDARD;
use base64::Engine as _;
use rand::{rngs::OsRng, CryptoRng, RngCore};

use crate::node::Ed25519Material;

const OPENSSH_HEADER: &str = "-----BEGIN OPENSSH PRIVATE KEY-----";
const OPENSSH_FOOTER: &str = "-----END OPENSSH PRIVATE KEY-----";

/// Encode an [`Ed25519Material`] instance as an OpenSSH private key.
pub fn encode_private_key(material: &Ed25519Material, comment: &str) -> String {
    let mut rng = OsRng;
    encode_private_key_with_rng(material, comment, &mut rng)
}

/// Encode an [`Ed25519Material`] instance as an OpenSSH private key using a custom RNG.
///
/// This is primarily exposed for deterministic unit testing.
pub fn encode_private_key_with_rng<R>(
    material: &Ed25519Material,
    comment: &str,
    rng: &mut R,
) -> String
where
    R: RngCore + CryptoRng,
{
    let mut document = Vec::new();
    document.extend_from_slice(b"openssh-key-v1\0");
    append_ssh_string(&mut document, b"none");
    append_ssh_string(&mut document, b"none");
    append_ssh_string(&mut document, &[]);
    document.extend_from_slice(&1u32.to_be_bytes());
    append_ssh_string(&mut document, &material.blob);

    let mut private_section = Vec::new();
    let check = rng.next_u32();
    private_section.extend_from_slice(&check.to_be_bytes());
    private_section.extend_from_slice(&check.to_be_bytes());
    append_ssh_string(&mut private_section, b"ssh-ed25519");
    append_ssh_string(&mut private_section, &material.public_key);
    append_ssh_string(&mut private_section, &material.private_key);
    append_ssh_string(&mut private_section, comment.as_bytes());
    append_padding(&mut private_section);

    append_ssh_string(&mut document, &private_section);

    let encoded = STANDARD.encode(&document);
    format_base64_document(&encoded)
}

/// Format the authorized key line, optionally appending a comment.
pub fn format_authorized_key(material: &Ed25519Material, comment: Option<&str>) -> String {
    let mut line = material.authorized_key.trim_end().to_string();
    if let Some(comment) = comment {
        if !comment.trim().is_empty() {
            line.push(' ');
            line.push_str(comment.trim());
        }
    }
    line.push('\n');
    line
}

fn append_ssh_string(buf: &mut Vec<u8>, data: &[u8]) {
    buf.extend_from_slice(&(data.len() as u32).to_be_bytes());
    buf.extend_from_slice(data);
}

fn append_padding(buf: &mut Vec<u8>) {
    let pad_len = (8 - (buf.len() % 8)) % 8;
    if pad_len == 0 {
        return;
    }
    for i in 1..=pad_len {
        buf.push(i as u8);
    }
}

fn format_base64_document(encoded: &str) -> String {
    let mut out = String::new();
    out.push_str(OPENSSH_HEADER);
    out.push('\n');
    for chunk in encoded.as_bytes().chunks(70) {
        out.push_str(&String::from_utf8_lossy(chunk));
        out.push('\n');
    }
    out.push_str(OPENSSH_FOOTER);
    out.push('\n');
    out
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::Node;
    use rand::{rngs::StdRng, SeedableRng};
    use std::io::{Cursor, Read};

    #[test]
    fn encode_private_key_renders_valid_document() {
        let seed = b"deterministic-seed-for-testing";
        let node = Node::master(seed).unwrap();
        let material = node.ed25519().unwrap();
        let mut rng = StdRng::seed_from_u64(42);
        let document = encode_private_key_with_rng(&material, "comment", &mut rng);

        assert!(document.starts_with(OPENSSH_HEADER));
        assert!(document.ends_with(&format!("{}\n", OPENSSH_FOOTER)));

        let encoded = document
            .lines()
            .filter(|line| !line.starts_with("-----"))
            .collect::<String>();
        let decoded = STANDARD.decode(encoded).unwrap();

        let mut cursor = Cursor::new(decoded.as_slice());
        let mut magic = [0u8; 15];
        cursor.read_exact(&mut magic).unwrap();
        assert_eq!(&magic, b"openssh-key-v1\0");

        let cipher = read_string(&mut cursor);
        assert_eq!(cipher, b"none");
        let kdf = read_string(&mut cursor);
        assert_eq!(kdf, b"none");
        let kdf_opts = read_string(&mut cursor);
        assert!(kdf_opts.is_empty());
        let key_count = read_u32(&mut cursor);
        assert_eq!(key_count, 1);
    }

    #[test]
    fn format_authorized_key_appends_comment() {
        let seed = b"another-deterministic-seed";
        let node = Node::master(seed).unwrap();
        let material = node.ed25519().unwrap();
        let line = format_authorized_key(&material, Some("example comment"));
        assert!(line.trim_end().ends_with("example comment"));
        let without = format_authorized_key(&material, None);
        assert!(!without.trim_end().ends_with("example comment"));
    }
    fn read_string(cursor: &mut Cursor<&[u8]>) -> Vec<u8> {
        let mut len_bytes = [0u8; 4];
        cursor.read_exact(&mut len_bytes).unwrap();
        let len = u32::from_be_bytes(len_bytes) as usize;
        let mut buf = vec![0u8; len];
        cursor.read_exact(&mut buf).unwrap();
        buf
    }

    fn read_u32(cursor: &mut Cursor<&[u8]>) -> u32 {
        let mut bytes = [0u8; 4];
        cursor.read_exact(&mut bytes).unwrap();
        u32::from_be_bytes(bytes)
    }
}
