use prost::Message;
use rocket::Request;
use rocket::data::{Data, FromData, Outcome, ToByteUnit};
use rocket::http::{ContentType, Status};
use std::ops::{Deref, DerefMut};

pub struct Proto<T>(pub T);

impl<T> Proto<T> {
    pub fn into_inner(self) -> T {
        self.0
    }
}

impl<T> Deref for Proto<T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl<T> DerefMut for Proto<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.0
    }
}

#[rocket::async_trait]
impl<'r, T: Message + Default> FromData<'r> for Proto<T> {
    type Error = String;

    async fn from_data(req: &'r Request<'_>, data: Data<'r>) -> Outcome<'r, Self> {
        if req.content_type() != Some(&ContentType::new("application", "protobuf"))
            && req.content_type() != Some(&ContentType::new("application", "octet-stream"))
        {
            return Outcome::Forward((data, Status::NotFound));
        }

        let limit = req.limits().get("protobuf").unwrap_or(1.mebibytes());
        let bytes = match data.open(limit).into_bytes().await {
            Ok(bytes) if bytes.is_complete() => bytes.into_inner(),
            Ok(_) => return Outcome::Error((Status::PayloadTooLarge, "Payload too large".into())),
            Err(e) => return Outcome::Error((Status::InternalServerError, e.to_string())),
        };

        match T::decode(&bytes[..]) {
            Ok(msg) => Outcome::Success(Proto(msg)),
            Err(e) => Outcome::Error((Status::UnprocessableEntity, e.to_string())),
        }
    }
}