use std::path::PathBuf;

use base64::{Engine as _, engine::general_purpose};
use clap::Parser;
use cosmian_kms_client::{
    ExportObjectParams, KmsClient,
    cosmian_kmip::kmip_2_1::{
        kmip_data_structures::KeyWrappingSpecification, kmip_types::CryptographicAlgorithm,
    },
    export_object,
    kmip_2_1::{kmip_attributes::Attributes, requests::create_symmetric_key_kmip_object},
    read_object_from_json_ttlv_file, write_kmip_object_to_file,
};
use cosmian_kms_crypto::crypto::{
    password_derivation::derive_key_from_password, wrap::wrap_object_with_key,
};

use crate::{
    actions::kms::{console, shared::SYMMETRIC_WRAPPING_KEY_SIZE},
    cli_bail,
    error::result::{KmsCliResult, KmsCliResultHelper},
};

/// Locally wrap a secret data or key in KMIP JSON TTLV format.
///
/// The secret data or key can be wrapped using either:
///  - a password derived into a symmetric key using Argon2
///  - symmetric key bytes in base64
///  - a key in the KMS (which will be exported first)
///  - a key in a KMIP JSON TTLV file
///
/// For the latter 2 cases, the key may be a symmetric key,
/// and RFC 5649 will be used, or a curve 25519 public key
/// and ECIES will be used.
#[derive(Parser, Default, Debug)]
#[clap(verbatim_doc_comment)]
pub struct WrapSecretDataOrKeyAction {
    /// The KMIP JSON TTLV input key file to wrap
    #[clap(required = true)]
    pub(crate) key_file_in: PathBuf,

    /// The KMIP JSON output file. When not specified, the input file is overwritten.
    #[clap(required = false)]
    pub(crate) key_file_out: Option<PathBuf>,

    /// A password to wrap the imported key.
    /// This password will be derived into an AES-256 symmetric key. For security reasons,
    /// a fresh salt is internally generated by `cosmian` and handled,
    /// and this final AES symmetric key will be displayed only once.
    #[clap(long = "wrap-password", short = 'p', required = false, group = "wrap")]
    pub(crate) wrap_password: Option<String>,

    /// A symmetric key as a base 64 string to wrap the imported key.
    #[clap(long = "wrap-key-b64", short = 'k', required = false, group = "wrap")]
    pub(crate) wrap_key_b64: Option<String>,

    /// The ID of a wrapping key in the KMS that will be exported and used to wrap the key.
    #[clap(long = "wrap-key-id", short = 'i', required = false, group = "wrap")]
    pub(crate) wrap_key_id: Option<String>,

    /// A wrapping key in a KMIP JSON TTLV file used to wrap the key.
    #[clap(long = "wrap-key-file", short = 'f', required = false, group = "wrap")]
    pub(crate) wrap_key_file: Option<PathBuf>,
}

impl WrapSecretDataOrKeyAction {
    /// Run the wrap key action.
    ///
    /// # Errors
    ///
    /// This function can return an error if:
    ///
    /// - The key file cannot be read.
    /// - The key is already wrapped and cannot be wrapped again.
    /// - The wrap key cannot be decoded from base64.
    /// - The wrap password cannot be derived into a symmetric key.
    /// - The wrap key cannot be exported from the KMS.
    /// - The wrap key file cannot be read.
    /// - The key block cannot be wrapped with the wrapping key.
    /// - The wrapped key object cannot be written to the output file.
    /// - The console output cannot be written.
    #[expect(clippy::print_stdout)]
    pub async fn run(&self, kms_rest_client: KmsClient) -> KmsCliResult<String> {
        // read the key file
        let mut object = read_object_from_json_ttlv_file(&self.key_file_in)?;

        // cannot wrap an already wrapped key
        if object.key_wrapping_data().is_some() {
            cli_bail!("cannot wrap an already wrapped key");
        }

        // cache the object type
        let object_type = object.object_type();

        // if the key must be wrapped, prepare the wrapping key
        let wrapping_key = if let Some(b64) = &self.wrap_key_b64 {
            let key_bytes = general_purpose::STANDARD
                .decode(b64)
                .with_context(|| "failed decoding the wrap key")?;
            create_symmetric_key_kmip_object(
                &key_bytes,
                &Attributes {
                    cryptographic_algorithm: Some(CryptographicAlgorithm::AES),
                    ..Default::default()
                },
            )?
        } else if let Some(password) = &self.wrap_password {
            let key_bytes = derive_key_from_password::<SYMMETRIC_WRAPPING_KEY_SIZE>(
                &[0_u8; 16],
                password.as_bytes(),
            )?;

            let symmetric_key_object = create_symmetric_key_kmip_object(
                key_bytes.as_ref(),
                &Attributes {
                    cryptographic_algorithm: Some(CryptographicAlgorithm::AES),
                    ..Default::default()
                },
            )?;

            // Print the wrapping key for user.
            println!(
                "Wrapping key: {}. This is the only time that this wrapping key will be printed.",
                general_purpose::STANDARD.encode(&*key_bytes)
            );
            symmetric_key_object
        } else if let Some(key_id) = &self.wrap_key_id {
            export_object(&kms_rest_client, key_id, ExportObjectParams::default())
                .await?
                .1
        } else if let Some(key_file) = &self.wrap_key_file {
            read_object_from_json_ttlv_file(key_file)?
        } else {
            cli_bail!("one of the wrapping options must be specified");
        };

        wrap_object_with_key(
            &mut object,
            &wrapping_key,
            &KeyWrappingSpecification::default(),
        )?;

        // set the output file path to the input file path if not specified
        let output_file = self
            .key_file_out
            .as_ref()
            .unwrap_or(&self.key_file_in)
            .clone();

        write_kmip_object_to_file(&object, &output_file)?;

        let stdout = format!(
            "The key of type {:?} in file {} was wrapped in file: {}",
            object_type,
            self.key_file_in.display(),
            output_file.display()
        );
        console::Stdout::new(&stdout).write()?;

        let (wrapping_key_bytes, _) = wrapping_key.key_block()?.key_bytes_and_attributes()?;

        Ok(general_purpose::STANDARD.encode(wrapping_key_bytes))
    }
}
