use core::{
    cell::UnsafeCell,
    future::poll_fn,
    ops::{AsyncFnOnce, Deref, DerefMut},
    sync::atomic::{AtomicBool, AtomicU32, Ordering},
    task::{Context, Poll},
};

use alloc::sync::Arc;
use crossbeam_queue::ArrayQueue;
use kernel_arch::task::Scheduler;
use libk_util::{sync::LockMethod, waker::QueueWaker};
use yggdrasil_abi::error::Error;

use crate::task::{sched::CpuQueue, thread::Thread};

struct ThreadedMutexInner {
    queue: ArrayQueue<Arc<Thread>>,
    lock: AtomicU32,
}

pub struct AsyncMutex<T> {
    value: UnsafeCell<T>,
    waker: QueueWaker,
    lock: AtomicBool,
}

pub struct AsyncMutexGuard<'a, T> {
    mutex: &'a AsyncMutex<T>,
}

pub struct Mutex<T> {
    value: UnsafeCell<T>,
    lock: ThreadedMutexInner,
}

pub struct MutexGuard<'a, T> {
    value: &'a UnsafeCell<T>,
    lock: &'a ThreadedMutexInner,
}

pub struct MappedAsyncMutexGuard<'a, U: ?Sized, T> {
    mutex: &'a AsyncMutex<T>,
    value: *mut U,
}

impl<T> AsyncMutex<T> {
    pub const fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            waker: QueueWaker::new(),
            lock: AtomicBool::new(false),
        }
    }

    pub fn try_lock(&self) -> bool {
        self.lock
            .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_ok()
    }

    pub fn poll_lock(&self, cx: &mut Context<'_>) -> Poll<AsyncMutexGuard<'_, T>> {
        self.waker.register(cx.waker());

        if self.try_lock() {
            self.waker.remove(cx.waker());
            return Poll::Ready(AsyncMutexGuard { mutex: self });
        }

        Poll::Pending
    }

    pub async fn lock(&self) -> AsyncMutexGuard<'_, T> {
        poll_fn(|cx| self.poll_lock(cx)).await
    }

    unsafe fn force_unlock(&self) {
        self.lock.store(false, Ordering::Release);
        self.waker.wake_one();
    }

    #[inline]
    pub fn get(&self) -> *mut T {
        self.value.get()
    }
}

unsafe impl<T: Sync> Send for AsyncMutex<T> {}
unsafe impl<T: Sync> Sync for AsyncMutex<T> {}

// impl<T> !Send for AsyncMutexGuard<'_, T> {}
unsafe impl<T: Sync> Send for AsyncMutexGuard<'_, T> {}
unsafe impl<T: Sync> Sync for AsyncMutexGuard<'_, T> {}

impl<'a, T> AsyncMutexGuard<'a, T> {
    pub async fn try_map_guard_async<
        U: 'a,
        E,
        F: AsyncFnOnce(&'a mut T) -> Result<&'a mut U, E>,
    >(
        self,
        mapper: F,
    ) -> Result<MappedAsyncMutexGuard<'a, U, T>, E> {
        let mutex = self.mutex;
        core::mem::forget(self);
        let value = mapper
            .async_call_once((unsafe { &mut *mutex.get() },))
            .await?;
        Ok(MappedAsyncMutexGuard { mutex, value })
    }
}

impl<T> Deref for AsyncMutexGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.mutex.get() }
    }
}

impl<T> DerefMut for AsyncMutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.mutex.get() }
    }
}

impl<T> Drop for AsyncMutexGuard<'_, T> {
    fn drop(&mut self) {
        unsafe {
            self.mutex.force_unlock();
        }
    }
}

impl<U: ?Sized, T> !Send for MappedAsyncMutexGuard<'_, U, T> {}
unsafe impl<U: Sync + ?Sized, T: Sync> Sync for MappedAsyncMutexGuard<'_, U, T> {}

impl<U: ?Sized, T> Deref for MappedAsyncMutexGuard<'_, U, T> {
    type Target = U;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.value }
    }
}

impl<U: ?Sized, T> DerefMut for MappedAsyncMutexGuard<'_, U, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.value }
    }
}

impl<U: ?Sized, T> Drop for MappedAsyncMutexGuard<'_, U, T> {
    fn drop(&mut self) {
        unsafe {
            self.mutex.force_unlock();
        }
    }
}

impl ThreadedMutexInner {
    const UNLOCKED: u32 = 0;
    const LOCKED: u32 = 1;

    pub fn new(queue_capacity: usize) -> Self {
        Self {
            queue: ArrayQueue::new(queue_capacity),
            lock: AtomicU32::new(Self::UNLOCKED),
        }
    }

    fn try_lock(&self) -> bool {
        self.lock
            .compare_exchange(
                Self::UNLOCKED,
                Self::LOCKED,
                Ordering::Acquire,
                Ordering::Relaxed,
            )
            .is_ok()
    }
}

impl<'q> LockMethod<'q> for ThreadedMutexInner {
    type Guard<'a>
        = ()
    where
        'a: 'q,
        Self: 'a;

    fn lock(&self) -> Result<(), Error> {
        // Fast-path

        loop {
            if self.try_lock() {
                return Ok(());
            }

            let thread = Thread::current();

            self.queue
                .push(thread.clone())
                .map_err(|_| Error::WouldBlock)?;

            thread.suspend()?;
        }
    }

    unsafe fn release(&self) {
        if self.lock.swap(Self::UNLOCKED, Ordering::Release) == Self::LOCKED {
            if let Some(t) = self.queue.pop() {
                t.enqueue();

                // Yield current thread to avoid congestion
                unsafe {
                    CpuQueue::local().yield_cpu();
                }
            }
        }
    }
}

impl<T> Mutex<T> {
    pub fn new(value: T) -> Self {
        Self {
            value: UnsafeCell::new(value),
            lock: ThreadedMutexInner::new(32),
        }
    }
}

impl<'q, T> LockMethod<'q> for Mutex<T> {
    type Guard<'a>
        = MutexGuard<'a, T>
    where
        'a: 'q,
        Self: 'a;

    fn lock(&'q self) -> Result<Self::Guard<'q>, Error> {
        self.lock.lock()?;
        Ok(MutexGuard {
            lock: &self.lock,
            value: &self.value,
        })
    }

    unsafe fn release(&self) {
        self.lock.release();
    }
}

unsafe impl<T: Send> Sync for Mutex<T> {}
unsafe impl<T: Send> Send for Mutex<T> {}

impl<T> Deref for MutexGuard<'_, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.value.get() }
    }
}

impl<T> DerefMut for MutexGuard<'_, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.value.get() }
    }
}

impl<T> Drop for MutexGuard<'_, T> {
    fn drop(&mut self) {
        unsafe { self.lock.release() }
    }
}
