/// Cryptography primitives for WASTE-rs.
///
/// Identity:   Ed25519 keypair (sign/verify peer identity)
/// Session:    X25519 ECDH → shared secret → ChaCha20-Poly1305 session key
/// Storage:    Keys persisted as JSON in the data directory
use anyhow::{Context, Result};
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine};
use chrono::Utc;
use proto::{PeerId, PeerInfo};
use ring::{
    aead::{self, BoundKey, Nonce, NonceSequence, CHACHA20_POLY1305},
    agreement::{self, EphemeralPrivateKey, UnparsedPublicKey, X25519},
    rand::{SecureRandom, SystemRandom},
    signature::{self, Ed25519KeyPair, KeyPair},
};
use serde::{Deserialize, Serialize};
use std::path::Path;
use tokio::fs;

// ── Identity keypair ──────────────────────────────────────────────────────────

#[derive(Clone)]
pub struct Identity {
    /// PKCS#8 Ed25519 keypair bytes (secret!)
    pkcs8: Vec<u8>,
    alias: String,
}

#[derive(Serialize, Deserialize)]
struct IdentityFile {
    pkcs8_base64: String,
    alias: String,
}

impl Identity {
    /// Load from disk, or generate a new keypair if none exists.
    pub async fn load_or_create(data_dir: &str, alias: &str) -> Result<Self> {
        let path = Path::new(data_dir).join("identity.json");
        if path.exists() {
            let raw = fs::read_to_string(&path)
                .await
                .context("reading identity file")?;
            let file: IdentityFile = serde_json::from_str(&raw)?;
            let pkcs8 = URL_SAFE_NO_PAD.decode(&file.pkcs8_base64)?;
            tracing::info!("Loaded existing identity from {:?}", path);
            Ok(Self { pkcs8, alias: file.alias })
        } else {
            tracing::info!("Generating new Ed25519 identity keypair");
            let rng = SystemRandom::new();
            let pkcs8 = Ed25519KeyPair::generate_pkcs8(&rng)
                .map_err(|e| anyhow::anyhow!("keygen failed: {:?}", e))?;
            let identity = Self {
                pkcs8: pkcs8.as_ref().to_vec(),
                alias: alias.to_string(),
            };
            let file = IdentityFile {
                pkcs8_base64: URL_SAFE_NO_PAD.encode(&identity.pkcs8),
                alias: alias.to_string(),
            };
            fs::create_dir_all(data_dir).await?;
            fs::write(&path, serde_json::to_string_pretty(&file)?).await?;
            tracing::info!("Saved new identity to {:?}", path);
            Ok(identity)
        }
    }

    fn keypair(&self) -> Result<Ed25519KeyPair> {
        Ed25519KeyPair::from_pkcs8(&self.pkcs8)
            .map_err(|e| anyhow::anyhow!("loading keypair: {:?}", e))
    }

    /// Our stable peer ID (base64url of the public key).
    pub fn peer_id(&self) -> PeerId {
        let kp = self.keypair().expect("identity keypair valid");
        PeerId(URL_SAFE_NO_PAD.encode(kp.public_key().as_ref()))
    }

    /// Full peer info for the Hello handshake.
    pub fn peer_info(&self) -> PeerInfo {
        PeerInfo {
            id: self.peer_id(),
            alias: self.alias.clone(),
            public_key: self.peer_id().0.clone(), // public key == peer id encoding
            created_at: Utc::now(),
        }
    }

    /// Sign arbitrary bytes. Returns base64url signature.
    pub fn sign(&self, data: &[u8]) -> Result<String> {
        let kp = self.keypair()?;
        let sig = kp.sign(data);
        Ok(URL_SAFE_NO_PAD.encode(sig.as_ref()))
    }

    /// Verify a signature from a peer (given their base64url public key).
    pub fn verify(public_key_b64: &str, data: &[u8], sig_b64: &str) -> Result<()> {
        let pk_bytes = URL_SAFE_NO_PAD.decode(public_key_b64)?;
        let sig_bytes = URL_SAFE_NO_PAD.decode(sig_b64)?;
        let peer_pk = signature::UnparsedPublicKey::new(&signature::ED25519, &pk_bytes);
        peer_pk
            .verify(data, &sig_bytes)
            .map_err(|_| anyhow::anyhow!("signature verification failed"))
    }
}

// ── X25519 ECDH session key derivation ───────────────────────────────────────

/// Generate an ephemeral X25519 keypair. Returns (private_key, public_key_b64).
pub fn generate_ephemeral() -> Result<(EphemeralPrivateKey, String)> {
    let rng = SystemRandom::new();
    let private = EphemeralPrivateKey::generate(&X25519, &rng)
        .map_err(|e| anyhow::anyhow!("X25519 keygen: {:?}", e))?;
    let public_b64 = URL_SAFE_NO_PAD.encode(
        private
            .compute_public_key()
            .map_err(|e| anyhow::anyhow!("pubkey: {:?}", e))?
            .as_ref(),
    );
    Ok((private, public_b64))
}

/// Complete ECDH: our ephemeral private key + their ephemeral public key → 32-byte shared secret.
pub fn ecdh(our_private: EphemeralPrivateKey, their_public_b64: &str) -> Result<Vec<u8>> {
    let their_bytes = URL_SAFE_NO_PAD.decode(their_public_b64)?;
    let their_pub = UnparsedPublicKey::new(&X25519, &their_bytes);
    agreement::agree_ephemeral(our_private, &their_pub, |key_material| {
        // In production: run through HKDF here
        key_material.to_vec()
    })
    .map_err(|e| anyhow::anyhow!("ECDH failed: {:?}", e))
}

// ── ChaCha20-Poly1305 AEAD ────────────────────────────────────────────────────

/// Encrypt plaintext with a 32-byte key. Returns (nonce_b64, ciphertext_b64).
pub fn encrypt(key: &[u8; 32], plaintext: &[u8]) -> Result<(String, String)> {
    let rng = SystemRandom::new();
    let mut nonce_bytes = [0u8; 12];
    rng.fill(&mut nonce_bytes)
        .map_err(|_| anyhow::anyhow!("RNG fill failed"))?;

    let unbound = aead::UnboundKey::new(&CHACHA20_POLY1305, key)
        .map_err(|e| anyhow::anyhow!("key: {:?}", e))?;
    let mut sealing = aead::SealingKey::new(unbound, OneNonce(Some(nonce_bytes)));

    let mut in_out = plaintext.to_vec();
    sealing
        .seal_in_place_append_tag(aead::Aad::empty(), &mut in_out)
        .map_err(|e| anyhow::anyhow!("encrypt: {:?}", e))?;

    Ok((
        URL_SAFE_NO_PAD.encode(&nonce_bytes),
        URL_SAFE_NO_PAD.encode(&in_out),
    ))
}

/// Decrypt ciphertext with a 32-byte key.
pub fn decrypt(key: &[u8; 32], nonce_b64: &str, ciphertext_b64: &str) -> Result<Vec<u8>> {
    let nonce_bytes: [u8; 12] = URL_SAFE_NO_PAD
        .decode(nonce_b64)?
        .try_into()
        .map_err(|_| anyhow::anyhow!("nonce length"))?;

    let unbound = aead::UnboundKey::new(&CHACHA20_POLY1305, key)
        .map_err(|e| anyhow::anyhow!("key: {:?}", e))?;
    let mut opening = aead::OpeningKey::new(unbound, OneNonce(Some(nonce_bytes)));

    let mut in_out = URL_SAFE_NO_PAD.decode(ciphertext_b64)?;
    let plaintext = opening
        .open_in_place(aead::Aad::empty(), &mut in_out)
        .map_err(|_| anyhow::anyhow!("decryption failed (bad key or tampered)"))?;

    Ok(plaintext.to_vec())
}

/// A single-use nonce holder for ring's NonceSequence trait.
struct OneNonce(Option<[u8; 12]>);

impl NonceSequence for OneNonce {
    fn advance(&mut self) -> std::result::Result<Nonce, ring::error::Unspecified> {
        self.0
            .take()
            .map(Nonce::assume_unique_for_key)
            .ok_or(ring::error::Unspecified)
    }
}