// SPDX-License-Identifier: GPL-2.0-or-later
//
// Wait Queue implementation with wait_until as the core primitive.
//
// The wait_until family returns Option<R>, allowing direct return of acquired resources.
// The wait_event family (returning bool) is implemented on top of wait_until for compatibility.

use core::{
    intrinsics::unlikely,
    marker::PhantomData,
    mem,
    sync::atomic::{AtomicU32, AtomicU8, Ordering},
};

use alloc::{
    boxed::Box,
    collections::{BTreeMap, VecDeque},
    rc::Rc,
    sync::{Arc, Weak},
    vec::Vec,
};
use log::warn;
use system_error::SystemError;

use crate::{
    arch::{ipc::signal::Signal, CurrentIrqArch},
    exception::InterruptArch,
    libs::mutex::MutexGuard,
    process::{ProcessControlBlock, ProcessManager, ProcessState},
    sched::{io_schedule, schedule, SchedMode},
    time::{
        timer::{clock, next_n_us_timer_jiffies, Timer},
        Duration, Instant,
    },
};

use super::spinlock::{SpinLock, SpinLockGuard};

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum WaitSignalMode {
    Uninterruptible,
    Interruptible,
    Killable,
}

const DEFAULT_WAIT_TAG: usize = usize::MAX;
const EMPTY_TAG_CACHE_LIMIT: usize = 8;

#[derive(Debug)]
struct InnerWaitQueue {
    dead: bool,
    waiters: BTreeMap<usize, VecDeque<Arc<Waker>>>,
    empty_tag_count: usize,
}

/// 等待队列：基于一次性 Waiter/Waker，避免唤醒丢失
#[derive(Debug)]
pub struct WaitQueue {
    inner: SpinLock<InnerWaitQueue>,
    num_waiters: AtomicU32,
}

/// 属于当前线程的等待者，不可跨线程共享
pub struct Waiter {
    waker: Arc<Waker>,
    _nosend: PhantomData<Rc<()>>,
}

/// 可跨 CPU/线程共享的唤醒器
#[derive(Debug)]
pub struct Waker {
    state: AtomicU8,
    tag: usize,
    target: Weak<ProcessControlBlock>,
}

#[allow(dead_code)]
impl WaitQueue {
    pub const fn default() -> Self {
        WaitQueue {
            inner: SpinLock::new(InnerWaitQueue::INIT),
            num_waiters: AtomicU32::new(0),
        }
    }

    pub fn register_waker(&self, waker: Arc<Waker>) -> Result<(), SystemError> {
        let mut guard = self.inner.lock_irqsave();
        if guard.dead {
            return Err(SystemError::ECHILD);
        }
        if waker.tag != DEFAULT_WAIT_TAG
            && guard
                .waiters
                .get(&waker.tag)
                .is_some_and(VecDeque::is_empty)
        {
            guard.empty_tag_count -= 1;
        }
        guard.waiters.entry(waker.tag).or_default().push_back(waker);
        self.num_waiters.fetch_add(1, Ordering::Release);
        Ok(())
    }

    pub fn remove_waker(&self, target: &Arc<Waker>) {
        let mut guard = self.inner.lock_irqsave();
        let (removed, became_empty) = if let Some(waiters) = guard.waiters.get_mut(&target.tag) {
            let before = waiters.len();
            waiters.retain(|w| !Arc::ptr_eq(w, target));
            let removed = before - waiters.len();
            (removed, removed > 0 && waiters.is_empty())
        } else {
            (0, false)
        };
        if removed > 0 {
            self.num_waiters
                .fetch_sub(removed as u32, Ordering::Release);
        }
        let evicted = became_empty
            .then(|| guard.cache_or_evict_empty_tag(target.tag))
            .flatten();
        drop(guard);
        drop(evicted);
    }

    /// Returns a transient, lockless snapshot of whether the queue is empty.
    ///
    /// This is only an advisory hint. A negative result must not be used to
    /// skip a correctness-critical, one-shot wakeup without an independent
    /// synchronization protocol.
    pub fn is_empty(&self) -> bool {
        self.num_waiters.load(Ordering::Acquire) == 0
    }

    // ==================== Core API: wait_until family ====================

    /// Waits until condition returns `Some(R)` (signals are ignored; use
    /// `wait_until_interruptible` for signal-aware waits).
    ///
    /// This is the core waiting primitive. The condition closure is called repeatedly
    /// until it returns `Some(_)`. The waker is registered BEFORE checking the condition
    /// to avoid missed wakeups.
    ///
    /// # Example
    /// ```ignore
    /// // Directly acquire a lock guard without race condition
    /// let guard = queue.wait_until(|| mutex.try_lock());
    /// ```
    #[track_caller]
    pub fn wait_until<F, R>(&self, cond: F) -> R
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Uninterruptible,
            None,
            None,
            None::<fn()>,
            false,
        )
        .unwrap()
    }

    /// 等待 IO 操作完成（不可中断）
    ///
    /// 与 wait_until 类似，但会正确标记进程在等待 IO，用于 iowait 统计
    #[track_caller]
    pub fn wait_until_io<F, R>(&self, cond: F) -> R
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Uninterruptible,
            None,
            None,
            None::<fn()>,
            true,
        )
        .unwrap()
    }

    /// Waits until condition returns `Some(R)` (interruptible by signals).
    ///
    /// Returns `Err(ERESTARTSYS)` if interrupted by a signal.
    #[track_caller]
    pub fn wait_until_interruptible<F, R>(&self, cond: F) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Interruptible,
            None,
            None,
            None::<fn()>,
            false,
        )
    }

    /// 等待 IO 操作完成（可被信号中断）
    ///
    /// 与 wait_until_interruptible 类似，但会正确标记进程在等待 IO
    #[track_caller]
    pub fn wait_until_io_interruptible<F, R>(&self, cond: F) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Interruptible,
            None,
            None,
            None::<fn()>,
            true,
        )
    }

    #[track_caller]
    pub fn wait_until_killable<F, R>(&self, cond: F) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Killable,
            None,
            None,
            None::<fn()>,
            false,
        )
    }

    /// Waits until condition returns `Some(R)` with timeout (interruptible).
    ///
    /// Returns:
    /// - `Ok(R)` if condition satisfied
    /// - `Err(ERESTARTSYS)` if interrupted by signal
    /// - `Err(EAGAIN_OR_EWOULDBLOCK)` if timeout
    #[track_caller]
    pub fn wait_until_timeout<F, R>(&self, cond: F, timeout: Duration) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Interruptible,
            Some(timeout),
            None,
            None::<fn()>,
            false,
        )
    }

    /// Waits with a timeout expressed in scheduler ticks.
    ///
    /// This preserves timeout ranges that cannot be represented by `Duration`'s
    /// microsecond storage and is useful for Linux socket timeout semantics.
    #[track_caller]
    pub fn wait_until_timeout_ticks<F, R>(
        &self,
        cond: F,
        timeout_ticks: u64,
    ) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
    {
        self.wait_until_impl(
            cond,
            WaitSignalMode::Interruptible,
            None,
            Some(timeout_ticks),
            None::<fn()>,
            false,
        )
    }

    /// Core implementation for all wait_until variants.
    ///
    /// Key design:
    /// - Create only ONE waiter/waker pair
    /// - Enqueue the waker BEFORE each condition check
    /// - This ensures no wakeup is lost between check and sleep
    fn wait_until_impl<F, R, B>(
        &self,
        cond: F,
        mode: WaitSignalMode,
        timeout: Option<Duration>,
        timeout_ticks: Option<u64>,
        before_sleep: Option<B>,
        is_io: bool,
    ) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
        B: FnMut(),
    {
        self.wait_until_impl_tagged(
            cond,
            mode,
            timeout,
            timeout_ticks,
            before_sleep,
            (is_io, DEFAULT_WAIT_TAG),
        )
    }

    fn wait_until_impl_tagged<F, R, B>(
        &self,
        mut cond: F,
        mode: WaitSignalMode,
        timeout: Option<Duration>,
        timeout_ticks: Option<u64>,
        mut before_sleep: Option<B>,
        wait_class: (bool, usize),
    ) -> Result<R, SystemError>
    where
        F: FnMut() -> Option<R>,
        B: FnMut(),
    {
        let (is_io, tag) = wait_class;
        // Fast path: check condition first
        if let Some(res) = cond() {
            return Ok(res);
        }

        // Compute one fixed deadline for the whole wait. Reusing it across wakeups
        // prevents timer re-registration or preemption from extending the budget.
        let deadline_ticks = timeout_ticks
            .map(|ticks| clock().saturating_add(ticks))
            .or_else(|| timeout.map(|duration| next_n_us_timer_jiffies(duration.total_micros())));

        // Create only ONE waiter/waker pair (key difference from old implementation)
        let (waiter, waker) = Waiter::new_pair_tagged(tag);

        fn cancel_or_timeout<F, R>(
            cond: &mut F,
            waker: &Arc<Waker>,
            err: SystemError,
        ) -> Result<R, SystemError>
        where
            F: FnMut() -> Option<R>,
        {
            // Close to prevent a concurrent wake from being lost.
            waker.close();
            if let Some(res) = cond() {
                return Ok(res);
            }
            Err(err)
        }

        loop {
            // Check timeout
            if deadline_ticks.is_some_and(|deadline| clock() >= deadline) {
                self.remove_waker(&waker);
                return cancel_or_timeout(&mut cond, &waker, SystemError::EAGAIN_OR_EWOULDBLOCK);
            }

            // Enqueue waker BEFORE checking condition (critical for correctness!)
            // This ensures that if condition becomes true after our check,
            // the subsequent wake_one() will find our waker in the queue.
            if self.register_waker(waker.clone()).is_err() {
                // Queue is dead, spin until condition is met
                loop {
                    if let Some(res) = cond() {
                        return Ok(res);
                    }
                    core::hint::spin_loop();
                }
            }

            // Check condition AFTER enqueuing
            if let Some(res) = cond() {
                // Condition satisfied, remove waker and return
                self.remove_waker(&waker);
                return Ok(res);
            }

            let has_pending = match mode {
                WaitSignalMode::Uninterruptible => false,
                WaitSignalMode::Interruptible => {
                    Signal::signal_pending_state(true, false, &ProcessManager::current_pcb())
                }
                WaitSignalMode::Killable => {
                    Signal::signal_pending_state(false, true, &ProcessManager::current_pcb())
                }
            };
            if has_pending {
                self.remove_waker(&waker);
                return cancel_or_timeout(&mut cond, &waker, SystemError::ERESTARTSYS);
            }

            // Execute before_sleep hook (e.g., release a lock)
            if let Some(ref mut hook) = before_sleep {
                hook();
            }

            // Setup timeout timer if needed
            let timer = if let Some(deadline) = deadline_ticks {
                if clock() >= deadline {
                    self.remove_waker(&waker);
                    return cancel_or_timeout(
                        &mut cond,
                        &waker,
                        SystemError::EAGAIN_OR_EWOULDBLOCK,
                    );
                }
                let t: Arc<Timer> = Timer::new(TimeoutWaker::new(waker.clone()), deadline);
                t.activate();
                Some(t)
            } else {
                None
            };

            // Wait to be woken
            let wait_result = match (mode, is_io) {
                (WaitSignalMode::Uninterruptible, false) => waiter.wait(false),
                (WaitSignalMode::Uninterruptible, true) => waiter.wait_io(false),
                (WaitSignalMode::Interruptible, false) => waiter.wait(true),
                (WaitSignalMode::Interruptible, true) => waiter.wait_io(true),
                (WaitSignalMode::Killable, _) => waiter.wait_killable(),
            };

            // Timer cleanup
            let was_timeout = timer.as_ref().map(|t| t.timeout()).unwrap_or(false);
            if !was_timeout {
                if let Some(t) = timer {
                    t.cancel();
                }
            }

            // Handle wait error (signal interruption)
            if let Err(e) = wait_result {
                self.remove_waker(&waker);
                return cancel_or_timeout(&mut cond, &waker, e);
            }

            // Note: We do NOT check condition here without re-enqueuing!
            // The loop will re-enqueue the waker and check condition again.
            // This is critical: if we checked condition here and it returned None,
            // and then someone called wake_one(), the wakeup would be lost because
            // our waker was already popped by the previous wake.
        }
    }

    // ==================== Compatibility API: wait_event family ====================
    // These are implemented on top of wait_until for backward compatibility.

    /// 可中断等待条件成立；`before_sleep` 在入队后、睡眠前执行（用于解锁等操作）
    pub fn wait_event_interruptible<F, B>(
        &self,
        mut cond: F,
        before_sleep: Option<B>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
        B: FnMut(),
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            None,
            None,
            before_sleep,
            false,
        )
    }

    /// Wait interruptibly while publishing a caller-defined class tag on the
    /// waiter. The tag lets a wakeup source select one eligible class without
    /// disturbing unrelated waiters on the same queue.
    pub fn wait_event_interruptible_tagged<F>(
        &self,
        tag: usize,
        mut cond: F,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        if tag == DEFAULT_WAIT_TAG {
            return Err(SystemError::EINVAL);
        }
        self.wait_until_impl_tagged(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            None,
            None,
            None::<fn()>,
            (false, tag),
        )
    }

    /// 可中断等待 IO 操作完成；正确标记为 IO 等待以用于 iowait 统计
    ///
    /// 用于块设备 IO、网络 IO 等场景。会设置 IN_IOWAIT 标志，使得 CPU 空闲时间
    /// 能够正确区分 iowait 和 pure idle。
    pub fn wait_event_io_interruptible<F, B>(
        &self,
        mut cond: F,
        before_sleep: Option<B>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
        B: FnMut(),
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            None,
            None,
            before_sleep,
            true,
        )
    }

    pub fn wait_event_killable<F, B>(
        &self,
        mut cond: F,
        before_sleep: Option<B>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
        B: FnMut(),
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Killable,
            None,
            None,
            before_sleep,
            false,
        )
    }

    /// 不可中断等待条件成立
    pub fn wait_event_uninterruptible<F, B>(
        &self,
        mut cond: F,
        before_sleep: Option<B>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
        B: FnMut(),
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Uninterruptible,
            None,
            None,
            before_sleep,
            false,
        )
    }

    /// 可中断等待条件成立，支持可选超时。
    ///
    /// - `timeout == None`：无限等待，直到条件成立或收到信号（返回 `ERESTARTSYS`）
    /// - `timeout == Some(d)`：超时返回 `EAGAIN_OR_EWOULDBLOCK`
    pub fn wait_event_interruptible_timeout<F>(
        &self,
        mut cond: F,
        timeout: Option<Duration>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            timeout,
            None,
            None::<fn()>,
            false,
        )
    }

    /// 可中断等待 IO 操作完成，支持可选超时
    ///
    /// 用于块设备、网络 socket 等 IO 操作。正确标记为 IO 等待以用于 iowait 统计。
    pub fn wait_event_io_interruptible_timeout<F>(
        &self,
        mut cond: F,
        timeout: Option<Duration>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            timeout,
            None,
            None::<fn()>,
            true,
        )
    }

    /// 不可中断等待条件成立，支持可选超时。
    ///
    /// - `timeout == None`：无限等待，直到条件成立。
    /// - `timeout == Some(d)`：超时返回 `EAGAIN_OR_EWOULDBLOCK`。
    pub fn wait_event_uninterruptible_timeout<F>(
        &self,
        mut cond: F,
        timeout: Option<Duration>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Uninterruptible,
            timeout,
            None,
            None::<fn()>,
            false,
        )
    }

    /// `wait_event_interruptible_timeout` 的扩展版本：提供 `before_sleep` 钩子。
    ///
    /// 典型用途：入队后、睡眠前释放锁（避免持锁睡眠）。
    pub fn wait_event_interruptible_timeout_with<F, B>(
        &self,
        mut cond: F,
        timeout: Option<Duration>,
        before_sleep: Option<B>,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
        B: FnMut(),
    {
        self.wait_until_impl(
            || if cond() { Some(()) } else { None },
            WaitSignalMode::Interruptible,
            timeout,
            None,
            before_sleep,
            false,
        )
    }

    // ==================== Wakeup API ====================

    /// 唤醒一个等待者
    pub fn wakeup(&self, _state: Option<ProcessState>) -> bool {
        self.wake_one()
    }

    pub fn wake_one(&self) -> bool {
        loop {
            let next = {
                let mut guard = self.inner.lock_irqsave();
                let tagged_waker = guard
                    .waiters
                    .iter_mut()
                    .find_map(|(tag, waiters)| waiters.pop_front().map(|waker| (*tag, waker)));
                let waker = tagged_waker.as_ref().map(|(_, waker)| waker);
                if waker.is_some() {
                    self.num_waiters.fetch_sub(1, Ordering::Release);
                }
                let evicted = tagged_waker
                    .as_ref()
                    .filter(|(tag, _)| guard.waiters.get(tag).is_some_and(VecDeque::is_empty))
                    .and_then(|(tag, _)| guard.cache_or_evict_empty_tag(*tag));
                drop(guard);
                drop(evicted);
                tagged_waker.map(|(_, waker)| waker)
            };

            let Some(waker) = next else { return false };
            if waker.wake() {
                return true;
            }
        }
    }

    /// Wake the first live waiter in an exact caller-defined class.
    pub fn wake_one_tagged(&self, tag: usize) -> bool {
        if tag == DEFAULT_WAIT_TAG {
            return false;
        }
        loop {
            let next = {
                let mut guard = self.inner.lock_irqsave();
                let waker = guard.waiters.get_mut(&tag).and_then(VecDeque::pop_front);
                let Some(waker) = waker else { return false };
                self.num_waiters.fetch_sub(1, Ordering::Release);
                let evicted = guard
                    .waiters
                    .get(&tag)
                    .is_some_and(VecDeque::is_empty)
                    .then(|| guard.cache_or_evict_empty_tag(tag))
                    .flatten();
                drop(guard);
                drop(evicted);
                waker
            };

            if next.wake() {
                return true;
            }
        }
    }

    /// 唤醒队列中全部等待者
    pub fn wakeup_all(&self, _state: Option<ProcessState>) {
        self.wake_all();
    }

    pub fn wake_all(&self) -> usize {
        let drained;
        {
            let mut guard = self.inner.lock_irqsave();
            drained = mem::take(&mut guard.waiters);
            guard.empty_tag_count = 0;
            self.num_waiters.store(0, Ordering::Release);
        }

        let mut woken = 0;
        for (_, wakers) in drained {
            for waker in wakers {
                if waker.wake() {
                    woken += 1;
                }
            }
        }
        woken
    }

    /// 标记等待队列失效，清空并唤醒剩余等待者
    pub fn mark_dead(&self) {
        let drained;
        {
            let mut guard = self.inner.lock_irqsave();
            guard.dead = true;
            drained = mem::take(&mut guard.waiters);
            guard.empty_tag_count = 0;
            self.num_waiters.store(0, Ordering::Release);
        }
        for (_, wakers) in drained {
            for waker in wakers {
                waker.wake();
                waker.close();
            }
        }
    }

    pub fn len(&self) -> usize {
        self.num_waiters.fetch_add(0, Ordering::Acquire) as usize
    }

    // ==================== Sleep with lock release ====================

    pub fn sleep_unlock_spinlock<T>(&self, to_unlock: SpinLockGuard<T>) -> Result<(), SystemError> {
        before_sleep_check(1);
        let mut to_unlock = Some(to_unlock);
        self.wait_event_interruptible(
            || false,
            Some(move || {
                if let Some(lock) = to_unlock.take() {
                    drop(lock);
                }
            }),
        )
    }

    pub fn sleep_unlock_mutex<T>(&self, to_unlock: MutexGuard<T>) -> Result<(), SystemError> {
        before_sleep_check(1);
        let mut to_unlock = Some(to_unlock);
        self.wait_event_interruptible(
            || false,
            Some(move || {
                if let Some(lock) = to_unlock.take() {
                    drop(lock);
                }
            }),
        )
    }

    pub fn sleep_uninterruptible_unlock_spinlock<T>(&self, to_unlock: SpinLockGuard<T>) {
        before_sleep_check(1);
        let mut to_unlock = Some(to_unlock);
        let _ = self.wait_event_uninterruptible(
            || false,
            Some(move || {
                if let Some(lock) = to_unlock.take() {
                    drop(lock);
                }
            }),
        );
    }

    pub fn sleep_uninterruptible_unlock_mutex<T>(&self, to_unlock: MutexGuard<T>) {
        before_sleep_check(1);
        let mut to_unlock = Some(to_unlock);
        let _ = self.wait_event_uninterruptible(
            || false,
            Some(move || {
                if let Some(lock) = to_unlock.take() {
                    drop(lock);
                }
            }),
        );
    }
}

impl InnerWaitQueue {
    pub const INIT: InnerWaitQueue = InnerWaitQueue {
        dead: false,
        waiters: BTreeMap::new(),
        empty_tag_count: 0,
    };

    fn cache_or_evict_empty_tag(&mut self, tag: usize) -> Option<VecDeque<Arc<Waker>>> {
        if tag == DEFAULT_WAIT_TAG {
            return None;
        }
        if self.empty_tag_count < EMPTY_TAG_CACHE_LIMIT {
            self.empty_tag_count += 1;
            None
        } else {
            self.waiters.remove(&tag)
        }
    }
}

impl Waiter {
    pub fn new_pair() -> (Self, Arc<Waker>) {
        Self::new_pair_tagged(DEFAULT_WAIT_TAG)
    }

    fn new_pair_tagged(tag: usize) -> (Self, Arc<Waker>) {
        let waker = Arc::new(Waker {
            state: AtomicU8::new(Waker::STATE_IDLE),
            tag,
            target: Arc::downgrade(&ProcessManager::current_pcb()),
        });
        let waiter = Waiter {
            waker: waker.clone(),
            _nosend: PhantomData,
        };
        (waiter, waker)
    }

    pub fn wait(&self, interruptible: bool) -> Result<(), SystemError> {
        block_current(self, interruptible)
    }

    pub fn wait_io(&self, interruptible: bool) -> Result<(), SystemError> {
        block_current_io(self, interruptible)
    }

    pub fn wait_killable(&self) -> Result<(), SystemError> {
        block_current_killable(self)
    }

    pub fn waker(&self) -> Arc<Waker> {
        self.waker.clone()
    }
}

impl Drop for Waiter {
    fn drop(&mut self) {
        self.waker.close();
    }
}

impl Waker {
    const STATE_IDLE: u8 = 0;
    const STATE_SLEEPING: u8 = 1;
    const STATE_NOTIFIED: u8 = 2;
    const STATE_CLOSED: u8 = 3;

    pub fn wake(&self) -> bool {
        loop {
            let state = self.state.load(Ordering::Acquire);
            match state {
                Self::STATE_CLOSED | Self::STATE_NOTIFIED => return false,
                Self::STATE_IDLE => {
                    if self
                        .state
                        .compare_exchange(
                            Self::STATE_IDLE,
                            Self::STATE_NOTIFIED,
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        if let Some(pcb) = self.target.upgrade() {
                            let _ = ProcessManager::wakeup(&pcb);
                        }
                        return true;
                    }
                }
                Self::STATE_SLEEPING => {
                    if self
                        .state
                        .compare_exchange(
                            Self::STATE_SLEEPING,
                            Self::STATE_NOTIFIED,
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        if let Some(pcb) = self.target.upgrade() {
                            let _ = ProcessManager::wakeup(&pcb);
                        }
                        return true;
                    }
                }
                _ => core::hint::spin_loop(),
            }
        }
    }

    #[inline]
    pub fn close(&self) {
        self.state.store(Self::STATE_CLOSED, Ordering::Release);
    }

    #[inline]
    pub fn is_closed(&self) -> bool {
        self.state.load(Ordering::Acquire) == Self::STATE_CLOSED
    }

    fn consume_notification(&self) -> bool {
        loop {
            let state = self.state.load(Ordering::Acquire);
            match state {
                Self::STATE_NOTIFIED => {
                    if self
                        .state
                        .compare_exchange(
                            Self::STATE_NOTIFIED,
                            Self::STATE_IDLE,
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        return true;
                    }
                }
                Self::STATE_CLOSED => return true,
                _ => return false,
            }
        }
    }

    fn prepare_sleep(&self) -> WakerSleepState {
        loop {
            let state = self.state.load(Ordering::Acquire);
            match state {
                Self::STATE_IDLE => {
                    if self
                        .state
                        .compare_exchange(
                            Self::STATE_IDLE,
                            Self::STATE_SLEEPING,
                            Ordering::AcqRel,
                            Ordering::Acquire,
                        )
                        .is_ok()
                    {
                        return WakerSleepState::Sleeping;
                    }
                }
                Self::STATE_SLEEPING => return WakerSleepState::Sleeping,
                Self::STATE_NOTIFIED => return WakerSleepState::Notified,
                Self::STATE_CLOSED => return WakerSleepState::Closed,
                _ => core::hint::spin_loop(),
            }
        }
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
enum WakerSleepState {
    Sleeping,
    Notified,
    Closed,
}

fn before_sleep_check(max_preempt: usize) {
    let pcb = ProcessManager::current_pcb();
    if unlikely(pcb.preempt_count() > max_preempt) {
        warn!(
            "Process {:?}: Try to sleep when preempt count is {}",
            pcb.raw_pid().data(),
            pcb.preempt_count()
        );
    }
}

/// 统一封装"标记阻塞 + 调度 + 信号检查"，避免各调用点重复逻辑
fn block_current(waiter: &Waiter, interruptible: bool) -> Result<(), SystemError> {
    block_current_impl(waiter, interruptible, false)
}

fn block_current_io(waiter: &Waiter, interruptible: bool) -> Result<(), SystemError> {
    block_current_impl(waiter, interruptible, true)
}

fn block_current_impl(
    waiter: &Waiter,
    interruptible: bool,
    is_io: bool,
) -> Result<(), SystemError> {
    loop {
        // 快路径：被提前唤醒
        if waiter.waker.consume_notification() {
            return Ok(());
        }

        let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };

        // Build a waiter/waker handshake to avoid lost wakeups.
        match waiter.waker.prepare_sleep() {
            WakerSleepState::Notified | WakerSleepState::Closed => {
                drop(irq_guard);
                return Ok(());
            }
            WakerSleepState::Sleeping => {}
        }

        ProcessManager::mark_sleep(interruptible)?;

        // Handle a wake racing between prepare_sleep and mark_sleep.
        if waiter.waker.consume_notification() {
            ProcessManager::undo_mark_sleep();
            drop(irq_guard);
            return Ok(());
        }
        drop(irq_guard);

        // 根据是否是 IO 等待选择不同的调度函数
        if is_io {
            io_schedule();
        } else {
            schedule(SchedMode::SM_NONE);
        }

        if interruptible
            && Signal::signal_pending_state(true, false, &ProcessManager::current_pcb())
        {
            return Err(SystemError::ERESTARTSYS);
        }
    }
}

fn block_current_killable(waiter: &Waiter) -> Result<(), SystemError> {
    loop {
        if waiter.waker.consume_notification() {
            return Ok(());
        }

        let irq_guard = unsafe { CurrentIrqArch::save_and_disable_irq() };

        match waiter.waker.prepare_sleep() {
            WakerSleepState::Notified | WakerSleepState::Closed => {
                drop(irq_guard);
                return Ok(());
            }
            WakerSleepState::Sleeping => {}
        }

        ProcessManager::mark_sleep(true)?;

        if waiter.waker.consume_notification() {
            ProcessManager::undo_mark_sleep();
            drop(irq_guard);
            return Ok(());
        }
        drop(irq_guard);

        schedule(SchedMode::SM_NONE);

        if Signal::signal_pending_state(false, true, &ProcessManager::current_pcb()) {
            return Err(SystemError::ERESTARTSYS);
        }
    }
}

/// 事件等待队列：按事件掩码唤醒
#[derive(Debug)]
pub struct EventWaitQueue {
    wait_list: SpinLock<Vec<(u64, Arc<Waker>)>>,
}

impl Default for EventWaitQueue {
    fn default() -> Self {
        Self::new()
    }
}

#[allow(dead_code)]
impl EventWaitQueue {
    pub fn new() -> Self {
        Self {
            wait_list: SpinLock::new(Default::default()),
        }
    }

    fn register_waker(&self, events: u64, waker: Arc<Waker>) {
        let mut entry = Some((events, waker));
        let mut spare: Vec<(u64, Arc<Waker>)> = Vec::new();

        loop {
            let needed_capacity;
            {
                let mut guard = self.wait_list.lock_irqsave();
                if guard.len() < guard.capacity() {
                    guard.push(entry.take().unwrap());
                    return;
                }

                needed_capacity = (guard.len() + 1).next_power_of_two().max(4);
                if spare.capacity() >= needed_capacity {
                    spare.extend(guard.drain(..));
                    spare = mem::replace(&mut *guard, spare);
                    continue;
                }
            }

            spare = Vec::with_capacity(needed_capacity);
        }
    }

    fn remove_waker(&self, target: &Arc<Waker>) {
        let mut guard = self.wait_list.lock_irqsave();
        guard.retain(|(_, waker)| !Arc::ptr_eq(waker, target));
    }

    pub fn wait_event_interruptible<F>(&self, events: u64, mut cond: F) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        before_sleep_check(0);
        if cond() {
            return Ok(());
        }

        let (waiter, waker) = Waiter::new_pair();
        loop {
            self.register_waker(events, waker.clone());

            if cond() {
                self.remove_waker(&waker);
                return Ok(());
            }

            if Signal::signal_pending_state(true, false, &ProcessManager::current_pcb()) {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(SystemError::ERESTARTSYS);
            }

            if let Err(err) = waiter.wait(true) {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(err);
            }
        }
    }

    pub fn wait_event_interruptible_timeout<F>(
        &self,
        events: u64,
        mut cond: F,
        timeout: Duration,
    ) -> Result<(), SystemError>
    where
        F: FnMut() -> bool,
    {
        before_sleep_check(0);
        if cond() {
            return Ok(());
        }

        let deadline = Instant::now() + timeout;
        let (waiter, waker) = Waiter::new_pair();
        loop {
            if Instant::now() >= deadline {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(SystemError::EAGAIN_OR_EWOULDBLOCK);
            }

            self.register_waker(events, waker.clone());

            if cond() {
                self.remove_waker(&waker);
                return Ok(());
            }

            if Signal::signal_pending_state(true, false, &ProcessManager::current_pcb()) {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(SystemError::ERESTARTSYS);
            }

            let remain = deadline
                .duration_since(Instant::now())
                .unwrap_or(Duration::ZERO);
            if remain == Duration::ZERO {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(SystemError::EAGAIN_OR_EWOULDBLOCK);
            }

            let timer: Arc<Timer> = Timer::new(
                TimeoutWaker::new(waker.clone()),
                next_n_us_timer_jiffies(remain.total_micros()),
            );
            timer.activate();

            let wait_result = waiter.wait(true);
            let was_timeout = timer.timeout();
            if !was_timeout {
                timer.cancel();
            }

            if was_timeout {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(SystemError::EAGAIN_OR_EWOULDBLOCK);
            }

            if let Err(err) = wait_result {
                self.remove_waker(&waker);
                waker.close();
                if cond() {
                    return Ok(());
                }
                return Err(err);
            }
        }
    }

    pub fn sleep(&self, events: u64) {
        before_sleep_check(0);
        let (waiter, waker) = Waiter::new_pair();
        self.register_waker(events, waker);
        let _ = waiter.wait(true);
    }

    pub fn sleep_unlock_spinlock<T>(&self, events: u64, to_unlock: SpinLockGuard<T>) {
        before_sleep_check(1);
        let (waiter, waker) = Waiter::new_pair();
        self.register_waker(events, waker);
        drop(to_unlock);
        let _ = waiter.wait(true);
    }

    pub fn wakeup_any(&self, events: u64) -> usize {
        let mut ret = 0;
        let mut guard = self.wait_list.lock_irqsave();
        guard.retain(|(es, waker)| {
            if *es & events > 0 {
                if waker.wake() {
                    ret += 1;
                }
                false
            } else {
                true
            }
        });
        ret
    }

    pub fn wakeup(&self, events: u64) -> usize {
        let mut ret = 0;
        let mut guard = self.wait_list.lock_irqsave();
        guard.retain(|(es, waker)| {
            if *es == events {
                if waker.wake() {
                    ret += 1;
                }
                false
            } else {
                true
            }
        });
        ret
    }

    pub fn wakeup_all(&self) {
        self.wakeup_any(u64::MAX);
    }
}

/// 通用的超时唤醒辅助结构
///
/// 用于定时器超时时唤醒等待队列中的 Waiter。
/// 相比直接唤醒 PCB，通过 Waker 唤醒可以：
/// 1. 与 Waiter::wait() 正确配合，避免竞态条件
/// 2. 使用原子状态机标记唤醒/睡眠状态
/// 3. 保持与等待队列机制的一致性
#[derive(Debug)]
pub struct TimeoutWaker {
    waker: Arc<Waker>,
}

impl TimeoutWaker {
    pub fn new(waker: Arc<Waker>) -> Box<Self> {
        Box::new(Self { waker })
    }
}

impl crate::time::timer::TimerFunction for TimeoutWaker {
    fn run(&mut self) -> Result<(), SystemError> {
        // 通过 Waker::wake() 唤醒，这样 Waiter::wait() 可以观察到
        // 注意：定时器唤醒必须通过 Waker::wake()，仅唤醒 PCB 是不够的
        self.waker.wake();
        Ok(())
    }
}
