use std::{
    future::Future,
    pin::Pin,
    sync::{
        atomic::{AtomicBool, Ordering},
        Arc,
    },
    task::{Context, Poll, Wake, Waker},
};

pub enum Either<L, R> {
    Left(L),
    Right(R),
}

pub fn select<L, R>(left: L, right: R) -> impl Future<Output = Either<L::Output, R::Output>> + Unpin
where
    L: Future + Unpin,
    R: Future + Unpin,
{
    Select {
        left: SelectState {
            woken: Arc::new(AtomicBool::new(true)),
            future: left,
        },
        right: SelectState {
            woken: Arc::new(AtomicBool::new(true)),
            future: right,
        },
    }
}

struct SelectWaker {
    woken: Arc<AtomicBool>,
    parent: Waker,
}

struct SelectState<T> {
    woken: Arc<AtomicBool>,
    future: T,
}

struct Select<L, R> {
    left: SelectState<L>,
    right: SelectState<R>,
}

impl<T: Future> SelectState<T> {
    fn check_for_poll(&self) -> bool {
        self.woken.swap(false, Ordering::AcqRel)
    }
}

impl Wake for SelectWaker {
    fn wake(self: Arc<Self>) {
        self.wake_by_ref();
    }

    fn wake_by_ref(self: &Arc<Self>) {
        self.woken.store(true, Ordering::Release);
        self.parent.wake_by_ref();
    }
}

impl<L, R> Future for Select<L, R>
where
    L: Future + Unpin,
    R: Future + Unpin,
{
    type Output = Either<L::Output, R::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let this = self.get_mut();

        if this.left.check_for_poll() {
            if let Poll::Ready(out) = Pin::new(&mut this.left.future).poll(cx) {
                return Poll::Ready(Either::Left(out));
            }
        }

        if this.right.check_for_poll() {
            if let Poll::Ready(out) = Pin::new(&mut this.right.future).poll(cx) {
                return Poll::Ready(Either::Right(out));
            }
        }

        Poll::Pending
    }
}