use crate::TryFromConn;
use std::{future::Future, marker::PhantomData, sync::Arc};
use trillium::{async_trait, Conn, Handler, Info, Status, Upgrade};
pub trait MutBorrowConn<'conn, ReturnType, Additional>: Send + Sync + 'conn {
type Fut: Future<Output = ReturnType> + Send + 'conn;
fn call(&self, conn: &'conn mut Conn, additional: Additional) -> Self::Fut;
}
impl<'conn, Fun, Fut, ReturnType, Additional> MutBorrowConn<'conn, ReturnType, Additional> for Fun
where
Fun: Fn(&'conn mut Conn, Additional) -> Fut + Send + Sync + 'conn,
Fut: Future<Output = ReturnType> + Send + 'conn,
{
type Fut = Fut;
fn call(&self, conn: &'conn mut Conn, additional: Additional) -> Fut {
self(conn, additional)
}
}
#[derive(Debug)]
pub struct ApiHandler<F, OutputHandler, TryFromConn>(
F,
PhantomData<OutputHandler>,
PhantomData<TryFromConn>,
);
impl<TryFromConnHandler, OutputHandler, Extracted>
ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
where
TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
OutputHandler: Handler,
Extracted: TryFromConn,
{
pub fn new(api_handler: TryFromConnHandler) -> Self {
Self::from(api_handler)
}
}
impl<TryFromConnHandler, OutputHandler, Extracted> From<TryFromConnHandler>
for ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
where
TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
OutputHandler: Handler,
Extracted: TryFromConn,
{
fn from(value: TryFromConnHandler) -> Self {
Self(value, PhantomData, PhantomData)
}
}
pub fn api<TryFromConnHandler, OutputHandler, Extracted>(
api_handler: TryFromConnHandler,
) -> ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
where
TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
Extracted: TryFromConn,
OutputHandler: Handler,
{
ApiHandler::from(api_handler)
}
#[async_trait]
impl<TryFromConnHandler, OutputHandler, Extracted> Handler
for ApiHandler<TryFromConnHandler, OutputHandler, Extracted>
where
TryFromConnHandler: for<'a> MutBorrowConn<'a, OutputHandler, Extracted>,
Extracted: TryFromConn,
Extracted::Error: Handler,
OutputHandler: Handler,
{
async fn run(&self, mut conn: Conn) -> Conn {
let mut output_handler: Result<OutputHandler, <Extracted as TryFromConn>::Error> =
match Extracted::try_from_conn(&mut conn).await {
Ok(extracted) => Ok(self.0.call(&mut conn, extracted).await),
Err(error_handler) => Err(error_handler),
};
if let Some(info) = conn.state_mut::<Info>() {
output_handler.init(info).await;
} else {
output_handler.init(&mut Info::default()).await;
}
let mut conn = output_handler.run(conn).await;
if conn.status().is_none() && conn.inner().response_body().is_some() {
conn.set_status(Status::Ok);
}
conn.with_state(OutputHandlerWrapper(
Arc::new(output_handler),
PhantomData::<Self>,
))
}
async fn before_send(&self, conn: Conn) -> Conn {
if let Some(OutputHandlerWrapper(handler, _)) = conn
.state::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
.cloned()
{
handler.before_send(conn).await
} else {
conn
}
}
fn has_upgrade(&self, upgrade: &Upgrade) -> bool {
upgrade
.state()
.get::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
.cloned()
.map_or(false, |OutputHandlerWrapper(handler, _)| {
handler.has_upgrade(upgrade)
})
}
async fn upgrade(&self, upgrade: Upgrade) {
if let Some(OutputHandlerWrapper(handler, _)) = upgrade
.state()
.get::<OutputHandlerWrapper<Self, OutputHandler, <Extracted as TryFromConn>::Error>>()
.cloned()
{
handler.upgrade(upgrade).await
}
}
}
struct OutputHandlerWrapper<TFC, OH, EH>(Arc<Result<OH, EH>>, PhantomData<TFC>);
impl<TFC, OH, EH> Clone for OutputHandlerWrapper<TFC, OH, EH> {
fn clone(&self) -> Self {
Self(Arc::clone(&self.0), self.1)
}
}