//! age.rs — AGE (rage) helpers + persistent key management for your mini-Redis.
//
// Features:
// - X25519 encryption/decryption (age style)
// - Ed25519 detached signatures + verification
// - Persistent named keys in DB (strings):
//      age:key:{name}       -> X25519 recipient (public encryption key, "age1...")
//      age:privkey:{name}   -> X25519 identity (secret encryption key, "AGE-SECRET-KEY-1...")
//      age:signpub:{name}   -> Ed25519 verify pubkey (public, used to verify signatures)
//      age:signpriv:{name}  -> Ed25519 signing secret key (private, used to sign)
// - Base64 wrapping for ciphertext/signature binary blobs.

use std::str::FromStr;

use secrecy::ExposeSecret;
use age::{Decryptor, Encryptor};
use age::x25519;

use ed25519_dalek::{Signature, Signer, Verifier, SigningKey, VerifyingKey};

use base64::{engine::general_purpose::STANDARD as B64, Engine as _};

use crate::protocol::Protocol;
use crate::server::Server;
use crate::error::DBError;

// ---------- Internal helpers ----------

#[derive(Debug)]
pub enum AgeWireError {
    ParseKey,
    Crypto(String),
    Utf8,
    SignatureLen,
    NotFound(&'static str),     // which kind of key was missing
    Storage(String),
}

impl AgeWireError {
    fn to_protocol(self) -> Protocol {
        match self {
            AgeWireError::ParseKey => Protocol::err("ERR age: invalid key"),
            AgeWireError::Crypto(e) => Protocol::err(&format!("ERR age: {e}")),
            AgeWireError::Utf8 => Protocol::err("ERR age: invalid UTF-8 plaintext"),
            AgeWireError::SignatureLen => Protocol::err("ERR age: bad signature length"),
            AgeWireError::NotFound(w) => Protocol::err(&format!("ERR age: missing {w}")),
            AgeWireError::Storage(e) => Protocol::err(&format!("ERR storage: {e}")),
        }
    }
}

fn parse_recipient(s: &str) -> Result<x25519::Recipient, AgeWireError> {
    x25519::Recipient::from_str(s).map_err(|_| AgeWireError::ParseKey)
}
fn parse_identity(s: &str) -> Result<x25519::Identity, AgeWireError> {
    x25519::Identity::from_str(s).map_err(|_| AgeWireError::ParseKey)
}
fn parse_ed25519_signing_key(s: &str) -> Result<SigningKey, AgeWireError> {
    // Parse base64-encoded signing key
    let bytes = B64.decode(s).map_err(|_| AgeWireError::ParseKey)?;
    if bytes.len() != 32 {
        return Err(AgeWireError::ParseKey);
    }
    let key_bytes: [u8; 32] = bytes.try_into().map_err(|_| AgeWireError::ParseKey)?;
    Ok(SigningKey::from_bytes(&key_bytes))
}
fn parse_ed25519_verifying_key(s: &str) -> Result<VerifyingKey, AgeWireError> {
    // Parse base64-encoded verifying key
    let bytes = B64.decode(s).map_err(|_| AgeWireError::ParseKey)?;
    if bytes.len() != 32 {
        return Err(AgeWireError::ParseKey);
    }
    let key_bytes: [u8; 32] = bytes.try_into().map_err(|_| AgeWireError::ParseKey)?;
    VerifyingKey::from_bytes(&key_bytes).map_err(|_| AgeWireError::ParseKey)
}

// ---------- Stateless crypto helpers (string in/out) ----------

pub fn gen_enc_keypair() -> (String, String) {
    let id = x25519::Identity::generate();
    let pk = id.to_public();
    (pk.to_string(), id.to_string().expose_secret().to_string()) // (recipient, identity)
}

pub fn gen_sign_keypair() -> (String, String) {
    use rand::RngCore;
    use rand::rngs::OsRng;
    
    // Generate random 32 bytes for the signing key
    let mut secret_bytes = [0u8; 32];
    OsRng.fill_bytes(&mut secret_bytes);
    
    let signing_key = SigningKey::from_bytes(&secret_bytes);
    let verifying_key = signing_key.verifying_key();
    
    // Encode as base64 for storage
    let signing_key_b64 = B64.encode(signing_key.to_bytes());
    let verifying_key_b64 = B64.encode(verifying_key.to_bytes());
    
    (verifying_key_b64, signing_key_b64) // (verify_pub, signing_secret)
}

/// Encrypt `msg` for `recipient_str` (X25519). Returns base64(ciphertext).
pub fn encrypt_b64(recipient_str: &str, msg: &str) -> Result<String, AgeWireError> {
    let recipient = parse_recipient(recipient_str)?;
    let enc = Encryptor::with_recipients(vec![Box::new(recipient)])
        .expect("failed to create encryptor"); // Handle Option<Encryptor>
    let mut out = Vec::new();
    {
        use std::io::Write;
        let mut w = enc.wrap_output(&mut out).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
        w.write_all(msg.as_bytes()).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
        w.finish().map_err(|e| AgeWireError::Crypto(e.to_string()))?;
    }
    Ok(B64.encode(out))
}

/// Decrypt base64(ciphertext) with `identity_str`. Returns plaintext String.
pub fn decrypt_b64(identity_str: &str, ct_b64: &str) -> Result<String, AgeWireError> {
    let id = parse_identity(identity_str)?;
    let ct = B64.decode(ct_b64.as_bytes()).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
    let dec = Decryptor::new(&ct[..]).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
    
    // The decrypt method returns a Result<StreamReader, DecryptError>
    let mut r = match dec {
        Decryptor::Recipients(d) => d.decrypt(std::iter::once(&id as &dyn age::Identity))
            .map_err(|e| AgeWireError::Crypto(e.to_string()))?,
        Decryptor::Passphrase(_) => return Err(AgeWireError::Crypto("Expected recipients, got passphrase".to_string())),
    };
    
    let mut pt = Vec::new();
    use std::io::Read;
    r.read_to_end(&mut pt).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
    String::from_utf8(pt).map_err(|_| AgeWireError::Utf8)
}

/// Sign bytes of `msg` (detached). Returns base64(signature bytes, 64 bytes).
pub fn sign_b64(signing_secret_str: &str, msg: &str) -> Result<String, AgeWireError> {
    let signing_key = parse_ed25519_signing_key(signing_secret_str)?;
    let sig = signing_key.sign(msg.as_bytes());
    Ok(B64.encode(sig.to_bytes()))
}

/// Verify detached signature (base64) for `msg` with pubkey.
pub fn verify_b64(verify_pub_str: &str, msg: &str, sig_b64: &str) -> Result<bool, AgeWireError> {
    let verifying_key = parse_ed25519_verifying_key(verify_pub_str)?;
    let sig_bytes = B64.decode(sig_b64.as_bytes()).map_err(|e| AgeWireError::Crypto(e.to_string()))?;
    if sig_bytes.len() != 64 {
        return Err(AgeWireError::SignatureLen);
    }
    let sig = Signature::from_bytes(sig_bytes[..].try_into().unwrap());
    Ok(verifying_key.verify(msg.as_bytes(), &sig).is_ok())
}

// ---------- Storage helpers ----------

fn sget(server: &Server, key: &str) -> Result<Option<String>, AgeWireError> {
    let st = server.current_storage().map_err(|e| AgeWireError::Storage(e.0))?;
    st.get(key).map_err(|e| AgeWireError::Storage(e.0))
}
fn sset(server: &Server, key: &str, val: &str) -> Result<(), AgeWireError> {
    let st = server.current_storage().map_err(|e| AgeWireError::Storage(e.0))?;
    st.set(key.to_string(), val.to_string()).map_err(|e| AgeWireError::Storage(e.0))
}

fn enc_pub_key_key(name: &str) -> String { format!("age:key:{name}") }
fn enc_priv_key_key(name: &str) -> String { format!("age:privkey:{name}") }
fn sign_pub_key_key(name: &str) -> String { format!("age:signpub:{name}") }
fn sign_priv_key_key(name: &str) -> String { format!("age:signpriv:{name}") }

// ---------- Command handlers (RESP Protocol) ----------
// Basic (stateless) ones kept for completeness

pub async fn cmd_age_genenc() -> Protocol {
    let (recip, ident) = gen_enc_keypair();
    Protocol::Array(vec![Protocol::BulkString(recip), Protocol::BulkString(ident)])
}

pub async fn cmd_age_gensign() -> Protocol {
    let (verify, secret) = gen_sign_keypair();
    Protocol::Array(vec![Protocol::BulkString(verify), Protocol::BulkString(secret)])
}

pub async fn cmd_age_encrypt(recipient: &str, message: &str) -> Protocol {
    match encrypt_b64(recipient, message) {
        Ok(b64) => Protocol::BulkString(b64),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_decrypt(identity: &str, ct_b64: &str) -> Protocol {
    match decrypt_b64(identity, ct_b64) {
        Ok(pt) => Protocol::BulkString(pt),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_sign(secret: &str, message: &str) -> Protocol {
    match sign_b64(secret, message) {
        Ok(b64sig) => Protocol::BulkString(b64sig),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_verify(verify_pub: &str, message: &str, sig_b64: &str) -> Protocol {
    match verify_b64(verify_pub, message, sig_b64) {
        Ok(true) => Protocol::SimpleString("1".to_string()),
        Ok(false) => Protocol::SimpleString("0".to_string()),
        Err(e) => e.to_protocol(),
    }
}

// ---------- NEW: Persistent, named-key commands ----------

pub async fn cmd_age_keygen(server: &Server, name: &str) -> Protocol {
    let (recip, ident) = gen_enc_keypair();
    if let Err(e) = sset(server, &enc_pub_key_key(name), &recip) { return e.to_protocol(); }
    if let Err(e) = sset(server, &enc_priv_key_key(name), &ident) { return e.to_protocol(); }
    Protocol::Array(vec![Protocol::BulkString(recip), Protocol::BulkString(ident)])
}

pub async fn cmd_age_signkeygen(server: &Server, name: &str) -> Protocol {
    let (verify, secret) = gen_sign_keypair();
    if let Err(e) = sset(server, &sign_pub_key_key(name), &verify) { return e.to_protocol(); }
    if let Err(e) = sset(server, &sign_priv_key_key(name), &secret) { return e.to_protocol(); }
    Protocol::Array(vec![Protocol::BulkString(verify), Protocol::BulkString(secret)])
}

pub async fn cmd_age_encrypt_name(server: &Server, name: &str, message: &str) -> Protocol {
    let recip = match sget(server, &enc_pub_key_key(name)) {
        Ok(Some(v)) => v,
        Ok(None) => return AgeWireError::NotFound("recipient (age:key:{name})").to_protocol(),
        Err(e) => return e.to_protocol(),
    };
    match encrypt_b64(&recip, message) {
        Ok(ct) => Protocol::BulkString(ct),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_decrypt_name(server: &Server, name: &str, ct_b64: &str) -> Protocol {
    let ident = match sget(server, &enc_priv_key_key(name)) {
        Ok(Some(v)) => v,
        Ok(None) => return AgeWireError::NotFound("identity (age:privkey:{name})").to_protocol(),
        Err(e) => return e.to_protocol(),
    };
    match decrypt_b64(&ident, ct_b64) {
        Ok(pt) => Protocol::BulkString(pt),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_sign_name(server: &Server, name: &str, message: &str) -> Protocol {
    let sec = match sget(server, &sign_priv_key_key(name)) {
        Ok(Some(v)) => v,
        Ok(None) => return AgeWireError::NotFound("signing secret (age:signpriv:{name})").to_protocol(),
        Err(e) => return e.to_protocol(),
    };
    match sign_b64(&sec, message) {
        Ok(sig) => Protocol::BulkString(sig),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_verify_name(server: &Server, name: &str, message: &str, sig_b64: &str) -> Protocol {
    let pubk = match sget(server, &sign_pub_key_key(name)) {
        Ok(Some(v)) => v,
        Ok(None) => return AgeWireError::NotFound("verify pubkey (age:signpub:{name})").to_protocol(),
        Err(e) => return e.to_protocol(),
    };
    match verify_b64(&pubk, message, sig_b64) {
        Ok(true) => Protocol::SimpleString("1".to_string()),
        Ok(false) => Protocol::SimpleString("0".to_string()),
        Err(e) => e.to_protocol(),
    }
}

pub async fn cmd_age_list(server: &Server) -> Protocol {
    // Returns 4 arrays: ["encpub", <names...>], ["encpriv", ...], ["signpub", ...], ["signpriv", ...]
    let st = match server.current_storage() { Ok(s) => s, Err(e) => return Protocol::err(&e.0) };

    let pull = |pat: &str, prefix: &str| -> Result<Vec<String>, DBError> {
        let keys = st.keys(pat)?;
        let mut names: Vec<String> = keys.into_iter()
            .filter_map(|k| k.strip_prefix(prefix).map(|x| x.to_string()))
            .collect();
        names.sort();
        Ok(names)
    };

    let encpub  = match pull("age:key:*",      "age:key:")      { Ok(v) => v, Err(e)=> return Protocol::err(&e.0) };
    let encpriv = match pull("age:privkey:*",  "age:privkey:")  { Ok(v) => v, Err(e)=> return Protocol::err(&e.0) };
    let signpub = match pull("age:signpub:*",  "age:signpub:")  { Ok(v) => v, Err(e)=> return Protocol::err(&e.0) };
    let signpriv= match pull("age:signpriv:*", "age:signpriv:") { Ok(v) => v, Err(e)=> return Protocol::err(&e.0) };

    let to_arr = |label: &str, v: Vec<String>| {
        let mut out = vec![Protocol::BulkString(label.to_string())];
        out.push(Protocol::Array(v.into_iter().map(Protocol::BulkString).collect()));
        Protocol::Array(out)
    };

    Protocol::Array(vec![
        to_arr("encpub", encpub),
        to_arr("encpriv", encpriv),
        to_arr("signpub", signpub),
        to_arr("signpriv", signpriv),
    ])
}