tina-core 0.0.2

//! tonic gRPC response
use std::fmt::Debug;

use bytes::BufMut;
use http::{HeaderMap, HeaderValue};
use http_body::Body;
use prost::Message;

use crate::{
    app_error_from,
    tina::{
        constant::Constants,
        data::{app_error::AppError, grpc::response_data::GrpcResData},
        grpc::{FromGrpcResponse, IntoGrpcResponse},
    },
};

/// Response Body
pub type BoxBody = http_body::combinators::UnsyncBoxBody<bytes::Bytes, AppError>;
/// Response
pub struct Response {
    pub(crate) inner: tonic::Response<BoxBody>,
}

impl Response {
    /// 构建
    pub fn new(body: BoxBody) -> Self {
        Self {
            inner: tonic::Response::new(body),
        }
    }
    /// 从parts构建
    pub fn from_parts<T: Message + Debug>(metadata: tonic::metadata::MetadataMap, data: T, extensions: tonic::Extensions) -> Self {
        let encoded_len = data.encoded_len();
        let mut buf = bytes::BytesMut::with_capacity(std::mem::size_of::<u8>() + std::mem::size_of::<u32>() + encoded_len);
        buf.put_u8(0);
        buf.put_u32(encoded_len as u32);
        data.encode(&mut buf).unwrap_or_else(|err| panic!("Encode buf failed, reason: {err:?}, data: {data:?}"));
        let body = buf.freeze();
        let body = http_body::Full::new(body).map_err(app_error_from!()).boxed_unsync();
        Self {
            inner: tonic::Response::from_parts(metadata, body, extensions),
        }
    }
}

// impl<T: Message + Debug> From<tonic::Response<T>> for Response {
//     fn from(value: tonic::Response<T>) -> Self {
//         let (metadata, data, extensions) = value.into_parts();
//         Self::from_parts(metadata, data, extensions)
//     }
// }

impl<T: Message + Debug, E: Message + Debug> From<tonic::Response<Result<T, E>>> for Response {
    fn from(value: tonic::Response<Result<T, E>>) -> Self {
        let (metadata, r, extensions) = value.into_parts();
        match r {
            Ok(data) => Self::from_parts(metadata, data, extensions),
            Err(err) => Self::from_parts(metadata, err, extensions),
        }
    }
}

impl Debug for Response {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.inner.fmt(f)
    }
}

#[async_trait]
impl IntoGrpcResponse for AppError {
    type Response = tonic::Response<AppError>;

    async fn into_grpc_response(self) -> Self::Response {
        let mut res_headers = HeaderMap::new();
        res_headers.insert(Constants::RESPONSE_ERROR_HEADER, HeaderValue::from_static("true"));
        res_headers.insert("content-type", HeaderValue::from_static("application/grpc"));
        // let status = tonic::Status::new(tonic::Code::Ok, "");
        let status = tonic::Status::from(&self);
        if let Err(err) = status.add_header(&mut res_headers) {
            tracing::error!("Add gRPC header failed: {err:?}");
        }
        let metadata = tonic::metadata::MetadataMap::from_headers(res_headers);
        tonic::Response::from_parts(metadata, self, tonic::Extensions::default())
    }
}

#[async_trait]
impl FromGrpcResponse for AppError {
    type Response = tonic::Response<AppError>;
    type Rejection = AppError;

    async fn from_grpc_response(res: Self::Response) -> Result<Self, Self::Rejection>
    where
        Self: Sized,
    {
        Ok(res.into_inner())
    }
}

#[async_trait]
impl<D> IntoGrpcResponse for GrpcResData<D>
where
    D: Message + Debug + Send + Sync + 'static,
{
    type Response = tonic::Response<D>;

    async fn into_grpc_response(self) -> Self::Response {
        let mut res_headers = self.metadata;
        res_headers.insert(Constants::RESPONSE_ERROR_HEADER, HeaderValue::from_static("false"));
        res_headers.insert("content-type", HeaderValue::from_static("application/grpc"));
        let status = tonic::Status::new(tonic::Code::Ok, "");
        if let Err(err) = status.add_header(&mut res_headers) {
            tracing::error!("Add gRPC header failed: {err:?}");
        }
        let metadata = tonic::metadata::MetadataMap::from_headers(res_headers);
        tonic::Response::from_parts(metadata, self.data, tonic::Extensions::default())
    }
}

#[async_trait]
impl<D> FromGrpcResponse for GrpcResData<D>
where
    D: Message + Default + Debug + Send + Sync + 'static,
{
    type Response = tonic::Response<D>;

    type Rejection = AppError;

    async fn from_grpc_response(res: Self::Response) -> Result<Self, Self::Rejection>
    where
        Self: Sized,
    {
        let (metadata, message, extensions) = res.into_parts();
        let mut metadata = metadata.into_headers();
        let mut extensions = extensions.into_http();
        let mut data = GrpcResData::new(message);
        {
            std::mem::swap(&mut data.metadata, &mut metadata);
        }
        {
            std::mem::swap(&mut data.extensions, &mut extensions);
        }
        Ok(data)
    }
}

#[async_trait]
impl<D, E, DM, EM> IntoGrpcResponse for Result<D, E>
where
    D: IntoGrpcResponse<Response = tonic::Response<DM>> + Debug + Send + Sync + 'static,
    E: IntoGrpcResponse<Response = tonic::Response<EM>> + Debug + Send + Sync + 'static,
    DM: Message + Debug + Send + Sync + 'static,
    EM: Message + Debug + Send + Sync + 'static,
{
    type Response = tonic::Response<Result<DM, EM>>;

    async fn into_grpc_response(self) -> Self::Response {
        match self {
            Ok(data) => {
                let res = data.into_grpc_response().await;
                let (metadata, message, extensions) = res.into_parts();
                tonic::Response::from_parts(metadata, Ok(message), extensions)
            }
            Err(err) => {
                let res = err.into_grpc_response().await;
                let (metadata, message, extensions) = res.into_parts();
                tonic::Response::from_parts(metadata, Err(message), extensions)
            }
        }
    }
}

#[async_trait]
impl<D, E, DM, EM> FromGrpcResponse for Result<D, E>
where
    D: FromGrpcResponse<Response = tonic::Response<DM>, Rejection = AppError> + Debug + Send + Sync + 'static,
    E: FromGrpcResponse<Response = tonic::Response<EM>, Rejection = AppError> + Debug + Send + Sync + 'static,
    DM: Message + Debug + Send + Sync + 'static,
    EM: Message + Debug + Send + Sync + 'static,
{
    type Response = tonic::Response<Result<DM, EM>>;
    type Rejection = E::Rejection;

    async fn from_grpc_response(res: Self::Response) -> Result<Self, Self::Rejection>
    where
        Self: Sized,
    {
        tracing::trace!("Result receive response: {:?}", res);
        let (metadata, r, extensions) = res.into_parts();
        match r {
            Ok(v) => {
                let res = tonic::Response::from_parts(metadata, v, extensions);
                match D::from_grpc_response(res).await {
                    Ok(v) => Ok(Ok(v)),
                    Err(err) => Err(err),
                }
            }
            Err(err) => {
                let res = tonic::Response::from_parts(metadata, err, extensions);
                match E::from_grpc_response(res).await {
                    Ok(v) => Ok(Err(v)),
                    Err(err) => Err(err),
                }
            }
        }
    }
}