use crate::WriteBufStorage;
use std::collections::VecDeque;
use std::marker::PhantomData;
use std::os::fd::AsRawFd;

use dambi::MmapSlab;
use dope_core::driver::{
    BackendToken, DriverLifecycle, DriverOps, FixedBuffers, PoolDriver, ProvidedRingOps,
};
use dope_core::{FdHandle, FixedFd};

use super::slot::{ConnFlags, CoreSlot};
pub(super) use crate::CodecLayer;
pub(crate) use crate::SlotId;
use dope_transport::Transport;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PushError {
    ConnClosed,
    GenerationMismatch,
    BufferFull,
}

pub(crate) enum RecvDispatch {
    Skipped,
    EofClose {
        slot_idx: usize,
    },
    Errno {
        slot_idx: usize,
        errno: i32,
        more: bool,
    },
    Chunk {
        slot_idx: usize,
        ptr: *const u8,
        len: usize,
        bid: u16,
        more: bool,
    },
    DrainOnly {
        slot_idx: usize,
        more: bool,
    },
}

pub(crate) enum WriteDispatch {
    Skipped,
    Ready { slot_idx: usize, result: i32 },
}

pub(crate) enum CodecAck {
    NoOp,
    Err,
    Progressed,
}

#[derive(Default, Clone, Copy)]
pub(super) struct RecvSlotState {
    pub(super) epoch_seq: u32,
    pub(super) epoch: u32,
    pub(super) armed: bool,
    pub(super) canceling: bool,
}

#[derive(Default, Clone, Copy)]
pub(super) struct WriteSlotState {
    pub(super) epoch_seq: u32,
    pub(super) in_flight_epoch: u32,
    pub(super) has_in_flight: bool,
}

impl WriteSlotState {
    #[inline(always)]
    pub(super) fn next_token<T: BackendToken>(&mut self, external: usize) -> (T, u32) {
        let token = T::next_for_slot(external, &mut self.epoch_seq);
        let (_, epoch) = token.parts();
        (token, epoch)
    }

    #[inline(always)]
    pub(super) fn mark_inflight(&mut self, epoch: u32) {
        self.has_in_flight = true;
        self.in_flight_epoch = epoch;
    }
}

#[derive(Default, Clone, Copy)]
pub(crate) struct ConnIoSlot {
    pub(super) recv: RecvSlotState,
    pub(super) write: WriteSlotState,
}

pub struct PoolCore<T: Transport, L: CodecLayer, U: Default + 'static, B: PoolDriver> {
    pub(crate) slots: MmapSlab<CoreSlot<L, U>>,
    pub(crate) free_conns: Vec<u16>,
    pub(crate) recv_rearm_pending: VecDeque<u16>,
    pub(crate) slot_id_base: usize,
    pub(super) _markers: PhantomData<(T, L, B)>,
}

impl<T: Transport, L: CodecLayer, U: Default + 'static, B: PoolDriver> PoolCore<T, L, U, B> {
    pub(crate) fn new(max_conn: usize, slot_id_base: usize) -> std::io::Result<Self> {
        let mut slots = MmapSlab::<CoreSlot<L, U>>::new_zeroed(max_conn)?;
        slots.prewarm();
        for slot in slots.iter_mut() {
            slot.codec_state.write(L::ConnState::default());
            slot.user.write(U::default());
        }

        let mut free_conns: Vec<u16> = (0..max_conn as u16).rev().collect();
        free_conns.shrink_to_fit();

        Ok(Self {
            slots,
            free_conns,
            recv_rearm_pending: VecDeque::with_capacity(max_conn),
            slot_id_base,
            _markers: PhantomData,
        })
    }

    #[inline(always)]
    pub(crate) fn internal_slot(&self, external: usize) -> Option<usize> {
        external
            .checked_sub(self.slot_id_base)
            .filter(|&i| i < self.slots.len())
    }

    #[cold]
    #[inline(never)]
    pub(crate) fn allocate_slot(
        &mut self,
        driver: &mut B::Driver,
        fd_handle: FdHandle,
    ) -> Option<(usize, SlotId)> {
        let fixed_idx = match driver.register_fixed(fd_handle.as_fd()) {
            Ok(Some(idx)) => idx,
            _ => return None,
        };
        let Some(slot) = self.free_conns.pop() else {
            driver.unregister_fixed(fixed_idx);
            return None;
        };
        let slot_idx = slot as usize;

        let fixed_fd = fd_handle.fixed_io(fixed_idx);
        std::mem::forget(fd_handle.into_owned());

        let next_gen = {
            let entry = &mut self.slots[slot_idx];
            entry.conn.reset_for_reuse();
            entry.conn.fd = fixed_fd;
            entry.conn.flags.set(ConnFlags::LIVE, true);
            entry.io.write.has_in_flight = false;

            entry.buf.pending_msghdr = unsafe { std::mem::zeroed() };
            entry.buf.pending_msghdr.msg_iov = entry.buf.pending_iovs.as_mut_ptr();
            entry.conn.generation
        };

        *self.slots[slot_idx].codec_state.get_mut() = L::ConnState::default();
        *self.slots[slot_idx].user.get_mut() = U::default();

        if !self.arm_recv(driver, slot_idx) {
            driver.unregister_fixed(fixed_idx);

            unsafe { libc::close(fixed_fd.as_fd().as_raw_fd()) };
            self.slots[slot_idx].conn.flags = ConnFlags::empty();
            self.free_conns.push(slot);
            return None;
        }

        Some((slot_idx, SlotId::from_parts(slot as u32, next_gen)))
    }

    pub(crate) fn handle_codec_inflight(&mut self, slot_idx: usize, result: i32) -> CodecAck {
        if !self.slots[slot_idx].conn.flags.contains(ConnFlags::LIVE) {
            let _ = self.consume_codec_inflight(slot_idx, usize::MAX);
            return CodecAck::NoOp;
        }
        if result < 0 {
            let _ = self.consume_codec_inflight(slot_idx, usize::MAX);
            return CodecAck::Err;
        }
        let n = result as usize;
        let _more = self.consume_codec_inflight(slot_idx, n);
        CodecAck::Progressed
    }

    pub(crate) fn resolve_live_slot(&self, conn_id: SlotId) -> Result<usize, PushError> {
        let slot_idx = conn_id.slot() as usize;
        if slot_idx >= self.slots.len() {
            return Err(PushError::ConnClosed);
        }
        let entry = &self.slots[slot_idx];
        if entry.conn.generation != conn_id.generation() {
            return Err(PushError::GenerationMismatch);
        }
        if !entry.conn.flags.contains(ConnFlags::LIVE) {
            return Err(PushError::ConnClosed);
        }
        Ok(slot_idx)
    }

    pub(crate) fn classify_write(
        &mut self,
        token: B::Token,
        result: i32,
        notif: bool,
    ) -> WriteDispatch {
        let (external, epoch) = token.parts();
        let Some(slot_idx) = self.internal_slot(external) else {
            return WriteDispatch::Skipped;
        };
        let write_slot = &mut self.slots[slot_idx].io.write;
        if !write_slot.has_in_flight || write_slot.in_flight_epoch != epoch {
            return WriteDispatch::Skipped;
        }
        if notif {
            return WriteDispatch::Skipped;
        }
        write_slot.has_in_flight = false;
        WriteDispatch::Ready { slot_idx, result }
    }

    pub(crate) fn classify_recv(
        &mut self,
        driver: &mut B::Driver,
        token: B::Token,
        result: i32,
        more: bool,
        bid: Option<u16>,
    ) -> RecvDispatch {
        let (external, epoch) = token.parts();
        let Some(slot_idx) = self.internal_slot(external) else {
            if let Some(b) = bid {
                driver.provided_ring().defer(b);
            }
            return RecvDispatch::Skipped;
        };
        if self.slots[slot_idx].io.recv.epoch != epoch {
            if let Some(b) = bid {
                driver.provided_ring().defer(b);
            }
            return RecvDispatch::Skipped;
        }
        let live = self.slots[slot_idx].conn.flags.contains(ConnFlags::LIVE);
        let canceling = self.slots[slot_idx].io.recv.canceling;
        if !live || canceling {
            if let Some(b) = bid {
                driver.provided_ring().defer(b);
            }
            if !more {
                self.slots[slot_idx].io.recv.armed = false;
            }
            return RecvDispatch::Skipped;
        }
        if result < 0 {
            if let Some(b) = bid {
                driver.provided_ring().defer(b);
            }
            return RecvDispatch::Errno {
                slot_idx,
                errno: -result,
                more,
            };
        }
        if result == 0 {
            if let Some(b) = bid {
                driver.provided_ring().defer(b);
            }
            self.slots[slot_idx].io.recv.armed = false;
            return RecvDispatch::EofClose { slot_idx };
        }
        let Some(b) = bid else {
            return RecvDispatch::DrainOnly { slot_idx, more };
        };
        let len = result as usize;
        let (ptr, _) = driver.provided_ring().ptr_len(b);
        RecvDispatch::Chunk {
            slot_idx,
            ptr,
            len,
            bid: b,
            more,
        }
    }

    pub(crate) fn arm_recv(&mut self, driver: &mut B::Driver, slot_idx: usize) -> bool {
        let buf_group = driver.buf_group();
        let external = self.slot_id_base + slot_idx;
        let entry = &mut self.slots[slot_idx];
        let token =
            <B::Token as BackendToken>::next_for_slot(external, &mut entry.io.recv.epoch_seq);
        let (_, epoch) = token.parts();
        let fd = entry.conn.fd;

        let armed = driver.arm_with_retry(|d| d.arm_recv_multi(token, fd, buf_group));
        let recv = &mut self.slots[slot_idx].io.recv;
        recv.epoch = epoch;
        recv.canceling = false;
        if armed {
            recv.armed = true;
            return true;
        }
        recv.armed = false;
        self.recv_rearm_pending.push_back(slot_idx as u16);
        false
    }

    #[inline(always)]
    pub(crate) fn flush_recv_rearm(&mut self, driver: &mut B::Driver) {
        if self.recv_rearm_pending.is_empty() {
            return;
        }
        self.flush_recv_rearm_slow(driver);
    }

    #[cold]
    #[inline(never)]
    fn flush_recv_rearm_slow(&mut self, driver: &mut B::Driver) {
        let n = self.recv_rearm_pending.len();
        for _ in 0..n {
            let Some(slot_idx) = self.recv_rearm_pending.pop_front() else {
                return;
            };
            let slot_idx = slot_idx as usize;
            let entry = &self.slots[slot_idx];
            if !entry.conn.flags.contains(ConnFlags::LIVE)
                || entry.io.recv.armed
                || entry.io.recv.canceling
            {
                continue;
            }
            let _ = self.arm_recv(driver, slot_idx);
        }
    }

    pub(crate) fn reserve_slot(&mut self) -> Option<(usize, SlotId)> {
        let slot = self.free_conns.pop()?;
        let slot_idx = slot as usize;
        let entry = &mut self.slots[slot_idx];
        entry.conn.reset_for_reuse();
        entry.io.write.has_in_flight = false;
        let next_gen = entry.conn.generation;

        *entry.codec_state.get_mut() = L::ConnState::default();
        *entry.user.get_mut() = U::default();

        Some((slot_idx, SlotId::from_parts(slot as u32, next_gen)))
    }

    pub(crate) fn activate_slot(
        &mut self,
        driver: &mut B::Driver,
        slot_idx: usize,
        fixed_fd: FixedFd,
    ) -> bool {
        {
            let entry = &mut self.slots[slot_idx];
            entry.conn.fd = fixed_fd;
            entry.conn.flags.set(ConnFlags::LIVE, true);

            entry.buf.pending_msghdr = unsafe { std::mem::zeroed() };
            entry.buf.pending_msghdr.msg_iov = entry.buf.pending_iovs.as_mut_ptr();
        }
        if !self.arm_recv(driver, slot_idx) {
            self.slots[slot_idx].conn.flags = ConnFlags::empty();
            return false;
        }
        true
    }

    pub(crate) fn release_reserved(&mut self, slot_idx: usize) {
        *self.slots[slot_idx].user.get_mut() = U::default();
        let entry = &mut self.slots[slot_idx];
        entry.conn.flags = ConnFlags::empty();
        self.free_conns.push(slot_idx as u16);
    }

    pub(crate) fn free_slot(&mut self, driver: &mut B::Driver, slot_idx: usize) {
        *self.slots[slot_idx].user.get_mut() = U::default();
        self.close_slot_core(driver, slot_idx);
    }

    fn close_slot_core(&mut self, driver: &mut B::Driver, slot_idx: usize) {
        let entry = &mut self.slots[slot_idx];
        if !entry.conn.flags.contains(ConnFlags::LIVE) {
            return;
        }
        let fixed_idx = entry.conn.fd.fixed_index();
        let fd_raw = std::os::fd::AsRawFd::as_raw_fd(&entry.conn.fd.as_fd());
        if entry.io.recv.armed {
            let token = <B::Token as BackendToken>::from_index_raw(
                self.slot_id_base + slot_idx,
                entry.io.recv.epoch,
            );
            entry.io.recv.canceling = true;
            driver.cancel_recv(token);
        } else {
            entry.io.recv.canceling = false;
        }
        entry.io.recv.armed = false;
        entry.io.write.has_in_flight = false;
        driver.unregister_fixed(fixed_idx);

        unsafe { libc::close(fd_raw) };
        entry.conn.flags = ConnFlags::empty();
        self.free_conns.push(slot_idx as u16);
    }

    pub(crate) fn try_submit_codec(&mut self, driver: &mut B::Driver, slot_idx: usize) {
        if self.slots[slot_idx].io.write.has_in_flight {
            return;
        }
        let Some((ptr, len)) = L::take_inflight(self.slots[slot_idx].codec_state.get_mut()) else {
            return;
        };
        let entry = &mut self.slots[slot_idx];
        let external = self.slot_id_base + slot_idx;
        let (token, epoch) = entry.io.write.next_token::<B::Token>(external);
        let conn_fd = entry.conn.fd;
        if driver.submit_send_tagged(token, conn_fd, ptr, len) {
            entry.io.write.mark_inflight(epoch);
        }
    }

    pub(crate) fn process_codec_inbound(
        &mut self,
        driver: &mut B::Driver,
        slot_idx: usize,
        src: &[u8],
    ) -> Option<Vec<u8>> {
        L::process_inbound(self.slots[slot_idx].codec_state.get_mut(), src);
        self.try_submit_codec(driver, slot_idx);
        if L::close_pending(self.slots[slot_idx].codec_state.get()) {
            self.slots[slot_idx]
                .conn
                .flags
                .set(ConnFlags::CLOSE_AFTER, true);
        }
        if !L::is_active(self.slots[slot_idx].codec_state.get()) {
            return None;
        }
        let plaintext = L::take_plaintext(self.slots[slot_idx].codec_state.get_mut());
        if plaintext.is_empty() {
            L::install_plaintext(self.slots[slot_idx].codec_state.get_mut(), plaintext);
            return None;
        }
        Some(plaintext)
    }

    pub(crate) fn install_plaintext(&mut self, slot_idx: usize, plaintext: Vec<u8>) {
        L::install_plaintext(self.slots[slot_idx].codec_state.get_mut(), plaintext);
    }

    pub(crate) fn encrypt_write_buf(&mut self, slot_idx: usize) {
        let head = self.slots[slot_idx].conn.write_head as usize;
        let inflight_end = self.slots[slot_idx].conn.write_inflight_end as usize;
        if head <= inflight_end {
            return;
        }
        let entry = &mut self.slots[slot_idx];
        let plaintext = &entry.buf.write_buf.as_mut_slice()[inflight_end..head];
        L::process_outbound(entry.codec_state.get_mut(), plaintext);
        entry.conn.write_head = 0;
        entry.conn.write_tail = 0;
        entry.conn.write_inflight_end = 0;
    }

    pub(crate) fn consume_codec_inflight(&mut self, slot_idx: usize, n: usize) -> bool {
        L::consume_inflight(self.slots[slot_idx].codec_state.get_mut(), n)
    }

    pub(crate) fn codec_idle(&self, slot_idx: usize) -> bool {
        let state = self.slots[slot_idx].codec_state.get();
        L::wire_slice(state).is_empty() && !L::has_inflight(state)
    }

    pub(crate) fn codec_close_pending(&self, slot_idx: usize) -> bool {
        L::close_pending(self.slots[slot_idx].codec_state.get())
    }
}

impl<T: Transport, L: CodecLayer, U: Default + 'static, B: PoolDriver> Drop
    for PoolCore<T, L, U, B>
{
    fn drop(&mut self) {
        for slot in self.slots.iter_mut() {
            unsafe {
                slot.codec_state.drop_in_place();
                slot.user.drop_in_place();
            }
        }
    }
}
