hirun 0.1.22

use super::{ActiveWorker, Attr, RawTask, TaskRef, WaitContext, Worker};
use crate::event::{Event, Scheduler, Timer};
use crate::Result;
use core::future::Future;
use core::ptr::NonNull;
use core::task::{Context, Waker};
use core::time::Duration;
use hioff::{container_of, container_of_mut};

/// 为Context<'_>提供增强接口
/// 注意：只对本crate的Waker生成的Context<'_>有效.
pub trait TaskContext {
    /// 设置当前task只在当前worker中调度，如果使用到了Scheduler，应该总是调用这个接口
    /// 异步函数中可以调用runtime::set_affinity.
    fn set_affinity(&mut self);
    /// 当前Future是否被提前终止，用于全局资源的清理，比如注册到sched()上的资源，应该确保在abort的时候取消注册
    /// 可能是JoinHandle::abort或者task内部调用Context::exit, Context::set_aborted所导致.
    /// 异步函数中可以调用runtime::task::aborted()获取.
    fn aborted(&mut self) -> bool;
    /// 修改abort标志，一般应用于Future的封装层，可能在封装层提前abort内部的Future.
    /// 比如Or/Deadline等就利用了这个接口
    /// 注意：不影响当前task的运行状态，如果需要退出当前task，则应调用exit.
    /// 异步函数中可以调用runtime::task::set_aborted
    fn set_aborted(&mut self, aborted: bool);
    /// 在task运行时主动退出，退出之后Future::poll会被再调度一次，再调度时aborted()返回true
    /// 异步函数中可以调用runtime::task::exit
    fn exit(&mut self);
    /// 主动执行task切换.
    /// 异步函数中可以调用runtime::yield_now.
    fn yield_now(&mut self);
}

#[allow(dead_code)]
pub(crate) trait RawTaskContext: TaskContext {
    fn worker_id(&mut self) -> u16;
    fn task(&mut self) -> &mut RawTask;
    fn task_ref(&mut self) -> TaskRef;
    // 后续可考虑放入TaskContext中.
    fn freeze_local(&mut self);
    fn unfreeze_local(&mut self);
    // 设置的值可以在同一个task内多个future之间共享.
    fn set_private(&mut self, data: *const ());
    fn get_private(&mut self) -> *const ();
    fn exited(&mut self) -> bool;
    fn cache(&mut self, task: TaskRef);
    fn spawn<T>(&mut self, future: T, attr: &Attr) -> Result<TaskRef>
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static;
    fn spawn_local<T>(&mut self, future: T, attr: &Attr) -> Result<TaskRef>
    where
        T: Future + 'static,
        T::Output: 'static;
}

pub(crate) struct TaskWaker<'a> {
    waker: Waker,
    task: NonNull<RawTask>,
    worker: &'a mut ActiveWorker,
    aborted: bool,
    exited: bool,
}

impl<'a> TaskWaker<'a> {
    pub(crate) fn new(task: &mut RawTask, worker: &'a mut ActiveWorker) -> Self {
        worker.set_current_task(Some(NonNull::from(&*task)));
        Self {
            waker: task.task_waker(),
            task: NonNull::from(task),
            worker,
            aborted: false,
            exited: false,
        }
    }

    fn from_ctx_mut(ctx: &mut Context<'a>) -> &'a mut Self {
        unsafe { container_of_mut!(ctx.waker(), Self, waker) }
    }

    fn from_ctx(ctx: &Context<'a>) -> &'a Self {
        unsafe { container_of!(ctx.waker(), Self, waker) }
    }

    pub(crate) fn waker(&self) -> &Waker {
        &self.waker
    }
}

impl Drop for TaskWaker<'_> {
    fn drop(&mut self) {
        self.worker.set_current_task(None);
    }
}

impl RawTaskContext for Context<'_> {
    fn task(&mut self) -> &mut RawTask {
        let this = TaskWaker::from_ctx_mut(self);
        unsafe { this.task.as_mut() }
    }

    fn task_ref(&mut self) -> TaskRef {
        self.task().task_ref()
    }

    fn unfreeze_local(&mut self) {
        self.task().status.unfreeze_local();
    }

    fn freeze_local(&mut self) {
        let id = self.worker_id();
        self.task().status.freeze_local(id);
    }

    fn set_private(&mut self, data: *const ()) {
        self.task().private = data;
    }

    fn get_private(&mut self) -> *const () {
        self.task().private
    }

    fn exited(&mut self) -> bool {
        let this = TaskWaker::from_ctx_mut(self);
        this.exited
    }

    fn worker_id(&mut self) -> u16 {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.worker_id()
    }

    fn cache(&mut self, task: TaskRef) {
        let this = TaskWaker::from_ctx_mut(self);
        task.cache(this.worker);
    }

    fn spawn<T>(&mut self, future: T, attr: &Attr) -> Result<TaskRef>
    where
        T: Future + Send + 'static,
        T::Output: Send + 'static,
    {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.spawn(future, attr)
    }

    fn spawn_local<T>(&mut self, future: T, attr: &Attr) -> Result<TaskRef>
    where
        T: Future + 'static,
        T::Output: 'static,
    {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.spawn_local(future, attr)
    }
}

impl TaskContext for Context<'_> {
    fn exit(&mut self) {
        let this = TaskWaker::from_ctx_mut(self);
        this.exited = true;
    }

    fn set_aborted(&mut self, aborted: bool) {
        let this = TaskWaker::from_ctx_mut(self);
        this.aborted = aborted;
    }

    fn aborted(&mut self) -> bool {
        let this = TaskWaker::from_ctx_mut(self);
        this.aborted
    }

    fn set_affinity(&mut self) {
        if self.task().status.get_local().is_none() {
            let id = self.worker_id();
            self.task().status.set_local(id);
        }
    }

    fn yield_now(&mut self) {
        self.task().status.yield_now();
    }
}

impl Scheduler for Context<'_> {
    unsafe fn add_event(&mut self, e: &Event, priority: i32) {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.add_event(e, priority)
    }

    unsafe fn del_event(&mut self, e: &Event) {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.del_event(e)
    }

    unsafe fn add_fd_event(&mut self, e: &Event, events: u32, fd: i32, flags: u16) -> Result<()> {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.add_fd_event(e, events, fd, flags)
    }

    unsafe fn mod_fd_event(&mut self, e: &Event, events: u32, fd: i32, flags: u16) -> Result<()> {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.mod_fd_event(e, events, fd, flags)
    }

    unsafe fn del_fd_event(&mut self, fd: i32) -> Result<()> {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.del_fd_event(fd)
    }

    unsafe fn set_timer(&mut self, t: &Timer, msecs: u32) {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.set_timer(t, msecs)
    }

    unsafe fn del_timer(&mut self, t: &Timer) {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.del_timer(t)
    }

    fn private_data(&self) -> *const () {
        let this = TaskWaker::from_ctx(self);
        this.worker.private_data()
    }

    fn set_private_data(&mut self, private: *const ()) {
        let this = TaskWaker::from_ctx_mut(self);
        this.worker.set_private_data(private);
    }

    fn stop(&self) {
        let this = TaskWaker::from_ctx(self);
        this.worker.stop()
    }

    fn now(&self) -> Duration {
        let this = TaskWaker::from_ctx(self);
        this.worker.now()
    }
}

impl WaitContext for Context<'_> {
    fn fd_wait(&mut self, fd: i32, index: usize, events: u32) -> Result<usize> {
        let this = TaskWaker::from_ctx_mut(self);
        let task = self.task();
        let old_size = task.fdset.fd_capacity();
        let index = task.fdset.fd_wait(
            fd,
            index,
            events,
            &task.event,
            this.worker,
        )?;
        if old_size == 0 {
            task.status.inc_ref();
            self.freeze_local();
        }
        Ok(index)
    }

    fn fd_awaked(&mut self, index: usize, events: u32) -> Option<u32> {
        let task = self.task();
        task.fdset.fd_awaked(index, events)
    }

    fn fd_abort(&mut self, index: usize) {
        let task = self.task();
        task.fdset.fd_abort(index)
    }

    fn fd_del(&mut self, fd: i32, index: usize) {
        let this = TaskWaker::from_ctx_mut(self);
        let task = self.task();
        task.fdset.fd_del(fd, index, this.worker);
    }
}