// SPDX-License-Identifier: GPL-3-0-or-later
// Copyright (c) 2025 Opinsys Oy
// Copyright (c) 2024-2025 Jarkko Sakkinen

use crate::{
    alg::AlgError,
    cli::Task,
    command::{AuthArgs, CommandError, InputArgs},
    io::read_file_input,
    task::{TaskAuth, TaskState},
};
use clap::Args;
use tpm2_crypto::tpm_make_name;
use tpm2_device::{with_device, TpmDevice};
use tpm2_protocol::{
    constant::TPM_MAX_COMMAND_SIZE,
    data::{Tpm2bName, Tpm2bPublic, TpmCc},
    frame::TpmLoadCommand,
    TpmHandle, TpmMarshal, TpmWriter,
};
use tpm2_tpmkey::TpmKey;
use tpm2_vtpm::{vtpm_policy_command_from_parts, VtpmHandle, VtpmHandleClass, VtpmPolicyCommand};

/// Loads a PEM or DER TPMKey file to cache.
#[derive(Args, Debug)]
#[command(verbatim_doc_comment)]
pub struct Load {
    #[clap(flatten)]
    pub auth_args: AuthArgs,

    #[clap(flatten)]
    pub input_args: InputArgs,
}

impl Task for Load {
    fn run(
        &self,
        task_state: &mut TaskState,
        writer: &mut dyn std::io::Write,
        _is_tty: bool,
    ) -> Result<(), CommandError> {
        with_device(
            task_state.device.clone(),
            |device| -> Result<(), CommandError> {
                let input_bytes = read_file_input(self.input_args.input.as_deref())?;
                if input_bytes.is_empty() {
                    return Ok(());
                }

                let tpm_key = TpmKey::from_pem(&input_bytes)
                    .or_else(|_| TpmKey::from_der(&input_bytes).map_err(AlgError::from))?;

                let parent_public = tpm_key
                    .parent_public()
                    .cloned()
                    .ok_or(CommandError::InvalidInput("parent missing".to_string()))?;

                let parent_handle_ref = Self::fetch_parent(task_state, device, &parent_public)?;

                let (parent_handle, policy_blob, name_alg, parent_empty_auth) =
                    task_state.resolve_policy(device, &parent_handle_ref)?;

                let (auths, policy_session_auth) = task_state.build_auth(
                    device,
                    &policy_blob,
                    name_alg,
                    parent_empty_auth,
                    &self.auth_args,
                )?;

                let run_load_result = Self::run_load(
                    task_state,
                    device,
                    parent_handle,
                    tpm_key.private(),
                    tpm_key.public(),
                    &auths,
                );

                if let Some(TaskAuth::Session(vhandle)) = policy_session_auth {
                    if let Err(e) = task_state.remove_session(device, vhandle) {
                        log::error!("vtpm:{vhandle:08x}: {e}");
                    }
                }

                let (object_handle, _, loaded_public) =
                    run_load_result.inspect_err(|e: &CommandError| {
                        log::debug!("run_load failed: {e}");
                    })?;

                let policy_blob = if let Some(policy) = &tpm_key.policy {
                    let mut buf = vec![0u8; TPM_MAX_COMMAND_SIZE as usize];
                    let len = {
                        let mut writer = TpmWriter::new(&mut buf);
                        let count = u32::try_from(policy.policy.len())?;
                        count.marshal(&mut writer).map_err(CommandError::Marshal)?;

                        for cmd in &policy.policy {
                            cmd.cc.marshal(&mut writer).map_err(CommandError::Marshal)?;

                            let body = cmd.body.clone();
                            let body_len = u32::try_from(body.len())?;

                            body_len
                                .marshal(&mut writer)
                                .map_err(CommandError::Marshal)?;
                            writer.write_bytes(&body).map_err(CommandError::Marshal)?;
                        }
                        writer.len()
                    };
                    buf.truncate(len);
                    let mut policy_vec: Vec<Box<dyn VtpmPolicyCommand>> = Vec::new();
                    for cmd in &policy.policy {
                        policy_vec.push(vtpm_policy_command_from_parts(cmd.cc, &cmd.body)?);
                    }
                    Some(policy_vec)
                } else {
                    None
                };

                let object_context = device.save_context(object_handle)?;
                let vhandle = task_state.cache.save_key(
                    object_context,
                    &loaded_public.inner,
                    &parent_public.inner,
                    tpm_key.empty_auth.unwrap_or_default(),
                    &policy_blob,
                )?;

                writeln!(writer, "vtpm:{vhandle:08x}")?;
                Ok(())
            },
        )
    }
}

impl Load {
    fn fetch_parent(
        task_state: &mut TaskState,
        device: &mut TpmDevice,
        parent_public: &Tpm2bPublic,
    ) -> Result<VtpmHandle, CommandError> {
        let name = tpm_make_name(&parent_public.inner)?;
        if let Some(phandle) = device.find_persistent(&name)? {
            return Ok(VtpmHandle::new(VtpmHandleClass::Tpm, phandle.0));
        }

        let vhandle_opt = task_state
            .cache
            .key_iter()
            .find(|(_, key)| key.public == parent_public.inner)
            .map(|(vhandle, _)| *vhandle);

        if let Some(vhandle) = vhandle_opt {
            return Ok(VtpmHandle::new(VtpmHandleClass::Vtpm, vhandle));
        }

        Err(CommandError::UnknownParent)
    }

    fn run_load(
        task_state: &mut TaskState,
        device: &mut TpmDevice,
        parent_handle: TpmHandle,
        in_private: &tpm2_protocol::data::Tpm2bPrivate,
        in_public: &Tpm2bPublic,
        auths: &[TaskAuth],
    ) -> Result<(TpmHandle, Tpm2bName, Tpm2bPublic), CommandError> {
        let cmd = TpmLoadCommand {
            in_private: *in_private,
            in_public: in_public.clone(),
            handles: [parent_handle],
        };

        let (resp, _) = task_state.execute(device, &cmd, auths)?;

        let resp = resp
            .Load()
            .map_err(|_| CommandError::ResponseMismatch(TpmCc::Load))?;

        task_state.track_handle(resp.handles[0])?;
        Ok((resp.handles[0], resp.name, in_public.clone()))
    }
}
