hirun 0.1.13

use crate::event::{Event, Scheduler, POLLIN, POLLONESHOT, POLLOUT};
use crate::runtime::{RawTaskContext, TaskContext, TaskRef};
use crate::{Error, Result};
use core::future::Future;
use core::marker::PhantomData;
use core::mem::ManuallyDrop;
use core::pin::Pin;
use core::task::{Context, Poll};
use hioff::container_of_mut;

pub(crate) struct WaitEvent<'a> {
    wait: bool,
    inner: WaitEventInner<'a>,
}

union WaitEventInner<'a> {
    once: ManuallyDrop<FdWaitOnce>,
    wait: ManuallyDrop<FdWait<'a>>,
}

impl Unpin for WaitEvent<'_> {}

unsafe impl Send for WaitEvent<'_> {}

impl WaitEvent<'_> {
    pub fn new(fd: i32, events: u32) -> Self {
        debug_assert!((events & (POLLIN | POLLOUT)) > 0);
        Self {
            wait: false,
            inner: WaitEventInner {
                once: ManuallyDrop::new(FdWaitOnce::new(fd, events)),
            },
        }
    }
}

impl Drop for WaitEvent<'_> {
    fn drop(&mut self) {
        if self.wait {
            unsafe { ManuallyDrop::drop(&mut self.inner.wait) };
        } else {
            unsafe { ManuallyDrop::drop(&mut self.inner.once) };
        }
    }
}

impl<'a> Future for WaitEvent<'a> {
    type Output = Result<u32>;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let private = ctx.get_private();
        if !private.is_null() {
            // 当前只有FdSetWait会调用ctx.set_private
            // 后续需要扩展此接口，允许设置多个key-value,避免多个场景下使用此机制带来的冲突
            if !self.wait {
                self.wait = true;
                let wait_ctx = unsafe { &mut *private.cast_mut().cast::<WaitContext>() };
                let once = unsafe { ManuallyDrop::take(&mut self.inner.once) };
                self.inner.wait = ManuallyDrop::new(FdWait::new(wait_ctx, once.fd, once.events));
            }
            let pinned = unsafe { Pin::new_unchecked(&mut *self.inner.wait) };
            Future::poll(pinned, ctx)
        } else {
            unsafe { &mut *self.inner.once }.poll(ctx)
        }
    }
}

// 只等待一次事件的发生, 接收到事件后就取消注册
struct FdWaitOnce {
    fd: i32,
    events: u32,
    recv: u32,
    event: Event,
    task: Option<TaskRef>,
}

impl FdWaitOnce {
    pub fn new(fd: i32, events: u32) -> Self {
        debug_assert!((events & (POLLIN | POLLOUT)) > 0);
        Self {
            fd,
            events,
            recv: 0,
            event: Event::new(Self::event_handle),
            task: None,
        }
    }
}

impl FdWaitOnce {
    fn add_event(&mut self, sched: &mut Scheduler) -> Poll<Result<u32>> {
        match unsafe { sched.add_fd_event(&self.event, self.events, self.fd) } {
            Ok(_) => Poll::Pending,
            Err(e) => Poll::Ready(Err(e)),
        }
    }

    fn event_handle(e: &Event, events: u32, sched: &mut Scheduler) {
        let this = unsafe { container_of_mut!(e, Self, event) };
        let _ = unsafe { sched.del_fd_event(&this.event, this.fd) };
        this.recv = events & this.events;
        let mut task = this.task.take().unwrap();
        task.status.unfreeze_local();
        task.fast_wake(sched);
    }

    fn abort(&mut self, ctx: &mut Context<'_>) {
        if let Some(_) = self.task.take() {
            let _ = unsafe { ctx.sched().del_fd_event(&self.event, self.fd) };
            ctx.unfreeze_local();
        }
    }

    fn poll(&mut self, ctx: &mut Context<'_>) -> Poll<Result<u32>> {
        if ctx.aborted() {
            self.abort(ctx);
            return Poll::Pending;
        }

        if self.recv > 0 {
            return Poll::Ready(Ok(self.recv));
        }

        if self.task.is_none() {
            self.task = Some(ctx.task_ref());
            ctx.freeze_local();
            self.add_event(ctx.sched())
        } else {
            Poll::Pending
        }
    }
}

struct WaitVTable {
    fd_wait: fn(this: *const (), fd: i32, events: u32, ctx: &mut Context<'_>) -> Result<usize>,
    fd_awaked: fn(this: *const (), index: usize, events: u32) -> Option<u32>,
    fd_abort: fn(this: *const (), index: usize, events: u32),
    fd_del: fn(this: *const (), fd: i32, ctx: &mut Context<'_>),
}

/// 在FdSet上下文中使用，提供单个Fd的异步Io事件监测功能，可结合Seq/Or/And监控多个fd的事件
pub(crate) struct WaitContext {
    vtbl: &'static WaitVTable,
    mark: PhantomData<*const WaitVTable>,
}

impl WaitContext {
    fn new(vtbl: &'static WaitVTable) -> Self {
        Self {
            vtbl,
            mark: PhantomData,
        }
    }

    fn fd_wait(&self, fd: i32, events: u32, ctx: &mut Context<'_>) -> Result<usize> {
        (self.vtbl.fd_wait)(self as *const _ as *const (), fd, events, ctx)
    }
    fn fd_awaked(&self, index: usize, events: u32) -> Option<u32> {
        (self.vtbl.fd_awaked)(self as *const _ as *const (), index, events)
    }
    fn fd_abort(&self, index: usize, events: u32) {
        (self.vtbl.fd_abort)(self as *const _ as *const (), index, events)
    }
    fn fd_del(&self, fd: i32, ctx: &mut Context<'_>) {
        (self.vtbl.fd_del)(self as *const _ as *const (), fd, ctx)
    }
}

#[allow(dead_code)]
impl WaitContext {
    /// 等待指定事件的发生，如果指定多个事件，那么任一事件发生则返回，具体事件由返回值指定.
    /// WaitContext只在FdSet的上下文中才出现，当fd错误或者超过FdSet的容量时将返回错误.
    pub fn wait(&self, fd: i32, events: u32) -> impl Future<Output = Result<u32>> + '_ {
        FdWait::new(self, fd, events)
    }

    /// 取消注册,
    /// FdSet有固定容量限制，如果任务内部的Fd不断在更新，此接口可删除不再需要的句柄.
    pub fn del(&self, fd: i32) -> impl Future<Output = ()> + '_ {
        FdDel::new(self, fd)
    }
}

struct FdState {
    event: Event,
    fd: i32,
    // 等待的事件
    events: u32,
    // 接收的事件
    recv: u32,
    // 在FdSet中的索引值
    index: u32,
    // 等待者计数, 0x00010001: 读写, 0x00010000: 只写, 0x00000001: 只读
    wait: u32,
}

impl FdState {
    fn set_fd(&mut self, fd: i32) {
        self.fd = fd;
        self.recv = 0;
        self.events = 0;
        self.wait = 0;
    }
}

/// 只支持Send, 不支持Sync, 在单线程下使用
// 利用FdSet可以方便的在一个异步任务中高效的处理多个Fd的异步IO事件
// 典型的应用场景为Tcp/Http Proxy, 在一个异步任务中同时在两个Fd上进行数据收发
// 用户调用WaitContext::wait(fd, events).await等待事件，结合seq/or/and等待多个事件
// N: 为最大的Fd数量，如果超过这个数量，WaitContext::wait操作将会返回失败.
pub(crate) struct FdSet<const N: usize> {
    ctx: WaitContext,
    // fdset未排序，顺序查找，大部分情况下就个位数fd，顺序查找满足要求
    fdset: [i32; N],
    stats: [FdState; N],
    idles: [u32; N],
    count: usize,
    waker: Option<TaskRef>,
}

unsafe impl<const N: usize> Send for FdSet<N> {}

impl<const N: usize> FdSet<N> {
    pub fn new() -> Self {
        Self {
            ctx: WaitContext::new(&Self::VTBL),
            stats: core::array::from_fn(|index| Self::new_fdstate(index)),
            idles: core::array::from_fn(|index| index as u32),
            fdset: [-1; N],
            count: 0,
            waker: None,
        }
    }

    pub fn wait<T: Future>(self, future: T) -> impl Future<Output = T::Output> {
        FdSetWait::new(self, future)
    }
}

impl<const N: usize> FdSet<N> {
    const VTBL: WaitVTable = WaitVTable {
        fd_wait: Self::vtbl_fd_wait,
        fd_awaked: Self::vtbl_fd_awaked,
        fd_abort: Self::vtbl_fd_abort,
        fd_del: Self::vtbl_fd_del,
    };
    fn vtbl_fd_wait(
        this: *const (),
        fd: i32,
        events: u32,
        ctx: &mut Context<'_>,
    ) -> Result<usize> {
        let this = unsafe { container_of_mut!(&*this.cast::<WaitContext>(), Self, ctx) };
        this.fd_wait(fd, events, ctx)
    }
    fn vtbl_fd_awaked(this: *const (), index: usize, events: u32) -> Option<u32> {
        let this = unsafe { container_of_mut!(&*this.cast::<WaitContext>(), Self, ctx) };
        this.fd_awaked(index, events)
    }
    fn vtbl_fd_abort(this: *const (), index: usize, events: u32) {
        let this = unsafe { container_of_mut!(&*this.cast::<WaitContext>(), Self, ctx) };
        this.fd_abort(index, events)
    }
    fn vtbl_fd_del(this: *const (), fd: i32, ctx: &mut Context<'_>) {
        let this = unsafe { container_of_mut!(&*this.cast::<WaitContext>(), Self, ctx) };
        this.fd_del(fd, ctx)
    }
}

impl<const N: usize> FdSet<N> {
    const RO_COUNTER: u32 = 0x0000_0001;
    const WO_COUNTER: u32 = 0x0001_0000;
    const WR_COUNTER: u32 = 0x0001_0001;
    const POLLMASK: u32 = POLLIN | POLLOUT;
    fn new_fdstate(index: usize) -> FdState {
        FdState {
            fd: -1,
            events: 0,
            recv: 0,
            event: Event::new(Self::event_handle),
            index: index as u32,
            wait: 0,
        }
    }
    fn event_handle(e: &Event, events: u32, sched: &mut Scheduler) {
        let stat = unsafe { container_of_mut!(e, FdState, event) };
        if stat.wait > 0 {
            stat.recv |= events & stat.events;
            let this = unsafe { container_of_mut!(stat, Self, stats[stat.index as usize]) };
            let task = this.waker.as_ref().unwrap().task_ref();
            task.fast_wake(sched);
        }
    }

    fn abort(&mut self, ctx: &mut Context<'_>) {
        if let Some(_) = self.waker.take() {
            for (index, fd) in self.fdset.iter().enumerate() {
                if *fd > -1 {
                    let stat = &self.stats[index];
                    if stat.events > 0 {
                        let _ = unsafe { ctx.sched().del_fd_event(&stat.event, *fd) };
                    }
                }
            }
            ctx.unfreeze_local();
        }
    }

    fn fd_index(&self, fd: i32) -> Option<usize> {
        for (fd_index, fd_fd) in self.fdset.iter().enumerate() {
            if *fd_fd == fd {
                return Some(fd_index);
            }
        }
        None
    }

    fn fd_del(&mut self, fd: i32, ctx: &mut Context<'_>) {
        let Some(index) = self.fd_index(fd) else {
            return;
        };
        let stat = &mut self.stats[index];
        let _ = unsafe { ctx.sched().del_fd_event(&stat.event, fd) };
        self.fdset[index] = -1;
        self.count -= 1;
        self.idles[self.count as usize] = index as u32;
    }

    fn fd_add(&mut self, fd: i32) -> Result<usize> {
        if let Some(index) = self.fd_index(fd) {
            return Ok(index);
        }
        if self.count < N {
            let index = self.idles[self.count as usize] as usize;
            self.count += 1;
            self.fdset[index] = fd;
            self.stats[index].set_fd(fd);
            Ok(index)
        } else {
            Err(Error::new(hierr::ERANGE))
        }
    }

    // TODO: 如果没有剩余空间考虑将最早无人等待的fd自动释放
    fn fd_wait(&mut self, fd: i32, events: u32, ctx: &mut Context<'_>) -> Result<usize> {
        debug_assert!((events & (POLLIN | POLLOUT)) > 0);
        let index = self.fd_add(fd)?;
        let stat = &mut self.stats[index];
        let new_events = stat.events | events;
        // 如果事件相同不需要重复注册，除非指定了ONESHOT模式
        if new_events != stat.events || (new_events & POLLONESHOT) > 0 {
            if stat.events == 0 {
                unsafe { ctx.sched().add_fd_event(&stat.event, new_events, stat.fd)? };
            } else {
                unsafe { ctx.sched().mod_fd_event(&stat.event, new_events, stat.fd)? };
            }
            stat.events = new_events;
        }
        match events & Self::POLLMASK {
            POLLIN => stat.wait += Self::RO_COUNTER,
            POLLOUT => stat.wait += Self::WO_COUNTER,
            Self::POLLMASK => stat.wait += Self::WR_COUNTER,
            _ => {},
        }
        // 此时强制清空stat.recv的对应事件, 因为业务层已经对应事件完全消耗完毕，
        // 如果不清空，会存在虚假唤醒情况
        stat.recv &= !events;
        if self.waker.is_none() {
            self.waker = Some(ctx.task_ref());
            ctx.freeze_local();
        }
        Ok(index)
    }

    fn fd_awaked(&mut self, index: usize, events: u32) -> Option<u32> {
        let stat = &mut self.stats[index];
        let events = stat.recv & events;
        match events & Self::POLLMASK {
            POLLIN => stat.wait -= Self::RO_COUNTER,
            POLLOUT => stat.wait -= Self::WO_COUNTER,
            Self::POLLMASK => stat.wait -= Self::WR_COUNTER,
            _ => return None,
        }
        Some(events)
    }
    fn fd_abort(&mut self, index: usize, events: u32) {
        let stat = &mut self.stats[index];
        match events & Self::POLLMASK {
            POLLIN => stat.wait -= Self::RO_COUNTER,
            POLLOUT => stat.wait -= Self::WO_COUNTER,
            Self::POLLMASK => stat.wait -= Self::WR_COUNTER,
            _ => {}
        }
    }
}

struct FdSetWait<const N: usize, T: Future> {
    fdset: FdSet<N>,
    future: T,
}

impl<const N: usize, T: Future> Unpin for FdSetWait<N, T> {}

impl<const N: usize, T: Future> FdSetWait<N, T> {
    fn new(fdset: FdSet<N>, future: T) -> Self {
        Self { fdset, future }
    }

    fn abort(&mut self, ctx: &mut Context<'_>) {
        let pinned = unsafe { Pin::new_unchecked(&mut self.future) };
        let _ = Future::poll(pinned, ctx);
        self.fdset.abort(ctx);
    }
}

impl<const N: usize, T: Future> Future for FdSetWait<N, T> {
    type Output = T::Output;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let old_private = ctx.get_private();
        // 这是为了兼容已有的AioFd::wait，从ctx.get_private获取到WaitContext
        ctx.set_private(&self.fdset.ctx as *const _ as *const ());
        if ctx.aborted() {
            self.abort(ctx);
            ctx.set_private(old_private);
            return Poll::Pending;
        }
        let pinned = unsafe { Pin::new_unchecked(&mut self.future) };
        let ret = Future::poll(pinned, ctx);
        ctx.set_private(old_private);
        match ret {
            Poll::Pending => Poll::Pending,
            Poll::Ready(val) => {
                // 异步任务正常结束的时候必须取消到所有注册的IO事件
                self.fdset.abort(ctx);
                Poll::Ready(val)
            }
        }
    }
}

struct FdWait<'a> {
    fd: i32,
    events: u32,
    index: usize,
    ctx: &'a WaitContext,
}

impl<'a> FdWait<'a> {
    fn new(ctx: &'a WaitContext, fd: i32, events: u32) -> Self {
        Self {
            ctx,
            fd,
            events,
            index: usize::MAX,
        }
    }
}

impl Future for FdWait<'_> {
    type Output = Result<u32>;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        if ctx.aborted() {
            self.ctx.fd_abort(self.index, self.events);
            return Poll::Pending;
        }
        if self.index == usize::MAX {
            match self.ctx.fd_wait(self.fd, self.events, ctx) {
                Ok(index) => self.index = index,
                Err(e) => return Poll::Ready(Err(e)),
            };
        }

        match self.ctx.fd_awaked(self.index, self.events) {
            Some(events) => Poll::Ready(Ok(events)),
            None => Poll::Pending,
        }
    }
}

struct FdDel<'a> {
    fd: i32,
    ctx: &'a WaitContext,
}

#[allow(dead_code)]
impl<'a> FdDel<'a> {
    fn new(ctx: &'a WaitContext, fd: i32) -> Self {
        Self { ctx, fd }
    }
}

impl Future for FdDel<'_> {
    type Output = ();
    fn poll(self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        let fd = self.fd;
        self.ctx.fd_del(fd, ctx);
        Poll::Ready(())
    }
}