mod small_slot;
mod table;

use std::{
    io,
    task::{Context, Poll},
};

use crate::sys::token::Completion;

use super::Driver;
use super::current::current_driver_mut;
pub(crate) use table::Requests;

#[repr(C, packed)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct RequestId {
    index: u16,
    epoch: u32,
    _wire_pad: u8,
}

impl RequestId {
    pub(crate) const EPOCH_MASK: u32 = 0x00ff_ffff;

    pub(crate) fn new(index: usize, epoch: u32) -> Self {
        Self {
            index: u16::try_from(index)
                .expect("dope invariant violated: request index exceeds u16 range"),
            epoch: epoch & Self::EPOCH_MASK,
            _wire_pad: 0,
        }
    }

    pub(in crate::driver) fn index(self) -> usize {
        self.index as usize
    }

    pub(in crate::driver) fn epoch(self) -> u32 {
        self.epoch
    }
}

enum RequestState<T> {
    Pending { data: Box<T>, ticket: RequestTicket },
    Done,
}

struct RequestTicket {
    index: RequestId,
}

impl RequestTicket {
    fn new(index: RequestId) -> Self {
        Self { index }
    }

    fn poll(&self, cx: &mut Context<'_>) -> Poll<Completion> {
        current_driver_mut().requests.poll(self.index, cx)
    }

    fn discard<T: 'static>(&self, data: &mut Option<T>) {
        current_driver_mut().requests.discard(self.index, data);
    }
}

pub struct Request<T: 'static> {
    state: RequestState<T>,
}

impl<T: 'static> Request<T> {
    pub(crate) fn submit_boxed(
        ctx: &mut Driver,
        data: Box<T>,
        sqe: crate::Sqe,
    ) -> Result<Self, (io::Error, Box<T>)> {
        match ctx.submit_sqe_op(sqe) {
            Ok(index) => Ok(Self {
                state: RequestState::Pending {
                    data,
                    ticket: RequestTicket::new(index),
                },
            }),
            Err(e) => Err((e, data)),
        }
    }
}

impl<T: Unpin + 'static> Request<T> {
    pub(crate) async fn await_submit(submit: Result<Self, (io::Error, T)>) -> (io::Result<u32>, T) {
        match submit {
            Ok(req) => {
                let (res, _, op) = req.await;
                (res, op)
            }
            Err((e, op)) => (Err(e), op),
        }
    }
}

impl<T: Unpin + 'static> std::future::Future for Request<T> {
    type Output = (io::Result<u32>, u32, T);

    fn poll(self: std::pin::Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let me = self.get_mut();
        match &me.state {
            RequestState::Pending { .. } => {}
            RequestState::Done => panic!("dope invariant violated: polled completed request"),
        }
        let RequestState::Pending { ticket, .. } = &me.state else {
            panic!("dope invariant violated: request poll state mismatch");
        };
        match ticket.poll(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(meta) => {
                let RequestState::Pending { data, .. } =
                    std::mem::replace(&mut me.state, RequestState::Done)
                else {
                    panic!("dope invariant violated: request completion state mismatch");
                };
                Poll::Ready((meta.result, meta.flags, *data))
            }
        }
    }
}

impl<T: 'static> Drop for Request<T> {
    fn drop(&mut self) {
        let (ticket, data) = match std::mem::replace(&mut self.state, RequestState::Done) {
            RequestState::Pending { ticket, data } => (ticket, data),
            RequestState::Done => return,
        };
        let mut data = Some(*data);
        ticket.discard(&mut data);
    }
}
