dope-session 0.2.3

use crate::WriteBufStorage;
use std::pin::Pin;
use std::task::{Context, Poll};

use dope_core::driver::PoolDriver;
use dope_core::profile::Profile;
use dope_transport::Transport;

use super::super::Pool;
use super::super::slot::CoreSlot;
use super::Server;
use super::state::SlotEgress;
use super::wake::WakeQueue;
use crate::SlotId;
use crate::protocol::server::{ResponseChunkStream, ServerProtocol};

impl<H: ServerProtocol, T: Transport, F: Profile, B: PoolDriver> Pool<Server<H, T>, F, B> {
    pub(super) fn encode_stream_chunks_to_buf(&mut self, slot_idx: usize) {
        let conn_id = {
            let conn = &self.core.slots[slot_idx].conn;
            SlotId::from_parts(slot_idx as u32, conn.generation)
        };
        let waker = WakeQueue::make_waker(conn_id);
        let mut cx = Context::from_waker(&waker);

        let CoreSlot {
            conn, buf, user, ..
        } = &mut self.core.slots[slot_idx];
        let user = user.get_mut();
        let SlotEgress::Stream(stream_state) = &mut user.egress else {
            return;
        };
        if stream_state.cur_chunk.is_some() {
            return;
        }
        let mut head = conn.write_head as usize;
        let cap = buf.write_buf.as_slice().len();

        loop {
            if stream_state.eof {
                break;
            }
            let polled: Poll<Option<dambi::Bytes>> = {
                let stream: Pin<&mut dyn ResponseChunkStream> = stream_state.stream.as_mut();
                stream.poll_chunk(&mut cx)
            };
            match polled {
                Poll::Pending => break,
                Poll::Ready(None) => {
                    stream_state.eof = true;
                    break;
                }
                Poll::Ready(Some(chunk)) => {
                    if chunk.is_empty() {
                        continue;
                    }
                    if cap.saturating_sub(head) < chunk.len() {
                        stream_state.cur_chunk = Some(chunk);
                        stream_state.cur_chunk_offset = 0;
                        break;
                    }
                    buf.write_buf.as_mut_slice()[head..head + chunk.len()].copy_from_slice(&chunk);
                    head += chunk.len();
                }
            }
        }

        conn.write_head = head.min(u16::MAX as usize) as u16;

        if stream_state.eof && stream_state.cur_chunk.is_none() {
            user.egress = SlotEgress::Idle;
        }
    }

    pub(super) fn submit_stream_write(&mut self, driver: &mut B::Driver, slot_idx: usize) {
        if self.core.slots[slot_idx].io.write.has_in_flight {
            return;
        }

        let needs_poll = {
            let user = self.core.slots[slot_idx].user.get();
            match &user.egress {
                SlotEgress::Stream(s) => !s.eof && s.cur_chunk.is_none(),
                _ => return,
            }
        };
        if needs_poll {
            let conn_id = {
                let conn = &self.core.slots[slot_idx].conn;
                SlotId::from_parts(slot_idx as u32, conn.generation)
            };
            let waker = WakeQueue::make_waker(conn_id);
            let mut cx = Context::from_waker(&waker);
            let user = self.core.slots[slot_idx].user.get_mut();
            let SlotEgress::Stream(stream_state) = &mut user.egress else {
                return;
            };
            let polled: Poll<Option<dambi::Bytes>> = {
                let stream: Pin<&mut dyn ResponseChunkStream> = stream_state.stream.as_mut();
                stream.poll_chunk(&mut cx)
            };
            match polled {
                Poll::Pending => return,
                Poll::Ready(None) => {
                    stream_state.eof = true;
                    user.egress = SlotEgress::Idle;
                    return;
                }
                Poll::Ready(Some(chunk)) => {
                    if chunk.is_empty() {
                        return;
                    }
                    stream_state.cur_chunk = Some(chunk);
                    stream_state.cur_chunk_offset = 0;
                }
            }
        }

        let (ptr, len) = {
            let user = self.core.slots[slot_idx].user.get();
            let SlotEgress::Stream(s) = &user.egress else {
                return;
            };
            let chunk = s.cur_chunk.as_ref().expect("just stashed");
            let off = s.cur_chunk_offset as usize;

            let ptr = unsafe { chunk.as_ptr().add(off) };
            let len = chunk.len() - off;
            (ptr, len)
        };
        let _ = self.driver_submit_send(driver, slot_idx, ptr, len as u32);
    }
}