use crate::util::{base64_encode, to_pem};
use ml_kem::kem::{Decapsulate, Encapsulate};
use ml_kem::{EncodedSizeUser, KemCore, MlKem1024, MlKem512, MlKem768};
use rust_util::XResult;
use std::convert::TryInto;
use crate::pivutil::ToStr;

#[derive(Debug, Clone, Copy)]
pub enum MlKemLen {
    Len512,
    Len768,
    Len1024,
}

impl ToStr for MlKemLen {
    fn to_str(&self) -> &str {
        match self {
            MlKemLen::Len512 => "mlkem512",
            MlKemLen::Len768 => "mlkem768",
            MlKemLen::Len1024 => "mlkem1024",
        }
    }
}

pub fn generate_mlkem_keypair(len: usize) -> XResult<(String, String, String, Vec<u8>, String)> {
    let (dk_private, ek_public) = match len {
        512 => generate_ml_kem_512(),
        768 => generate_ml_kem_768(),
        1024 => generate_ml_kem_1024(),
        _ => return simple_error!("Invalid ML-KEM={}", len),
    };
    let secret_key_der_base64 = base64_encode(&dk_private);
    let secret_key_pem = to_pem(&format!("ML-KEM-{} PRIVATE KEY", len), &dk_private);
    let public_key_pem = to_pem(&format!("ML-KEM-{} PUBLIC KEY", len), &ek_public);
    let public_key_der = ek_public;
    let jwk_ec_key = "".to_string();
    Ok((
        secret_key_der_base64,
        secret_key_pem,
        public_key_pem,
        public_key_der,
        jwk_ec_key,
    ))
}

pub fn try_parse_decapsulate_key_private_then_decapsulate(
    key_bytes: &[u8],
    ciphertext_bytes: &[u8],
) -> XResult<(MlKemLen, Vec<u8>)> {
    if let Ok(shared_secret) =
        parse_decapsulate_key_512_private_then_decapsulate(key_bytes, ciphertext_bytes)
    {
        return Ok((MlKemLen::Len512, shared_secret.to_vec()));
    }
    if let Ok(shared_secret) =
        parse_decapsulate_key_768_private_then_decapsulate(key_bytes, ciphertext_bytes)
    {
        return Ok((MlKemLen::Len768, shared_secret.to_vec()));
    }
    if let Ok(shared_secret) =
        parse_decapsulate_key_1024_private_then_decapsulate(key_bytes, ciphertext_bytes)
    {
        return Ok((MlKemLen::Len1024, shared_secret.to_vec()));
    }
    simple_error!("Invalid decapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
}

pub fn try_parse_decapsulate_key_private_get_encapsulate(
    key_bytes: &[u8],
) -> XResult<(MlKemLen, Vec<u8>)> {
    if let Ok(encapsulate_key) = parse_decapsulate_key_512_private_get_encapsulate(key_bytes) {
        return Ok((MlKemLen::Len512, encapsulate_key));
    }
    if let Ok(encapsulate_key) = parse_decapsulate_key_768_private_get_encapsulate(key_bytes) {
        return Ok((MlKemLen::Len768, encapsulate_key));
    }
    if let Ok(encapsulate_key) = parse_decapsulate_key_1024_private_get_encapsulate(key_bytes) {
        return Ok((MlKemLen::Len1024, encapsulate_key));
    }
    simple_error!("Invalid decapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
}

pub fn generate_ml_kem_512() -> (Vec<u8>, Vec<u8>) {
    let mut rng = rand::thread_rng();
    let (dk_private, ek_public) = <MlKem512 as KemCore>::generate(&mut rng);
    (
        dk_private.as_bytes().0.to_vec(),
        ek_public.as_bytes().0.to_vec(),
    )
}

pub fn generate_ml_kem_768() -> (Vec<u8>, Vec<u8>) {
    let mut rng = rand::thread_rng();
    let (dk_private, ek_public) = <MlKem768 as KemCore>::generate(&mut rng);
    (
        dk_private.as_bytes().0.to_vec(),
        ek_public.as_bytes().0.to_vec(),
    )
}

pub fn generate_ml_kem_1024() -> (Vec<u8>, Vec<u8>) {
    let mut rng = rand::thread_rng();
    let (dk_private, ek_public) = <MlKem1024 as KemCore>::generate(&mut rng);
    (
        dk_private.as_bytes().0.to_vec(),
        ek_public.as_bytes().0.to_vec(),
    )
}

pub fn try_parse_encapsulation_key_public_then_encapsulate(bytes: &[u8]) -> XResult<(MlKemLen, Vec<u8>, Vec<u8>)> {
    if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_512_public_then_encapsulate(bytes) {
        return Ok((MlKemLen::Len512, ciphertext, shared_key));
    }
    if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_768_public_then_encapsulate(bytes) {
        return Ok((MlKemLen::Len768, ciphertext, shared_key));
    }
    if let Ok((ciphertext, shared_key)) = parse_encapsulation_key_1024_public_then_encapsulate(bytes) {
        return Ok((MlKemLen::Len1024, ciphertext, shared_key));
    }
    simple_error!("Invalid encapsulation key, only allow MK-KEM-512, ML-KEM-768, ML-KEM-1024")
}

pub fn parse_encapsulation_key_512_public_then_encapsulate(
    bytes: &[u8],
) -> XResult<(Vec<u8>, Vec<u8>)> {
    let ek = <MlKem512 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
        bytes.try_into(),
        "Parse encapsulation key 512 failed: {}"
    ));
    let (ciphertext, shared_key) = opt_result!(
        ek.encapsulate(&mut rand::thread_rng()),
        "Encapsulation key 512 encapsulate failed: {:?}"
    );
    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
}

pub fn parse_encapsulation_key_768_public_then_encapsulate(
    bytes: &[u8],
) -> XResult<(Vec<u8>, Vec<u8>)> {
    let ek = <MlKem768 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
        bytes.try_into(),
        "Parse encapsulation key 768 failed: {}"
    ));
    let (ciphertext, shared_key) = opt_result!(
        ek.encapsulate(&mut rand::thread_rng()),
        "Encapsulation key 768 encapsulate failed: {:?}"
    );
    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
}

pub fn parse_encapsulation_key_1024_public_then_encapsulate(
    bytes: &[u8],
) -> XResult<(Vec<u8>, Vec<u8>)> {
    let ek = <MlKem1024 as KemCore>::EncapsulationKey::from_bytes(&opt_result!(
        bytes.try_into(),
        "Parse encapsulation key 1024 failed: {}"
    ));
    let (ciphertext, shared_key) = opt_result!(
        ek.encapsulate(&mut rand::thread_rng()),
        "Encapsulation key 1024 encapsulate failed: {:?}"
    );
    Ok((ciphertext.0.to_vec(), shared_key.0.to_vec()))
}

pub fn parse_decapsulate_key_512_private_then_decapsulate(
    key_bytes: &[u8],
    ciphertext_bytes: &[u8],
) -> XResult<Vec<u8>> {
    let dk = <MlKem512 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 512 failed: {}"
    ));
    let shared_key = opt_result!(
        dk.decapsulate(opt_result!(
            ciphertext_bytes.try_into(),
            "Parse encoded ciphertext 512 failed: {}"
        )),
        "Decapsulation key 512 decapsulate failed: {:?}"
    );
    Ok(shared_key.0.to_vec())
}

pub fn parse_decapsulate_key_768_private_then_decapsulate(
    key_bytes: &[u8],
    ciphertext_bytes: &[u8],
) -> XResult<Vec<u8>> {
    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 768 failed: {}"
    ));
    let shared_key = opt_result!(
        dk.decapsulate(opt_result!(
            ciphertext_bytes.try_into(),
            "Parse encoded ciphertext 768 failed: {}"
        )),
        "Decapsulation key 768 decapsulate failed: {:?}"
    );
    Ok(shared_key.0.to_vec())
}

pub fn parse_decapsulate_key_1024_private_then_decapsulate(
    key_bytes: &[u8],
    ciphertext_bytes: &[u8],
) -> XResult<Vec<u8>> {
    let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 1024 failed: {}"
    ));
    let shared_key = opt_result!(
        dk.decapsulate(opt_result!(
            ciphertext_bytes.try_into(),
            "Parse encoded ciphertext 1024 failed: {}"
        )),
        "Decapsulation key 1024 decapsulate failed: {:?}"
    );
    Ok(shared_key.0.to_vec())
}

pub fn parse_decapsulate_key_512_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
    let dk = <MlKem512 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 512 failed: {}"
    ));
    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
}

pub fn parse_decapsulate_key_768_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
    let dk = <MlKem768 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 768 failed: {}"
    ));
    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
}

pub fn parse_decapsulate_key_1024_private_get_encapsulate(key_bytes: &[u8]) -> XResult<Vec<u8>> {
    let dk = <MlKem1024 as KemCore>::DecapsulationKey::from_bytes(&opt_result!(
        key_bytes.try_into(),
        "Parse decapsulation key 1024 failed: {}"
    ));
    Ok(dk.encapsulation_key().as_bytes().0.to_vec())
}