use async_trait::async_trait;
use dagger_sdk::Container;
use std::{future::Future, pin::Pin, sync::Arc};

pub type DynMiddleware = Arc<dyn DaggerMiddleware + Send + Sync>;

#[async_trait]
pub trait DaggerMiddleware {
    async fn handle(
        &self,
        container: dagger_sdk::Container,
    ) -> eyre::Result<dagger_sdk::Container> {
        Ok(container)
    }
}

pub struct DaggerMiddlewareFn<F>
where
    F: Fn(Container) -> Pin<Box<dyn Future<Output = eyre::Result<Container>> + Send>>,
{
    pub func: F,
}

pub fn middleware<F>(func: F) -> Box<DaggerMiddlewareFn<F>>
where
    F: Fn(Container) -> Pin<Box<dyn Future<Output = eyre::Result<Container>> + Send>>,
{
    Box::new(DaggerMiddlewareFn { func })
}

#[async_trait]
impl<F> DaggerMiddleware for DaggerMiddlewareFn<F>
where
    F: Fn(Container) -> Pin<Box<dyn Future<Output = eyre::Result<Container>> + Send>> + Send + Sync,
{
    async fn handle(&self, container: Container) -> eyre::Result<Container> {
        // Call the closure stored in the struct
        (self.func)(container).await
    }
}

#[cfg(test)]
mod test {
    use futures::FutureExt;

    use super::*;

    #[tokio::test]
    async fn can_add_middleware() -> eyre::Result<()> {
        middleware(|c| async move { Ok(c) }.boxed());

        Ok(())
    }
}