use core::{
    marker::PhantomData,
    task::{RawWaker, RawWakerVTable, Waker},
};

use alloc::{collections::VecDeque, sync::Weak};
use futures_util::task::ArcWake;

use crate::sync::IrqSafeSpinlock;

/// Async/await primitive to suspend and wake up tasks waiting on some shared resource
pub struct QueueWaker {
    queue: IrqSafeSpinlock<VecDeque<Waker>>,
}

pub struct WakeWeak<T: ArcWake> {
    _pd: PhantomData<T>,
}

impl<T: ArcWake> WakeWeak<T> {
    pub const VTABLE: RawWakerVTable = RawWakerVTable::new(
        Self::clone_waker,
        Self::wake,
        Self::wake_by_ref,
        Self::drop_waker,
    );

    unsafe fn drop_waker(ptr: *const ()) {
        let weak = Weak::from_raw(ptr as *const T);

        // Explicitly drop the thing
        drop(weak);
    }

    unsafe fn wake(ptr: *const ()) {
        let weak = Weak::from_raw(ptr as *const T);

        if let Some(strong) = weak.upgrade() {
            strong.wake();
        }

        // Weak gets dropped
    }

    unsafe fn wake_by_ref(ptr: *const ()) {
        let weak = Weak::from_raw(ptr as *const T);

        if let Some(strong) = weak.upgrade() {
            strong.wake();
        }

        // Weak doesn't get dropped
        core::mem::forget(weak);
    }

    unsafe fn clone_waker(ptr: *const ()) -> RawWaker {
        let weak = Weak::from_raw(ptr as *const T);

        let waker = Self::weak_waker(weak.clone());

        // Prevent decrement of the refcount
        core::mem::forget(weak);

        waker
    }

    pub fn weak_waker(weak: Weak<T>) -> RawWaker {
        let raw = Weak::into_raw(weak);

        RawWaker::new(raw as *const (), &Self::VTABLE)
    }
}

impl QueueWaker {
    /// Constructs an empty [QueueWaker]
    pub const fn new() -> Self {
        Self {
            queue: IrqSafeSpinlock::new(VecDeque::new()),
        }
    }

    /// Registers a [Waker] reference to be waken up by this [QueueWaker]
    pub fn register(&self, waker: &Waker) {
        let mut queue = self.queue.lock();

        if queue.iter().any(|other| other.will_wake(waker)) {
            return;
        }

        queue.push_back(waker.clone());
    }

    /// Removes a [Waker] reference from this [QueueWaker]
    pub fn remove(&self, waker: &Waker) -> bool {
        let mut queue = self.queue.lock();
        let mut index = 0;
        let mut removed = false;

        while index < queue.len() {
            if queue[index].will_wake(waker) {
                removed = true;
                queue.remove(index);
            }
            index += 1;
        }

        removed
    }

    /// Wakes up up to `limit` tasks waiting on this queue
    pub fn wake_some(&self, limit: usize) -> usize {
        let mut queue = self.queue.lock();
        let mut count = 0;

        while count < limit
            && let Some(item) = queue.pop_front()
        {
            item.wake();
            count += 1;
        }

        count
    }

    /// Wakes up a single task waiting on this queue
    pub fn wake_one(&self) -> bool {
        self.wake_some(1) != 0
    }

    /// Wakes up all tasks waiting on this queue
    pub fn wake_all(&self) -> usize {
        self.wake_some(usize::MAX)
    }
}
