/*
 * Copyright 2025 Yubico AB
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

use std::fs;
use yubihsmrs::object::{ObjectAlgorithm, ObjectDescriptor, ObjectHandle, ObjectType};
use yubihsmrs::Session;
use crate::traits::ui_traits::YubihsmUi;
use crate::traits::command_traits::Command;
use crate::traits::operation_traits::YubihsmOperations;
use crate::ui::helper_operations::{delete_objects, display_menu_headers, display_wrapkey_shares, generate_object, list_objects};
use crate::ui::helper_operations::{display_object_properties, get_new_spec_table, get_script_input_data};
use crate::ui::device_menu::DeviceMenu;
use crate::ui::asym_menu::AsymmetricMenu;
use crate::common::error::MgmError;
use crate::common::types::{NewObjectSpec, SelectionItem};
use crate::common::validators::{aes_key_validator, pem_private_rsa_file_validator, pem_public_rsa_file_validator};
use crate::common::algorithms::MgmAlgorithm;
use crate::common::util::get_delegated_capabilities;
use crate::hsm_operations::asym::AsymmetricOperations;
use crate::hsm_operations::sym::SymmetricOperations;
use crate::hsm_operations::wrap::{WrapCommand, WrapKeyType, WrapOperations, WrapOpSpec, WrapType};
use crate::ui::helper_io::{get_path, get_pem_from_file, write_bytes_to_file};
use crate::script::script_recorder::SessionRecorder;
use crate::script::script_types;
use crate::script::script_types::{RecordableObjectSpec, RecordedOperation, MaskLevel};

static WRAP_HEADER: &str = "Wrap keys";

pub struct WrapMenu<T: YubihsmUi + Clone> {
    ui: T,
}

impl<T: YubihsmUi + Clone> WrapMenu<T> {

    pub fn new(interface: T) -> Self {
        WrapMenu { ui: interface  }
    }
    
    pub fn exec_command(&self, session: &Session, recorder: &Option<SessionRecorder>,  authkey: &ObjectDescriptor) -> Result<(), MgmError> {
        loop {
            display_menu_headers(&self.ui, &[crate::MAIN_HEADER, WRAP_HEADER],
                                 "Wrap key operations allow you to manage and use wrap keys keys stored on the YubiHSM")?;

            let cmd = self.ui.select_command(&WrapCommand::authorized_commands(authkey))?;
            display_menu_headers(&self.ui, &[crate::MAIN_HEADER, WRAP_HEADER, cmd.label()], cmd.description())?;

            let result = match cmd {
                WrapCommand::List => list_objects(&self.ui, &WrapOperations, session),
                WrapCommand::GetKeyProperties => display_object_properties(&self.ui, &WrapOperations, session),
                WrapCommand::Generate => generate_object(&self.ui, recorder, &WrapOperations, session, authkey, ObjectType::WrapKey),
                WrapCommand::Import => self.import(session, recorder, authkey),
                WrapCommand::Delete => delete_objects(&self.ui, recorder, &WrapOperations, session, &WrapOperations.get_all_objects(session)?),
                WrapCommand::GetPublicKey => AsymmetricMenu::new(self.ui.clone()).get_public_key(session, ObjectType::WrapKey),
                WrapCommand::ExportWrapped => self.export_wrapped(session, recorder, authkey),
                WrapCommand::ImportWrapped => self.import_wrapped(session, recorder, authkey),
                WrapCommand::GetRandom => DeviceMenu::new(self.ui.clone()).get_random(session),
                WrapCommand::Exit => std::process::exit(0),
            };

            if let Err(err) = result {
                self.ui.display_error_message(err.to_string().as_str());
            }
        }
    }

    pub fn import(&self, session: &Session, recorder: &Option<SessionRecorder>,  authkey: &ObjectDescriptor) -> Result<(), MgmError> {
        if self.ui.get_confirmation("Re-create from shares?")? {
            self.import_from_shares(session, recorder)
        } else {
            self.import_full_key(session, recorder, authkey)
        }
    }

    fn import_full_key(&self, session: &Session, recorder: &Option<SessionRecorder>,  authkey: &ObjectDescriptor) -> Result<(), MgmError> {
        let mut new_key = NewObjectSpec::default();


        let mut input = self.ui.get_string_input(
            "Enter wrap key in HEX format or path to PEM file containing RSA key:", true, None, None)?;
        loop {
            if aes_key_validator(&input).is_ok() || pem_private_rsa_file_validator(&input).is_ok() || pem_public_rsa_file_validator(&input).is_ok() {
                break;
            }
            self.ui.display_error_message("Input is neither valid AES key in HEX format nor valid path to a file containing RSA key in PEM format");
            input = self.ui.get_string_input(
                "Try again or press ESC to return to menu:", true, None, None)?;
        }

        if aes_key_validator(&input).is_ok() {
            self.ui.display_info_message("Detected HEX string. Parsing as AES wrap key...");
            new_key.object_type = ObjectType::WrapKey;
            new_key.data.push(hex::decode(input.clone())?);
            new_key.algorithm = WrapOperations::get_algorithm_from_keylen(new_key.data[0].len())?;
        } else if pem_private_rsa_file_validator(&input).is_ok() || pem_public_rsa_file_validator(&input).is_ok() {
            self.ui.display_info_message("Detected PEM file with private RSA key. Parsing as RSA key...");
            let pem = get_pem_from_file(&input)?;
            if pem.len() > 1 {
                self.ui.display_warning("Multiple PEM objects found in file. Only the first one is read");
            }
            let pem = pem[0].to_owned();
            let (_type, _algo, _bytes) = AsymmetricOperations::parse_asym_pem(pem)?;
            new_key.data.push(_bytes);
            new_key.algorithm = _algo;
            match _type {
                ObjectType::AsymmetricKey => new_key.object_type = ObjectType::WrapKey,
                ObjectType::PublicKey => new_key.object_type = ObjectType::PublicWrapKey,
                _ => unreachable!()
            }
        } else {
            unreachable!();
        }
        let key_type = WrapOperations::get_wrapkey_type(&new_key.object_type, &new_key.algorithm)?;

        new_key.id = self.ui.get_new_object_id(0)?;
        new_key.label = self.ui.get_object_label("")?;
        new_key.domains = self.ui.select_object_domains(authkey.domains())?;
        new_key.capabilities = self.ui.select_object_capabilities(
            &WrapOperations.get_applicable_capabilities(authkey, Some(new_key.object_type), Some(new_key.algorithm))?,
            Some("Select object capabilities"))?;
        new_key.delegated_capabilities = self.ui.select_object_capabilities(
            &get_delegated_capabilities(authkey),
            Some("Select delegated capabilities"))?;

        if !self.ui.get_note_confirmation(
            "Import wrap key with:",
            &get_new_spec_table(&new_key))? {
            self.ui.display_info_message("Object is not imported");
            return Ok(());
        }

        let progress = self.ui.start_progress(Some("Importing key..."));
        new_key.id = WrapOperations.import(session, &new_key)?;
        self.ui.stop_progress(progress, None);
        self.ui.display_success_message(format!("Imported {} object with ID 0x{:04x} into the YubiHSM", new_key.object_type, new_key.id).as_str());

        let mut n_shares = 0;
        let mut n_threshold = 0;
        if key_type == WrapKeyType::Aes {
            self.ui.display_info_message("Split wrap key? Note that the wrap key is already imported into the YubiHSM2. Key split is done outside the device");
            if self.ui.get_confirmation("Split wrap key?")? {
                n_shares = self.ui.get_split_aes_n_shares("Enter the number of shares to create:")?;
                n_threshold = self.ui.get_split_aes_m_threshold("Enter the number of shares necessary to re-create the key:", n_shares)?;
                let split_key = WrapOperations::split_wrap_key(&new_key, n_threshold, n_shares)?;
                display_wrapkey_shares(&self.ui, split_key.shares_data)?;
            }
        }
        self.record_import_wrapkey(recorder, &new_key, if AsymmetricOperations::is_rsa_key_algorithm(&new_key.algorithm) { Some(input) } else {None} , n_threshold, n_shares)?;

        Ok(())
    }

    pub fn import_from_shares(&self, session: &Session, recorder: &Option<SessionRecorder>) -> Result<(), MgmError> {
        let shares = self.recover_wrapkey_shares()?;
        let mut new_key = WrapOperations::get_wrapkey_from_shares(shares)?;
        new_key.label = self.ui.get_object_label("")?;

        if !self.ui.get_note_confirmation("Import wrap key with:", &get_new_spec_table(&new_key))? {
            self.ui.display_info_message("Key is not imported");
            return Ok(());
        }

        let progress = self.ui.start_progress(Some("Importing key..."));
        new_key.id = WrapOperations.import(session, &new_key)?;
        self.ui.stop_progress(progress, None);
        self.ui.display_success_message(format!("Imported wrap key with ID 0x{:04x} on the device", new_key.id).as_str());

        self.record_import_wrapkey(recorder, &new_key, None, 0, 0)?;

        Ok(())
    }

    fn export_wrapped(&self, session: &Session, recorder: &Option<SessionRecorder>, authkey: &ObjectDescriptor) -> Result<(), MgmError> {
        let wrapkeys = WrapOperations::get_wrapping_keys(session, authkey)?;
        let wrapkey = self.ui.select_one_object(
            &wrapkeys,
            Some("Select the wrapping key to use for exporting objects:"))?;
        let wrapkey_type = WrapOperations::get_wrapkey_type(wrapkey.object_type(), wrapkey.algorithm())?;

        let wrap_type = match wrapkey_type {
            WrapKeyType::Aes => WrapType::Object,
            WrapKeyType::RsaPublic => {
                self.ui.select_one_item(
                    &SelectionItem::get_items(&[WrapType::Object, WrapType::Key]),
                    Some(&WrapType::Object),
                    Some("Select type of wrapping:"))?
            }
            _ => unreachable!()
        };
        let mut wrap_op = WrapOpSpec {
            wrapkey_id: wrapkey.object_id(),
            wrapkey_type,
            wrap_type,
            include_ed_seed: false,
            aes_algorithm: None,
            oaep_algorithm: None,
            mgf1_algorithm: None,
        };

        let exportable_objects = WrapOperations::get_exportable_objects(session, &wrapkey, wrap_type)?;
        if exportable_objects.is_empty() {
            self.ui.display_info_message("No objects available for export by this user");
            return Ok(());
        }
        let export_objects = self.ui.select_multiple_objects(
            &exportable_objects,
            false,
            Some("Select objects to export"))?;
        if exportable_objects.iter().any(|x| x.algorithm() == &ObjectAlgorithm::Ed25519) {
            wrap_op.include_ed_seed = self.ui.get_confirmation("Include Ed25519 seed in the wrapped export? (required for importing Ed25519 keys)")?;
        };
        let export_objects = export_objects.iter().map(|obj| ObjectHandle { object_id: obj.object_id(), object_type: *obj.object_type() }).collect();

        if wrapkey_type == WrapKeyType::RsaPublic {
            wrap_op.aes_algorithm = Some(self.ui.select_algorithm(
                &SymmetricOperations.get_generation_algorithms(),
                Some(ObjectAlgorithm::Aes256),
                Some("Select AES algorithm to use for wrapping"))?);
            wrap_op.oaep_algorithm = Some(self.ui.select_algorithm(
                &MgmAlgorithm::RSA_OAEP_ALGORITHMS,
                Some(ObjectAlgorithm::RsaOaepSha256),
                Some("Select OAEP algorithm to use for wrapping"))?);
            wrap_op.mgf1_algorithm = Some(self.ui.select_algorithm(
                &MgmAlgorithm::MGF1_ALGORITHMS,
                Some(ObjectAlgorithm::Mgf1Sha256),
                Some("Select MGF1 algorithm to use for wrapping"))?);
        }

        let dir = get_path(&self.ui,
                                "Enter path to backup directory:",
                                true,
                                ".")?;

        let wrapped_objects = WrapOperations::export_wrapped(session, &wrap_op, &export_objects)?;

        for object in &wrapped_objects {
            if object.error.is_some() {
                self.ui.display_warning(format!("Failed to wrap {} with ID 0x{:04x}: {}. Skipping...", object.object_type, object.object_id, object.error.as_ref().unwrap()).as_str());
                continue;
            }
            let filename = format!("{}/0x{:04x}-{}.yhw", dir, object.object_id, object.object_type);
            write_bytes_to_file(&self.ui,&object.wrapped_data, filename.as_str())?;
        }

        if let Some(rec) = recorder {
            let d = if rec.mask == MaskLevel::All {
                script_types::PROMPT.to_string()
            } else {
                dir
            };

            rec.record(RecordedOperation::ExportWrapped {
                wrap_spec: wrap_op,
                objects: export_objects,
                destination_directory: d,
            })?;
        }

        Ok(())
    }

    pub fn import_wrapped(&self, session: &Session, recorder: &Option<SessionRecorder>, authkey: &ObjectDescriptor) -> Result<(), MgmError> {
        let filepath = self.ui.get_path_input(
            "Enter absolute path to wrapped object file:",
            true,
            None,
            Some("Files containing wrapped YubiHSM objects usually have the file extension .yhw"))?;
        self.import_wrapped_from_file(session, recorder, authkey, &filepath)
    }

    pub fn import_wrapped_from_file(&self, session: &Session, recorder: &Option<SessionRecorder>, authkey: &ObjectDescriptor, filepath: &str) -> Result<(), MgmError> {
        let wrapped = fs::read(&filepath)?;

        let wrapkeys = WrapOperations::get_unwrapping_keys(session, authkey)?;
        let wrapkey = self.ui.select_one_object(
            &wrapkeys,
            Some("Select the unwrapping key to use for importing objects:"))?;
        let wrapkey_type = WrapOperations::get_wrapkey_type(wrapkey.object_type(), wrapkey.algorithm())?;

        let mut wrap_op = WrapOpSpec {
            wrapkey_id: wrapkey.object_id(),
            wrapkey_type,
            wrap_type: WrapType::Object,
            include_ed_seed: false,
            aes_algorithm: None,
            oaep_algorithm: None,
            mgf1_algorithm: None,
        };
        if wrapkey_type == WrapKeyType::Rsa {
            wrap_op.oaep_algorithm = Some(self.ui.select_algorithm(
                &MgmAlgorithm::RSA_OAEP_ALGORITHMS,
                Some(ObjectAlgorithm::RsaOaepSha256),
                Some("Select OAEP algorithm to use for wrapping"))?);
            wrap_op.mgf1_algorithm = Some(self.ui.select_algorithm(
                &MgmAlgorithm::MGF1_ALGORITHMS,
                Some(ObjectAlgorithm::Mgf1Sha256),
                Some("Select MGF1 algorithm to use for wrapping"))?);
        }

        let res = WrapOperations::import_wrapped(session, &wrap_op, &wrapped, None);
        let handle = match res {
            Ok(h) => {
                self.record_import_wrapped(recorder, &wrap_op, &filepath, None)?;
                h
            },
            Err(e) => {
                if wrapkey_type == WrapKeyType::Rsa {
                    self.ui.display_info_message(format!("Failed to unwrap as object: {}. Trying as key data...", e).as_str());
                    wrap_op.wrap_type = WrapType::Key;

                    let algo = self.ui.select_algorithm(
                        &WrapOperations::get_rsa_unwrapped_key_algorithms(),
                        None,
                        Some("Select wrapped key algorithm"))?;
                    let caps = if SymmetricOperations::is_aes_algorithm(&algo) {
                        SymmetricOperations.get_applicable_capabilities(&wrapkey, None, None)?
                    } else {
                        AsymmetricOperations.get_applicable_capabilities(&wrapkey, None, Some(algo))?
                    };

                    let mut new_key = NewObjectSpec::default();
                    new_key.algorithm = algo;
                    new_key.object_type = if SymmetricOperations::is_aes_algorithm(&algo) {
                        ObjectType::SymmetricKey
                    } else {
                        ObjectType::AsymmetricKey
                    };
                    new_key.id = self.ui.get_new_object_id(0)?;
                    new_key.label = self.ui.get_object_label("")?;
                    new_key.domains = self.ui.select_object_domains(authkey.domains())?;
                    new_key.capabilities = self.ui.select_object_capabilities(
                        &caps,
                        Some("Select object capabilities"))?;

                    let handle = WrapOperations::import_wrapped(session, &wrap_op, &wrapped, Some(new_key.clone()))?;
                    self.record_import_wrapped(recorder, &wrap_op, &filepath, Some(&new_key))?;
                    handle

                } else {
                    return Err(e)
                }
            }
        };

        self.ui.display_success_message(format!("Successfully imported object {}, with ID 0x{:04x}", handle.object_type, handle.object_id).as_str());

        Ok(())
    }

    fn recover_wrapkey_shares(&self) -> Result<Vec<String>, MgmError> {
        self.ui.display_warning("Note that the wrap key will be recreated outside the YubiHSM before importing it in its whole into the device");
        self.ui.get_string_input("Press any key to recreate wrap key from shares", false, None, None)?;

        let n_shares = self.ui.get_split_aes_n_shares(
            "Enter the number of shares to re-create the AES wrap key:")?;
        let mut shares_vec = Vec::new();

        self.ui.clear_screen();
        shares_vec.push(self.ui.get_split_aes_share("Enter share number 1:", None)?);

        if n_shares > 1 {
            for i in 2..=n_shares {
                self.ui.clear_screen();
                shares_vec.push(self.ui.get_split_aes_share(format!("Enter share number {}:", i).as_str(), Some(shares_vec[0].len() as u8))?);
            }
        }
        self.ui.clear_screen();

        self.ui.display_info_message(format!("{} shares have been registered", n_shares).as_str());

        Ok(shares_vec)
    }

    fn record_import_wrapkey(&self, recorder: &Option<SessionRecorder>, spec: &NewObjectSpec, filename: Option<String>, n_threshold: u8, n_shares: u8) -> Result<(), MgmError> {
        if let Some(rec) = recorder {
            let value = get_script_input_data(&rec.mask, spec, filename)?;
            rec.record(RecordedOperation::ImportWrapKey { spec: RecordableObjectSpec::from(spec), value, n_threshold, n_shares })?;
        }
        Ok(())
    }

    fn record_import_wrapped(&self, recorder: &Option<SessionRecorder>, wrapping_spec: &WrapOpSpec, filepath: &str, new_key_spec: Option<&NewObjectSpec>) -> Result<(), MgmError> {
        if let Some(rec) = recorder {
            let fp = if rec.mask == MaskLevel::All {
                script_types::PROMPT
            } else {
                filepath
            };

            rec.record(RecordedOperation::ImportWrapped {
                wrap_spec: wrapping_spec.clone(),
                wrapped_filepath: fp.to_string(),
                new_key_spec: new_key_spec.map(RecordableObjectSpec::from),
            })?;
        }
        Ok(())
    }
}