#![feature(pin, arbitrary_self_types, futures_api, async_await, await_macro)]

use std::future::Future;
use std::marker::{PhantomData, Unpin};
use std::mem::PinMut;
use std::task::{Context, Poll};

fn main() {

}

pub async fn sequential_or(fut1: impl Future<Output = bool>, fut2: impl Future<Output = bool>) -> bool {
    await!(fut1) || await!(fut2)
}

pub async fn parallel_or(fut1: impl Future<Output = bool>, fut2: impl Future<Output = bool>) -> bool {
    let fut1 = map_fut(fut1, |b| if b {
        Ok(())
    } else {
        Err(())
    });
    let fut2 = map_fut(fut2, |b| if b {
        Ok(())
    } else {
        Err(())
    });
    await!(parallel_or_else(fut1, fut2)).is_ok()
}

pub async fn ident_fut<T>(x: T) -> T {
    x
}

pub async fn map_fut<T, U>(fut: impl Future<Output = T>, f: impl FnOnce(T) -> U) -> U {
    f(await!(fut))
}

#[derive(Debug, Clone)]
pub struct Forever<T>(PhantomData<T>);

pub fn forever<T>() -> Forever<T> {
    Forever(PhantomData)
}

impl<T> Future for Forever<T> {
    type Output = T;
    fn poll(self: PinMut<Self>, cx: &mut Context) -> Poll<Self::Output> {
        cx.local_waker().wake();
        Poll::Pending
    }
}

#[derive(Debug, Clone)]
pub struct WaitN<F: Future> {
    count: usize,  // unpinned
    fut: F,  // pinned
}

pub fn wait_n<F: Future>(count: usize, fut: F) -> WaitN<F> {
    WaitN {
        count,
        fut,
    }
}

impl<F: Future> WaitN<F> {
    fn count_mut<'a>(self: PinMut<'a, Self>) -> &'a mut usize {
        unsafe { &mut PinMut::get_mut(self).count }
    }
    fn fut_pinmut<'a>(self: PinMut<'a, Self>) -> PinMut<'a, F> {
        unsafe { PinMut::map(self, |this| &mut this.fut) }
    }
}

impl<F: Future> Future for WaitN<F> {
    type Output = F::Output;
    fn poll(mut self: PinMut<Self>, cx: &mut Context) -> Poll<Self::Output> {
        if self.count <= 0 {
            self.reborrow().fut_pinmut().poll(cx)
        } else {
            *self.count_mut() -= 1;
            cx.local_waker().wake();
            Poll::Pending
        }
    }
}

#[derive(Debug, Clone)]
pub struct ParallelOrElse<T, E1, E2, F1, F2>
where
    F1: Future<Output = Result<T, E1>>,
    F2: Future<Output = Result<T, E2>>,
{
    fut1: F1,  // pinned
    fut2: F2,  // pinned
    error1: Option<E1>,  // unpinned
    error2: Option<E2>,  // unpinned
}

pub fn parallel_or_else<T, E1, E2, F1, F2>(fut1: F1, fut2: F2) -> ParallelOrElse<T, E1, E2, F1, F2>
where
    F1: Future<Output = Result<T, E1>>,
    F2: Future<Output = Result<T, E2>>,
{
    ParallelOrElse {
        fut1,
        fut2,
        error1: None,
        error2: None,
    }
}

impl<T, E1, E2, F1, F2> ParallelOrElse<T, E1, E2, F1, F2>
where
    F1: Future<Output = Result<T, E1>>,
    F2: Future<Output = Result<T, E2>>,
{
    fn fut1_pinmut<'a>(self: PinMut<'a, Self>) -> PinMut<'a, F1> {
        unsafe { PinMut::map(self, |this| &mut this.fut1) }
    }
    fn fut2_pinmut<'a>(self: PinMut<'a, Self>) -> PinMut<'a, F2> {
        unsafe { PinMut::map(self, |this| &mut this.fut2) }
    }
    fn error1_mut<'a>(self: PinMut<'a, Self>) -> &'a mut Option<E1> {
        unsafe { &mut PinMut::get_mut(self).error1 }
    }
    fn error2_mut<'a>(self: PinMut<'a, Self>) -> &'a mut Option<E2> {
        unsafe { &mut PinMut::get_mut(self).error2 }
    }
}

impl<T, E1, E2, F1, F2> Future for ParallelOrElse<T, E1, E2, F1, F2>
where
    F1: Future<Output = Result<T, E1>>,
    F2: Future<Output = Result<T, E2>>,
{
    type Output = Result<T, (E1, E2)>;
    fn poll(mut self: PinMut<Self>, cx: &mut Context) -> Poll<Self::Output> {
        if self.error1.is_none() {
            *self.reborrow().error1_mut() = match self.reborrow().fut1_pinmut().poll(cx) {
                Poll::Ready(Ok(x)) => return Poll::Ready(Ok(x)),
                Poll::Ready(Err(e)) => Some(e),
                Poll::Pending => None,
            };
        }
        if self.error2.is_none() {
            *self.reborrow().error2_mut() = match self.reborrow().fut2_pinmut().poll(cx) {
                Poll::Ready(Ok(x)) => return Poll::Ready(Ok(x)),
                Poll::Ready(Err(e)) => Some(e),
                Poll::Pending => None,
            };
        }
        if self.error1.is_some() && self.error2.is_some() {
            return Poll::Ready(Err((self.reborrow().error1_mut().take().unwrap(), self.reborrow().error2_mut().take().unwrap())));
        }
        Poll::Pending
    }
}

impl<T, E1, E2, F1, F2> Unpin for ParallelOrElse<T, E1, E2, F1, F2>
where
    F1: Future<Output = Result<T, E1>> + Unpin,
    F2: Future<Output = Result<T, E2>> + Unpin,
{}