hirun 0.1.1

use super::{const_buf, mut_buf, Fd};
use crate::runtime::{sched_from_waker, task_from_waker, TaskRef};
use crate::{
    event::{Event, Scheduler, POLLET, POLLIN, POLLOUT},
    Error,
};
use core::future::Future;
use core::ops::Deref;
use core::pin::Pin;
use core::ptr::NonNull;
use core::task::{Context, Poll};
use hioff::container_of_mut;

#[repr(C)]
pub struct AioFd<'a> {
    pub(crate) fd: &'a Fd,
    event: Event,
    events: u32,
    task: Option<TaskRef>,
    sched: Option<NonNull<Scheduler>>,
}

// Fd支持Sync，当前实现的Scheduler支持Sync，可以跨线程del_fd_event
unsafe impl Send for AioFd<'_> {}
unsafe impl Sync for AioFd<'_> {}

impl<'a> AioFd<'a> {
    pub fn new(fd: &'a Fd) -> Self {
        Self {
            fd,
            event: Event::new(Self::event_handle),
            task: None,
            sched: None,
            events: 0,
        }
    }
}

impl AioFd<'_> {
    pub async fn wait(&mut self, events: u32) -> Result<(), Error> {
        FdWait::new(self, events).await
    }

    pub async fn read_all(&mut self, mut buf: &mut [u8]) -> Result<usize, Error> {
        let mut recved = 0;
        loop {
            let ret = unsafe { libc::read(self.fd.fd, mut_buf(buf), buf.len()) };
            // 频繁调用场景，但rust无法指定分支预测功能
            #[allow(clippy::comparison_chain)]
            if ret > 0 {
                let n = ret as usize;
                recved += n;
                if n == buf.len() {
                    return Ok(recved);
                }
                buf = &mut buf[n..];
            } else if ret == 0 {
                return Ok(recved);
            } else {
                let e = Error::last_error();
                if e.errno == libc::EAGAIN {
                    self.wait(POLLIN).await?;
                } else if e.errno != libc::EINTR {
                    return Err(e);
                }
            }
        }
    }

    pub async fn write_all(&mut self, mut buf: &[u8]) -> Result<usize, Error> {
        let mut sended = 0;
        loop {
            let ret = unsafe { libc::write(self.fd.fd, const_buf(buf), buf.len()) };
            if ret >= 0 {
                let n = ret as usize;
                sended += n;
                if n == buf.len() {
                    return Ok(sended);
                }
                buf = &buf[n..];
            } else {
                let e = Error::last_error();
                if e.errno == libc::EAGAIN {
                    self.wait(POLLOUT).await?;
                } else if e.errno != libc::EINTR {
                    return Err(e);
                }
            }
        }
    }

    pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize, Error> {
        loop {
            let ret = unsafe { libc::read(self.fd.fd, mut_buf(buf), buf.len()) };
            if ret >= 0 {
                return Ok(ret as usize);
            } else {
                let e = Error::last_error();
                if e.errno == libc::EAGAIN {
                    self.wait(POLLIN).await?;
                } else if e.errno != libc::EINTR {
                    return Err(e);
                }
            }
        }
    }

    pub async fn write(&mut self, buf: &[u8]) -> Result<usize, Error> {
        loop {
            let ret = unsafe { libc::write(self.fd.fd, const_buf(buf), buf.len()) };
            if ret >= 0 {
                return Ok(ret as usize);
            } else {
                let e = Error::last_error();
                if e.errno == libc::EAGAIN {
                    self.wait(POLLOUT).await?;
                } else if e.errno != libc::EINTR {
                    return Err(e);
                }
            }
        }
    }
}

impl AioFd<'_> {
    fn event_handle(e: &Event, _events: u32, sched: &mut Scheduler) {
        let this = unsafe { container_of_mut!(e, Self, event) };
        if let Some(task) = this.task.take() {
            task.fast_wake(sched);
        }
    }

    fn set_task(&mut self, ctx: &mut Context<'_>) {
        let task = TaskRef::new(unsafe { task_from_waker(ctx.waker()) });
        task.inc_ref();
        debug_assert!(self.task.is_none());
        self.task = Some(task);
    }

    fn add(&mut self, events: u32, ctx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        if self.events == events {
            self.set_task(ctx);
            return Poll::Pending;
        }
        let ret = if let Some(mut sched) = self.sched {
            unsafe { sched.as_mut() }.mod_fd_event(&self.event, events | POLLET, self.fd.fd)
        } else {
            let sched = unsafe { sched_from_waker(ctx.waker()).as_mut() };
            let events = events | POLLET;
            let ret = sched.add_fd_event(&self.event, events, self.fd.fd);
            if ret.is_ok() {
                self.sched = Some(NonNull::from(sched));
            }
            ret
        };
        match ret {
            Ok(_) => {
                self.events = events;
                self.set_task(ctx);
                Poll::Pending
            }
            Err(e) => Poll::Ready(Err(e)),
        }
    }
}

impl Drop for AioFd<'_> {
    fn drop(&mut self) {
        if let Some(sched) = self.sched {
            let _ = unsafe { sched.as_ref() }.del_fd_event(self.fd.fd);
        }
    }
}

impl Deref for AioFd<'_> {
    type Target = Fd;
    fn deref(&self) -> &Self::Target {
        self.fd
    }
}

struct FdWait<'a, 'b> {
    aio: &'a mut AioFd<'b>,
    events: u32,
}

impl<'a, 'b> FdWait<'a, 'b> {
    fn new(aio: &'a mut AioFd<'b>, events: u32) -> Self {
        Self { aio, events }
    }
}

impl Future for FdWait<'_, '_> {
    type Output = Result<(), Error>;
    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
        if self.events > 0 {
            let events = self.events;
            self.events = 0;
            self.aio.add(events, ctx)
        } else {
            Poll::Ready(Ok(()))
        }
    }
}