// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{connection::Owner, units::Rates, Checkpoints, Connection, Result, Timer, Trace};
use core::task::{Context, Poll};
use futures::ready;

mod thread;
pub(crate) mod timer;

use thread::Thread;

#[derive(Debug)]
pub struct Driver<'a, C: Connection> {
    pub connection: C,
    local_thread: Thread<'a>,
    local_rates: Rates,
    peer_streams: Vec<(Poll<()>, Thread<'a>)>,
    peer_rates: Rates,
    can_accept: bool,
    is_finished: bool,
}

impl<'a, C: Connection> Driver<'a, C> {
    pub fn new(scenario: &'a crate::scenario::Connection, connection: C) -> Self {
        Self {
            connection,
            local_thread: Thread::new(&scenario.ops, Owner::Local),
            local_rates: Default::default(),
            peer_streams: scenario
                .peer_streams
                .iter()
                .map(|ops| (Poll::Pending, Thread::new(ops, Owner::Remote)))
                .collect(),
            peer_rates: Default::default(),
            can_accept: true,
            is_finished: false,
        }
    }

    pub async fn run<T: Trace, Ch: Checkpoints, Ti: Timer>(
        mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        timer: &mut Ti,
    ) -> Result<C> {
        futures::future::poll_fn(|cx| self.poll_with_timer(trace, checkpoints, timer, cx)).await?;
        Ok(self.connection)
    }

    #[inline]
    pub fn poll_with_timer<T: Trace, Ch: Checkpoints, Ti: Timer>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        timer: &mut Ti,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        let now = timer.now();
        let res = self.poll(trace, checkpoints, now, cx);

        if let Some(target) = timer::Provider::next_expiration(&self) {
            // update the timer with the next expiration
            let _ = timer.poll(target, cx);
        };

        res
    }

    #[inline]
    pub fn poll<T: Trace, Ch: Checkpoints>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: timer::Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        trace.enter_connection(self.connection.id());

        if self.is_finished {
            return self.connection.poll_finish(cx);
        }

        let mut poll_accept = false;
        let mut all_ready = true;

        trace.enter(now, 0, 0);
        let result = self.local_thread.poll(
            &mut self.connection,
            trace,
            checkpoints,
            &mut self.local_rates,
            now,
            cx,
        );
        trace.exit(now);

        match result {
            Poll::Ready(Ok(_)) => {}
            Poll::Ready(Err(err)) => return Err(err).into(),
            Poll::Pending => all_ready = false,
        }

        for (idx, (accepted, thread)) in self.peer_streams.iter_mut().enumerate() {
            // if we're still waiting to accept this stream move on
            if accepted.is_pending() {
                all_ready = false;
                poll_accept = self.can_accept;
                continue;
            }

            trace.enter(now, 1, idx);
            let result = thread.poll(
                &mut self.connection,
                trace,
                checkpoints,
                &mut self.peer_rates,
                now,
                cx,
            );
            trace.exit(now);

            match result {
                Poll::Ready(Ok(_)) => {}
                Poll::Ready(Err(err)) => return Err(err).into(),
                Poll::Pending => all_ready = false,
            }
        }

        if poll_accept {
            match self.connection.poll_accept_stream(cx) {
                Poll::Ready(Ok(Some(id))) => {
                    trace.accept(now, id);
                    if let Some((accepted, _)) = self.peer_streams.get_mut(id as usize) {
                        *accepted = Poll::Ready(());
                        cx.waker().wake_by_ref();
                    } else {
                        todo!("return a not found error")
                    }
                }
                Poll::Ready(Ok(None)) => self.can_accept = false,
                Poll::Ready(Err(err)) => return Err(err).into(),
                Poll::Pending => all_ready = false,
            }
        }

        all_ready &= !timer::Provider::is_armed(&self);

        if all_ready {
            self.is_finished = true;
            self.connection.poll_finish(cx)
        } else {
            ready!(self.connection.poll_progress(cx))?;
            Poll::Pending
        }
    }
}

impl<'a, C: Connection> crate::client::Connection for Driver<'a, C> {
    #[inline]
    fn id(&self) -> u64 {
        self.connection.id()
    }

    #[inline]
    fn poll<T, Ch>(
        &mut self,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: timer::Timestamp,
        cx: &mut Context<'_>,
    ) -> Poll<Result<()>>
    where
        T: Trace,
        Ch: Checkpoints,
    {
        Self::poll(self, trace, checkpoints, now, cx)
    }
}

impl<'a, C: Connection> timer::Provider for Driver<'a, C> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.local_thread.timers(query)?;
        for (_, thread) in self.peer_streams.iter() {
            thread.timers(query)?;
        }
        Ok(())
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{connection::Owner, Result};
use core::{
    pin::Pin,
    task::{Context, Poll},
};
use once_cell::sync::Lazy;
use std::mem::MaybeUninit;
use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};

const READ_BUFFER_SIZE: usize = 100_000;
const SEND_BUFFER_SIZE: usize = 100_000_000;

static SEND_BUFFER: Lazy<Vec<u8>> = Lazy::new(|| vec![42; SEND_BUFFER_SIZE]);

#[derive(Debug)]
pub struct Connection<T: AsyncRead + AsyncWrite> {
    id: u64,
    inner: Pin<Box<T>>,
    stream_opened: bool,
}

impl<T: AsyncRead + AsyncWrite> Connection<T> {
    pub fn new(id: u64, inner: Pin<Box<T>>) -> Self {
        // force an allocation up-front
        let _ = &*SEND_BUFFER;

        Self {
            id,
            inner,
            stream_opened: false,
        }
    }

    fn open_stream(&mut self) -> Result<()> {
        if self.stream_opened {
            return Err("cannot open more than one duplex stream at a time".into());
        }

        self.stream_opened = true;
        Ok(())
    }

    fn close_stream(&mut self) -> Result<()> {
        if !self.stream_opened {
            return Err("attempted to close the stream which wasn't opened".into());
        }

        self.stream_opened = false;
        Ok(())
    }
}

impl<T: AsyncRead + AsyncWrite> super::Connection for Connection<T> {
    fn id(&self) -> u64 {
        self.id
    }

    fn poll_open_bidirectional_stream(&mut self, _: u64, _: &mut Context) -> Poll<Result<()>> {
        self.open_stream().into()
    }

    fn poll_open_send_stream(&mut self, _: u64, _: &mut Context) -> Poll<Result<()>> {
        self.open_stream().into()
    }

    fn poll_accept_stream(&mut self, _: &mut Context) -> Poll<Result<Option<u64>>> {
        let id: u64 = 0;
        match self.open_stream() {
            Ok(()) => Ok(Some(id)).into(),
            Err(err) => Err(err).into(),
        }
    }

    fn poll_send(
        &mut self,
        _owner: Owner,
        _id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        let mut sent: u64 = 0;
        while sent < bytes {
            let to_send = (bytes - sent) as usize;
            let to_send = to_send.min(SEND_BUFFER_SIZE);
            let to_send = &SEND_BUFFER[0..to_send];
            match self.inner.as_mut().poll_write(cx, to_send) {
                Poll::Ready(result) => {
                    sent += result? as u64;
                }
                Poll::Pending => {
                    break;
                }
            }
        }

        if sent == 0 {
            return Poll::Pending;
        }

        if sent == bytes {
            if let Poll::Ready(res) = self.inner.as_mut().poll_flush(cx) {
                res?;
            }
        }

        Ok(sent).into()
    }

    fn poll_receive(
        &mut self,
        _owner: Owner,
        _id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        let mut buf: [MaybeUninit<u8>; READ_BUFFER_SIZE] =
            unsafe { MaybeUninit::uninit().assume_init() };

        let mut received: u64 = 0;
        while received < bytes {
            let mut buf = ReadBuf::uninit(&mut buf);

            match self.inner.as_mut().poll_read(cx, &mut buf) {
                Poll::Ready(_) => {
                    if buf.filled().is_empty() && self.stream_opened {
                        self.close_stream()?;
                        return Ok(0).into();
                    }
                    received += buf.filled().len() as u64;
                }
                Poll::Pending => {
                    break;
                }
            }
        }

        if received == 0 {
            return Poll::Pending;
        }

        Ok(received).into()
    }

    fn poll_send_finish(&mut self, _: Owner, _: u64, _: &mut Context) -> Poll<Result<()>> {
        Ok(()).into()
    }

    fn poll_receive_finish(&mut self, _: Owner, _: u64, _: &mut Context) -> Poll<Result<()>> {
        Ok(()).into()
    }

    fn poll_progress(&mut self, _: &mut Context) -> Poll<Result<()>> {
        Ok(()).into()
    }

    fn poll_finish(&mut self, _cx: &mut Context) -> Poll<Result<()>> {
        if self.stream_opened {
            self.close_stream()?;
        }
        Ok(()).into()
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use core::task::Poll;

#[derive(Clone, Debug, Default)]
pub struct IdPrefixReader {
    bytes: [u8; core::mem::size_of::<u64>()],
    cursor: u8,
}

impl IdPrefixReader {
    pub fn remaining(&mut self) -> &mut [u8] {
        &mut self.bytes[self.cursor as usize..]
    }

    pub fn on_read(&mut self, len: usize) -> Poll<u64> {
        let cursor = self.cursor as usize + len;
        if cursor >= self.bytes.len() {
            Poll::Ready(u64::from_be_bytes(self.bytes))
        } else {
            self.cursor = cursor as u8;
            Poll::Pending
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

pub type Result<T, E = Error> = core::result::Result<T, E>;
pub type Error = Box<dyn std::error::Error>;

mod checkpoints;
pub mod client;
mod connection;
mod driver;
#[cfg(any(test, feature = "testing"))]
pub mod testing;

pub mod duplex;
pub mod helper;
pub mod multiplex;
pub mod operation;
#[cfg(feature = "s2n-quic")]
pub mod s2n_quic;
pub mod scenario;
pub mod stats;
pub mod timer;
pub mod trace;
pub mod units;

pub use checkpoints::Checkpoints;
pub use client::Driver as Client;
pub use connection::Connection;
pub use driver::Driver;
pub use timer::Timer;
pub use trace::Trace;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{connection::Owner, Result};
use bytes::Buf;
use core::{
    pin::Pin,
    task::{Context, Poll},
};
use futures::ready;
use std::collections::{hash_map::Entry, HashMap, HashSet, VecDeque};
use tokio::io::{AsyncRead, AsyncWrite};

mod buffer;
mod frame;
mod stream;

use buffer::{ReadBuffer, WriteBuffer};
use frame::Frame;
use stream::{ReceiveStream, SendStream, Stream};

const DEFAULT_STREAM_WINDOW: u64 = 256000;
const DEFAULT_MAX_STREAMS: u64 = 100;

#[derive(Clone, Debug)]
pub struct Config {
    pub stream_window: u64,
    pub max_streams: u64,
    pub max_stream_data_thresh: u64,
    pub max_stream_frame_len: u32,
    pub max_write_queue_len: usize,
    pub peer_max_streams: u64,
    pub peer_max_stream_data: u64,
}

impl Default for Config {
    fn default() -> Self {
        Self {
            stream_window: DEFAULT_STREAM_WINDOW,
            max_streams: DEFAULT_MAX_STREAMS,
            max_stream_data_thresh: DEFAULT_STREAM_WINDOW / 2,
            max_stream_frame_len: (1 << 15),
            max_write_queue_len: 250,
            peer_max_streams: DEFAULT_MAX_STREAMS,
            peer_max_stream_data: DEFAULT_STREAM_WINDOW,
        }
    }
}

#[derive(Debug)]
pub struct Connection<T: AsyncRead + AsyncWrite> {
    id: u64,
    inner: Pin<Box<T>>,
    rx_open: bool,
    tx_open: bool,
    tx_flushing: bool,
    frame: Option<frame::Frame>,
    read_buf: ReadBuffer,
    decoder: frame::Decoder,
    write_buf: WriteBuffer,
    stream_controller: stream::Controller,
    streams: [HashMap<u64, Stream>; 2],
    max_stream_data: HashSet<(Owner, u64)>,
    pending_accept: VecDeque<u64>,
    config: Config,
}

impl<T: AsyncRead + AsyncWrite> Connection<T> {
    pub fn new(id: u64, inner: Pin<Box<T>>, config: Config) -> Self {
        let mut write_buf = WriteBuffer::default();

        if config.stream_window != DEFAULT_STREAM_WINDOW {
            write_buf.push(Frame::InitialMaxStreamData {
                up_to: config.stream_window,
            });
        }

        if config.max_streams != DEFAULT_MAX_STREAMS {
            write_buf.push(Frame::MaxStreams {
                up_to: config.max_streams,
            });
        }

        Self {
            id,
            inner,
            rx_open: true,
            tx_open: true,
            tx_flushing: false,
            frame: None,
            read_buf: Default::default(),
            decoder: Default::default(),
            write_buf,
            stream_controller: stream::Controller::new(config.max_streams, config.peer_max_streams),
            streams: [Default::default(), Default::default()],
            max_stream_data: Default::default(),
            pending_accept: Default::default(),
            config,
        }
    }

    fn flush_write_buffer(&mut self, cx: &mut Context) -> Result<()> {
        if !self.tx_open {
            return if self.write_buf.is_empty() {
                Ok(())
            } else {
                Err("stream was closed with pending data".into())
            };
        }

        let mut total_len = 0;

        if self.inner.as_ref().is_write_vectored() {
            while !self.write_buf.is_empty() {
                let chunks = self.write_buf.chunks();

                match self.inner.as_mut().poll_write_vectored(cx, &chunks) {
                    Poll::Ready(result) => {
                        let len = result?;
                        self.write_buf.advance(len);
                        total_len += len;
                    }
                    Poll::Pending => {
                        // we don't need to flush since this call does that
                        self.tx_flushing = false;
                    }
                }
            }
        } else {
            while let Some(mut chunk) = self.write_buf.pop_front() {
                match self.inner.as_mut().poll_write(cx, &chunk) {
                    Poll::Ready(result) => {
                        let len = result?;
                        chunk.advance(len);
                        total_len += len;

                        if !chunk.is_empty() {
                            self.write_buf.push_front(chunk);
                        }

                        if len == 0 {
                            return if self.write_buf.is_empty() {
                                self.tx_open = false;
                                Ok(())
                            } else {
                                Err("stream was closed with pending data".into())
                            };
                        }
                    }
                    Poll::Pending => {
                        self.write_buf.push_front(chunk);
                        // we don't need to flush since this call does that
                        self.tx_flushing = false;
                        break;
                    }
                }
            }
        }

        if total_len > 0 {
            self.write_buf.notify(cx);

            // if we don't have anything buffered, tell the socket
            // to flush to the operating system
            self.tx_flushing = self.write_buf.is_empty();
        }

        if self.tx_flushing {
            if let Poll::Ready(res) = self.inner.as_mut().poll_flush(cx) {
                self.tx_flushing = false;
                res?;
            }
        }

        Ok(())
    }

    fn flush_read_buffer(&mut self, cx: &mut Context) -> Result<()> {
        loop {
            if self.frame.is_none() {
                if let Some(frame) = self.decoder.decode(&mut self.read_buf)? {
                    self.frame = Some(frame);
                }
            }

            match self.dispatch_frame(cx) {
                Poll::Ready(result) => result?,
                Poll::Pending => return Ok(()),
            }
        }
    }

    fn fill_read_buffer(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        // don't fill the buffer if we have a pending frame
        if self.frame.is_some() {
            return Poll::Pending;
        }

        // the socket is closed
        if !self.rx_open {
            return Ok(()).into();
        }

        let rx_open = &mut self.rx_open;
        let inner = self.inner.as_mut();

        ready!(self.read_buf.read(|buf| {
            ready!(inner.poll_read(cx, buf))?;

            // the socket returned a 0 write
            if buf.filled().is_empty() {
                *rx_open = false;
            }

            Ok(()).into()
        }))?;

        Ok(()).into()
    }

    fn dispatch_frame(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        use frame::Frame::*;

        match self.frame.take() {
            Some(StreamOpen { id, bidirectional }) => {
                // TODO make sure the peer hasn't opened too many
                let mut stream = Stream {
                    rx: Some(ReceiveStream::new(self.config.stream_window)),
                    tx: None,
                };
                if bidirectional {
                    stream.tx = Some(SendStream::new(self.config.peer_max_stream_data));
                };
                self.streams[Owner::Remote].insert(id, stream);
                self.pending_accept.push_back(id);
                cx.waker().wake_by_ref();
            }
            Some(StreamData {
                id,
                owner,
                mut data,
            }) => {
                if let Some(rx) = self.streams[owner]
                    .get_mut(&id)
                    .and_then(|stream| stream.rx.as_mut())
                {
                    let len = data.len() as u64;
                    let len = rx.buffer(len)?;
                    data.advance(len as _);

                    if !data.is_empty() {
                        self.frame = Some(StreamData { id, owner, data });
                        return Poll::Pending;
                    }
                }
            }
            Some(MaxStreams { up_to }) => {
                self.stream_controller.max_streams(up_to);
            }
            Some(MaxStreamData { id, owner, up_to }) => {
                if let Some(tx) = self.streams[owner]
                    .get_mut(&id)
                    .and_then(|stream| stream.tx.as_mut())
                {
                    tx.max_data(up_to);
                }
            }
            Some(StreamFinish { id, owner }) => {
                if let Some(rx) = self.streams[owner]
                    .get_mut(&id)
                    .ok_or("invalid stream")?
                    .rx
                    .as_mut()
                {
                    rx.finish();
                }
            }
            Some(InitialMaxStreamData { up_to }) => {
                self.config.peer_max_stream_data = up_to;
            }
            None => return Poll::Pending,
        }

        Ok(()).into()
    }
}

impl<T: AsyncRead + AsyncWrite> super::Connection for Connection<T> {
    fn id(&self) -> u64 {
        self.id
    }

    fn poll_open_bidirectional_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>> {
        ready!(self.stream_controller.open(cx));

        self.write_buf.push(Frame::StreamOpen {
            id,
            bidirectional: true,
        });

        let stream = Stream {
            rx: Some(ReceiveStream::new(self.config.stream_window)),
            tx: Some(SendStream::new(self.config.peer_max_stream_data)),
        };
        self.streams[Owner::Local].insert(id, stream);

        Ok(()).into()
    }

    fn poll_open_send_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>> {
        ready!(self.stream_controller.open(cx));

        self.write_buf.push(Frame::StreamOpen {
            id,
            bidirectional: false,
        });

        let stream = Stream {
            rx: None,
            tx: Some(SendStream::new(self.config.peer_max_stream_data)),
        };
        self.streams[Owner::Local].insert(id, stream);

        Ok(()).into()
    }

    fn poll_accept_stream(&mut self, _cx: &mut Context) -> Poll<Result<Option<u64>>> {
        if !self.rx_open {
            return Ok(None).into();
        }

        if let Some(id) = self.pending_accept.pop_front() {
            Ok(Some(id)).into()
        } else {
            Poll::Pending
        }
    }

    fn poll_send(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        if !self.write_buf.request_push(self.config.max_write_queue_len) {
            return Poll::Pending;
        }

        let stream = self.streams[owner]
            .get_mut(&id)
            .ok_or("missing stream")?
            .tx
            .as_mut()
            .ok_or("missing tx stream")?;

        let allowed_bytes = bytes.min(self.config.max_stream_frame_len as _);

        if let Some(data) = stream.send(allowed_bytes, cx) {
            let len = data.len() as u64;
            self.write_buf.push(frame::Frame::StreamData {
                id,
                owner: !owner,
                data,
            });
            Ok(len).into()
        } else {
            Poll::Pending
        }
    }

    fn poll_receive(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        let stream = self.streams[owner]
            .get_mut(&id)
            .ok_or("missing stream")?
            .rx
            .as_mut()
            .ok_or("missing rx stream")?;

        let len = ready!(stream.receive(bytes, cx))?;

        if stream.receive_window() < self.config.stream_window / 2 {
            self.max_stream_data.insert((owner, id));
        }

        Ok(len).into()
    }

    fn poll_send_finish(&mut self, owner: Owner, id: u64, _cx: &mut Context) -> Poll<Result<()>> {
        if !self.write_buf.request_push(self.config.max_write_queue_len) {
            return Poll::Pending;
        }

        if let Entry::Occupied(mut entry) = self.streams[owner].entry(id) {
            let stream = entry.get_mut();

            if stream.tx.take().is_some() {
                self.write_buf
                    .push(frame::Frame::StreamFinish { id, owner: !owner });
            }

            if stream.rx.is_none() {
                entry.remove();
                self.stream_controller.close();
            }
        }

        Ok(()).into()
    }

    fn poll_receive_finish(&mut self, owner: Owner, id: u64, cx: &mut Context) -> Poll<Result<()>> {
        if let Entry::Occupied(mut entry) = self.streams[owner].entry(id) {
            let stream = entry.get_mut();

            if let Some(rx) = stream.rx.as_mut() {
                ready!(rx.poll_finish(cx));
            }
            stream.rx = None;

            if stream.tx.is_none() {
                entry.remove();
                self.stream_controller.close();
            }
        }

        Ok(()).into()
    }

    fn poll_progress(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        loop {
            if let Some(up_to) = self.stream_controller.transmit() {
                self.write_buf.push_priority(Frame::MaxStreams { up_to });
            }

            for (owner, id) in self.max_stream_data.drain() {
                if let Some(stream) = self.streams[owner].get_mut(&id).and_then(|s| s.rx.as_mut()) {
                    let up_to = stream.credit(self.config.stream_window);
                    self.write_buf.push_priority(Frame::MaxStreamData {
                        id,
                        owner: !owner,
                        up_to,
                    });
                }
            }

            self.flush_write_buffer(cx)?;
            self.flush_read_buffer(cx)?;
            ready!(self.fill_read_buffer(cx))?;

            // the connection is done
            if !self.rx_open && self.frame.is_none() {
                return Ok(()).into();
            }
        }
    }

    fn poll_finish(&mut self, cx: &mut Context) -> Poll<Result<()>> {
        if self.tx_open {
            self.flush_write_buffer(cx)?;

            // wait to shutdown the socket until we have written everything
            if !self.write_buf.is_empty() {
                return Poll::Pending;
            }

            // notify the peer we're not writing anything anymore
            // https://docs.rs/tokio/latest/tokio/io/trait.AsyncWrite.html#tymethod.poll_shutdown
            // > Invocation of a shutdown implies an invocation of flush. Once
            // > this method returns Ready it implies that a flush successfully
            // > happened before the shutdown happened. That is, callers donâ€™t need
            // > to call flush before calling shutdown. They can rely that by
            // > calling shutdown any pending buffered data will be written out.
            ready!(self.inner.as_mut().poll_shutdown(cx))?;
            self.tx_open = false;
        }

        loop {
            // work to read all of the remaining data
            self.flush_read_buffer(cx)?;
            ready!(self.fill_read_buffer(cx))?;

            // the connection is done
            if !self.rx_open && self.frame.is_none() {
                return Ok(()).into();
            }
        }
    }
}

#[derive(Debug, Default)]
struct Blocked(Option<core::task::Waker>);

impl Blocked {
    pub fn block(&mut self, cx: &mut Context) {
        self.0 = Some(cx.waker().clone());
    }

    pub fn unblock(&mut self) {
        if let Some(waker) = self.0.take() {
            waker.wake();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{scenario::Scenario, testing, timer, trace::MemoryLogger, units::*, Driver};
    use futures_test::task::new_count_waker;
    use insta::assert_display_snapshot;
    use std::collections::HashSet;

    fn test(config: Config, scenario: &Scenario) -> (MemoryLogger, MemoryLogger) {
        let traces = &scenario.traces;

        let (client, server) = testing::Connection::pair(10000);

        let mut client = {
            let scenario = &scenario.clients[0].connections[0];
            let conn = Box::pin(client);
            let conn = super::Connection::new(0, conn, config.clone());
            Driver::new(scenario, conn)
        };
        let mut client_trace = MemoryLogger::new(traces.clone());
        let mut client_checkpoints = HashSet::new();
        let mut client_timer = timer::Testing::default();

        let mut server = {
            let scenario = &scenario.servers[0].connections[0];
            let conn = Box::pin(server);
            let conn = super::Connection::new(1, conn, config);
            Driver::new(scenario, conn)
        };
        let mut server_trace = MemoryLogger::new(traces.clone());
        let mut server_checkpoints = HashSet::new();
        let mut server_timer = timer::Testing::default();

        let (waker, count) = new_count_waker();
        let mut prev_count = 0;
        let mut cx = core::task::Context::from_waker(&waker);

        loop {
            let c = client.poll_with_timer(
                &mut client_trace,
                &mut client_checkpoints,
                &mut client_timer,
                &mut cx,
            );
            let s = server.poll_with_timer(
                &mut server_trace,
                &mut server_checkpoints,
                &mut server_timer,
                &mut cx,
            );

            match (c, s) {
                (Poll::Ready(Ok(())), Poll::Ready(Ok(()))) => break,
                (Poll::Ready(Err(e)), _) | (_, Poll::Ready(Err(e))) => panic!("{e}"),
                _ => {
                    let current_count = count.get();
                    if current_count > prev_count {
                        prev_count = current_count;
                        continue;
                    }

                    if client_timer.advance_pair(&mut server_timer).is_none() {
                        eprintln!("the timer did not advance!");
                        eprintln!("server trace:");
                        eprintln!("{}", server_trace.as_str().unwrap());
                        eprintln!("{server:#?}");
                        eprintln!("client trace:");
                        eprintln!("{}", client_trace.as_str().unwrap());
                        eprintln!("{client:#?}");
                        panic!("test is deadlocked");
                    }
                }
            }
        }

        (client_trace, server_trace)
    }

    macro_rules! test {
        ($name:ident, $config:expr, $builder:expr) => {
            #[test]
            fn $name() -> crate::Result<()> {
                let scenario = Scenario::build(|scenario| {
                    let server = scenario.create_server();

                    scenario.create_client(|client| {
                        client.connect_to(server, $builder);
                    });
                });

                let (client_trace, server_trace) = test($config, &scenario);

                assert_display_snapshot!(
                    concat!(stringify!($name), "__client"),
                    client_trace.as_str().unwrap()
                );
                assert_display_snapshot!(
                    concat!(stringify!($name), "__server"),
                    server_trace.as_str().unwrap()
                );

                Ok(())
            }
        };
    }

    test!(single_stream, Config::default(), |conn| {
        conn.open_send_stream(
            |local| {
                local.send(1.kilobytes());
            },
            |remote| {
                remote.receive(1.kilobytes());
            },
        );
    });

    test!(single_slow_send_stream, Config::default(), |conn| {
        conn.open_send_stream(
            |local| {
                local.set_send_rate(100.bytes() / 50.millis());
                local.send(1.kilobytes());
            },
            |remote| {
                remote.receive(1.kilobytes());
            },
        );
    });

    test!(single_slow_recv_stream, Config::default(), |conn| {
        conn.open_send_stream(
            |local| {
                local.send(1.kilobytes());
            },
            |remote| {
                remote.set_receive_rate(100.bytes() / 50.millis());
                remote.receive(1.kilobytes());
            },
        );
    });

    test!(
        low_stream_window,
        Config {
            stream_window: 50,
            ..Config::default()
        },
        |conn| {
            conn.open_send_stream(
                |local| {
                    local.set_send_rate(100.bytes() / 50.millis());
                    local.send(1.kilobytes());
                },
                |remote| {
                    remote.receive(1.kilobytes());
                },
            );
        }
    );

    test!(
        low_max_streams,
        Config {
            max_streams: 2,
            peer_max_streams: 2,
            ..Config::default()
        },
        |conn| {
            conn.scope(|scope| {
                for _ in 0..4 {
                    scope.spawn(|conn| {
                        conn.open_send_stream(
                            |local| {
                                local.set_send_rate(500.bytes() / 50.millis());
                                local.send(1.kilobytes());
                            },
                            |remote| {
                                remote.receive(1.kilobytes());
                            },
                        );
                    });
                }
            });
        }
    );

    test!(multiple_streams, Config::default(), |conn| {
        conn.scope(|scope| {
            for _ in 0..3 {
                scope.spawn(|conn| {
                    conn.open_send_stream(
                        |local| {
                            local.set_send_rate(100.bytes() / 50.millis());
                            local.send(1.kilobytes());
                        },
                        |remote| {
                            remote.receive(1.kilobytes());
                        },
                    );
                });
            }
        });
    });

    test!(checkpoint, Config::default(), |conn| {
        let (park, unpark) = conn.checkpoint();
        conn.scope(|scope| {
            scope.spawn(|conn| {
                conn.open_send_stream(
                    |local| {
                        local.set_send_rate(100.bytes() / 50.millis());
                        local.send(1.kilobytes());
                        local.unpark(unpark);
                        local.send(1.kilobytes());
                    },
                    |remote| {
                        remote.receive(2.kilobytes());
                    },
                );
            });
            scope.spawn(|conn| {
                conn.open_send_stream(
                    |local| {
                        local.park(park);
                        local.set_send_rate(100.bytes() / 50.millis());
                        local.send(1.kilobytes());
                    },
                    |remote| {
                        remote.receive(1.kilobytes());
                    },
                );
            });
        });
    });
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::units::{duration_format, Byte, Duration, Rate};
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum Connection {
    /// Pause for the specified duration before processing the next op
    Sleep {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Open a bidirectional stream with an identifier
    OpenBidirectionalStream { stream_id: u64 },
    /// Open a unidirectional stream with an identifier
    OpenSendStream { stream_id: u64 },
    /// Send a specific amount of data over the stream id
    Send { stream_id: u64, bytes: Byte },
    /// Finish sending data on the stream
    SendFinish { stream_id: u64 },
    /// Sets the send rate for a stream
    SendRate {
        stream_id: u64,
        #[serde(flatten)]
        rate: Rate,
    },
    /// Send a specific amount of data over the stream id
    Receive { stream_id: u64, bytes: Byte },
    /// Receives all of the data on the stream until it is finished
    ReceiveAll { stream_id: u64 },
    /// Finish receiving data on the stream
    ReceiveFinish { stream_id: u64 },
    /// Sets the receive rate for a stream
    ReceiveRate {
        stream_id: u64,
        #[serde(flatten)]
        rate: Rate,
    },
    /// Parks the current thread and waits for the checkpoint to be unparked
    Park { checkpoint: u64 },
    /// Notifies the parked checkpoint that it can continue
    Unpark { checkpoint: u64 },
    /// Emit a trace event
    Trace { trace_id: u64 },
    /// Profiles the time it takes to perform the contained operations
    Profile {
        trace_id: u64,
        operations: Vec<Connection>,
    },
    /// Evaluates the contained operations for the specified value
    Iterate {
        value: IterateValue,
        operations: Vec<Connection>,
        #[serde(skip_serializing_if = "Option::is_none")]
        trace_id: Option<u64>,
    },
    /// Perform operations concurrently
    Scope { threads: Vec<Vec<Connection>> },
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum IterateValue {
    /// Iterate for the specified duration
    Time {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Iterate for the specified count
    Count { amount: u64 },
}

impl IterateValue {
    pub(crate) fn is_zero(&self) -> bool {
        match self {
            Self::Time { amount } => *amount == Duration::ZERO,
            Self::Count { amount } => *amount == 0,
        }
    }
}

impl From<Duration> for IterateValue {
    fn from(amount: Duration) -> Self {
        Self::Time { amount }
    }
}

impl From<u64> for IterateValue {
    fn from(amount: u64) -> Self {
        Self::Count { amount }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum Client {
    /// Pause for the specified duration before processing the next op
    Sleep {
        #[serde(with = "duration_format", rename = "timeout_ms")]
        timeout: Duration,
    },
    /// Open a connection with an identifier
    Connect {
        server_id: u64,
        #[serde(skip_serializing_if = "Option::is_none", default)]
        router_id: Option<u64>,
        server_connection_id: u64,
        client_connection_id: u64,
    },
    /// Parks the current thread and waits for the checkpoint to be unparked
    Park { checkpoint: u64 },
    /// Notifies the parked checkpoint that it can continue
    Unpark { checkpoint: u64 },
    /// Emit a trace event
    Trace { trace_id: u64 },
    /// Perform operations concurrently
    Scope { threads: Vec<Vec<Client>> },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize, Hash)]
#[serde(rename_all = "snake_case")]
pub enum Router {
    /// Pause for the specified duration before processing the next op
    Sleep {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },

    /// Set the number of packets that can be buffered server->client
    ServerBufferCount { packet_count: u32 },
    /// Set the chance of a server->client packet being dropped
    ServerDropRate { packet_count: u32 },
    /// Set the chance of a server->client packet being reordered
    ServerReorderRate { packet_count: u32 },
    /// Set the chance of a server->client packet being corrupted
    ServerCorruptRate { packet_count: u32 },
    /// Set the amount of delay for server->client packets
    ServerDelay {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Set the amount of jitter for server->client packets
    ServerJitter {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Set the server->client MTU
    ServerMtu { mtu: u16 },

    /// Set the number of packets that can be buffered server->client
    ClientBufferCount { packet_count: u32 },
    /// Set the chance of a client->server packet being dropped
    ClientDropRate { packet_count: u32 },
    /// Set the chance of a client->server packet being reordered
    ClientReorderRate { packet_count: u32 },
    /// Set the chance of a client->server packet being corrupted
    ClientCorruptRate { packet_count: u32 },
    /// Set the amount of delay for client->server packets
    ClientDelay {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Set the amount of jitter for client->server packets
    ClientJitter {
        #[serde(with = "duration_format", rename = "amount_ms")]
        amount: Duration,
    },
    /// Set the client->server MTU
    ClientMtu { mtu: u16 },

    /// Set the chance of a port being rebound
    ClientRebindPortRate { packet_count: u32 },
    /// Set the chance of an IP being rebound
    ClientRebindAddressRate { packet_count: u32 },

    /// Rebinds all of the ports and/or addresses currently being used
    RebindAll { ports: bool, addresses: bool },

    /// Emit a trace event
    Trace { trace_id: u64 },
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{connection::Owner, helper::IdPrefixReader, scenario, Result};
use bytes::Bytes;
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
};
use futures::ready;
use s2n_quic::{
    connection,
    stream::{LocalStream, PeerStream, SplittableStream},
};
use s2n_quic_core::stream::testing::Data;
use std::{
    collections::{hash_map::Entry, HashMap},
    sync::Arc,
};

fn stream_error(err: s2n_quic::stream::Error) -> Result<()> {
    if let s2n_quic::stream::Error::StreamReset { error, .. } = err {
        if *error == 0 {
            return Ok(());
        }
    }

    if let s2n_quic::stream::Error::ConnectionError { error, .. } = err {
        return conn_error(error);
    }

    Err(err.into())
}

fn conn_error(err: s2n_quic::connection::Error) -> Result<()> {
    if let s2n_quic::connection::Error::Application { error, .. } = err {
        if *error == 0 {
            return Ok(());
        }
    }

    Err(err.into())
}

impl<'a> crate::client::Client<'a> for s2n_quic::Client {
    type Connect = Connect<'a>;
    type Connection = crate::Driver<'a, Connection>;

    fn connect(
        &mut self,
        addr: std::net::SocketAddr,
        server_name: &str,
        _server_conn_id: u64,
        scenario: &'a Arc<scenario::Connection>,
    ) -> Self::Connect {
        let connect = s2n_quic::client::Connect::new(addr).with_server_name(server_name);
        let attempt = s2n_quic::Client::connect(self, connect);
        Connect { attempt, scenario }
    }
}

pub struct Connect<'a> {
    attempt: s2n_quic::client::ConnectionAttempt,
    scenario: &'a scenario::Connection,
}

impl<'a> Future for Connect<'a> {
    type Output = Result<crate::Driver<'a, Connection>>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let conn = ready!(Pin::new(&mut self.attempt).poll(cx))?;
        let conn = Connection::new(conn);
        let conn = crate::Driver::new(self.scenario, conn);
        Ok(conn).into()
    }
}

pub struct Connection {
    conn: s2n_quic::Connection,
    streams: [HashMap<u64, Stream>; 2],
    opened_streams: HashMap<u64, (Bytes, LocalStream)>,
    unidentified_peer_stream: Option<(IdPrefixReader, PeerStream)>,
}

impl From<s2n_quic::Connection> for Connection {
    fn from(conn: s2n_quic::Connection) -> Self {
        Self::new(conn)
    }
}

impl Connection {
    pub fn new(connection: s2n_quic::Connection) -> Self {
        Self {
            conn: connection,
            streams: [HashMap::new(), HashMap::new()],
            opened_streams: HashMap::new(),
            unidentified_peer_stream: Default::default(),
        }
    }

    pub fn into_inner(self) -> s2n_quic::Connection {
        self.conn
    }

    fn open_local_stream<
        F: FnOnce(&mut s2n_quic::Connection, &mut Context) -> Poll<Result<S, connection::Error>>,
        S: Into<LocalStream>,
    >(
        &mut self,
        id: u64,
        open: F,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        // the stream has already been opened and is waiting to send the prefix
        if let Entry::Occupied(mut entry) = self.opened_streams.entry(id) {
            let (prefix, stream) = entry.get_mut();
            return match stream.poll_send(prefix, cx) {
                Poll::Ready(Ok(_)) => {
                    let (_, stream) = entry.remove();
                    let stream = Stream::new(stream);
                    self.streams[Owner::Local].insert(id, stream);
                    Poll::Ready(Ok(()))
                }
                Poll::Ready(Err(err)) => {
                    entry.remove();
                    Poll::Ready(stream_error(err))
                }
                Poll::Pending => Poll::Pending,
            };
        }

        let mut stream = ready!(open(&mut self.conn, cx))?.into();

        let mut prefix = Bytes::copy_from_slice(&id.to_be_bytes());

        match stream.poll_send(&mut prefix, cx) {
            Poll::Ready(Ok(_)) => {
                let stream = Stream::new(stream);
                self.streams[Owner::Local].insert(id, stream);
                Poll::Ready(Ok(()))
            }
            Poll::Ready(Err(err)) => Poll::Ready(stream_error(err)),
            Poll::Pending => {
                self.opened_streams.insert(id, (prefix, stream));
                Poll::Pending
            }
        }
    }
}

impl super::Connection for Connection {
    fn id(&self) -> u64 {
        self.conn.id()
    }

    fn poll_open_bidirectional_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>> {
        self.open_local_stream(id, |conn, cx| conn.poll_open_bidirectional_stream(cx), cx)
    }

    fn poll_open_send_stream(&mut self, id: u64, cx: &mut Context) -> Poll<Result<()>> {
        self.open_local_stream(id, |conn, cx| conn.poll_open_send_stream(cx), cx)
    }

    fn poll_accept_stream(&mut self, cx: &mut Context) -> Poll<Result<Option<u64>>> {
        loop {
            if let Some((id, stream)) = self.unidentified_peer_stream.as_mut() {
                let len = ready!(futures::io::AsyncRead::poll_read(
                    Pin::new(stream),
                    cx,
                    id.remaining()
                ))?;
                let id = ready!(id.on_read(len));

                let (_, stream) = self.unidentified_peer_stream.take().unwrap();
                let stream = Stream::new(stream);
                self.streams[Owner::Remote].insert(id, stream);
                return Poll::Ready(Ok(Some(id)));
            }

            let stream = ready!(self.conn.poll_accept(cx));

            if let Ok(Some(stream)) = stream {
                self.unidentified_peer_stream = Some((Default::default(), stream));
            } else {
                return Poll::Ready(Ok(None));
            };
        }
    }

    fn poll_send(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        self.streams[owner]
            .get_mut(&id)
            .unwrap()
            .tx
            .as_mut()
            .unwrap()
            .poll_send(bytes, cx)
    }

    fn poll_receive(
        &mut self,
        owner: Owner,
        id: u64,
        bytes: u64,
        cx: &mut Context,
    ) -> Poll<Result<u64>> {
        self.streams[owner]
            .get_mut(&id)
            .unwrap()
            .rx
            .as_mut()
            .unwrap()
            .poll_receive(bytes, cx)
    }

    fn poll_send_finish(&mut self, owner: Owner, id: u64, _cx: &mut Context) -> Poll<Result<()>> {
        if let Entry::Occupied(mut entry) = self.streams[owner].entry(id) {
            let stream = entry.get_mut();
            if let Some(mut stream) = stream.tx.take() {
                stream.inner.finish().or_else(stream_error)?;
            }

            if stream.rx.is_none() {
                entry.remove();
            }
        }

        Poll::Ready(Ok(()))
    }

    fn poll_receive_finish(
        &mut self,
        owner: Owner,
        id: u64,
        _cx: &mut Context,
    ) -> Poll<Result<()>> {
        if let Entry::Occupied(mut entry) = self.streams[owner].entry(id) {
            let stream = entry.get_mut();
            if let Some(mut stream) = stream.rx.take() {
                let _ = stream.inner.stop_sending(0u8.into());
            }

            if stream.tx.is_none() {
                entry.remove();
            }
        }

        Poll::Ready(Ok(()))
    }
}

macro_rules! chunks {
    () => {
        [
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
            Bytes::new(),
        ]
    };
}

struct Stream {
    rx: Option<ReceiveStream>,
    tx: Option<SendStream>,
}

impl Stream {
    fn new(stream: impl SplittableStream) -> Self {
        let (rx, tx) = stream.split();
        let rx = rx.map(ReceiveStream::new);
        let tx = tx.map(SendStream::new);
        Self { rx, tx }
    }
}

struct ReceiveStream {
    inner: s2n_quic::stream::ReceiveStream,
    buffered: u64,
    is_open: bool,
}

impl ReceiveStream {
    fn new(inner: s2n_quic::stream::ReceiveStream) -> Self {
        Self {
            inner,
            buffered: 0,
            is_open: true,
        }
    }

    fn poll_receive(&mut self, bytes: u64, cx: &mut Context) -> Poll<Result<u64>> {
        if !self.is_open && self.buffered == 0 {
            return Ok(0).into();
        }

        while self.buffered <= bytes && self.is_open {
            let mut chunks = chunks!();

            if let Poll::Ready(res) = self.inner.poll_receive_vectored(&mut chunks, cx) {
                let (count, is_open) = res?;
                self.is_open &= is_open;

                for chunk in &chunks[..count] {
                    self.buffered += chunk.len() as u64;
                }
            } else {
                break;
            }
        }

        let received_len = bytes.min(self.buffered);
        self.buffered -= received_len;

        if !self.is_open && received_len == 0 {
            return Ok(0).into();
        }

        if received_len == 0 {
            Poll::Pending
        } else {
            Ok(received_len).into()
        }
    }
}

struct SendStream {
    inner: s2n_quic::stream::SendStream,
    data: Data,
}

impl SendStream {
    fn new(inner: s2n_quic::stream::SendStream) -> Self {
        Self {
            inner,
            data: Data::new(u64::MAX),
        }
    }

    fn poll_send(&mut self, mut bytes: u64, cx: &mut Context) -> Poll<Result<u64>> {
        if bytes == 0 {
            return Ok(0).into();
        }

        let mut len = 0;
        let mut data = self.data;

        while bytes > 0 {
            let mut chunks = chunks!();

            let count = data.send(bytes as usize, &mut chunks).unwrap();
            let initial_len: u64 = chunks.iter().map(|chunk| chunk.len() as u64).sum();

            let count = if let Poll::Ready(count) =
                self.inner.poll_send_vectored(&mut chunks[..count], cx)?
            {
                count
            } else {
                break;
            };

            if count == chunks.len() {
                len += initial_len;
                bytes -= initial_len;
                continue;
            }

            let remaining_len: u64 = chunks[count..].iter().map(|chunk| chunk.len() as u64).sum();

            len += initial_len - remaining_len;

            break;
        }

        if len == 0 {
            return Poll::Pending;
        }

        self.data.seek_forward(len as usize);

        Poll::Ready(Ok(len))
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{operation as op, Result};
use serde::{Deserialize, Serialize};
use std::{collections::BTreeMap, path::Path, sync::Arc};

pub mod builder;
mod id;

pub use builder::Builder;
pub use id::Id;

#[derive(Clone, Debug, Default, Deserialize, Serialize, Hash)]
pub struct Scenario {
    pub id: Id,
    pub clients: Vec<Arc<Client>>,
    pub servers: Vec<Arc<Server>>,
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub routers: Vec<Arc<Router>>,
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub traces: Arc<Vec<String>>,
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub certificates: Vec<Arc<Certificate>>,
}

impl Scenario {
    pub fn build<F: FnOnce(&mut builder::Builder)>(f: F) -> Self {
        let mut builder = builder::Builder::new();
        f(&mut builder);
        builder.finish()
    }

    pub fn open(path: &Path) -> Result<Self> {
        let file = std::fs::File::open(path)?;
        let mut file = std::io::BufReader::new(file);
        let scenario = serde_json::from_reader(&mut file)?;
        Ok(scenario)
    }

    pub fn write<W: std::io::Write>(&self, out: &mut W) -> std::io::Result<()> {
        serde_json::to_writer(out, self)?;
        Ok(())
    }
}

#[derive(Clone, Debug, Default, Deserialize, Serialize, Hash)]
pub struct Client {
    #[serde(skip_serializing_if = "String::is_empty", default)]
    pub name: String,
    pub scenario: Vec<op::Client>,
    pub connections: Vec<Arc<Connection>>,
    #[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
    pub configuration: BTreeMap<String, String>,
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub certificate_authorities: Vec<u64>,
}

#[derive(Clone, Debug, Default, Deserialize, Serialize, Hash)]
pub struct Server {
    #[serde(skip_serializing_if = "String::is_empty", default)]
    pub name: String,
    pub connections: Vec<Arc<Connection>>,
    #[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
    pub configuration: BTreeMap<String, String>,
    pub private_key: u64,
    pub certificate: u64,
    pub certificate_authority: u64,
}

impl Server {
    pub fn on_server_name(&self, server_name: &str) -> Result<&Arc<Connection>> {
        let (conn_idx, _) = server_name.split_once('.').ok_or("invalid hostname")?;
        let conn_idx: usize = conn_idx.parse()?;
        let conn = self.connections.get(conn_idx).ok_or("invalid connection")?;
        Ok(conn)
    }
}

#[derive(Clone, Debug, Deserialize, Serialize, Hash)]
pub struct Connection {
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub ops: Vec<op::Connection>,
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub peer_streams: Vec<Vec<op::Connection>>,
}

#[derive(Clone, Debug, Deserialize, Serialize, Hash)]
pub struct Router {
    #[serde(skip_serializing_if = "String::is_empty", default)]
    pub name: String,
    pub scenario: Vec<op::Router>,
    #[serde(skip_serializing_if = "BTreeMap::is_empty", default)]
    pub configuration: BTreeMap<String, String>,
}

#[derive(Clone, Debug, Deserialize, Serialize, Hash)]
pub struct Certificate {
    #[serde(skip_serializing_if = "String::is_empty", default)]
    pub pem: String,

    #[serde(skip_serializing_if = "Vec::is_empty", with = "pkcs12_format", default)]
    pub pkcs12: Vec<u8>,
}

pub(crate) mod pkcs12_format {
    use serde::{self, Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        if serializer.is_human_readable() {
            let out = base64::encode(bytes);
            serializer.serialize_str(&out)
        } else {
            serializer.serialize_bytes(bytes)
        }
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
    where
        D: Deserializer<'de>,
    {
        if deserializer.is_human_readable() {
            let s = String::deserialize(deserializer)?;
            let out = base64::decode(s).map_err(serde::de::Error::custom)?;
            Ok(out)
        } else {
            Vec::deserialize(deserializer)
        }
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::units::duration_format;
use core::time::Duration;
use serde::{Deserialize, Serialize};
use std::{collections::HashMap, time::SystemTime};

pub trait Print: Serialize {
    fn print(&self) -> crate::Result<()> {
        use std::io::Write;
        let out = std::io::stdout();
        let mut out = out.lock();
        serde_json::to_writer(&mut out, self)?;
        writeln!(out)?;
        Ok(())
    }
}

impl<T: Serialize> Print for T {}

#[derive(Clone, Debug, Serialize, Deserialize)]
pub struct Initialize {
    pub pid: u64,
    pub driver: String,
    pub scenario: String,
    pub start_time: SystemTime,
    #[serde(default, skip_serializing_if = "is_default")]
    pub traces: Vec<String>,
}

impl Default for Initialize {
    fn default() -> Self {
        Self {
            pid: 0,
            driver: Default::default(),
            scenario: Default::default(),
            start_time: std::time::SystemTime::now(),
            traces: vec![],
        }
    }
}

#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Stats {
    #[serde(rename = "t", with = "duration_format")]
    pub time: Duration,
    #[serde(default, skip_serializing_if = "is_default")]
    pub cpu: f32,
    #[serde(default, skip_serializing_if = "is_default")]
    pub cycles: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub instructions: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub branches: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub context_switches: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub memory: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub virtual_memory: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub syscalls: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub connections: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub accept: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub allocs: Stat,
    #[serde(default, skip_serializing_if = "is_default")]
    pub reallocs: Stat,
    #[serde(default, skip_serializing_if = "is_default")]
    pub deallocs: Stat,
    #[serde(default, skip_serializing_if = "is_default")]
    pub connect_time: Stat,
    #[serde(default, skip_serializing_if = "is_default")]
    pub send: HashMap<StreamId, Stat>,
    #[serde(default, skip_serializing_if = "is_default")]
    pub receive: HashMap<StreamId, Stat>,
    #[serde(default, skip_serializing_if = "is_default")]
    pub profiles: HashMap<u64, Histogram>,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Stat {
    #[serde(default, skip_serializing_if = "is_default")]
    pub count: u64,
    #[serde(default, skip_serializing_if = "is_default")]
    pub total: u64,
}

impl Stat {
    pub fn average(&self) -> f64 {
        self.total as f64 / self.count as f64
    }
}

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Histogram {
    #[serde(default, skip_serializing_if = "is_default")]
    pub stat: Stat,
    #[serde(default, skip_serializing_if = "is_default")]
    pub buckets: Vec<Bucket>,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct Bucket {
    pub lower: u64,
    pub upper: u64,
    pub count: u64,
}

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct StreamId {
    pub connection_id: u64,
    pub id: u64,
}

impl<'de> Deserialize<'de> for StreamId {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: serde::de::Deserializer<'de>,
    {
        let s = String::deserialize(deserializer)?;
        let (connection_id, id) = s.split_once(':').unwrap();
        let connection_id = connection_id.parse().unwrap();
        let id = id.parse().unwrap();
        Ok(Self { connection_id, id })
    }
}

impl Serialize for StreamId {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: serde::Serializer,
    {
        use std::io::{Cursor, Write};

        let mut out = [0u8; 1024];
        let mut f = Cursor::new(&mut out[..]);
        write!(f, "{}:{}", self.connection_id, self.id).unwrap();
        let len = f.position() as usize;

        let out = &out[..len];
        let out = unsafe { core::str::from_utf8_unchecked(out) };

        serializer.serialize_str(out)
    }
}

#[inline(always)]
fn is_default<T: Default + PartialEq>(v: &T) -> bool {
    &T::default() == v
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::{
    io,
    pin::Pin,
    sync::{Arc, Mutex},
    task::{Context, Poll, Waker},
};
use tokio::io::{AsyncRead, AsyncWrite};

/// Simple testing connection with a single bidirectional stream
#[derive(Debug)]
pub struct Connection {
    local: Arc<Mutex<Stream>>,
    peer: Arc<Mutex<Stream>>,
}

impl Connection {
    /// Creates a pair of connections with the specified buffer limit
    pub fn pair(max_buffer: usize) -> (Connection, Connection) {
        let client: Arc<Mutex<Stream>> = Arc::new(Mutex::new(Stream::new(max_buffer)));
        let server: Arc<Mutex<Stream>> = Arc::new(Mutex::new(Stream::new(max_buffer)));
        let client_conn = Connection {
            local: client.clone(),
            peer: server.clone(),
        };
        let server_conn = Connection {
            local: server,
            peer: client,
        };
        (client_conn, server_conn)
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        if let Ok(mut stream) = self.local.lock() {
            stream.close();
        }
        if let Ok(mut stream) = self.peer.lock() {
            stream.close();
        }
    }
}

impl AsyncRead for Connection {
    fn poll_read(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        let mut stream = self.local.lock().unwrap();
        Pin::new(&mut *stream).poll_read(cx, buf)
    }
}

impl AsyncWrite for Connection {
    fn poll_write(
        self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        let mut stream = self.peer.lock().unwrap();
        Pin::new(&mut *stream).poll_write(cx, buf)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        let mut stream = self.peer.lock().unwrap();
        Pin::new(&mut *stream).poll_flush(cx)
    }

    fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        let mut stream = self.peer.lock().unwrap();
        Pin::new(&mut *stream).poll_shutdown(cx)
    }
}

#[derive(Debug, Default)]
struct Stream {
    buffer: Vec<u8>,
    read_waker: Option<Waker>,
    write_waker: Option<Waker>,
    is_closed: bool,
    max_buffer: usize,
}

impl Stream {
    fn new(max_buffer: usize) -> Self {
        Self {
            max_buffer,
            ..Default::default()
        }
    }

    fn close(&mut self) {
        self.is_closed = true;
        self.wake_writer();
        self.wake_reader();
    }

    fn wake_reader(&mut self) {
        if let Some(waker) = self.read_waker.take() {
            waker.wake();
        }
    }

    fn wake_writer(&mut self) {
        if let Some(waker) = self.write_waker.take() {
            waker.wake();
        }
    }

    fn ensure_open(&self) -> io::Result<()> {
        if self.is_closed {
            Err(io::Error::new(io::ErrorKind::ConnectionReset, "closed"))
        } else {
            Ok(())
        }
    }
}

impl AsyncRead for Stream {
    fn poll_read(
        mut self: Pin<&mut Self>,
        context: &mut Context<'_>,
        buf: &mut tokio::io::ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        if self.buffer.is_empty() {
            // do an empty write
            if self.is_closed {
                return Ok(()).into();
            }

            self.read_waker = Some(context.waker().clone());
            return Poll::Pending;
        }

        let len = self.buffer.len().min(buf.remaining());
        buf.put_slice(&self.buffer[..len]);
        self.buffer.drain(..len);

        // let the peer know it can write more
        self.wake_writer();

        Poll::Ready(Ok(()))
    }
}

impl AsyncWrite for Stream {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<Result<usize, io::Error>> {
        self.ensure_open()?;

        let len = (self.max_buffer - self.buffer.len()).min(buf.len());

        if len == 0 {
            self.write_waker = Some(cx.waker().clone());
            return Poll::Pending;
        }

        self.buffer.extend_from_slice(&buf[..len]);

        // let the peer know it can read more
        self.wake_reader();

        Poll::Ready(Ok(len))
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
        self.ensure_open()?;

        if self.buffer.is_empty() {
            Poll::Ready(Ok(()))
        } else {
            self.write_waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }

    fn poll_shutdown(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<(), io::Error>> {
        if !self.buffer.is_empty() {
            self.write_waker = Some(cx.waker().clone());
            return Poll::Pending;
        }

        self.is_closed = true;

        // let the peer know we're closed now
        self.wake_reader();

        Poll::Ready(Ok(()))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::io::{AsyncReadExt, AsyncWriteExt};

    #[tokio::test]
    async fn read_write_shutdown() -> io::Result<()> {
        let (mut client, mut server) = Connection::pair(100);

        tokio::spawn(async move {
            client.write_all(b"hello!").await?;
            client.shutdown().await?;
            io::Result::Ok(())
        });

        let mut buf = vec![];
        server.read_to_end(&mut buf).await?;

        assert_eq!(buf, b"hello!");

        Ok(())
    }

    #[tokio::test]
    async fn read_write_drop() -> io::Result<()> {
        let (mut client, mut server) = Connection::pair(100);

        tokio::spawn(async move {
            client.write_all(b"hello!").await?;
            // let the drop handle the connection shutdown
            io::Result::Ok(())
        });

        let mut buf = vec![];
        server.read_to_end(&mut buf).await?;

        assert_eq!(buf, b"hello!");

        Ok(())
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use core::task::{Context, Poll};
pub use s2n_quic_core::time::{
    timer::{Provider, Query, Result},
    Timestamp,
};

pub trait Timer {
    fn now(&self) -> Timestamp;
    fn poll(&mut self, target: Timestamp, cx: &mut Context) -> Poll<()>;
}

#[cfg(any(test, feature = "testing"))]
mod testing;
mod tokio;

pub use self::tokio::Timer as Tokio;
#[cfg(any(test, feature = "testing"))]
pub use testing::Timer as Testing;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use crate::{operation as op, timer::Timestamp, units::*};
use core::fmt;
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc,
};

mod usdt;
pub use self::usdt::Usdt;

pub trait Trace {
    #[inline(always)]
    fn enter_connection(&mut self, id: u64) {
        let _ = id;
    }

    #[inline(always)]
    fn exec(&mut self, now: Timestamp, op: &op::Connection) {
        let _ = now;
        let _ = op;
    }

    #[inline(always)]
    fn exec_client(&mut self, now: Timestamp, op: &op::Client) {
        let _ = now;
        let _ = op;
    }

    #[inline(always)]
    fn enter(&mut self, now: Timestamp, scope: u64, thread: usize) {
        let _ = now;
        let _ = scope;
        let _ = thread;
    }

    #[inline(always)]
    fn exit(&mut self, now: Timestamp) {
        let _ = now;
    }

    #[inline(always)]
    fn send(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        let _ = now;
        let _ = stream_id;
        let _ = len;
    }

    #[inline(always)]
    fn send_finish(&mut self, now: Timestamp, stream_id: u64) {
        let _ = now;
        let _ = stream_id;
    }

    #[inline(always)]
    fn receive(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        let _ = now;
        let _ = stream_id;
        let _ = len;
    }

    #[inline(always)]
    fn receive_finish(&mut self, now: Timestamp, stream_id: u64) {
        let _ = now;
        let _ = stream_id;
    }

    #[inline(always)]
    fn accept(&mut self, now: Timestamp, stream_id: u64) {
        let _ = now;
        let _ = stream_id;
    }

    #[inline(always)]
    fn open(&mut self, now: Timestamp, stream_id: u64) {
        let _ = now;
        let _ = stream_id;
    }

    #[inline(always)]
    fn trace(&mut self, now: Timestamp, id: u64) {
        let _ = now;
        let _ = id;
    }

    #[inline(always)]
    fn profile(&mut self, now: Timestamp, id: u64, time: Duration) {
        let _ = now;
        let _ = id;
        let _ = time;
    }

    #[inline(always)]
    fn park(&mut self, now: Timestamp, id: u64) {
        let _ = now;
        let _ = id;
    }

    #[inline(always)]
    fn unpark(&mut self, now: Timestamp, id: u64) {
        let _ = now;
        let _ = id;
    }

    #[inline(always)]
    fn connect(&mut self, now: Timestamp, connection_id: u64, time: Duration) {
        let _ = now;
        let _ = connection_id;
        let _ = time;
    }
}

impl<A: Trace, B: Trace> Trace for (A, B) {
    #[inline(always)]
    fn enter_connection(&mut self, id: u64) {
        self.0.enter_connection(id);
        self.1.enter_connection(id);
    }

    #[inline(always)]
    fn exec(&mut self, now: Timestamp, op: &op::Connection) {
        self.0.exec(now, op);
        self.1.exec(now, op);
    }

    #[inline(always)]
    fn exec_client(&mut self, now: Timestamp, op: &op::Client) {
        self.0.exec_client(now, op);
        self.1.exec_client(now, op);
    }

    #[inline(always)]
    fn enter(&mut self, now: Timestamp, scope: u64, thread: usize) {
        self.0.enter(now, scope, thread);
        self.1.enter(now, scope, thread);
    }

    #[inline(always)]
    fn exit(&mut self, now: Timestamp) {
        self.0.exit(now);
        self.1.exit(now);
    }

    #[inline(always)]
    fn send(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        self.0.send(now, stream_id, len);
        self.1.send(now, stream_id, len);
    }

    #[inline(always)]
    fn send_finish(&mut self, now: Timestamp, stream_id: u64) {
        self.0.send_finish(now, stream_id);
        self.1.send_finish(now, stream_id);
    }

    #[inline(always)]
    fn receive(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        self.0.receive(now, stream_id, len);
        self.1.receive(now, stream_id, len);
    }

    #[inline(always)]
    fn receive_finish(&mut self, now: Timestamp, stream_id: u64) {
        self.0.receive_finish(now, stream_id);
        self.1.receive_finish(now, stream_id);
    }

    #[inline(always)]
    fn accept(&mut self, now: Timestamp, stream_id: u64) {
        self.0.accept(now, stream_id);
        self.1.accept(now, stream_id)
    }

    #[inline(always)]
    fn open(&mut self, now: Timestamp, stream_id: u64) {
        self.0.open(now, stream_id);
        self.1.open(now, stream_id)
    }

    #[inline(always)]
    fn trace(&mut self, now: Timestamp, id: u64) {
        self.0.trace(now, id);
        self.1.trace(now, id)
    }

    #[inline(always)]
    fn profile(&mut self, now: Timestamp, id: u64, time: Duration) {
        self.0.profile(now, id, time);
        self.1.profile(now, id, time)
    }

    #[inline]
    fn park(&mut self, now: Timestamp, id: u64) {
        self.0.park(now, id);
        self.1.park(now, id);
    }

    #[inline]
    fn unpark(&mut self, now: Timestamp, id: u64) {
        self.0.unpark(now, id);
        self.1.unpark(now, id);
    }

    #[inline(always)]
    fn connect(&mut self, now: Timestamp, connection_id: u64, time: Duration) {
        self.0.connect(now, connection_id, time);
        self.1.connect(now, connection_id, time);
    }
}

impl<T: Trace> Trace for Option<T> {
    #[inline(always)]
    fn enter_connection(&mut self, id: u64) {
        if let Some(t) = self.as_mut() {
            t.enter_connection(id);
        }
    }

    #[inline(always)]
    fn exec(&mut self, now: Timestamp, op: &op::Connection) {
        if let Some(t) = self.as_mut() {
            t.exec(now, op);
        }
    }

    #[inline(always)]
    fn exec_client(&mut self, now: Timestamp, op: &op::Client) {
        if let Some(t) = self.as_mut() {
            t.exec_client(now, op);
        }
    }

    #[inline(always)]
    fn enter(&mut self, now: Timestamp, scope: u64, thread: usize) {
        if let Some(t) = self.as_mut() {
            t.enter(now, scope, thread);
        }
    }

    #[inline(always)]
    fn exit(&mut self, now: Timestamp) {
        if let Some(t) = self.as_mut() {
            t.exit(now);
        }
    }

    #[inline(always)]
    fn send(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        if let Some(t) = self.as_mut() {
            t.send(now, stream_id, len);
        }
    }

    #[inline(always)]
    fn send_finish(&mut self, now: Timestamp, stream_id: u64) {
        if let Some(t) = self.as_mut() {
            t.send_finish(now, stream_id);
        }
    }

    #[inline(always)]
    fn receive(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        if let Some(t) = self.as_mut() {
            t.receive(now, stream_id, len);
        }
    }

    #[inline(always)]
    fn receive_finish(&mut self, now: Timestamp, stream_id: u64) {
        if let Some(t) = self.as_mut() {
            t.receive_finish(now, stream_id);
        }
    }

    #[inline(always)]
    fn accept(&mut self, now: Timestamp, stream_id: u64) {
        if let Some(t) = self.as_mut() {
            t.accept(now, stream_id);
        }
    }

    #[inline(always)]
    fn open(&mut self, now: Timestamp, stream_id: u64) {
        if let Some(t) = self.as_mut() {
            t.open(now, stream_id);
        }
    }

    #[inline(always)]
    fn trace(&mut self, now: Timestamp, id: u64) {
        if let Some(t) = self.as_mut() {
            t.trace(now, id);
        }
    }

    #[inline(always)]
    fn profile(&mut self, now: Timestamp, id: u64, time: Duration) {
        if let Some(t) = self.as_mut() {
            t.profile(now, id, time);
        }
    }

    #[inline]
    fn park(&mut self, now: Timestamp, id: u64) {
        if let Some(t) = self.as_mut() {
            t.park(now, id);
        }
    }

    #[inline]
    fn unpark(&mut self, now: Timestamp, id: u64) {
        if let Some(t) = self.as_mut() {
            t.unpark(now, id);
        }
    }

    #[inline]
    fn connect(&mut self, now: Timestamp, connection_id: u64, time: Duration) {
        if let Some(t) = self.as_mut() {
            t.connect(now, connection_id, time);
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct Disabled(());

impl Trace for Disabled {}

#[derive(Debug)]
pub struct Logger<O: Output> {
    id: u64,
    traces: Arc<Vec<String>>,
    scope: Vec<(u64, usize)>,
    output: O,
    verbose: bool,
}

pub type MemoryLogger = Logger<std::io::Cursor<Vec<u8>>>;
pub type StdioLogger = Logger<std::io::BufWriter<std::io::Stdout>>;

impl<O: Output> Logger<O> {
    pub fn new(traces: Arc<Vec<String>>) -> Self {
        Self {
            id: 0,
            traces,
            scope: vec![],
            output: O::new(),
            verbose: false,
        }
    }

    pub fn verbose(&mut self, enabled: bool) {
        self.verbose = enabled;
    }

    fn log(&mut self, now: Timestamp, v: impl fmt::Display) {
        self.output.log(self.id, &self.scope, self.verbose, now, v)
    }
}

impl MemoryLogger {
    pub fn as_str(&self) -> Option<&str> {
        core::str::from_utf8(self.output.get_ref()).ok()
    }
}

impl Clone for StdioLogger {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            traces: self.traces.clone(),
            scope: vec![],
            output: Output::new(),
            verbose: self.verbose,
        }
    }
}

pub trait Output {
    type Io: std::io::Write;

    fn new() -> Self;

    fn write<F: FnOnce(&mut Self::Io) -> std::io::Result<()>>(
        &mut self,
        f: F,
    ) -> std::io::Result<()>;

    fn log(
        &mut self,
        id: u64,
        scope: &[(u64, usize)],
        verbose: bool,
        now: Timestamp,
        v: impl fmt::Display,
    ) {
        let _ = self.write(|out| {
            use std::io::Write;
            write!(out, "{now} ")?;

            write!(out, "[{id}")?;
            if verbose {
                for (scope, thread) in scope.iter() {
                    write!(out, ":{scope}.{thread}")?;
                }
            }
            write!(out, "] ")?;

            writeln!(out, "{v}")?;
            Ok(())
        });
    }
}

impl Output for std::io::BufWriter<std::io::Stdout> {
    type Io = Self;

    fn new() -> Self {
        std::io::BufWriter::new(std::io::stdout())
    }

    fn write<F: FnOnce(&mut Self::Io) -> std::io::Result<()>>(
        &mut self,
        f: F,
    ) -> std::io::Result<()> {
        f(self)?;
        use std::io::Write;
        self.flush()?;
        Ok(())
    }
}

impl Output for std::io::Cursor<Vec<u8>> {
    type Io = Self;

    fn new() -> Self {
        std::io::Cursor::new(vec![])
    }

    fn write<F: FnOnce(&mut Self::Io) -> std::io::Result<()>>(
        &mut self,
        f: F,
    ) -> std::io::Result<()> {
        f(self)?;
        Ok(())
    }
}

impl<O: Output> Trace for Logger<O> {
    #[inline(always)]
    fn enter_connection(&mut self, id: u64) {
        self.id = id;
    }

    #[inline(always)]
    fn exec(&mut self, now: Timestamp, op: &op::Connection) {
        if self.verbose {
            self.log(now, format_args!("exec: {op:?}"));
        }
    }

    #[inline(always)]
    fn exec_client(&mut self, now: Timestamp, op: &op::Client) {
        if self.verbose {
            self.log(now, format_args!("exec: {op:?}"));
        }
    }

    #[inline(always)]
    fn enter(&mut self, _now: Timestamp, scope: u64, thread: usize) {
        self.scope.push((scope, thread));
    }

    #[inline(always)]
    fn exit(&mut self, _now: Timestamp) {
        self.scope.pop();
    }

    #[inline(always)]
    fn send(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        self.log(now, format_args!("send[{stream_id}]={len}"));
    }

    #[inline(always)]
    fn send_finish(&mut self, now: Timestamp, stream_id: u64) {
        self.log(now, format_args!("sfin[{stream_id}]"));
    }

    #[inline(always)]
    fn receive(&mut self, now: Timestamp, stream_id: u64, len: u64) {
        self.log(now, format_args!("recv[{stream_id}]={len}"));
    }

    #[inline(always)]
    fn receive_finish(&mut self, now: Timestamp, stream_id: u64) {
        self.log(now, format_args!("rfin[{stream_id}]"));
    }

    #[inline(always)]
    fn accept(&mut self, now: Timestamp, stream_id: u64) {
        self.log(now, format_args!("acpt[{stream_id}]"));
    }

    #[inline(always)]
    fn open(&mut self, now: Timestamp, stream_id: u64) {
        self.log(now, format_args!("open[{stream_id}]"));
    }

    #[inline(always)]
    fn trace(&mut self, now: Timestamp, id: u64) {
        if let Some(msg) = self.traces.get(id as usize).filter(|_| self.verbose) {
            let id = self.id;
            let scope = &self.scope;
            let verbose = self.verbose;
            self.output
                .log(id, scope, verbose, now, format_args!("trce[{msg}]"));
        } else {
            self.log(now, format_args!("trce[{id}]"));
        }
    }

    #[inline(always)]
    fn profile(&mut self, now: Timestamp, id: u64, time: Duration) {
        if let Some(msg) = self.traces.get(id as usize).filter(|_| self.verbose) {
            let id = self.id;
            let scope = &self.scope;
            let verbose = self.verbose;
            self.output.log(
                id,
                scope,
                verbose,
                now,
                format_args!("prof[{msg}]={time:?}"),
            );
        } else {
            self.log(now, format_args!("prof[{id}]={time:?}"));
        }
    }

    #[inline(always)]
    fn park(&mut self, now: Timestamp, id: u64) {
        self.log(now, format_args!("park[{id}]"));
    }

    #[inline(always)]
    fn unpark(&mut self, now: Timestamp, id: u64) {
        self.log(now, format_args!("uprk[{id}]"));
    }

    fn connect(&mut self, now: Timestamp, connection_id: u64, time: Duration) {
        self.log(
            now,
            format_args!("conn[{connection_id}]={:?}us", time.as_micros()),
        );
    }
}

#[derive(Clone, Debug, Default)]
pub struct Throughput(Arc<ThroughputInner>);

impl Throughput {
    pub fn reporter(&self, freq: Duration) {
        let handle = self.clone();
        tokio::spawn(async move {
            let start = tokio::time::Instant::now();
            let mut prev = start;
            while !handle.0.is_open.fetch_or(false, Ordering::Relaxed) {
                tokio::time::sleep(freq).await;
                prev = handle.report(start, prev);
            }
            handle.report(start, prev);
        });
    }

    fn report(
        &self,
        start: tokio::time::Instant,
        prev: tokio::time::Instant,
    ) -> tokio::time::Instant {
        let now = tokio::time::Instant::now();
        let elapsed = now - prev;
        let ts = unsafe { Timestamp::from_duration(now - start) };
        let v = self.0.results.take();
        eprintln!("{ts} {}", v / elapsed);
        now
    }
}

#[derive(Debug, Default)]
struct ThroughputInner {
    results: ThroughputResults,
    is_open: Arc<AtomicBool>,
}

impl Trace for Throughput {
    fn send(&mut self, _now: Timestamp, _stream_id: u64, len: u64) {
        self.0.results.tx.fetch_add(len, Ordering::Relaxed);
    }

    fn receive(&mut self, _now: Timestamp, _stream_id: u64, len: u64) {
        self.0.results.rx.fetch_add(len, Ordering::Relaxed);
    }
}

impl Drop for ThroughputInner {
    fn drop(&mut self) {
        self.is_open.store(false, Ordering::Relaxed);
    }
}

#[derive(Clone, Debug, Default)]
pub struct ThroughputResults<Counter = Arc<AtomicU64>> {
    rx: Counter,
    tx: Counter,
}

impl ThroughputResults {
    pub fn take(&self) -> ThroughputResults<Byte> {
        ThroughputResults {
            rx: self.rx.swap(0, Ordering::Relaxed).bytes(),
            tx: self.tx.swap(0, Ordering::Relaxed).bytes(),
        }
    }
}

impl core::ops::Div<Duration> for ThroughputResults<Byte> {
    type Output = ThroughputResults<Rate>;

    fn div(self, duration: Duration) -> ThroughputResults<Rate> {
        ThroughputResults {
            rx: self.rx / duration,
            tx: self.tx / duration,
        }
    }
}

impl fmt::Display for ThroughputResults<Rate> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "throughput: rx={:#} tx={:#}", self.rx, self.tx)
    }
}


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

mod byte;
mod duration;
mod rate;

pub use byte::*;
pub use duration::*;
pub use rate::*;


// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use super::{AddressMap, Client, Connection};
use crate::{
    driver::timer::{self, Timer, Timestamp},
    operation as op, scenario, Checkpoints, Result, Trace,
};
use core::{
    future::Future,
    task::{Context, Poll},
};
use futures::ready;

pub struct Thread<'a, C: Client<'a>> {
    scenario: &'a scenario::Client,
    ops: &'a [op::Client],
    index: usize,
    op: Option<Op<'a, C>>,
    timer: Timer,
}

impl<'a, C: Client<'a>> Thread<'a, C> {
    pub fn new(scenario: &'a scenario::Client, ops: &'a [op::Client]) -> Self {
        Self {
            scenario,
            ops,
            index: 0,
            op: None,
            timer: Timer::default(),
        }
    }

    #[inline]
    pub(super) fn poll<T: Trace, Ch: Checkpoints>(
        &mut self,
        client: &mut C,
        address_map: &AddressMap,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        loop {
            while self.op.is_none() {
                if let Some(op) = self.ops.get(self.index) {
                    self.index += 1;
                    self.on_op(client, address_map, op, trace, checkpoints, now, cx);
                } else {
                    // we are all done processing the operations
                    return Poll::Ready(Ok(()));
                }
            }

            ready!(self.poll_op(client, address_map, trace, checkpoints, now, cx))?;
            self.op = None;
        }
    }

    #[allow(clippy::too_many_arguments)]
    #[inline]
    fn on_op<T: Trace, Ch: Checkpoints>(
        &mut self,
        client: &mut C,
        addresses: &AddressMap,
        op: &'a op::Client,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: Timestamp,
        cx: &mut Context,
    ) {
        trace.exec_client(now, op);
        use op::Client::*;
        match op {
            Sleep { timeout } => {
                self.timer.sleep(now, *timeout);
                self.op = Some(Op::Sleep);
            }
            Connect {
                server_id,
                router_id,
                server_connection_id,
                client_connection_id,
            } => {
                let addr = if let Some(router_id) = router_id {
                    addresses.router(*router_id, *server_id)
                } else {
                    addresses.server(*server_id)
                };
                let hostname = addresses.hostname(*server_connection_id);
                let ops = &self.scenario.connections[*client_connection_id as usize];
                let connect = client.connect(addr, hostname, *server_connection_id, ops);
                self.op = Some(Op::Connect {
                    connect,
                    id: *client_connection_id,
                    start: now,
                });
            }
            Park { checkpoint } => {
                trace.park(now, *checkpoint);
                self.op = Some(Op::Wait {
                    checkpoint: *checkpoint,
                });
            }
            Unpark { checkpoint } => {
                checkpoints.unpark(*checkpoint, cx);
            }
            Trace { trace_id } => {
                trace.trace(now, *trace_id);
            }
            Scope { threads } => {
                if !threads.is_empty() {
                    let threads = threads
                        .iter()
                        .map(|thread| Thread::new(self.scenario, thread))
                        .collect();
                    self.op = Some(Op::Scope { threads });
                }
            }
        }
    }

    #[inline]
    fn poll_op<T: Trace, Ch: Checkpoints>(
        &mut self,
        client: &mut C,
        addresses: &AddressMap,
        trace: &mut T,
        checkpoints: &mut Ch,
        now: Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        match self.op.as_mut().unwrap() {
            Op::Sleep => {
                ready!(self.timer.poll(now));
            }
            Op::Connect { connect, id, start } => {
                let connect = core::pin::Pin::new(connect);
                let connection = ready!(connect.poll(cx))?;
                let time = now - *start;
                trace.enter_connection(connection.id());
                trace.connect(now, *id, time);
                self.op = Some(Op::Connection { connection });
                return self.poll_op(client, addresses, trace, checkpoints, now, cx);
            }
            Op::Connection { connection } => {
                ready!(connection.poll(trace, checkpoints, now, cx))?;
            }
            Op::Wait { checkpoint } => {
                ready!(checkpoints.park(*checkpoint));
                trace.unpark(now, *checkpoint);
            }
            Op::Scope { threads } => {
                let mut all_ready = true;
                let op_idx = self.index;
                for (idx, thread) in threads.iter_mut().enumerate() {
                    trace.enter(now, op_idx as _, idx);
                    let result = thread.poll(client, addresses, trace, checkpoints, now, cx);
                    trace.exit(now);
                    match result {
                        Poll::Ready(Ok(_)) => {}
                        Poll::Ready(Err(err)) => return Err(err).into(),
                        Poll::Pending => all_ready = false,
                    }
                }
                if !all_ready {
                    return Poll::Pending;
                }
            }
        }

        // clear the timer for the next operation
        self.timer.cancel();
        Ok(()).into()
    }
}

impl<'a, C: Client<'a>> timer::Provider for Thread<'a, C> {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.timer.timers(query)?;
        match &self.op {
            Some(Op::Connection { connection }) => {
                connection.timers(query)?;
            }
            Some(Op::Scope { threads }) => {
                for thread in threads {
                    thread.timers(query)?;
                }
            }
            _ => {}
        }
        Ok(())
    }
}

enum Op<'a, C: Client<'a>> {
    Sleep,
    Connect {
        connect: C::Connect,
        id: u64,
        start: Timestamp,
    },
    Connection {
        connection: C::Connection,
    },
    Wait {
        checkpoint: u64,
    },
    Scope {
        threads: Vec<Thread<'a, C>>,
    },
}