pub mod uring;

use std::future::Future;
use std::os::unix::io::FromRawFd;
use std::{io, net, ops, time};

use futures::executor;
use futures::task::LocalSpawn;

use nix::sys::socket;

pub fn run<C, F>(c: C) -> anyhow::Result<()>
where
    C: FnOnce(Context) -> F,
    F: Future<Output = ()> + 'static,
{
    let mut runner = Runner::new(4096)?;
    runner.spawn(c);
    runner.run()
}

pub struct Runner {
    uring: uring::Uring,
    pool: executor::LocalPool,
}

impl Runner {
    pub fn new(size: u32) -> anyhow::Result<Self> {
        let uring = uring::Uring::new(size)?;
        let pool = executor::LocalPool::new();

        Ok(Self { uring, pool })
    }

    pub fn spawn<C, F>(&mut self, c: C)
    where
        C: FnOnce(Context) -> F,
        F: Future<Output = ()> + 'static,
    {
        let context = Context::new(self.uring.clone(), self.pool.spawner());
        let f = c(context);

        self.pool
            .spawner()
            .spawn_local_obj(Box::pin(f).into())
            .expect("failed to spawn");
    }

    pub fn run(&mut self) -> anyhow::Result<()> {
        loop {
            self.pool.run_until_stalled();

            // Submit any pending tasks and wait for a result.
            let count = self.uring.submit_and_wait()?;
            if count == 0 {
                anyhow::bail!("no more pending IO"); // TODO return ok if the pool is empty
            }
        }
    }
}

#[derive(Clone)]
pub struct Context {
    uring: uring::Uring,
    spawner: executor::LocalSpawner,
    // TODO cancel support
    // TODO timeout support
    // TODO chain support
}

impl Context {
    pub fn new(uring: uring::Uring, spawner: executor::LocalSpawner) -> Self {
        Self { uring, spawner }
    }

    pub fn spawn<C, F>(&mut self, c: C)
    where
        C: FnOnce(Context) -> F,
        F: Future<Output = ()> + 'static,
    {
        let f = c(self.clone());

        self.spawner
            .spawn_local_obj(Box::pin(f).into())
            .expect("failed to spawn"); // TODO handle
    }

    /// Accepts a new connection on the given listener.
    pub async fn accept(
        &mut self,
        socket: &mut net::TcpListener,
    ) -> io::Result<net::TcpStream> {
        let submission = uring::Submission::Accept { socket };
        let task = self.uring.run(submission, uring::Flags::empty());

        match task.await? {
            uring::Completion::Accept { stream } => Ok(stream),
            _ => unreachable!(),
        }
    }

    /// Returns a buffer of at least the given size.
    pub fn buffer(&mut self, size: usize) -> Box<[u8]> {
        // TODO use fixed buffers when available
        vec![0; size].into_boxed_slice()
    }

    pub async fn connect<A>(&mut self, addr: A) -> anyhow::Result<net::TcpStream>
    where
        A: net::ToSocketAddrs,
    {
        let fd = socket::socket(
            socket::AddressFamily::Inet, // TODO support ipv6
            socket::SockType::Stream,
            socket::SockFlag::empty(),
            socket::SockProtocol::Tcp,
        )?;

        let socket = unsafe { net::TcpStream::from_raw_fd(fd) };

        let addr = addr.to_socket_addrs()?.next().unwrap();

        let submission = uring::Submission::Connect { socket: socket, addr: addr };
        let task = self.uring.run(submission, uring::Flags::empty());

        match task.await? {
            uring::Completion::Connect { stream } => Ok(stream),
            _ => unreachable!(),
        }
    }

    /// Wrapper around net::TcpListener::bind.
    pub fn listen<A>(&mut self, addr: A) -> io::Result<net::TcpListener>
    where
        A: net::ToSocketAddrs,
    {
        // TODO register the fd with io_uring
        net::TcpListener::bind(addr)
    }

    /// Reads at least one byte into the buffer range and returns the size. If the stream is
    /// closed, returns zero.
    pub async fn read(&mut self, stream: &mut net::TcpStream, buffer: &mut [u8]) -> io::Result<usize> {
        let submission = uring::Submission::Read {
            stream,
            buffer,
        };

        let task = self.uring.run(submission, uring::Flags::empty());

        match task.await? {
            uring::Completion::Read { size } => Ok(size),
            _ => unreachable!(),
        }
    }

    /// Reads until the entire buffer range is populated or the stream is closed.
    pub async fn read_full(
        &mut self,
        stream: &mut net::TcpStream,
        buffer: &mut [u8],
    ) -> io::Result<usize>
    {
        // TODO
        self.read(stream, buffer).await
    }

    pub fn read_close(&self, stream: &mut net::TcpStream) -> io::Result<()> {
        stream.shutdown(net::Shutdown::Read)
    }

    /// Writes at least one byte within the buffer range to the socket. Returns the amount written.
    pub async fn write(
        &mut self,
        stream: &mut net::TcpStream,
        buffer: &[u8],
    ) -> io::Result<usize>
    {
        let submission = uring::Submission::Write {
            stream,
            buffer,
        };

        let task = self.uring.run(submission, uring::Flags::empty());

        match task.await? {
            uring::Completion::Write { size } => Ok(size),
            _ => unreachable!(),
        }
    }

    /// Writes at least one byte within the buffer range to the socket before running the next task. Returns the amount written.
    pub async fn write_then(
        &mut self,
        stream: &mut net::TcpStream,
        buffer: &[u8],
    ) -> io::Result<usize>
    {

        let submission = uring::Submission::Write{stream, buffer};
        let task = self.uring.run(submission, uring::Flags::IO_LINK);

        match task.await? {
            uring::Completion::Write { size } => Ok(size),
            _ => unreachable!(),
        }
    }

    /// Writes the entire buffer range to the socket.
    pub async fn write_full(
        &mut self,
        stream: &mut net::TcpStream,
        buffer: &[u8],
    ) -> io::Result<()>
    {
        // TODO
        self.write(stream, buffer).await?;
        Ok(())
    }

    pub fn write_close(&self, stream: &mut net::TcpStream) -> io::Result<()> {
        stream.shutdown(net::Shutdown::Write)
    }

    /// Cancels all active tasks.
    pub fn cancel(&mut self) {
        // TODO implement
    }

    /// Creates a new context that will be cancelled after the given duration.
    pub fn timeout(&mut self, _expires: time::Duration) -> Self {
        // TODO implement
        self.clone()
    }
}

fn _range_bounds<R>(range: R, min: usize, max: usize) -> ops::Range<usize>
where
    R: ops::RangeBounds<usize>,
{
    let start = match range.start_bound() {
        ops::Bound::Included(n) => *n,
        ops::Bound::Excluded(n) => n + 1,
        ops::Bound::Unbounded => min,
    };

    let end = match range.end_bound() {
        ops::Bound::Included(n) => n + 1,
        ops::Bound::Excluded(n) => *n,
        ops::Bound::Unbounded => max,
    };

    start..end
}