use detee_sgx::{prelude::*, HRaTlsConfigBuilder};
use detee_shared::sgx::pb::dtpm_proto::DtpmGetConfigReq;
use hyper_rustls::HttpsConnectorBuilder;
use rustls::ClientConfig;
use std::sync::{Arc, RwLock};
use tonic::transport::{Channel, Endpoint};

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

use crate::{config::Config, utils::hratls_url_and_mr_enclave_from_app_id};

#[derive(thiserror::Error, Debug)]
pub enum Error {
    #[error("Failed to connect to the brain: {0}")]
    BrainConnection(#[from] tonic::transport::Error),
    #[error("Received error from dtpm: {}", _0.message())]
    ResponseStatus(#[from] tonic::Status),
    #[error("Hex: {0}")]
    HexDecode(#[from] hex::FromHexError),
    #[error("Disk access error: {0}")]
    DiskAccess(#[from] std::io::Error),
    #[error("HRatls: {0}")]
    SgxHRatls(#[from] detee_sgx::error::SgxError),
    #[error("DtpmConfig: {0}")]
    DtpmConfig(String),
}

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

pub async fn connect_dtpm_grpc_client(
    hratls_uri: String,
    package_mr_enclave: Option<[u8; 32]>,
) -> Result<DtpmConfigManagerClient<Channel>> {
    let private_key_pem = Config::get_hratls_private_key().0;
    let mr_signer = vec![Config::get_mr_signer()];

    let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();

    let hratls_config = Arc::new(RwLock::new(
        HRaTlsConfig::new()
            .allow_instance_measurement(InstanceMeasurement::new().with_mrsigners(mr_signer))
            .with_hratls_private_key_pem(private_key_pem),
    ));

    if let Some(mr_enclave) = package_mr_enclave {
        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)?.connect_with_connector(connector).await?;

    Ok(DtpmConfigManagerClient::new(channel))
}

pub async fn attest_and_send_config(loaded_config: DtpmConfig, uuid: &str) -> Result<()> {
    let config_data = Some(DtpmConfigData::from(loaded_config));
    let req_data = DtpmSetConfigReq { config_data, ..Default::default() };

    log::trace!("Decoded the configuration... {:?}", req_data);

    let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(uuid).await;
    log::info!("hratls uri: {}\nmr_enclave: {:?}", &hratls_uri, &mr_enclave);

    let client = connect_dtpm_grpc_client(hratls_uri, mr_enclave).await?;

    let response = client
        .max_decoding_message_size(10240000)
        .set_config(tonic::Request::new(req_data))
        .await?;

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

    Ok(())
}

pub async fn get_config_from_enclave(uuid: &str) -> Result<DtpmConfig> {
    let (hratls_uri, mr_enclave) = hratls_url_and_mr_enclave_from_app_id(uuid).await;
    log::info!("hratls uri: {}\nmr_enclave: {:?}", &hratls_uri, &mr_enclave);

    let client = connect_dtpm_grpc_client(hratls_uri, None).await?;

    let mgr_config_pb = client
        .max_decoding_message_size(10240000)
        .get_config(tonic::Request::new(DtpmGetConfigReq { empty: None }))
        .await?
        .into_inner();

    let config: DtpmConfig = mgr_config_pb
        .config_data
        .ok_or(Error::DtpmConfig("config data not found".to_string()))?
        .into();
    Ok(config)
}