use prost::Message;
use prost_reflect::{DynamicMessage, MessageDescriptor};
use tonic::{
    Status,
    codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder},
};

pub struct JsonCodec {
    request_descriptor: MessageDescriptor,
    response_descriptor: MessageDescriptor,
}

impl JsonCodec {
    pub fn new(
        request_descriptor: MessageDescriptor,
        response_descriptor: MessageDescriptor,
    ) -> Self {
        Self {
            request_descriptor,
            response_descriptor,
        }
    }
}

impl Codec for JsonCodec {
    type Encode = serde_json::Value;
    type Decode = serde_json::Value;

    type Encoder = JsonEncoder;
    type Decoder = JsonDecoder;

    fn encoder(&mut self) -> Self::Encoder {
        JsonEncoder(self.request_descriptor.clone())
    }

    fn decoder(&mut self) -> Self::Decoder {
        JsonDecoder(self.response_descriptor.clone())
    }
}

pub struct JsonEncoder(MessageDescriptor);

impl Encoder for JsonEncoder {
    type Item = serde_json::Value;
    type Error = Status;

    fn encode(&mut self, item: Self::Item, dst: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
        let message = DynamicMessage::deserialize(self.0.clone(), item)
            .map_err(|error| Status::invalid_argument(error.to_string()))?;

        message.encode_raw(dst);
        Ok(())
    }
}

pub struct JsonDecoder(MessageDescriptor);

impl Decoder for JsonDecoder {
    type Item = serde_json::Value;
    type Error = Status;

    fn decode(&mut self, src: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
        let mut message = DynamicMessage::new(self.0.clone());
        message
            .merge(src)
            .map_err(|error| Status::internal(error.to_string()))?;

        serde_json::to_value(&message)
            .map(Some)
            .map_err(|error| Status::internal(error.to_string()))
    }
}