use core::{
    future::poll_fn,
    task::{Context, Poll},
    time::Duration,
};

use alloc::collections::BTreeMap;

use futures_util::{future::BoxFuture, FutureExt};
use libk_util::sync::LockMethod;
use yggdrasil_abi::{error::Error, io::RawFd};

use crate::{
    task::{runtime, sync::Mutex},
    vfs::{FileReadiness, FileRef},
};

/// Poll channel implementation. Allows blocking until a file descriptor signals an event or a
/// timeout is reached.
pub struct FdPoll {
    fds: Mutex<BTreeMap<RawFd, FileRef>>,
}

impl FdPoll {
    /// Creates a new poll channel
    pub fn new() -> Self {
        Self {
            fds: Mutex::new(BTreeMap::new()),
        }
    }

    /// Adds a file descriptor of interest to the poll channel
    pub fn add(&self, fd: RawFd, file: FileRef) {
        self.fds.lock().unwrap().insert(fd, file);
    }

    pub fn remove(&self, fd: RawFd) {
        self.fds.lock().unwrap().remove(&fd);
    }

    /// Polls the channel once, returning either a file descriptor or timeout
    pub async fn wait(
        &self,
        timeout: Option<Duration>,
        retain: bool,
    ) -> Option<(RawFd, Result<(), Error>)> {
        let mut timeout = timeout.map(|t| runtime::sleep(t).boxed());
        poll_fn(|cx| self.poll_once(cx, &mut timeout, retain)).await
    }

    fn poll_once(
        &self,
        cx: &mut Context<'_>,
        timeout: &mut Option<BoxFuture<()>>,
        retain: bool,
    ) -> Poll<Option<(RawFd, Result<(), Error>)>> {
        if let Some(timeout) = timeout.as_mut()
            && timeout.as_mut().poll(cx).is_ready()
        {
            // Timeout
            return Poll::Ready(None);
        }

        let mut lock = self.fds.lock().unwrap();
        for (&fd, file) in lock.iter() {
            if let Poll::Ready(result) = file.poll_read(cx) {
                if !retain {
                    lock.remove(&fd);
                }
                return Poll::Ready(Some((fd, result)));
            }
        }

        // Wait
        Poll::Pending
    }
}

impl FileReadiness for FdPoll {
    fn poll_read(&self, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        self.poll_once(cx, &mut None, true).map(|_| Ok(()))
    }
}
