use base64::{Engine, engine::general_purpose::STANDARD};
use bip39::Mnemonic;
use ed25519_dalek::SigningKey;
use ring::aead::{AES_256_GCM, Aad, LessSafeKey, NONCE_LEN, UnboundKey};
use ring::pbkdf2;
use serde::{Deserialize, Serialize};
use std::{path::Path, str::FromStr};

#[derive(Serialize, Deserialize)]
pub struct EncryptedMnemonic {
    pub salt: String,
    pub iv: String,
    pub data: String,
}

/// Recover ed25519 signing key from seed
pub fn recover_signing_key(seed: [u8; 64]) -> Result<SigningKey, Box<dyn std::error::Error>> {
    let mut secret_bytes = [0u8; 32];
    secret_bytes.copy_from_slice(&seed[..32]);

    let signing_key = SigningKey::from_bytes(&secret_bytes);
    Ok(signing_key)
}

/// Encrypt a mnemonic string with a password and write it to `path`.
///
/// Uses AES-256-GCM with a PBKDF2-derived key (10 000 iterations, SHA-256).
/// The output file is a JSON-serialized `EncryptedMnemonic`.
pub fn save_encrypted_mnemonic<P: AsRef<Path>>(
    path: P,
    mnemonic_str: &str,
    password: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let salt: [u8; 16] = rand::random();
    let mut derived_key = [0u8; 32];

    pbkdf2::derive(
        pbkdf2::PBKDF2_HMAC_SHA256,
        std::num::NonZeroU32::new(10000).unwrap(),
        &salt,
        password.as_bytes(),
        &mut derived_key,
    );

    let aes_key =
        UnboundKey::new(&AES_256_GCM, &derived_key).map_err(|_| "Failed to create AES key")?;
    let key = LessSafeKey::new(aes_key);

    let iv: [u8; NONCE_LEN] = rand::random();

    let mut in_out = mnemonic_str.as_bytes().to_vec();
    let aad = Aad::empty();
    key.seal_in_place_append_tag(
        ring::aead::Nonce::assume_unique_for_key(iv),
        aad,
        &mut in_out,
    )
    .map_err(|_| "Failed to encrypt mnemonic")?;

    let encrypted_mnemonic = EncryptedMnemonic {
        salt: STANDARD.encode(salt),
        iv: STANDARD.encode(iv),
        data: STANDARD.encode(&in_out),
    };

    // Ensure parent directory exists
    if let Some(parent) = path.as_ref().parent() {
        std::fs::create_dir_all(parent)?;
    }

    let json = serde_json::to_string_pretty(&encrypted_mnemonic)?;
    std::fs::write(path, json)?;

    Ok(())
}

/// Recover mnemonic from encrypted storage
pub fn recover_mnemonic<P: AsRef<Path>>(
    path: P,
    password: &str,
) -> Result<Mnemonic, Box<dyn std::error::Error>> {
    let enc_mnemonic: EncryptedMnemonic = serde_json::from_str(&std::fs::read_to_string(path)?)?;

    let salt = STANDARD.decode(&enc_mnemonic.salt)?;
    let iv_bytes = STANDARD.decode(&enc_mnemonic.iv)?;
    let mut ciphertext = STANDARD.decode(&enc_mnemonic.data)?;

    let mut aes_key = [0u8; 32];
    ring::pbkdf2::derive(
        ring::pbkdf2::PBKDF2_HMAC_SHA256,
        std::num::NonZeroU32::new(10000).unwrap(),
        &salt,
        password.as_bytes(),
        &mut aes_key,
    );

    let unbound_key =
        UnboundKey::new(&AES_256_GCM, &aes_key).map_err(|_| "Failed to create AES key")?;
    let key = LessSafeKey::new(unbound_key);

    let mut iv = [0u8; NONCE_LEN];
    iv.copy_from_slice(&iv_bytes);

    let aad = Aad::empty();
    let plaintext = key
        .open_in_place(
            ring::aead::Nonce::assume_unique_for_key(iv),
            aad,
            &mut ciphertext,
        )
        .map_err(|_| "Decryption failed - wrong password?")?;

    let mnemonic_str = String::from_utf8(plaintext.to_vec())?;
    let mnemonic = Mnemonic::from_str(&mnemonic_str)?;

    Ok(mnemonic)
}
