// SPDX-License-Identifier: Apache-2.0

use detee_sgx::prelude::*;
use detee_sgx::HRaTlsConfigBuilder;
use detee_shared::common_proto::Empty;
use detee_shared::sgx::pb::dtpm_proto::DtpmGetConfigRes;
use detee_shared::sgx::types::dtpm::FileEntry;
use hyper_rustls::HttpsConnectorBuilder;
use rustls::ClientConfig;
use std::sync::{Arc, RwLock};
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tonic::codec::CompressionEncoding;
use tonic::transport::{Channel, Endpoint};

use detee_shared::sgx::pb::dtpm_proto::dtpm_config_manager_client::DtpmConfigManagerClient;
use detee_shared::sgx::pb::dtpm_proto::{DtpmSetConfigReq, FileEntry as FileEntryPb};
use detee_shared::sgx::types::dtpm::DtpmConfig;

use crate::config::Config;
use crate::sgx::utils::hratls_url_and_mr_enclave_from_app_id;

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Failed to connect to DTPM: {0}")]
    DtpmConnection(#[from] tonic::transport::Error),
    #[error("Received error from DTPM: {}", _0.message())]
    ResponseStatus(#[from] tonic::Status),
    #[error("Disk access error: {0}")]
    DiskAccess(#[from] std::io::Error),
    #[error("HRatls: {0}")]
    SgxHRatls(#[from] detee_sgx::error::SgxError),
    #[error("Dtpm: {0}")]
    Dtpm(String),
    #[error(transparent)]
    ConfigError(#[from] crate::config::Error),
}

type Result<T> = std::result::Result<T, Error>;

pub async fn dtpm_client(
    hratls_uri: &str,
    mr_enclave: &[u8; 32],
) -> Result<DtpmConfigManagerClient<Channel>> {
    let private_key_pem = Config::get_hratls_private_key()?;

    let hratls_config =
        Arc::new(RwLock::new(HRaTlsConfig::new().with_hratls_private_key_pem(private_key_pem)));

    hratls_config.write().unwrap().allow_more_instance_measurement(
        InstanceMeasurement::new().with_mrenclaves(vec![*mr_enclave]),
    );

    let client_tls_config = ClientConfig::from_hratls_config(hratls_config.clone())?;
    let connector = HttpsConnectorBuilder::new()
        .with_tls_config(client_tls_config)
        .https_only()
        .enable_http2()
        .build();

    let channel =
        Endpoint::from_shared(hratls_uri.to_string())?.connect_with_connector(connector).await?;

    Ok(DtpmConfigManagerClient::new(channel).send_compressed(CompressionEncoding::Zstd))
}

pub async fn update_config(app_id: &str, config: DtpmConfig) -> Result<()> {
    let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(app_id).await?;
    let mr_enclave = mr_enclave.expect("App contract does not have a mr_enclave");

    let dtpm_client = dtpm_client(&hratls_uri, &mr_enclave).await?;

    upload_files_pb(config.filesystems.clone(), &dtpm_client).await?;
    let req = DtpmSetConfigReq { config_data: Some(config.into()), ..Default::default() };
    set_config_pb(req, &dtpm_client).await?;
    Ok(())
}

pub async fn get_config(app_id: &str) -> Result<DtpmConfig> {
    let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(app_id).await?;
    let mr_enclave = mr_enclave.expect("App contract does not have a mr_enclave");

    let dtpm_client = dtpm_client(&hratls_uri, &mr_enclave).await?;
    let config_res = get_config_pb(&dtpm_client).await?;
    let config: DtpmConfig =
        config_res.config_data.ok_or(Error::Dtpm("config data not found".to_string()))?.into();
    Ok(config)
}

pub async fn upload_files_pb(
    reqs: Vec<FileEntry>,
    client: &DtpmConfigManagerClient<Channel>,
) -> Result<()> {
    let (tx, rx) = mpsc::channel(6);
    tokio::spawn(async move {
        for file in reqs {
            let file_pb: FileEntryPb = file.into();
            let _ = tx.send(file_pb).await;
        }
    });

    let fs_stream = ReceiverStream::new(rx);

    let _ = client.clone().upload_files(tonic::Request::new(fs_stream)).await?;

    Ok(())
}

pub(crate) async fn set_config_pb(
    req: DtpmSetConfigReq,
    client: &DtpmConfigManagerClient<Channel>,
) -> Result<()> {
    let response = client.clone().set_config(tonic::Request::new(req)).await?;

    log::trace!("Received respose from the server...{:?}", response.into_inner());

    Ok(())
}

pub(crate) async fn get_config_pb(
    client: &DtpmConfigManagerClient<Channel>,
) -> Result<DtpmGetConfigRes> {
    let mgr_config_pb =
        client.clone().get_config(tonic::Request::new(Empty {})).await?.into_inner();
    Ok(mgr_config_pb)
}