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

use crate::sys::token::Completion;
use dambi::{LiveKey, Slab};

use super::RequestId;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) struct WakerKey(usize);

impl WakerKey {
    fn as_usize(self) -> usize {
        self.0
    }
}

pub(crate) struct WakerSlots {
    slots: Vec<Option<Waker>>,
}

impl WakerSlots {
    pub(crate) fn with_capacity(min_cap: usize) -> Self {
        Self {
            slots: vec![None; min_cap],
        }
    }

    pub(crate) fn ensure_index(&mut self, index: usize) {
        if index >= self.slots.len() {
            self.slots.resize_with(index + 1, || None);
        }
    }

    pub(crate) fn put(&mut self, index: usize, waker: &Waker) -> WakerKey {
        debug_assert!(
            index < self.slots.len(),
            "dope invariant violated: waker index out of bounds"
        );
        let slot = &mut self.slots[index];
        if let Some(old) = slot {
            if !old.will_wake(waker) {
                *old = waker.clone();
            }
        } else {
            *slot = Some(waker.clone());
        }
        WakerKey(index)
    }

    pub(crate) fn update(&mut self, key: WakerKey, waker: &Waker) {
        let index = key.as_usize();
        debug_assert!(
            index < self.slots.len(),
            "dope invariant violated: waker index out of bounds"
        );
        let Some(old) = self.slots[index].as_mut() else {
            panic!("dope invariant violated: waiting request missing waker");
        };
        if !old.will_wake(waker) {
            *old = waker.clone();
        }
    }

    pub(crate) fn take(&mut self, key: WakerKey) -> Waker {
        let index = key.as_usize();
        debug_assert!(
            index < self.slots.len(),
            "dope invariant violated: waker index out of bounds"
        );
        self.slots[index]
            .take()
            .unwrap_or_else(|| panic!("dope invariant violated: waiting request missing waker"))
    }

    pub(crate) fn clear_existing(&mut self, index: usize) {
        debug_assert!(
            index < self.slots.len(),
            "dope invariant violated: waker index out of bounds"
        );
        let _ = self.slots[index].take();
    }
}

const DROPPED_INLINE_CAP: usize = 256;

type DroppedSlot = super::small_slot::SmallSlot<DROPPED_INLINE_CAP>;

pub(crate) struct DroppedData {
    slots: Vec<DroppedSlot>,
}

impl DroppedData {
    pub(in crate::driver) fn with_capacity(min_cap: usize) -> Self {
        let mut slots = Vec::with_capacity(min_cap);
        slots.resize_with(min_cap, DroppedSlot::new);
        Self { slots }
    }

    fn ensure_index(&mut self, index: usize) {
        if index < self.slots.len() {
            return;
        }
        self.slots.resize_with(index + 1, DroppedSlot::new);
    }

    pub(in crate::driver) fn put<T: 'static>(&mut self, index: usize, value: T) {
        self.ensure_index(index);
        self.slots[index].put(value);
    }

    pub(in crate::driver) fn clear(&mut self, index: usize) {
        if index >= self.slots.len() {
            return;
        }
        self.slots[index].clear();
    }
}

pub(in crate::driver) enum RequestEntry {
    Submitted,
    Waiting(WakerKey),
    Dropped,
    Completed(Completion),
}

impl RequestEntry {
    fn submitted() -> Self {
        Self::Submitted
    }

    fn complete(&mut self, completion: Completion) -> Result<Option<WakerKey>, ()> {
        match self {
            Self::Submitted => {
                *self = Self::Completed(completion);
                Ok(None)
            }
            Self::Waiting(key) => {
                let key = *key;
                *self = Self::Completed(completion);
                Ok(Some(key))
            }
            Self::Dropped => Err(()),
            Self::Completed(_) => panic!("dope invariant violated: request completion repeated"),
        }
    }

    fn discard<T: 'static>(
        &mut self,
        dropped: &mut DroppedData,
        index: usize,
        data: &mut Option<T>,
    ) {
        match self {
            Self::Submitted | Self::Waiting(_) => {
                if let Some(value) = data.take() {
                    dropped.put(index, value);
                } else {
                    dropped.clear(index);
                }
                *self = Self::Dropped;
            }
            Self::Dropped | Self::Completed(_) => {}
        }
    }
}

pub(crate) struct Requests {
    slots: Slab<RequestEntry>,
    epochs: Vec<u32>,
    wakers: WakerSlots,
    dropped: DroppedData,
}

impl Requests {
    pub(in crate::driver) fn with_capacity(min_cap: usize) -> Self {
        Self {
            slots: Slab::with_capacity(min_cap),
            epochs: Vec::with_capacity(min_cap),
            wakers: WakerSlots::with_capacity(min_cap),
            dropped: DroppedData::with_capacity(min_cap),
        }
    }

    fn dropped() -> Completion {
        Completion {
            result: Err(io::Error::other("op dropped")),
            flags: 0,
        }
    }

    fn live_key(&self, id: RequestId) -> LiveKey {
        let index = id.index();
        let epoch = id.epoch();
        let slot_epoch = self
            .epochs
            .get(index)
            .copied()
            .unwrap_or_else(|| panic!("dope invariant violated: request epoch missing"));
        assert!(
            slot_epoch == epoch,
            "dope invariant violated: request epoch mismatch",
        );
        self.slots
            .live_key(index)
            .unwrap_or_else(|| panic!("dope invariant violated: request slot missing"))
    }

    fn remove(&mut self, key: LiveKey) -> RequestEntry {
        let index = key.as_usize();
        self.dropped.clear(index);
        self.slots.remove_live(key)
    }

    fn remove_completed(&mut self, id: RequestId) -> Completion {
        match self.remove(self.live_key(id)) {
            RequestEntry::Completed(completion) => completion,
            _ => panic!("dope invariant violated: request remove lifecycle mismatch"),
        }
    }

    fn ensure_index(&mut self, index: usize) {
        self.wakers.ensure_index(index);
        if self.epochs.len() <= index {
            self.epochs.resize(index + 1, 0);
        }
    }

    fn next_epoch(&mut self, index: usize) -> u32 {
        self.ensure_index(index);
        let epoch = self.epochs[index].wrapping_add(1) & RequestId::EPOCH_MASK;
        let epoch = epoch.max(1);
        self.epochs[index] = epoch;
        epoch
    }

    fn insert_entry(&mut self, entry: RequestEntry) -> RequestId {
        let index = self.slots.insert(entry);
        let epoch = self.next_epoch(index);
        self.dropped.clear(index);
        RequestId::new(index, epoch)
    }

    pub(in crate::driver) fn insert(&mut self) -> RequestId {
        self.insert_entry(RequestEntry::submitted())
    }

    pub(in crate::driver) fn rollback(&mut self, index: RequestId) {
        let index = index.index();
        self.dropped.clear(index);
        let _ = self.slots.remove(index);
        self.wakers.clear_existing(index);
    }

    pub(in crate::driver) fn poll(
        &mut self,
        index: RequestId,
        cx: &mut Context<'_>,
    ) -> Poll<Completion> {
        let key = self.live_key(index);
        let slot = index.index();
        match self.slots.get_live(key) {
            RequestEntry::Submitted => {
                let waker = self.wakers.put(slot, cx.waker());
                *self.slots.get_live(key) = RequestEntry::Waiting(waker);
                Poll::Pending
            }
            RequestEntry::Waiting(waker) => {
                self.wakers.update(*waker, cx.waker());
                Poll::Pending
            }
            RequestEntry::Completed(_) => Poll::Ready(self.remove_completed(index)),
            RequestEntry::Dropped => Poll::Ready(Self::dropped()),
        }
    }

    pub(in crate::driver) fn discard<T: 'static>(
        &mut self,
        index: RequestId,
        data: &mut Option<T>,
    ) {
        let key = self.live_key(index);
        let slot = index.index();
        match self.slots.get_live(key) {
            RequestEntry::Submitted => {
                let (slots, dropped) = (&mut self.slots, &mut self.dropped);
                slots.get_live(key).discard(dropped, slot, data);
            }
            RequestEntry::Waiting(waker) => {
                let waker = *waker;
                let (slots, dropped) = (&mut self.slots, &mut self.dropped);
                slots.get_live(key).discard(dropped, slot, data);
                let _ = self.wakers.take(waker);
            }
            RequestEntry::Completed(_) => match self.remove(key) {
                RequestEntry::Completed(_) => {}
                _ => panic!("dope invariant violated: request remove lifecycle mismatch"),
            },
            RequestEntry::Dropped => {}
        }
    }

    pub(in crate::driver) fn on_request(&mut self, id: RequestId, result: i32, flags: u32) {
        if io_uring::cqueue::notif(flags) {
            return;
        }
        let index = id.index();
        let Some(entry) = self.slots.get_mut(index) else {
            return;
        };
        if self.epochs.get(index).copied() != Some(id.epoch()) {
            return;
        }
        let completion = Completion {
            result: if result >= 0 {
                Ok(result as u32)
            } else {
                Err(io::Error::from_raw_os_error(-result))
            },
            flags,
        };

        match entry.complete(completion) {
            Ok(Some(waker)) => {
                self.wakers.take(waker).wake();
            }
            Ok(None) => {}
            Err(()) => {
                self.wakers.clear_existing(index);
                match self.remove(self.live_key(id)) {
                    RequestEntry::Dropped => {}
                    _ => panic!("dope invariant violated: request remove lifecycle mismatch"),
                }
            }
        }
    }
}
