pub(crate) mod poll;
pub(crate) mod recv;
pub(crate) mod state;
pub(crate) mod wake;
pub(crate) mod write;
pub(crate) mod write_stream;

pub(crate) use self::state::ServerSlotState;

use self::wake::WakeQueue;
use super::Direction;
use super::Pool;
use super::core::PoolCore;
use crate::protocol::server::{LongFuture, ServerProtocol};
use dope_core::FdHandle;
use dope_core::driver::{
    Backend, BackendToken, DriverConfig as _, DriverLifecycle, DriverOps, PoolDriver,
};
use dope_core::profile::Profile;
use dope_runtime::runtime::Runtime;
use dope_transport::{Tcp, Transport};
use std::collections::VecDeque;
use std::io;
use std::marker::PhantomData;
use std::net::SocketAddr;
use std::sync::Arc;
use std::sync::atomic::AtomicBool;

const ACCEPT_PENDING_INIT_CAP: usize = 1024;

#[derive(Default)]
pub(crate) struct AcceptSlot {
    pub(super) epoch_seq: u32,
    pub(super) epoch: u32,
    pub(super) armed: bool,
    pub(super) canceling: bool,
    pub(super) pending_rearm: bool,
}

pub struct Server<H: ServerProtocol, T: Transport = Tcp> {
    pub(crate) protocol: H,
    pub(crate) listener_fd: FdHandle,
    pub(crate) accept_slot: AcceptSlot,
    pub(crate) accept_pending: VecDeque<FdHandle>,
    pub(crate) socket_opts: dope_transport::config::StreamSocketConfig,
    pub(crate) accept_slot_id: usize,
    transport: PhantomData<T>,
}

impl<H: ServerProtocol, T: Transport> Server<H, T> {
    #[inline(always)]
    pub fn protocol(&self) -> &H {
        &self.protocol
    }

    #[inline(always)]
    pub fn protocol_mut(&mut self) -> &mut H {
        &mut self.protocol
    }

    #[inline(always)]
    pub fn listener_fd(&self) -> &FdHandle {
        &self.listener_fd
    }

    #[cold]
    #[inline(never)]
    pub(crate) fn accept_arm<B: PoolDriver>(&mut self, driver: &mut B::Driver) {
        if self.accept_slot.armed || self.accept_slot.canceling {
            self.accept_slot.pending_rearm = false;
            return;
        }
        let token = <B::Token as BackendToken>::next_for_slot(
            self.accept_slot_id,
            &mut self.accept_slot.epoch_seq,
        );
        let (_, epoch) = token.parts();
        let fd = self.listener_fd.as_fd();
        let armed = driver.arm_with_retry(|d| d.arm_accept_multi(token, fd));
        self.accept_slot.epoch = epoch;
        if armed {
            self.accept_slot.armed = true;
            self.accept_slot.pending_rearm = false;
        } else {
            self.accept_slot.pending_rearm = true;
        }
    }

    #[inline(always)]
    pub(crate) fn accept_flush_rearm<B: PoolDriver>(&mut self, driver: &mut B::Driver) {
        if self.accept_slot.pending_rearm {
            self.accept_arm::<B>(driver);
        }
    }

    pub(crate) fn on_accept_event<B: PoolDriver>(
        &mut self,
        token: B::Token,
        result: i32,
        more: bool,
    ) {
        let (_, epoch) = token.parts();
        if self.accept_slot.epoch != epoch {
            return;
        }

        if result >= 0 {
            let fd = dope_core::take_fd(result);
            self.accept_pending.push_back(fd);
            if !more {
                self.accept_slot.armed = false;
                self.accept_slot.pending_rearm = !self.accept_slot.canceling;
            }
            return;
        }

        let errno = -result;
        if errno == libc::EINVAL {
            panic!("dope invariant violated: accept-multi unsupported on required uring platform");
        }
        if Tcp::is_accept_retryable(errno) {
            if !more {
                self.accept_slot.armed = false;
                self.accept_slot.pending_rearm = !self.accept_slot.canceling;
            }
            return;
        }
        if errno == libc::ECANCELED {
            self.accept_slot.armed = false;
            self.accept_slot.pending_rearm = !self.accept_slot.canceling;
            self.accept_slot.canceling = false;
            return;
        }
        self.accept_slot.armed = false;
        self.accept_slot.pending_rearm = !self.accept_slot.canceling;
    }

    #[inline(always)]
    pub(crate) fn has_pending(&self) -> bool {
        WakeQueue::is_pending() || !self.accept_pending.is_empty() || self.accept_slot.pending_rearm
    }
}

impl<H: ServerProtocol, T: Transport, B: PoolDriver> Direction<B> for Server<H, T> {
    type Transport = T;
    type SlotUser = ServerSlotState<H>;
    type CodecLayer = H::CodecLayer;

    #[inline(always)]
    fn has_pending(&self) -> bool {
        Self::has_pending(self)
    }

    #[inline(always)]
    fn tick(
        &mut self,
        core: &mut PoolCore<Self::Transport, Self::CodecLayer, Self::SlotUser, B>,
        driver: &mut B::Driver,
    ) {
        while let Some(fd) = self.accept_pending.pop_front() {
            let _ = T::apply_stream_socket_opts(&fd, &self.socket_opts);
            let Some((slot_idx, conn_id)) = core.allocate_slot(driver, fd) else {
                continue;
            };
            let user = core.slots[slot_idx].user.get_mut();
            if let Some(future) = self.protocol.on_register(conn_id, &mut user.protocol_state) {
                user.long_future = Some(LongFuture {
                    generation: conn_id.generation(),
                    future,
                });
                WakeQueue::push(conn_id);
            }
        }
        self.accept_flush_rearm::<B>(driver);
        core.flush_recv_rearm(driver);
    }

    #[inline(always)]
    fn on_accept_event(&mut self, token: B::Token, result: i32, more: bool) {
        Self::on_accept_event::<B>(self, token, result, more)
    }
}

#[derive(Clone, Debug)]
pub struct TcpServerConfig {
    pub max_conn: usize,
    pub bind: SocketAddr,
    pub backlog: i32,
    pub socket_opts: dope_transport::config::StreamSocketConfig,
    pub shutdown: Option<Arc<AtomicBool>>,
    pub cpu_id: Option<u16>,
}

impl<H, F, B> Pool<Server<H, Tcp>, F, B>
where
    H: ServerProtocol,
    F: Profile,
    B: PoolDriver,
{
    pub fn build(protocol: H, cfg: TcpServerConfig) -> io::Result<Runtime<Self>> {
        let shutdown = cfg.shutdown.clone();
        let max_conn = cfg.max_conn;
        let cpu_id = cfg.cpu_id;
        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);

        let driver_cfg =
            <<B as Backend>::Config as dope_core::driver::DriverConfig>::for_tcp_profile(
                F::RING_ENTRIES,
                max_conn,
                crate::pool::slot::INGRESS_BUF_CAP,
            )
            .with_defer_taskrun(F::DEFER_TASKRUN)
            .with_cpu_id(cpu_id);
        let mut driver = <B as Backend>::new_driver(driver_cfg)?;
        let reactor = Self::build_raw(protocol, cfg, &mut driver, 0)?;
        Ok(Runtime::new(reactor, driver).with_shutdown(shutdown))
    }

    pub fn build_raw(
        protocol: H,
        cfg: TcpServerConfig,
        driver: &mut B::Driver,
        slot_id_base: usize,
    ) -> io::Result<Self> {
        let TcpServerConfig {
            max_conn,
            bind,
            backlog,
            socket_opts,
            shutdown: _,
            cpu_id: _,
        } = cfg;
        assert!(max_conn > 0 && max_conn <= u16::MAX as usize);

        WakeQueue::reserve(max_conn);

        let listener_cfg = dope_transport::config::ListenerSocketConfig {
            reuseport: dope_transport::config::SocketToggle::Enabled,
            fastopen_backlog: None,
        };

        B::set_current(driver as *mut _);

        let listener_result = Tcp::bind_listener(bind, backlog, &listener_cfg);

        B::clear_current();

        let listener_fd = listener_result?;
        let core =
            PoolCore::<Tcp, H::CodecLayer, ServerSlotState<H>, B>::new(max_conn, slot_id_base)?;
        let accept_slot_id = slot_id_base + max_conn;

        let direction = Server {
            protocol,
            listener_fd,
            accept_slot: AcceptSlot::default(),
            accept_pending: VecDeque::with_capacity(ACCEPT_PENDING_INIT_CAP),
            socket_opts,
            accept_slot_id,
            transport: PhantomData,
        };

        let mut reactor = Self {
            direction,
            core,
            profile: PhantomData,
        };
        reactor.direction.accept_arm::<B>(driver);
        if !reactor.direction.accept_slot.armed && !reactor.direction.accept_slot.pending_rearm {
            return Err(io::Error::other("dope-net: failed to arm multishot accept"));
        }
        Ok(reactor)
    }
}

impl<H, T, F, B> Pool<Server<H, T>, F, B>
where
    H: ServerProtocol,
    T: Transport,
    F: Profile,
    B: PoolDriver,
{
    pub fn slot_id_range(&self) -> std::ops::Range<usize> {
        let base = self.core.slot_id_base;
        base..(self.direction.accept_slot_id + 1)
    }

    pub fn conn_slot_count(&self) -> usize {
        self.core.slots.len()
    }

    pub fn protocol(&self) -> &H {
        &self.direction.protocol
    }

    pub fn protocol_mut(&mut self) -> &mut H {
        &mut self.direction.protocol
    }

    pub fn listener_fd(&self) -> &FdHandle {
        &self.direction.listener_fd
    }
}

impl<H, T, F, B> super::PoolTick<B> for Pool<Server<H, T>, F, B>
where
    H: ServerProtocol,
    T: Transport,
    F: Profile,
    B: PoolDriver,
{
    #[inline(always)]
    fn tick(&mut self, driver: &mut B::Driver) {
        self.direction.tick(&mut self.core, driver);
        self.drain_ready_queue(driver);
    }
}
