//! Seed storage: encrypted at rest with Argon2id + ChaCha20-Poly1305.

use std::fs;
use std::io::{self, Write};
use std::path::PathBuf;

use argon2::Argon2;
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Nonce,
};
use rand::RngCore;
use warzone_protocol::identity::Seed;
use zeroize::Zeroize;

/// Magic bytes to identify encrypted seed files.
const MAGIC: &[u8; 4] = b"WZS1";
/// Salt length for Argon2.
const SALT_LEN: usize = 16;
/// Nonce length for ChaCha20-Poly1305.
const NONCE_LEN: usize = 12;

/// Get the warzone data directory. Respects WARZONE_HOME env var,
/// falls back to ~/.warzone.
pub fn data_dir() -> PathBuf {
    if let Ok(wz) = std::env::var("WARZONE_HOME") {
        PathBuf::from(wz)
    } else {
        let home = std::env::var("HOME").unwrap_or_else(|_| ".".into());
        PathBuf::from(home).join(".warzone")
    }
}

fn seed_path() -> PathBuf {
    data_dir().join("identity.seed")
}

/// Derive a 32-byte encryption key from a passphrase using Argon2id.
fn derive_key(passphrase: &[u8], salt: &[u8]) -> [u8; 32] {
    let mut key = [0u8; 32];
    Argon2::default()
        .hash_password_into(passphrase, salt, &mut key)
        .expect("Argon2 should not fail with valid params");
    key
}

/// Prompt for a passphrase (hidden input).
fn prompt_passphrase(prompt: &str) -> String {
    eprint!("{}", prompt);
    io::stderr().flush().unwrap();
    let mut pass = String::new();
    // Try to disable echo. If that fails (e.g. piped input), just read normally.
    #[cfg(unix)]
    {
        use std::os::unix::io::AsRawFd;
        let fd = io::stdin().as_raw_fd();
        let mut termios = unsafe {
            let mut t = std::mem::zeroed();
            libc::tcgetattr(fd, &mut t);
            t
        };
        let old = termios;
        termios.c_lflag &= !libc::ECHO;
        unsafe { libc::tcsetattr(fd, libc::TCSANOW, &termios) };
        io::stdin().read_line(&mut pass).unwrap();
        unsafe { libc::tcsetattr(fd, libc::TCSANOW, &old) };
        eprintln!();
    }
    #[cfg(not(unix))]
    {
        io::stdin().read_line(&mut pass).unwrap();
    }
    pass.trim().to_string()
}

/// Save seed encrypted with a passphrase.
pub fn save_seed(seed: &Seed) -> anyhow::Result<()> {
    let path = seed_path();
    if let Some(parent) = path.parent() {
        fs::create_dir_all(parent)?;
    }

    let passphrase = prompt_passphrase("Set passphrase (empty for no encryption): ");

    if passphrase.is_empty() {
        // Plaintext (legacy, for testing)
        fs::write(&path, &seed.0)?;
    } else {
        let confirm = prompt_passphrase("Confirm passphrase: ");
        if passphrase != confirm {
            anyhow::bail!("Passphrases don't match");
        }

        let mut salt = [0u8; SALT_LEN];
        rand::rngs::OsRng.fill_bytes(&mut salt);

        let mut key = derive_key(passphrase.as_bytes(), &salt);

        let cipher = ChaCha20Poly1305::new((&key).into());
        let mut nonce_bytes = [0u8; NONCE_LEN];
        rand::rngs::OsRng.fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let ciphertext = cipher
            .encrypt(nonce, seed.0.as_slice())
            .map_err(|_| anyhow::anyhow!("encryption failed"))?;

        // File format: MAGIC(4) + salt(16) + nonce(12) + ciphertext(32+16=48)
        let mut file_data = Vec::with_capacity(4 + SALT_LEN + NONCE_LEN + ciphertext.len());
        file_data.extend_from_slice(MAGIC);
        file_data.extend_from_slice(&salt);
        file_data.extend_from_slice(&nonce_bytes);
        file_data.extend_from_slice(&ciphertext);

        fs::write(&path, &file_data)?;
        key.zeroize();
    }

    #[cfg(unix)]
    {
        use std::os::unix::fs::PermissionsExt;
        fs::set_permissions(&path, fs::Permissions::from_mode(0o600))?;
    }

    Ok(())
}

/// Load raw seed bytes (for deriving eth address etc).
pub fn load_seed_raw() -> anyhow::Result<[u8; 32]> {
    let seed = load_seed()?;
    Ok(seed.0)
}

/// Load seed, decrypting if necessary.
pub fn load_seed() -> anyhow::Result<Seed> {
    let path = seed_path();
    let bytes = fs::read(&path)
        .map_err(|_| anyhow::anyhow!("No identity found. Run `warzone init` first."))?;

    // Check if encrypted
    if bytes.len() >= 4 && &bytes[..4] == MAGIC {
        // Encrypted format
        if bytes.len() < 4 + SALT_LEN + NONCE_LEN + 48 {
            anyhow::bail!("Corrupted encrypted seed file");
        }

        let salt = &bytes[4..4 + SALT_LEN];
        let nonce_bytes = &bytes[4 + SALT_LEN..4 + SALT_LEN + NONCE_LEN];
        let ciphertext = &bytes[4 + SALT_LEN + NONCE_LEN..];

        let passphrase = prompt_passphrase("Passphrase: ");
        let mut key = derive_key(passphrase.as_bytes(), salt);

        let cipher = ChaCha20Poly1305::new((&key).into());
        let nonce = Nonce::from_slice(nonce_bytes);

        let plaintext = cipher
            .decrypt(nonce, ciphertext)
            .map_err(|_| anyhow::anyhow!("Wrong passphrase"))?;

        key.zeroize();

        if plaintext.len() != 32 {
            anyhow::bail!("Corrupted seed data");
        }
        let mut seed_bytes = [0u8; 32];
        seed_bytes.copy_from_slice(&plaintext);
        Ok(Seed::from_bytes(seed_bytes))
    } else if bytes.len() == 32 {
        // Legacy plaintext
        let mut seed_bytes = [0u8; 32];
        seed_bytes.copy_from_slice(&bytes);
        Ok(Seed::from_bytes(seed_bytes))
    } else {
        anyhow::bail!("Corrupted seed file (unknown format)")
    }
}