//! Python bindings via PyO3
//!
//! Exposes all cryptographic operations to Python.

use pyo3::prelude::*;
use pyo3::exceptions::PyValueError;

use crate::error::CryptoError;
use crate::keys::{hkdf, hmac};
use crate::symmetric::aes_gcm;
use crate::hash::chain::{sha256, sha256_hex};
use crate::pqc::{kyber, dilithium, sphincs, SecurityLevel};

/// Convert CryptoError to Python exception
impl From<CryptoError> for PyErr {
    fn from(err: CryptoError) -> PyErr {
        PyValueError::new_err(err.to_string())
    }
}

/// HKDF key derivation
#[pyfunction]
#[pyo3(signature = (secret, salt=None, info=None, length=32))]
fn derive_key(
    secret: &[u8],
    salt: Option<&[u8]>,
    info: Option<&[u8]>,
    length: usize,
) -> PyResult<Vec<u8>> {
    let info_bytes = info.unwrap_or(b"");
    hkdf::derive_key(secret, salt, info_bytes, length)
        .map_err(|e| e.into())
}

/// HMAC-SHA256
#[pyfunction]
fn hmac_sha256(key: &[u8], message: &[u8]) -> Vec<u8> {
    hmac::hmac_sha256(key, message).to_vec()
}

/// Verify HMAC-SHA256
#[pyfunction]
fn verify_hmac(key: &[u8], message: &[u8], tag: &[u8]) -> bool {
    hmac::verify_hmac(key, message, tag)
}

/// AES-256-GCM encrypt
#[pyfunction]
fn encrypt_aes_gcm(
    key: &[u8],
    nonce: &[u8],
    plaintext: &[u8],
    aad: &[u8],
) -> PyResult<Vec<u8>> {
    aes_gcm::encrypt_aes_gcm_slice(key, nonce, plaintext, aad)
        .map_err(|e| e.into())
}

/// AES-256-GCM decrypt
#[pyfunction]
fn decrypt_aes_gcm(
    key: &[u8],
    nonce: &[u8],
    ciphertext: &[u8],
    aad: &[u8],
) -> PyResult<Vec<u8>> {
    aes_gcm::decrypt_aes_gcm_slice(key, nonce, ciphertext, aad)
        .map_err(|e| e.into())
}

/// AES-256-GCM encrypt with auto nonce
#[pyfunction]
fn encrypt_aes_gcm_auto(key: &[u8], plaintext: &[u8], aad: &[u8]) -> PyResult<Vec<u8>> {
    if key.len() != 32 {
        return Err(PyValueError::new_err("Key must be 32 bytes"));
    }
    let key_arr: [u8; 32] = key.try_into().unwrap();
    aes_gcm::encrypt_aes_gcm_with_nonce(&key_arr, plaintext, aad)
        .map_err(|e| e.into())
}

/// AES-256-GCM decrypt with embedded nonce
#[pyfunction]
fn decrypt_aes_gcm_auto(key: &[u8], data: &[u8], aad: &[u8]) -> PyResult<Vec<u8>> {
    if key.len() != 32 {
        return Err(PyValueError::new_err("Key must be 32 bytes"));
    }
    let key_arr: [u8; 32] = key.try_into().unwrap();
    aes_gcm::decrypt_aes_gcm_with_nonce(&key_arr, data, aad)
        .map_err(|e| e.into())
}

/// SHA-256 hash
#[pyfunction]
fn py_sha256(data: &[u8]) -> Vec<u8> {
    sha256(data).to_vec()
}

/// SHA-256 hash as hex string
#[pyfunction]
fn py_sha256_hex(data: &[u8]) -> String {
    sha256_hex(data)
}

/// Generate random bytes
#[pyfunction]
fn random_bytes(length: usize) -> Vec<u8> {
    crate::random_bytes(length)
}

/// Generate random 32-byte key
#[pyfunction]
fn random_key() -> Vec<u8> {
    crate::random_key().to_vec()
}

/// Generate random 12-byte nonce
#[pyfunction]
fn random_nonce() -> Vec<u8> {
    crate::random_nonce().to_vec()
}

// PQC Functions

fn parse_security_level(level: &str) -> PyResult<SecurityLevel> {
    match level.to_lowercase().as_str() {
        "level1" | "1" | "128" => Ok(SecurityLevel::Level1),
        "level3" | "3" | "192" => Ok(SecurityLevel::Level3),
        "level5" | "5" | "256" => Ok(SecurityLevel::Level5),
        _ => Err(PyValueError::new_err(format!("Unknown security level: {}", level))),
    }
}

/// Generate Kyber key pair
#[pyfunction]
#[pyo3(signature = (level="level3"))]
fn kyber_keygen(level: &str) -> PyResult<(Vec<u8>, Vec<u8>)> {
    let sec_level = parse_security_level(level)?;
    let keypair = kyber::kyber_keygen(sec_level)
        .map_err(PyErr::from)?;
    Ok((keypair.public_key.clone(), keypair.secret_key().to_vec()))
}

/// Kyber encapsulate
#[pyfunction]
fn kyber_encapsulate(public_key: &[u8]) -> PyResult<(Vec<u8>, Vec<u8>)> {
    let (ciphertext, shared) = kyber::kyber_encapsulate(public_key)
        .map_err(PyErr::from)?;
    Ok((ciphertext, shared.to_vec()))
}

/// Kyber decapsulate
#[pyfunction]
fn kyber_decapsulate(ciphertext: &[u8], secret_key: &[u8]) -> PyResult<Vec<u8>> {
    let shared = kyber::kyber_decapsulate(ciphertext, secret_key)
        .map_err(PyErr::from)?;
    Ok(shared.to_vec())
}

/// Generate Dilithium key pair
#[pyfunction]
#[pyo3(signature = (level="level3"))]
fn dilithium_keygen(level: &str) -> PyResult<(Vec<u8>, Vec<u8>)> {
    let sec_level = parse_security_level(level)?;
    let keypair = dilithium::dilithium_keygen(sec_level)
        .map_err(PyErr::from)?;
    Ok((keypair.public_key.clone(), keypair.secret_key().to_vec()))
}

/// Dilithium sign
#[pyfunction]
fn dilithium_sign(secret_key: &[u8], message: &[u8]) -> PyResult<Vec<u8>> {
    dilithium::dilithium_sign(secret_key, message)
        .map_err(PyErr::from)
}

/// Dilithium verify
#[pyfunction]
fn dilithium_verify(public_key: &[u8], message: &[u8], signature: &[u8]) -> PyResult<bool> {
    dilithium::dilithium_verify(public_key, message, signature)
        .map_err(PyErr::from)
}

fn parse_sphincs_variant(variant: &str) -> PyResult<sphincs::SphincsVariant> {
    match variant.to_lowercase().as_str() {
        "sha2-128f" | "sha2_128f" => Ok(sphincs::SphincsVariant::Sha2_128f),
        "sha2-128s" | "sha2_128s" => Ok(sphincs::SphincsVariant::Sha2_128s),
        "sha2-192f" | "sha2_192f" => Ok(sphincs::SphincsVariant::Sha2_192f),
        "sha2-192s" | "sha2_192s" => Ok(sphincs::SphincsVariant::Sha2_192s),
        "sha2-256f" | "sha2_256f" => Ok(sphincs::SphincsVariant::Sha2_256f),
        "sha2-256s" | "sha2_256s" => Ok(sphincs::SphincsVariant::Sha2_256s),
        "shake-128f" | "shake_128f" => Ok(sphincs::SphincsVariant::Shake_128f),
        "shake-128s" | "shake_128s" => Ok(sphincs::SphincsVariant::Shake_128s),
        "shake-192f" | "shake_192f" => Ok(sphincs::SphincsVariant::Shake_192f),
        "shake-192s" | "shake_192s" => Ok(sphincs::SphincsVariant::Shake_192s),
        "shake-256f" | "shake_256f" => Ok(sphincs::SphincsVariant::Shake_256f),
        "shake-256s" | "shake_256s" => Ok(sphincs::SphincsVariant::Shake_256s),
        _ => Err(PyValueError::new_err(format!("Unknown SPHINCS+ variant: {}", variant))),
    }
}

/// Generate SPHINCS+ key pair
#[pyfunction]
#[pyo3(signature = (variant="sha2-192f"))]
fn sphincs_keygen(variant: &str) -> PyResult<(Vec<u8>, Vec<u8>)> {
    let var = parse_sphincs_variant(variant)?;
    let keypair = sphincs::sphincs_keygen(var)
        .map_err(PyErr::from)?;
    Ok((keypair.public_key.clone(), keypair.secret_key().to_vec()))
}

/// SPHINCS+ sign
#[pyfunction]
#[pyo3(signature = (secret_key, message, variant="sha2-192f"))]
fn sphincs_sign(secret_key: &[u8], message: &[u8], variant: &str) -> PyResult<Vec<u8>> {
    let var = parse_sphincs_variant(variant)?;
    sphincs::sphincs_sign_with_variant(secret_key, message, var)
        .map_err(PyErr::from)
}

/// SPHINCS+ verify
#[pyfunction]
#[pyo3(signature = (public_key, message, signature, variant="sha2-192f"))]
fn sphincs_verify(
    public_key: &[u8],
    message: &[u8],
    signature: &[u8],
    variant: &str,
) -> PyResult<bool> {
    let var = parse_sphincs_variant(variant)?;
    sphincs::sphincs_verify_with_variant(public_key, message, signature, var)
        .map_err(PyErr::from)
}

/// Python module definition
#[pymodule]
fn pqc_py(m: &Bound<'_, PyModule>) -> PyResult<()> {
    // Version
    m.add("VERSION", crate::VERSION)?;

    // HKDF
    m.add_function(wrap_pyfunction!(derive_key, m)?)?;

    // HMAC
    m.add_function(wrap_pyfunction!(hmac_sha256, m)?)?;
    m.add_function(wrap_pyfunction!(verify_hmac, m)?)?;

    // AES-GCM
    m.add_function(wrap_pyfunction!(encrypt_aes_gcm, m)?)?;
    m.add_function(wrap_pyfunction!(decrypt_aes_gcm, m)?)?;
    m.add_function(wrap_pyfunction!(encrypt_aes_gcm_auto, m)?)?;
    m.add_function(wrap_pyfunction!(decrypt_aes_gcm_auto, m)?)?;

    // Hashing
    m.add_function(wrap_pyfunction!(py_sha256, m)?)?;
    m.add_function(wrap_pyfunction!(py_sha256_hex, m)?)?;

    // Random
    m.add_function(wrap_pyfunction!(random_bytes, m)?)?;
    m.add_function(wrap_pyfunction!(random_key, m)?)?;
    m.add_function(wrap_pyfunction!(random_nonce, m)?)?;

    // Kyber (ML-KEM)
    m.add_function(wrap_pyfunction!(kyber_keygen, m)?)?;
    m.add_function(wrap_pyfunction!(kyber_encapsulate, m)?)?;
    m.add_function(wrap_pyfunction!(kyber_decapsulate, m)?)?;

    // Dilithium (ML-DSA)
    m.add_function(wrap_pyfunction!(dilithium_keygen, m)?)?;
    m.add_function(wrap_pyfunction!(dilithium_sign, m)?)?;
    m.add_function(wrap_pyfunction!(dilithium_verify, m)?)?;

    // SPHINCS+ (SLH-DSA)
    m.add_function(wrap_pyfunction!(sphincs_keygen, m)?)?;
    m.add_function(wrap_pyfunction!(sphincs_sign, m)?)?;
    m.add_function(wrap_pyfunction!(sphincs_verify, m)?)?;

    Ok(())
}
