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

use super::timer::{self, Timer, Timestamp};
use crate::{
    connection::Owner,
    operation as op,
    units::{Byte, Rate, Rates},
    Checkpoints, Connection, Result, Trace,
};
use core::task::{Context, Poll};
use futures::ready;

#[derive(Debug)]
pub struct Thread<'a> {
    ops: &'a [op::Connection],
    index: usize,
    op: Option<Op<'a>>,
    timer: Timer,
    owner: Owner,
}

impl<'a> Thread<'a> {
    #[inline]
    pub fn new(ops: &'a [op::Connection], owner: Owner) -> Self {
        Self {
            ops,
            index: 0,
            op: None,
            timer: Timer::default(),
            owner,
        }
    }

    #[inline]
    pub(crate) fn poll<C: Connection, T: Trace, Ch: Checkpoints>(
        &mut self,
        conn: &mut C,
        trace: &mut T,
        checkpoints: &mut Ch,
        rates: &mut Rates,
        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(op, trace, checkpoints, rates, now, cx);
                } else {
                    // we are all done processing the operations
                    return Poll::Ready(Ok(()));
                }
            }

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

    fn reset(&mut self, cx: &mut Context) {
        self.index = 0;
        self.op = None;
        cx.waker().wake_by_ref();
    }

    #[inline]
    fn on_op<T: Trace, Ch: Checkpoints>(
        &mut self,
        op: &'a op::Connection,
        trace: &mut T,
        checkpoints: &mut Ch,
        rates: &mut Rates,
        now: Timestamp,
        cx: &mut Context,
    ) {
        trace.exec(now, op);
        use op::{Connection::*, IterateValue};
        match op {
            Sleep { amount } => {
                self.timer.sleep(now, *amount);
                self.op = Some(Op::Sleep);
            }
            OpenBidirectionalStream { stream_id } => {
                self.op = Some(Op::OpenBidirectionalStream { id: *stream_id });
            }
            OpenSendStream { stream_id } => {
                self.op = Some(Op::OpenSendStream { id: *stream_id });
            }
            Send { stream_id, bytes } => {
                self.op = Some(Op::Send {
                    id: *stream_id,
                    remaining: *bytes,
                    rate: rates.send.get(stream_id).cloned(),
                });
            }
            SendFinish { stream_id } => {
                self.op = Some(Op::SendFinish { id: *stream_id });
            }
            SendRate { stream_id, rate } => {
                rates.send.insert(*stream_id, *rate);
            }
            Receive { stream_id, bytes } => {
                self.op = Some(Op::Receive {
                    id: *stream_id,
                    remaining: *bytes,
                    rate: rates.receive.get(stream_id).cloned(),
                });
            }
            ReceiveAll { stream_id } => {
                self.op = Some(Op::ReceiveAll {
                    id: *stream_id,
                    rate: rates.receive.get(stream_id).cloned(),
                });
            }
            ReceiveFinish { stream_id } => {
                self.op = Some(Op::ReceiveFinish { id: *stream_id });
            }
            ReceiveRate { stream_id, rate } => {
                rates.receive.insert(*stream_id, *rate);
            }
            Trace { trace_id } => {
                trace.trace(now, *trace_id);
            }
            Profile {
                trace_id,
                operations,
            } => {
                self.op = Some(Op::Profile {
                    start: now,
                    trace_id: *trace_id,
                    thread: Box::new(Thread::new(operations, self.owner)),
                });
            }
            Iterate {
                value,
                operations,
                trace_id,
            } => {
                let start = now;
                let trace_id = *trace_id;
                let thread = Box::new(Thread::new(operations, self.owner));

                self.op = Some(match value {
                    IterateValue::Count { amount } => Op::Iterate {
                        start,
                        count: *amount,
                        thread,
                        trace_id,
                    },
                    IterateValue::Time { amount } => {
                        self.timer.sleep(now, *amount);
                        Op::IterateFor {
                            start,
                            thread,
                            trace_id,
                        }
                    }
                });
            }
            Park { checkpoint } => {
                trace.park(now, *checkpoint);
                self.op = Some(Op::Wait {
                    checkpoint: *checkpoint,
                });
            }
            Unpark { checkpoint } => {
                // notify the checkpoint that it can make progress
                checkpoints.unpark(*checkpoint, cx);
            }
            Scope { threads } => {
                if !threads.is_empty() {
                    let threads = threads
                        .iter()
                        .map(|thread| Thread::new(thread, self.owner))
                        .collect();
                    self.op = Some(Op::Scope { threads });
                }
            }
        }
    }

    #[inline]
    fn poll_op<C: Connection, T: Trace, Ch: Checkpoints>(
        &mut self,
        conn: &mut C,
        trace: &mut T,
        checkpoints: &mut Ch,
        rates: &mut Rates,
        now: Timestamp,
        cx: &mut Context,
    ) -> Poll<Result<()>> {
        let owner = self.owner;
        match self.op.as_mut().unwrap() {
            Op::Sleep => {
                ready!(self.timer.poll(now));
            }
            Op::OpenBidirectionalStream { id } => {
                ready!(conn.poll_open_bidirectional_stream(*id, cx))?;
                trace.open(now, *id);
            }
            Op::OpenSendStream { id } => {
                ready!(conn.poll_open_send_stream(*id, cx))?;
                trace.open(now, *id);
            }
            Op::Send {
                id,
                remaining,
                rate,
            } => {
                return self.timer.transfer(remaining, rate, now, cx, |bytes, cx| {
                    let amount = ready!(conn.poll_send(owner, *id, *bytes, cx))?;
                    trace.send(now, *id, amount);
                    Ok(amount).into()
                })
            }
            Op::SendFinish { id } => {
                ready!(conn.poll_send_finish(owner, *id, cx))?;
                trace.send_finish(now, *id);
            }
            Op::Receive {
                id,
                remaining,
                rate,
            } => {
                return self.timer.transfer(remaining, rate, now, cx, |bytes, cx| {
                    let amount = ready!(conn.poll_receive(owner, *id, *bytes, cx))?;
                    trace.receive(now, *id, amount);
                    Ok(amount).into()
                })
            }
            Op::ReceiveAll { id, rate } => {
                let mut remaining = Byte::MAX;
                return self
                    .timer
                    .transfer(&mut remaining, rate, now, cx, |bytes, cx| {
                        let amount = ready!(conn.poll_receive(owner, *id, *bytes, cx))?;
                        trace.receive(now, *id, amount);
                        Ok(amount).into()
                    });
            }
            Op::ReceiveFinish { id } => {
                ready!(conn.poll_receive_finish(owner, *id, cx))?;
                trace.receive_finish(now, *id);
            }
            Op::Wait { checkpoint } => {
                ready!(checkpoints.park(*checkpoint));
                trace.unpark(now, *checkpoint);
            }
            Op::Iterate {
                start,
                count,
                thread,
                trace_id,
            } => {
                ready!(thread.poll(conn, trace, checkpoints, rates, now, cx))?;

                if let Some(trace_id) = trace_id {
                    let time = now.saturating_duration_since(*start);
                    trace.profile(now, *trace_id, time);
                    *start = now;
                }

                if let Some(next) = count.checked_sub(1) {
                    if next > 0 {
                        *count = next;
                        thread.reset(cx);
                        return Poll::Pending;
                    }
                }
            }
            Op::IterateFor {
                start,
                thread,
                trace_id,
            } => {
                if let Poll::Ready(res) = thread.poll(conn, trace, checkpoints, rates, now, cx) {
                    res?;
                    thread.reset(cx);

                    if let Some(trace_id) = trace_id {
                        let time = now.saturating_duration_since(*start);
                        trace.profile(now, *trace_id, time);
                        *start = now;
                    }
                }

                ready!(self.timer.poll(now));
            }
            Op::Profile {
                trace_id,
                start,
                thread,
            } => {
                ready!(thread.poll(conn, trace, checkpoints, rates, now, cx))?;
                let time = now.saturating_duration_since(*start);
                trace.profile(now, *trace_id, time);
            }
            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(conn, trace, checkpoints, 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 !all_ready {
                    return Poll::Pending;
                }
            }
        }

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

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

#[derive(Debug)]
enum Op<'a> {
    Sleep,
    OpenBidirectionalStream {
        id: u64,
    },
    OpenSendStream {
        id: u64,
    },
    Send {
        id: u64,
        remaining: Byte,
        rate: Option<Rate>,
    },
    SendFinish {
        id: u64,
    },
    Receive {
        id: u64,
        remaining: Byte,
        rate: Option<Rate>,
    },
    ReceiveAll {
        id: u64,
        rate: Option<Rate>,
    },
    ReceiveFinish {
        id: u64,
    },
    Wait {
        checkpoint: u64,
    },
    Iterate {
        start: Timestamp,
        count: u64,
        thread: Box<Thread<'a>>,
        trace_id: Option<u64>,
    },
    IterateFor {
        start: Timestamp,
        thread: Box<Thread<'a>>,
        trace_id: Option<u64>,
    },
    Profile {
        trace_id: u64,
        start: Timestamp,
        thread: Box<Thread<'a>>,
    },
    Scope {
        threads: Vec<Thread<'a>>,
    },
}


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

use crate::units::{Byte, ByteExt, Rate};
use core::{
    task::{Context, Poll},
    time::Duration,
};
use futures::ready;
pub use s2n_quic_core::time::{
    timer::{self, Provider, Query, Result},
    Timestamp,
};

#[derive(Debug, Default)]
pub struct Timer {
    sleep: s2n_quic_core::time::Timer,
    window: Byte,
}

impl timer::Provider for Timer {
    #[inline]
    fn timers<Q: timer::Query>(&self, query: &mut Q) -> timer::Result {
        self.sleep.timers(query)
    }
}

impl Timer {
    pub fn poll(&mut self, now: Timestamp) -> Poll<()> {
        if !self.sleep.is_armed() {
            Poll::Ready(())
        } else {
            self.sleep.poll_expiration(now)
        }
    }

    pub fn cancel(&mut self) {
        self.sleep.cancel();
    }

    pub fn sleep(&mut self, now: Timestamp, duration: Duration) {
        self.sleep.set(now + duration);
    }

    pub fn transfer<F: FnMut(Byte, &mut Context) -> Poll<crate::Result<u64>>>(
        &mut self,
        remaining: &mut Byte,
        rate: &Option<Rate>,
        now: Timestamp,
        cx: &mut Context,
        mut f: F,
    ) -> Poll<crate::Result<()>> {
        loop {
            let amount = if let Some(Rate { bytes, period }) = rate.as_ref() {
                if self.poll(now).is_ready() {
                    self.window = *bytes.min(remaining);
                    self.sleep(now, *period);
                }

                if self.window == Byte::default() {
                    return Poll::Pending;
                }

                let amount = ready!(f(self.window, cx))?.bytes();

                self.window -= amount;

                amount
            } else {
                ready!(f(*remaining, cx))?.bytes()
            };

            // if the transfer returned 0 bytes that means it's done
            if amount == 0.bytes() {
                *remaining = Byte::default();
            } else {
                *remaining -= amount.bytes();
            }

            if *remaining == Byte::default() {
                self.window = Byte::default();
                return Poll::Ready(Ok(()));
            }
        }
    }
}


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

use super::Frame;
use crate::Result;
use bytes::{Buf, BufMut, Bytes, BytesMut};
use core::{
    mem::MaybeUninit,
    task::{Context, Poll},
};
use futures::ready;
use std::{collections::VecDeque, io::IoSlice};
use tokio::io::ReadBuf;

const CHUNK_CAPACITY: usize = 4096;

#[derive(Debug, Default)]
pub struct ReadBuffer {
    len: usize,
    tail: BytesMut,
    head: VecDeque<BytesMut>,
}

impl ReadBuffer {
    pub fn read<F: FnOnce(&mut ReadBuf) -> Poll<Result<()>>>(
        &mut self,
        read: F,
    ) -> Poll<Result<()>> {
        let capacity = self.tail.capacity();

        if capacity == 0 {
            self.tail.reserve(CHUNK_CAPACITY);
        } else if capacity < 32 {
            let tail = core::mem::replace(&mut self.tail, BytesMut::with_capacity(CHUNK_CAPACITY));
            self.head.push_back(tail);
        }

        let buf = self.tail.chunk_mut();
        let buf = unsafe { &mut *(buf as *mut _ as *mut [MaybeUninit<u8>]) };
        let mut buf = ReadBuf::uninit(buf);
        ready!(read(&mut buf))?;
        let len = buf.filled().len();

        if len == 0 {
            return Ok(()).into();
        }

        unsafe {
            self.tail.advance_mut(len);
            self.len += len;
        }

        Ok(()).into()
    }

    fn check_consistency(&self) {
        if cfg!(debug_assertions) {
            let mut actual_len = self.tail.len();
            for chunk in self.head.iter() {
                actual_len += chunk.len();
            }
            assert_eq!(actual_len, self.remaining(), "{self:?}");
        }
    }
}

impl Buf for ReadBuffer {
    fn remaining(&self) -> usize {
        self.len
    }

    fn chunk(&self) -> &[u8] {
        for chunk in self.head.iter() {
            if !chunk.is_empty() {
                return chunk.chunk();
            }
        }

        self.tail.chunk()
    }

    fn advance(&mut self, mut cnt: usize) {
        while cnt > 0 {
            let len = if let Some(front) = self.head.front_mut() {
                let len = front.len().min(cnt);
                front.advance(len);
                if front.is_empty() {
                    let _ = self.head.pop_front();
                }
                len
            } else {
                self.tail.advance(cnt);
                cnt
            };

            cnt -= len;
            self.len -= len;
        }

        self.check_consistency();
    }

    fn copy_to_bytes(&mut self, len: usize) -> Bytes {
        while let Some(mut front) = self.head.pop_front() {
            if front.is_empty() {
                continue;
            }

            self.len -= len;

            if front.len() == len {
                self.check_consistency();
                return front.freeze();
            }

            let out = front.split_to(len);
            self.head.push_front(front);

            self.check_consistency();
            return out.freeze();
        }

        let out = self.tail.split_to(len).freeze();

        self.len -= len;
        self.check_consistency();

        out
    }
}

#[derive(Debug, Default)]
pub struct WriteBuffer {
    unfilled: BytesMut,
    queue: VecDeque<Bytes>,
    push_interest: bool,
}

impl WriteBuffer {
    pub fn request_push(&mut self, max_queue_len: usize) -> bool {
        let can_push = max_queue_len > self.queue.len();
        self.push_interest |= !can_push;
        can_push
    }

    pub fn push(&mut self, frame: Frame) {
        let capacity = self.unfilled.capacity();

        if capacity == 0 {
            self.unfilled.reserve(CHUNK_CAPACITY);
        } else if capacity < 32 {
            self.unfilled = BytesMut::with_capacity(CHUNK_CAPACITY);
        }

        frame.write_header(&mut self.unfilled);
        self.queue.push_back(self.unfilled.split().freeze());

        if let Some(data) = frame.body() {
            self.queue.push_back(data);
        }
    }

    pub fn push_priority(&mut self, frame: Frame) {
        let capacity = self.unfilled.capacity();

        const CAPACITY: usize = 4096;

        if capacity == 0 {
            self.unfilled.reserve(CAPACITY);
        } else if capacity < 32 {
            self.unfilled = BytesMut::with_capacity(CAPACITY);
        }

        frame.write_header(&mut self.unfilled);
        let header = self.unfilled.split().freeze();

        // push the body first
        if let Some(data) = frame.body() {
            self.queue.push_front(data);
        }

        self.queue.push_front(header);
    }

    pub fn is_empty(&self) -> bool {
        self.queue.is_empty()
    }

    pub fn pop_front(&mut self) -> Option<Bytes> {
        self.queue.pop_front()
    }

    pub fn push_front(&mut self, chunk: Bytes) {
        self.queue.push_front(chunk);
    }

    pub fn notify(&mut self, cx: &mut Context) {
        if self.push_interest {
            self.push_interest = false;
            cx.waker().wake_by_ref();
        }
    }

    pub fn advance(&mut self, mut len: usize) {
        while let Some(mut chunk) = self.queue.pop_front() {
            let next_len = len.saturating_sub(chunk.len());

            if next_len > 0 {
                len = next_len;
                continue;
            }

            if chunk.len() > len {
                chunk.advance(len);
                self.queue.push_front(chunk);
            }

            return;
        }
    }

    pub fn chunks(&self) -> Vec<IoSlice> {
        self.queue.iter().map(|v| IoSlice::new(v)).collect()
    }
}


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

use crate::{connection::Owner, Result};
use bytes::{Buf, BufMut, Bytes};
use core::mem::size_of;
use enum_primitive_derive::Primitive;
use num_traits::FromPrimitive;
use std::io;

#[derive(Clone, Copy, Debug, Primitive)]
#[repr(u8)]
enum Tag {
    StreamOpen = 0,
    StreamData = 1,
    StreamFinish = 2,
    MaxStreams = 3,
    MaxStreamData = 4,
    InitialMaxStreamData = 5,
}

#[derive(Clone, Debug)]
pub enum Frame {
    StreamOpen { id: u64, bidirectional: bool },
    StreamData { id: u64, owner: Owner, data: Bytes },
    StreamFinish { id: u64, owner: Owner },
    MaxStreams { up_to: u64 },
    MaxStreamData { id: u64, owner: Owner, up_to: u64 },
    InitialMaxStreamData { up_to: u64 },
}

impl Frame {
    pub fn write_header<B: BufMut>(&self, buf: &mut B) {
        match self {
            Self::StreamOpen { id, bidirectional } => {
                buf.put_u8(Tag::StreamOpen as _);
                buf.put_u64(*id);
                buf.put_u8(*bidirectional as u8);
            }
            Self::StreamData { id, owner, data } => {
                buf.put_u8(Tag::StreamData as _);
                buf.put_u64(*id);
                buf.put_u8(*owner as _);
                buf.put_u32(data.len() as _);
            }
            Self::StreamFinish { id, owner } => {
                buf.put_u8(Tag::StreamFinish as _);
                buf.put_u64(*id);
                buf.put_u8(*owner as _);
            }
            Self::MaxStreams { up_to } => {
                buf.put_u8(Tag::MaxStreams as _);
                buf.put_u64(*up_to);
            }
            Self::MaxStreamData { id, owner, up_to } => {
                buf.put_u8(Tag::MaxStreamData as _);
                buf.put_u64(*id);
                buf.put_u8(*owner as _);
                buf.put_u64(*up_to);
            }
            Self::InitialMaxStreamData { up_to } => {
                buf.put_u8(Tag::InitialMaxStreamData as _);
                buf.put_u64(*up_to);
            }
        }
    }

    pub fn body(self) -> Option<Bytes> {
        if let Self::StreamData { data, .. } = self {
            Some(data)
        } else {
            None
        }
    }
}

#[derive(Debug, Default)]
pub struct Decoder {
    tag: Option<Tag>,
    stream: Option<(u64, Owner, usize)>,
}

impl Decoder {
    pub fn decode<B: Buf>(&mut self, buf: &mut B) -> Result<Option<Frame>> {
        if let Some((id, owner, len)) = self.stream.take() {
            return self.decode_stream(buf, id, owner, len);
        }

        if let Some(tag) = self.tag.take() {
            return self.decode_frame(buf, tag);
        }

        if buf.remaining() == 0 {
            return Ok(None);
        }

        let tag = Tag::from_u8(buf.get_u8())
            .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "invalid frame tag"))?;

        self.decode_frame(buf, tag)
    }

    fn decode_frame<B: Buf>(&mut self, buf: &mut B, tag: Tag) -> Result<Option<Frame>> {
        match tag {
            Tag::StreamOpen => {
                if buf.remaining() < (size_of::<u64>() + size_of::<u8>()) {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let id = buf.get_u64();
                let bidirectional = buf.get_u8() != 0;

                Ok(Some(Frame::StreamOpen { id, bidirectional }))
            }
            Tag::StreamData => {
                if buf.remaining() < (size_of::<u64>() + size_of::<u8>() + size_of::<u32>()) {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let id = buf.get_u64();
                let owner = Self::decode_owner(buf)?;
                let len = buf.get_u32();
                self.decode_stream(buf, id, owner, len as _)
            }
            Tag::StreamFinish => {
                if buf.remaining() < (size_of::<u64>() + size_of::<u8>()) {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let id = buf.get_u64();
                let owner = Self::decode_owner(buf)?;

                Ok(Some(Frame::StreamFinish { id, owner }))
            }
            Tag::MaxStreams => {
                if buf.remaining() < size_of::<u64>() {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let up_to = buf.get_u64();

                Ok(Some(Frame::MaxStreams { up_to }))
            }
            Tag::MaxStreamData => {
                if buf.remaining() < (size_of::<u64>() + size_of::<u8>() + size_of::<u64>()) {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let id = buf.get_u64();
                let owner = Self::decode_owner(buf)?;
                let up_to = buf.get_u64();

                Ok(Some(Frame::MaxStreamData { id, owner, up_to }))
            }
            Tag::InitialMaxStreamData => {
                if buf.remaining() < (size_of::<u64>()) {
                    self.tag = Some(tag);
                    return Ok(None);
                }

                let up_to = buf.get_u64();

                Ok(Some(Frame::InitialMaxStreamData { up_to }))
            }
        }
    }

    fn decode_owner<B: Buf>(buf: &mut B) -> Result<Owner> {
        let owner = buf.get_u8();
        Ok(match owner {
            0 => Owner::Local,
            1 => Owner::Remote,
            _ => return Err(io::Error::new(io::ErrorKind::InvalidData, "invalid owner id").into()),
        })
    }

    fn decode_stream<B: Buf>(
        &mut self,
        buf: &mut B,
        id: u64,
        owner: Owner,
        len: usize,
    ) -> Result<Option<Frame>> {
        if len == 0 {
            return self.decode(buf);
        }

        let chunk_len = buf.chunk().len();

        if chunk_len == 0 {
            self.stream = Some((id, owner, len));
            return Ok(None);
        }

        Ok(if chunk_len >= len {
            let data = buf.copy_to_bytes(len);
            Some(Frame::StreamData { id, owner, data })
        } else {
            let data = buf.copy_to_bytes(chunk_len);
            self.stream = Some((id, owner, len - chunk_len));
            Some(Frame::StreamData { id, owner, data })
        })
    }
}


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

use super::Blocked;
use crate::Result;
use bytes::Bytes;
use core::task::{Context, Poll};
use s2n_quic_core::stream::testing::Data;

#[derive(Debug)]
pub struct Stream {
    pub rx: Option<ReceiveStream>,
    pub tx: Option<SendStream>,
}

#[derive(Debug)]
pub struct ReceiveStream {
    received_offset: u64,
    buffered_offset: u64,
    window_offset: u64,
    is_finished: bool,
    blocked: Blocked,
}

impl ReceiveStream {
    pub fn new(window_offset: u64) -> Self {
        Self {
            received_offset: 0,
            buffered_offset: 0,
            window_offset,
            is_finished: false,
            blocked: Default::default(),
        }
    }

    pub fn buffer(&mut self, len: u64) -> Result<u64> {
        if self.is_finished {
            return Err("stream is already finished".into());
        }

        let len = (self.window_offset - self.buffered_offset).min(len);

        if len == 0 {
            return Ok(0);
        }

        self.buffered_offset += len;

        self.blocked.unblock();

        Ok(len)
    }

    pub fn receive(&mut self, len: u64, cx: &mut Context) -> Poll<Result<u64>> {
        let len = (self.buffered_offset - self.received_offset).min(len);

        if len == 0 {
            return if self.is_finished {
                Ok(0).into()
            } else {
                self.blocked.block(cx);
                Poll::Pending
            };
        }

        self.received_offset += len;

        Ok(len).into()
    }

    pub fn finish(&mut self) {
        self.is_finished = true;
        self.blocked.unblock();
    }

    pub fn receive_window(&self) -> u64 {
        self.window_offset - self.received_offset
    }

    pub fn credit(&mut self, credits: u64) -> u64 {
        self.window_offset += credits;
        self.window_offset
    }

    pub fn poll_finish(&mut self, cx: &mut Context) -> Poll<()> {
        if self.is_finished {
            Poll::Ready(())
        } else {
            self.blocked.block(cx);
            Poll::Pending
        }
    }
}

#[derive(Debug, Default)]
pub struct SendStream {
    max_data: u64,
    data: Data,
    blocked: Blocked,
}

impl SendStream {
    pub fn new(max_data: u64) -> Self {
        Self {
            max_data,
            data: Data::new(u64::MAX),
            blocked: Default::default(),
        }
    }

    pub fn max_data(&mut self, max_data: u64) {
        self.max_data = max_data.max(self.max_data);
        self.blocked.unblock()
    }

    pub fn send(&mut self, len: u64, cx: &mut Context) -> Option<Bytes> {
        let window = self.max_data - self.data.offset();
        let len = len.min(window) as usize;

        if len == 0 {
            self.blocked.block(cx);
            return None;
        }

        self.data.send_one(len)
    }
}

#[derive(Debug, Default)]
pub struct Controller {
    open_offset: u64,
    peer_window_offset: u64,
    local_window_offset: u64,
    transmitted_window_offset: u64,
    blocked: Blocked,
}

impl Controller {
    pub fn new(local_window: u64, peer_window: u64) -> Self {
        let local_window = local_window.max(1);
        let peer_window = peer_window.max(1);
        Self {
            open_offset: 0,
            peer_window_offset: local_window,
            local_window_offset: peer_window,
            transmitted_window_offset: local_window.min(peer_window),
            blocked: Default::default(),
        }
    }

    pub fn open(&mut self, cx: &mut Context) -> Poll<()> {
        if self.capacity() == 0 {
            self.blocked.block(cx);
            return Poll::Pending;
        }

        self.open_offset += 1;

        Poll::Ready(())
    }

    pub fn close(&mut self) {
        self.local_window_offset += 1;
        if self.capacity() > 0 {
            self.blocked.unblock();
        }
    }

    pub fn transmit(&mut self) -> Option<u64> {
        // send a max streams if we are using more than half capacity
        if self.local_window_offset >= self.transmitted_window_offset * 2 {
            self.transmitted_window_offset = self.local_window_offset;
            Some(self.local_window_offset)
        } else {
            None
        }
    }

    pub fn max_streams(&mut self, up_to: u64) {
        self.peer_window_offset = up_to;
        if self.capacity() > 0 {
            self.blocked.unblock()
        }
    }

    pub fn capacity(&self) -> u64 {
        self.local_window_offset.min(self.peer_window_offset) - self.open_offset
    }
}


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

use std::sync::Arc;

macro_rules! sleep {
    () => {
        pub fn sleep(&mut self, amount: core::time::Duration) -> &mut Self {
            self.ops
                .push(crate::operation::Connection::Sleep { amount });
            self
        }
    };
}

macro_rules! trace {
    () => {
        pub fn trace(&mut self, name: &str) -> &mut Self {
            let trace_id = self.state.trace(name);
            self.ops
                .push(crate::operation::Connection::Trace { trace_id });
            self
        }

        pub fn profile<F: FnOnce(&mut Self)>(&mut self, name: &str, f: F) -> &mut Self {
            let trace_id = self.state.trace(name);
            let mut builder = self.child_scope();
            f(&mut builder);
            let operations = builder.finish_scope();

            self.ops.push(crate::operation::Connection::Profile {
                trace_id,
                operations,
            });

            self
        }
    };
}

macro_rules! iterate {
    () => {
        pub fn iterate<I: Into<crate::operation::IterateValue>, F: FnOnce(&mut Self)>(
            &mut self,
            count: I,
            f: F,
        ) -> &mut Self {
            let mut builder = self.child_scope();
            f(&mut builder);
            let mut operations = builder.finish_scope();

            let count = count.into();

            if operations.is_empty() || count.is_zero() {
                return self;
            }

            let mut trace_id = None;

            // optimize out nested iterate/profile statements
            if operations.len() == 1 {
                if let Some(crate::operation::Connection::Profile {
                    trace_id: child_id,
                    operations: child,
                }) = operations.get_mut(0)
                {
                    trace_id = Some(*child_id);
                    operations = core::mem::take(child);
                }
            }

            self.ops.push(crate::operation::Connection::Iterate {
                value: count,
                operations,
                trace_id,
            });

            self
        }
    };
}

#[macro_use]
pub mod checkpoint;

pub mod certificate;
pub mod client;
pub mod connection;
pub mod scope;
pub mod server;
mod state;
pub mod stream;

pub use client::Client;
pub use connection::Connection;
pub use scope::Scope;
pub use server::Server;
pub use stream::Stream;

use state::State;

#[derive(Debug)]
pub struct Builder {
    state: State,
}

impl Builder {
    pub(super) fn new() -> Self {
        Self {
            state: Default::default(),
        }
    }

    pub fn create_ca(&mut self) -> certificate::Authority {
        self.create_ca_with(|_| {})
    }

    pub fn create_ca_with<F: FnOnce(&mut certificate::AuthorityBuilder)>(
        &mut self,
        f: F,
    ) -> certificate::Authority {
        self.state.create_ca_with(f)
    }

    pub fn create_server(&mut self) -> Server {
        self.create_server_with(|_| {})
    }

    pub fn create_server_with<F: FnOnce(&mut server::Builder)>(&mut self, f: F) -> Server {
        Server::new(self.state.clone(), f)
    }

    pub fn create_client<F: FnOnce(&mut client::Builder)>(&mut self, f: F) {
        let id = self.state.clients.push(super::Client {
            name: String::new(),
            scenario: vec![],
            connections: vec![],
            configuration: Default::default(),
            certificate_authorities: vec![],
        }) as u64;

        let mut builder = client::Builder::new(id, self.state.clone());
        f(&mut builder);

        let client = &mut self.state.clients.borrow_mut()[id as usize];

        client.scenario = builder.finish();
    }

    pub(super) fn finish(self) -> super::Scenario {
        let clients = self
            .state
            .clients
            .take()
            .into_iter()
            .map(|mut client| {
                client.certificate_authorities.sort_unstable();

                Arc::new(client)
            })
            .collect();
        let servers = self
            .state
            .servers
            .take()
            .into_iter()
            .map(Arc::new)
            .collect();
        let mut traces = self.state.trace.take().into_iter().collect::<Vec<_>>();
        traces.sort_by(|(_, a), (_, b)| a.cmp(b));
        let traces = Arc::new(traces.into_iter().map(|(value, _)| value).collect());
        let certificates = self.state.certificates.take();

        let mut scenario = super::Scenario {
            id: Default::default(),
            clients,
            servers,
            // TODO implement router builder
            routers: vec![],
            traces,
            certificates: vec![],
        };

        let mut hash = crate::scenario::Id::hasher();
        core::hash::Hash::hash(&scenario, &mut hash);
        core::hash::Hash::hash(&certificates, &mut hash);
        scenario.id = hash.finish();

        scenario.certificates = certificate::Certificate::build_all(certificates, &scenario.id);

        scenario
    }
}

pub trait Endpoint {
    type Peer: Endpoint;
}

impl Endpoint for Client {
    type Peer = Server;
}

impl Endpoint for Server {
    type Peer = Client;
}

#[derive(Debug)]
pub struct Local;

#[derive(Debug)]
pub struct Remote;

#[cfg(test)]
mod tests;


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

use core::fmt;
use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Default, PartialEq, Eq, Serialize, Deserialize, Hash)]
pub struct Id(String);

impl Id {
    pub(crate) fn hasher() -> Hasher {
        Hasher::default()
    }
}

impl fmt::Display for Id {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        self.0.fmt(f)
    }
}

#[derive(Debug, Default)]
pub struct Hasher {
    hash: sha2::Sha256,
}

impl Hasher {
    pub fn finish(self) -> Id {
        use sha2::Digest;
        let hash = self.hash.finalize();
        let out = base64::encode_config(hash, base64::URL_SAFE_NO_PAD);
        // '_' is not allowed in DNS names
        let out = out.replace('_', "-");
        Id(out)
    }
}

impl core::hash::Hasher for Hasher {
    fn write(&mut self, bytes: &[u8]) {
        use sha2::Digest;
        self.hash.update(bytes)
    }

    fn finish(&self) -> u64 {
        unimplemented!()
    }
}


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

use crate::scenario;
use rcgen::SignatureAlgorithm;
use std::{collections::HashMap, sync::Arc};

static DEFAULT_ALG: &SignatureAlgorithm = &rcgen::PKCS_ECDSA_P256_SHA256;

#[derive(Clone, Debug, Hash)]
pub(crate) enum Certificate {
    Authority {
        alg: &'static SignatureAlgorithm,
    },
    PrivateKey {
        alg: &'static SignatureAlgorithm,
        authority: u64,
        intermediates: Vec<&'static SignatureAlgorithm>,
    },
    // Placeholder for a public cert
    Public,
}

fn create_ca(domain: &str, name: String, alg: &'static SignatureAlgorithm) -> rcgen::Certificate {
    use rcgen::{
        BasicConstraints, Certificate, CertificateParams, DistinguishedName, DnType, IsCa,
        KeyUsagePurpose,
    };

    let mut params = CertificateParams::new(vec![domain.to_string()]);
    params.alg = alg;
    params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained);
    params.distinguished_name = DistinguishedName::new();
    params.distinguished_name.push(DnType::CountryName, "US");
    params.distinguished_name.push(DnType::CommonName, name);
    params.key_usages = vec![
        KeyUsagePurpose::DigitalSignature,
        KeyUsagePurpose::KeyCertSign,
        KeyUsagePurpose::CrlSign,
    ];
    Certificate::from_params(params).unwrap()
}

fn create_cert(domain: &str, name: String, alg: &'static SignatureAlgorithm) -> rcgen::Certificate {
    use rcgen::{
        Certificate, CertificateParams, DistinguishedName, DnType, ExtendedKeyUsagePurpose,
        KeyUsagePurpose,
    };

    let mut params = CertificateParams::new(vec![domain.to_string(), format!("*.{domain}")]);
    params.alg = alg;
    params.use_authority_key_identifier_extension = true;
    params.distinguished_name = DistinguishedName::new();
    params.distinguished_name.push(DnType::CountryName, "US");
    params.distinguished_name.push(DnType::CommonName, name);
    params.key_usages = vec![
        KeyUsagePurpose::DigitalSignature,
        KeyUsagePurpose::KeyEncipherment,
    ];
    params.extended_key_usages = vec![
        ExtendedKeyUsagePurpose::ServerAuth,
        ExtendedKeyUsagePurpose::ClientAuth,
    ];
    Certificate::from_params(params).unwrap()
}

impl Certificate {
    pub(crate) fn build_all(
        certs: Vec<Self>,
        id: &scenario::Id,
    ) -> Vec<Arc<scenario::Certificate>> {
        let mut cas = HashMap::new();
        let mut ias = HashMap::new();
        let mut out = vec![];

        let domain = format!("{id}.net");
        for (cert_idx, cert) in certs.into_iter().enumerate() {
            match cert {
                Self::Authority { alg } => {
                    let cert = create_ca(&domain, format!("netbench CA {cert_idx}"), alg);
                    let pem = cert.serialize_pem().unwrap();

                    out.push(Arc::new(scenario::Certificate {
                        pem,
                        pkcs12: vec![],
                    }));

                    cas.insert(cert_idx as u64, cert);
                }
                Self::PrivateKey {
                    alg,
                    authority,
                    intermediates,
                } => {
                    // create any IAs we need
                    for (idx, alg) in intermediates.iter().copied().enumerate() {
                        ias.entry((authority, idx, alg)).or_insert_with(|| {
                            create_ca(&domain, format!("netbench IA {authority} {idx}"), alg)
                        });
                    }

                    // create a reverse chain of authorities that need to sign this cert
                    let ca = cas.get(&authority).unwrap();
                    let authorities = intermediates
                        .iter()
                        .enumerate()
                        .rev()
                        .map(|(idx, alg)| ias.get(&(authority, idx, alg)).unwrap())
                        .chain(Some(ca));

                    let cert = create_cert(&domain, format!("netbench Leaf {cert_idx}"), alg);
                    let mut chain = String::new();
                    let private_key = cert.serialize_private_key_pem();

                    let mut current_cert = &cert;
                    for authority in authorities {
                        let public = current_cert.serialize_pem_with_signer(authority).unwrap();
                        chain.push_str(&public);
                        current_cert = authority;
                    }

                    let pkcs12 = {
                        let public = openssl::x509::X509::from_pem(chain.as_bytes()).unwrap();
                        let key = openssl::pkey::PKey::private_key_from_pem(private_key.as_bytes())
                            .unwrap();
                        openssl::pkcs12::Pkcs12::builder()
                            .pkey(&key)
                            .cert(&public)
                            .build2("")
                            .unwrap()
                            .to_der()
                            .unwrap()
                    };

                    out.push(Arc::new(scenario::Certificate {
                        pem: private_key,
                        pkcs12,
                    }));
                    out.push(Arc::new(scenario::Certificate {
                        pem: chain,
                        pkcs12: vec![],
                    }));
                }
                Self::Public => {
                    // noop - handled by private key
                }
            }
        }
        out
    }
}

#[derive(Clone, Debug)]
pub struct Authority {
    id: u64,
    state: super::State,
}

impl Authority {
    pub(crate) fn new<F: FnOnce(&mut AuthorityBuilder)>(state: super::State, f: F) -> Self {
        let mut builder = AuthorityBuilder { alg: DEFAULT_ALG };
        f(&mut builder);

        let id = state
            .certificates
            .push(Certificate::Authority { alg: builder.alg }) as u64;

        Self { id, state }
    }

    pub fn key_pair(&self) -> KeyPair {
        self.key_pair_with(|_| {})
    }

    pub fn key_pair_with<F: FnOnce(&mut KeyPairBuilder)>(&self, f: F) -> KeyPair {
        let mut builder = KeyPairBuilder {
            authority: self.id,
            intermediates: vec![],
            alg: DEFAULT_ALG,
        };

        f(&mut builder);

        let KeyPairBuilder {
            authority,
            intermediates,
            alg,
        } = builder;

        let private_key = self.state.certificates.push(Certificate::PrivateKey {
            alg,
            authority,
            intermediates,
        }) as u64;
        let certificate = self.state.certificates.push(Certificate::Public) as u64;

        KeyPair {
            private_key,
            certificate,
            authority,
        }
    }
}

#[derive(Debug)]
pub struct AuthorityBuilder {
    alg: &'static SignatureAlgorithm,
}

macro_rules! authority {
    ($(($alg:ident, $lower:ident)),* $(,)?) => {
        impl AuthorityBuilder {
            $(
                pub fn $lower(&mut self) -> &mut Self {
                    self.alg = &rcgen::$alg;
                    self
                }
            )*
        }
    };
}

authority!(
    (PKCS_ECDSA_P256_SHA256, ecdsa),
    (PKCS_ED25519, ed25519),
    (PKCS_RSA_SHA256, rsa_256),
    (PKCS_RSA_SHA384, rsa_384),
    (PKCS_RSA_SHA512, rsa_512),
);

#[derive(Copy, Clone, Debug)]
pub struct KeyPair {
    pub(crate) authority: u64,
    pub(crate) private_key: u64,
    pub(crate) certificate: u64,
}

#[derive(Debug)]
pub struct KeyPairBuilder {
    authority: u64,
    intermediates: Vec<&'static SignatureAlgorithm>,
    alg: &'static SignatureAlgorithm,
}

impl KeyPairBuilder {
    pub fn push_ia(&mut self) -> &mut Self {
        self.push_ia_with(|_| {})
    }

    pub fn push_ia_with<F: FnOnce(&mut AuthorityBuilder)>(&mut self, f: F) -> &mut Self {
        let mut builder = AuthorityBuilder { alg: self.alg };
        f(&mut builder);
        self.intermediates.push(builder.alg);
        self
    }
}


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

use core::marker::PhantomData;

#[derive(Debug)]
pub struct Checkpoint<Endpoint, Location, Op> {
    pub(crate) id: u64,
    endpoint: PhantomData<Endpoint>,
    location: PhantomData<Location>,
    op: PhantomData<Op>,
}

impl<Endpoint, Location, Op> Checkpoint<Endpoint, Location, Op> {
    pub(crate) fn new(id: u64) -> Self {
        Self {
            id,
            endpoint: PhantomData,
            location: PhantomData,
            op: PhantomData,
        }
    }
}

#[derive(Debug)]
pub struct Park;

#[derive(Debug)]
pub struct Unpark;

macro_rules! sync {
    ($endpoint:ty, $location:ty) => {
        pub fn park(
            &mut self,
            checkpoint: crate::scenario::builder::checkpoint::Checkpoint<
                $endpoint,
                $location,
                crate::scenario::builder::checkpoint::Park,
            >,
        ) -> &mut Self {
            self.ops.push(crate::operation::Connection::Park {
                checkpoint: checkpoint.id,
            });
            self
        }

        pub fn unpark(
            &mut self,
            checkpoint: crate::scenario::builder::checkpoint::Checkpoint<
                $endpoint,
                $location,
                crate::scenario::builder::checkpoint::Unpark,
            >,
        ) -> &mut Self {
            self.ops.push(crate::operation::Connection::Unpark {
                checkpoint: checkpoint.id,
            });
            self
        }
    };
}


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

use super::{
    checkpoint::{Checkpoint, Park, Unpark},
    connection::{self, Connect, Connection},
    Local,
};
use crate::operation as op;
use core::marker::PhantomData;

#[derive(Debug)]
pub struct Builder {
    id: u64,
    state: super::State,
    ops: Vec<op::Client>,
}

impl Builder {
    pub(crate) fn new(id: u64, state: super::State) -> Self {
        Self {
            id,
            state,
            ops: vec![],
        }
    }

    pub fn connect_to<F: FnOnce(&mut connection::Builder<Client>), To: Connect<Client>>(
        &mut self,
        to: To,
        f: F,
    ) -> Connection<Client> {
        let mut builder = connection::Builder::new(self.state.connection());
        f(&mut builder);

        let template = builder.finish();
        let connection = Connection {
            endpoint_id: self.id,
            state: self.state.clone(),
            template,
            endpoint: PhantomData,
        };

        let op = to.connect_to(&connection);
        self.ops.push(op);

        connection
    }

    pub fn scope<F: FnOnce(&mut Scope)>(&mut self, f: F) -> &mut Self {
        let mut scope = Scope::new(self.id, self.state.clone());
        f(&mut scope);

        let threads = scope.finish();

        if threads.is_empty() {
            // no-op
        } else if threads.len() == 1 {
            // only a single thread was spawned, which is the same as not spawning it
            self.ops.extend(threads.into_iter().flatten());
        } else {
            self.ops.push(op::Client::Scope { threads });
        }

        self
    }

    pub fn checkpoint(
        &mut self,
    ) -> (
        Checkpoint<Client, Local, Park>,
        Checkpoint<Client, Local, Unpark>,
    ) {
        self.state.checkpoint()
    }

    pub(crate) fn finish(self) -> Vec<op::Client> {
        self.ops
    }
}

pub struct Scope {
    id: u64,
    state: super::State,
    threads: Vec<Vec<op::Client>>,
}

impl Scope {
    pub(crate) fn new(id: u64, state: super::State) -> Self {
        Self {
            id,
            state,
            threads: vec![],
        }
    }

    pub fn spawn<F: FnOnce(&mut Builder)>(&mut self, f: F) -> &mut Self {
        let mut builder = Builder::new(self.id, self.state.clone());
        f(&mut builder);
        self.threads.push(builder.finish());
        self
    }

    pub(crate) fn finish(self) -> Vec<Vec<op::Client>> {
        self.threads
    }
}

#[derive(Debug)]
pub struct Client {}


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

use super::{
    checkpoint::{Checkpoint, Park, Unpark},
    state::{IdPool, RefVec},
    stream::{ReceiveStream, SendStream, Stream},
    Client, Endpoint, Local, Remote, Scope, Server,
};
use crate::operation as op;
use core::marker::PhantomData;
use std::sync::Arc;

#[derive(Clone, Debug, Default)]
pub struct State {
    peer_streams: RefVec<Vec<op::Connection>>,
    stream: IdPool,
    scenario: super::State,
}

impl State {
    pub(crate) fn new(scenario: super::State) -> Self {
        Self {
            scenario,
            ..Default::default()
        }
    }
}

impl core::ops::Deref for State {
    type Target = super::State;

    fn deref(&self) -> &Self::Target {
        &self.scenario
    }
}

#[derive(Debug)]
pub struct Builder<Endpoint> {
    ops: Vec<op::Connection>,
    state: State,
    endpoint: PhantomData<Endpoint>,
}

impl<E: Endpoint> Builder<E> {
    pub(crate) fn new(state: State) -> Self {
        Self {
            ops: vec![],
            state,
            endpoint: PhantomData,
        }
    }

    fn child_scope(&self) -> Self {
        Self::new(self.state.clone())
    }

    pub fn checkpoint<Location>(
        &self,
    ) -> (
        Checkpoint<E, Location, Park>,
        Checkpoint<E, Location, Unpark>,
    ) {
        self.state.checkpoint()
    }

    sync!(E, Local);
    sleep!();
    trace!();
    iterate!();

    pub fn scope<F: FnOnce(&mut Scope<E>)>(&mut self, f: F) -> &mut Self {
        let mut scope = Scope::new(self.state.clone());
        f(&mut scope);

        let threads = scope.finish();

        if threads.is_empty() {
            // no-op
        } else if threads.len() == 1 {
            // only a single thread was spawned, which is the same as not spawning it
            self.ops.extend(threads.into_iter().flatten());
        } else {
            self.ops.push(op::Connection::Scope { threads });
        }

        self
    }

    pub fn concurrently<A: FnOnce(&mut Builder<E>), B: FnOnce(&mut Builder<E>)>(
        &mut self,
        a: A,
        b: B,
    ) -> &mut Self {
        self.scope(|scope| {
            scope.spawn(a);
            scope.spawn(b);
        })
    }

    pub fn open_bidirectional_stream<
        L: FnOnce(&mut Stream<E, Local>),
        R: FnOnce(&mut Stream<E::Peer, Remote>),
    >(
        &mut self,
        local: L,
        remote: R,
    ) -> &mut Self {
        let id = self.state.stream.next_id();
        let mut local_stream = Stream::new(id, self.state.clone());
        let mut remote_stream = Stream::new(id, self.state.clone());

        local(&mut local_stream);
        remote(&mut remote_stream);

        self.ops
            .push(op::Connection::OpenBidirectionalStream { stream_id: id });
        self.ops.extend(local_stream.finish());

        debug_assert_eq!(id, self.state.peer_streams.len() as u64);
        self.state.peer_streams.push(remote_stream.finish());

        self
    }

    pub fn open_send_stream<
        L: FnOnce(&mut SendStream<E, Local>),
        R: FnOnce(&mut ReceiveStream<E::Peer, Remote>),
    >(
        &mut self,
        local: L,
        remote: R,
    ) -> &mut Self {
        let id = self.state.stream.next_id();
        let mut local_stream = SendStream::new(id, self.state.clone());
        let mut remote_stream = ReceiveStream::new(id, self.state.clone());

        local(&mut local_stream);
        remote(&mut remote_stream);

        self.ops
            .push(op::Connection::OpenSendStream { stream_id: id });
        self.ops.extend(local_stream.finish());

        debug_assert_eq!(id, self.state.peer_streams.len() as u64);
        self.state.peer_streams.push(remote_stream.finish());

        self
    }

    pub(crate) fn finish(self) -> crate::scenario::Connection {
        let peer_streams = self.state.peer_streams.take();
        let ops = self.ops;
        crate::scenario::Connection { ops, peer_streams }
    }

    pub(crate) fn finish_scope(self) -> Vec<op::Connection> {
        self.ops
    }
}

#[derive(Debug)]
pub struct Connection<Endpoint> {
    pub(crate) state: super::State,
    pub(crate) endpoint_id: u64,
    pub(crate) template: crate::scenario::Connection,
    pub(crate) endpoint: PhantomData<Endpoint>,
}

pub trait Connect<Endpoint> {
    fn connect_to(&self, handle: &Connection<Endpoint>) -> op::Client;
}

impl Connect<Client> for Connection<Server> {
    fn connect_to(&self, handle: &Connection<Client>) -> op::Client {
        let server_id = self.endpoint_id;
        let server = &mut self.state.servers.borrow_mut()[server_id as usize];

        fn push(
            connections: &mut Vec<Arc<crate::scenario::Connection>>,
            ops: &Vec<op::Connection>,
            peer_streams: &Vec<Vec<op::Connection>>,
        ) -> u64 {
            // try to dedupe the connection operations if one exists
            for (id, prev) in connections.iter().enumerate() {
                if &prev.ops == ops && &prev.peer_streams == peer_streams {
                    return id as u64;
                }
            }

            let id = connections.len() as u64;

            connections.push(Arc::new(crate::scenario::Connection {
                ops: ops.clone(),
                peer_streams: peer_streams.clone(),
            }));

            id
        }

        let server_connection_id = push(
            &mut server.connections,
            &self.template.ops,
            &handle.template.peer_streams,
        );

        let certificate_authority = server.certificate_authority;

        let client = &mut self.state.clients.borrow_mut()[handle.endpoint_id as usize];
        let client_connection_id = push(
            &mut client.connections,
            &handle.template.ops,
            &self.template.peer_streams,
        );

        if !client
            .certificate_authorities
            .contains(&certificate_authority)
        {
            client.certificate_authorities.push(certificate_authority);
        }

        op::Client::Connect {
            server_id,
            router_id: None,
            server_connection_id,
            client_connection_id,
        }
    }
}

impl Connect<Client> for Server {
    fn connect_to(&self, handle: &Connection<Client>) -> op::Client {
        self.with(|_| {
            // empty instructions
        })
        .connect_to(handle)
    }
}

impl Connect<Client> for &Server {
    fn connect_to(&self, handle: &Connection<Client>) -> op::Client {
        self.with(|_| {
            // empty instructions
        })
        .connect_to(handle)
    }
}


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

use super::{connection, Endpoint};
use crate::operation as op;
use core::marker::PhantomData;

pub struct Scope<Endpoint> {
    state: connection::State,
    threads: Vec<Vec<op::Connection>>,
    endpoint: PhantomData<Endpoint>,
}

impl<E: Endpoint> Scope<E> {
    pub(crate) fn new(state: connection::State) -> Self {
        Self {
            state,
            threads: vec![],
            endpoint: PhantomData,
        }
    }

    pub fn spawn<F: FnOnce(&mut connection::Builder<E>)>(&mut self, f: F) -> &mut Self {
        let mut builder = connection::Builder::new(self.state.clone());
        f(&mut builder);
        self.threads.push(builder.finish_scope());
        self
    }

    pub(crate) fn finish(self) -> Vec<Vec<op::Connection>> {
        self.threads
    }
}


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

use super::{
    certificate::KeyPair,
    connection::{self, Connection},
};
use core::marker::PhantomData;

#[derive(Debug)]
pub struct Server {
    id: u64,
    state: super::State,
}

pub struct Builder {
    state: super::State,
    key_pair: Option<KeyPair>,
}

impl Builder {
    pub fn set_cert(&mut self, key_pair: KeyPair) -> &mut Self {
        self.key_pair = Some(key_pair);
        self
    }
}

impl Server {
    pub(crate) fn new<F: FnOnce(&mut Builder)>(state: super::State, f: F) -> Self {
        let mut builder = Builder {
            state,
            key_pair: None,
        };
        f(&mut builder);
        let state = builder.state;

        let mut server = crate::scenario::Server::default();

        let key_pair = builder.key_pair.unwrap_or_else(|| state.default_key_pair());
        server.certificate = key_pair.certificate;
        server.private_key = key_pair.private_key;
        server.certificate_authority = key_pair.authority;

        let id = state.servers.push(server) as u64;

        Self { id, state }
    }

    pub(crate) fn with<F: FnOnce(&mut connection::Builder<Server>)>(
        &self,
        f: F,
    ) -> Connection<Server> {
        let mut builder = connection::Builder::new(self.state.connection());
        f(&mut builder);

        Connection {
            endpoint_id: self.id,
            state: self.state.clone(),
            template: builder.finish(),
            endpoint: PhantomData,
        }
    }
}


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

use super::{
    certificate::{self, KeyPair},
    checkpoint::{Checkpoint, Park, Unpark},
    connection,
};
use crate::scenario::{Client, Server};
use core::fmt;
use std::{
    cell::{Cell, RefCell, RefMut},
    collections::HashMap,
    rc::Rc,
};

#[derive(Clone, Debug, Default)]
pub struct State {
    pub(crate) checkpoint: IdPool,
    pub(crate) servers: RefVec<Server>,
    pub(crate) clients: RefVec<Client>,
    pub(crate) trace: RefMap<String, u64>,
    pub(crate) certificates: RefVec<certificate::Certificate>,
    default_key_pair: Rc<Cell<Option<KeyPair>>>,
}

impl State {
    pub fn checkpoint<Endpoint, Location>(
        &self,
    ) -> (
        Checkpoint<Endpoint, Location, Park>,
        Checkpoint<Endpoint, Location, Unpark>,
    ) {
        let id = self.checkpoint.next_id();
        (Checkpoint::new(id), Checkpoint::new(id))
    }

    pub fn connection(&self) -> connection::State {
        connection::State::new(self.clone())
    }

    pub fn trace(&self, name: &str) -> u64 {
        if let Some(v) = self.trace.0.borrow().get(name) {
            return *v;
        }

        let mut map = self.trace.0.borrow_mut();

        let id = map.len() as u64;
        map.insert(name.to_string(), id);
        id
    }

    pub fn create_ca(&self) -> certificate::Authority {
        self.create_ca_with(|_| {})
    }

    pub fn create_ca_with<F: FnOnce(&mut certificate::AuthorityBuilder)>(
        &self,
        f: F,
    ) -> certificate::Authority {
        certificate::Authority::new(self.clone(), f)
    }

    pub fn default_key_pair(&self) -> KeyPair {
        if let Some(key) = self.default_key_pair.get() {
            key
        } else {
            let ca = self.create_ca();
            let pair = ca.key_pair();
            self.default_key_pair.set(Some(pair));
            pair
        }
    }
}

#[derive(Clone, Default)]
pub struct IdPool(Rc<RefCell<u64>>);

impl fmt::Debug for IdPool {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "IdPool({})", *self.0.borrow())
    }
}

impl IdPool {
    pub fn next_id(&self) -> u64 {
        let mut next_id = self.0.borrow_mut();
        let id = *next_id;
        *next_id += 1;
        id
    }
}

#[derive(Clone, Debug)]
pub struct RefMap<Key, Value>(Rc<RefCell<HashMap<Key, Value>>>);

impl<Key, Value> Default for RefMap<Key, Value> {
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<Key: core::hash::Hash + Eq, Value> RefMap<Key, Value> {
    #[allow(dead_code)]
    pub fn insert(&self, key: Key, value: Value) {
        self.0.borrow_mut().insert(key, value);
    }

    pub fn take(&self) -> HashMap<Key, Value> {
        core::mem::take(&mut self.0.borrow_mut())
    }
}

#[derive(Clone, Debug)]
pub struct RefVec<Value>(Rc<RefCell<Vec<Value>>>);

impl<Value> Default for RefVec<Value> {
    fn default() -> Self {
        Self(Default::default())
    }
}

impl<Value> RefVec<Value> {
    pub fn len(&self) -> usize {
        self.0.borrow().len()
    }

    pub fn push(&self, value: Value) -> usize {
        let mut v = self.0.borrow_mut();
        let len = v.len();
        v.push(value);
        len
    }

    pub fn take(&self) -> Vec<Value> {
        core::mem::take(&mut self.0.borrow_mut())
    }

    pub fn borrow_mut(&self) -> RefMut<Vec<Value>> {
        self.0.borrow_mut()
    }
}


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

use super::connection;
use crate::{
    operation as op,
    units::{Byte, Rate},
};
use core::marker::PhantomData;

macro_rules! send_stream {
    () => {
        pub fn send(&mut self, bytes: Byte) -> &mut Self {
            self.ops.push(crate::operation::Connection::Send {
                stream_id: self.id,
                bytes,
            });
            self
        }

        pub fn set_send_rate(&mut self, rate: Rate) -> &mut Self {
            self.ops.push(crate::operation::Connection::SendRate {
                stream_id: self.id,
                rate,
            });
            self
        }
    };
}

macro_rules! receive_stream {
    () => {
        pub fn receive(&mut self, bytes: Byte) -> &mut Self {
            self.ops.push(crate::operation::Connection::Receive {
                stream_id: self.id,
                bytes,
            });
            self
        }

        pub fn set_receive_rate(&mut self, rate: Rate) -> &mut Self {
            self.ops.push(crate::operation::Connection::ReceiveRate {
                stream_id: self.id,
                rate,
            });
            self
        }

        pub fn receive_all(&mut self) -> &mut Self {
            self.ops
                .push(crate::operation::Connection::ReceiveAll { stream_id: self.id });
            self
        }
    };
}

pub struct Stream<Endpoint, Location> {
    id: u64,
    ops: Vec<op::Connection>,
    state: connection::State,
    endpoint: PhantomData<Endpoint>,
    location: PhantomData<Location>,
}

impl<Endpoint, Location> Stream<Endpoint, Location> {
    send_stream!();
    receive_stream!();
    sync!(Endpoint, Location);
    sleep!();
    trace!();
    iterate!();

    pub(crate) fn new(id: u64, state: connection::State) -> Self {
        Self {
            id,
            ops: vec![],
            state,
            endpoint: PhantomData,
            location: PhantomData,
        }
    }

    fn child_scope(&self) -> Self {
        Self::new(self.id, self.state.clone())
    }

    fn finish_scope(self) -> Vec<op::Connection> {
        self.ops
    }

    pub fn concurrently<
        S: FnOnce(&mut SendStream<Endpoint, Location>),
        R: FnOnce(&mut ReceiveStream<Endpoint, Location>),
    >(
        &mut self,
        send: S,
        receive: R,
    ) -> &mut Self {
        let mut send_stream = SendStream::new(self.id, self.state.clone());
        let mut receive_stream = ReceiveStream::new(self.id, self.state.clone());
        send(&mut send_stream);
        receive(&mut receive_stream);
        let threads = vec![send_stream.ops, receive_stream.ops];
        self.ops.push(op::Connection::Scope { threads });
        self
    }

    pub(crate) fn finish(mut self) -> Vec<op::Connection> {
        let stream_id = self.id;
        self.ops.push(op::Connection::SendFinish { stream_id });
        self.ops.push(op::Connection::ReceiveFinish { stream_id });
        self.ops
    }
}

pub struct SendStream<Endpoint, Location> {
    id: u64,
    ops: Vec<op::Connection>,
    state: connection::State,
    endpoint: PhantomData<Endpoint>,
    location: PhantomData<Location>,
}

impl<Endpoint, Location> SendStream<Endpoint, Location> {
    send_stream!();
    sync!(Endpoint, Location);
    sleep!();
    trace!();
    iterate!();

    pub(crate) fn new(id: u64, state: connection::State) -> Self {
        Self {
            id,
            ops: vec![],
            state,
            endpoint: PhantomData,
            location: PhantomData,
        }
    }

    fn child_scope(&self) -> Self {
        Self::new(self.id, self.state.clone())
    }

    fn finish_scope(self) -> Vec<op::Connection> {
        self.ops
    }

    pub(crate) fn finish(mut self) -> Vec<op::Connection> {
        let stream_id = self.id;
        self.ops.push(op::Connection::SendFinish { stream_id });
        self.ops
    }
}

pub struct ReceiveStream<Endpoint, Location> {
    id: u64,
    ops: Vec<op::Connection>,
    state: connection::State,
    endpoint: PhantomData<Endpoint>,
    location: PhantomData<Location>,
}

impl<Endpoint, Location> ReceiveStream<Endpoint, Location> {
    receive_stream!();
    sync!(Endpoint, Location);
    sleep!();
    trace!();
    iterate!();

    pub(crate) fn new(id: u64, state: connection::State) -> Self {
        Self {
            id,
            ops: vec![],
            state,
            endpoint: PhantomData,
            location: PhantomData,
        }
    }

    fn child_scope(&self) -> Self {
        Self::new(self.id, self.state.clone())
    }

    fn finish_scope(self) -> Vec<op::Connection> {
        self.ops
    }

    pub(crate) fn finish(mut self) -> Vec<op::Connection> {
        let stream_id = self.id;
        self.ops.push(op::Connection::ReceiveFinish { stream_id });
        self.ops
    }
}


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

use crate::{
    scenario::{Builder, Scenario},
    units::*,
};
use insta::assert_json_snapshot;

fn scenario<F: FnOnce(&mut Builder)>(f: F) -> Scenario {
    let mut scenario = Scenario::build(f);
    // hash traversal order was changed somewhere between 1.53 and 1.58 so
    // we won't compare the id for now
    scenario.id = Default::default();
    // cert keys are not deterministic
    scenario.certificates.clear();
    scenario
}

macro_rules! scenario_test {
    ($name:ident, $builder:expr) => {
        #[test]
        fn $name() {
            assert_json_snapshot!(scenario($builder));
        }
    };
}

scenario_test!(simple, |scenario| {
    let server = scenario.create_server();

    scenario.create_client(|client| {
        client.connect_to(server, |conn| {
            conn.open_send_stream(
                |local| {
                    local.set_send_rate(1024.bytes() / 50.millis());
                    local.send(1.megabytes());
                },
                |peer| {
                    peer.set_receive_rate(1024.bytes() / 50.millis());
                    peer.receive(1.megabytes());
                },
            );
        });
    });
});

scenario_test!(conn_checkpoints, |scenario| {
    let server = scenario.create_server();

    scenario.create_client(|client| {
        client.connect_to(server, |conn| {
            let (cp1_rx, cp1_tx) = conn.checkpoint();

            conn.concurrently(
                |conn| {
                    conn.open_send_stream(
                        |local| {
                            local.set_send_rate(10.kilobytes() / 50.millis());
                            local.send(1.megabytes() / 2);
                            local.unpark(cp1_tx);
                            local.send(1.megabytes() / 2);
                        },
                        |peer| {
                            peer.set_receive_rate(10.kilobytes() / 50.millis());
                            peer.receive(1.megabytes());
                        },
                    );
                },
                |conn| {
                    conn.open_send_stream(
                        |local| {
                            local.park(cp1_rx);
                            local.set_send_rate(1024.bytes() / 50.millis());
                            local.send(1.megabytes());
                        },
                        |peer| {
                            peer.set_receive_rate(1024.bytes() / 50.millis());
                            peer.receive(1.megabytes());
                        },
                    );
                },
            );
        });
    });
});

scenario_test!(linked_streams, |scenario| {
    let server = scenario.create_server();

    scenario.create_client(|client| {
        client.connect_to(server, |conn| {
            let (b_park, b_unpark) = conn.checkpoint();

            conn.concurrently(
                |conn| {
                    conn.open_bidirectional_stream(
                        |local| {
                            local.concurrently(
                                |sender| {
                                    sender.set_send_rate(1024.bytes() / 50.millis());
                                    sender.send(1.megabytes());
                                },
                                |receiver| {
                                    receiver.receive_all();
                                },
                            );

                            local.unpark(b_unpark);
                        },
                        |peer| {
                            peer.sleep(100.millis());

                            peer.set_receive_rate(1024.bytes() / 50.millis());
                            peer.receive(100.kilobytes());

                            peer.set_receive_rate(10.bytes() / 50.millis());
                            peer.receive_all();

                            peer.send(2.megabytes());
                        },
                    );
                },
                |conn| {
                    conn.open_send_stream(
                        |local| {
                            local.park(b_park);
                            local.send(1.megabytes());
                        },
                        |peer| {
                            peer.receive(1.megabytes());
                        },
                    );
                },
            );
        });
    });
});

scenario_test!(custom_cert, |scenario| {
    let ca = scenario.create_ca();

    let server = scenario.create_server_with(|server| {
        server.set_cert(ca.key_pair());
    });

    scenario.create_client(|client| {
        client.connect_to(server, |conn| {
            conn.open_send_stream(
                |local| {
                    local.set_send_rate(1024.bytes() / 50.millis());
                    local.send(1.megabytes());
                },
                |peer| {
                    peer.set_receive_rate(1024.bytes() / 50.millis());
                    peer.receive(1.megabytes());
                },
            );
        });
    });
});

scenario_test!(long_chain_cert, |scenario| {
    let ca = scenario.create_ca_with(|ca| {
        ca.ecdsa();
    });

    let server = scenario.create_server_with(|server| {
        let key = ca.key_pair_with(|key| {
            key.push_ia();
            key.push_ia_with(|ia| {
                ia.ed25519();
            });
        });
        server.set_cert(key);
    });

    scenario.create_client(|client| {
        client.connect_to(server, |conn| {
            conn.open_send_stream(
                |local| {
                    local.set_send_rate(1024.bytes() / 50.millis());
                    local.send(1.megabytes());
                },
                |peer| {
                    peer.set_receive_rate(1024.bytes() / 50.millis());
                    peer.receive(1.megabytes());
                },
            );
        });
    });
});


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

use super::Timestamp;
use core::{
    task::{Context, Poll, Waker},
    time::Duration,
};

#[derive(Debug)]
pub struct Timer {
    now: Timestamp,
    target: Option<Timestamp>,
    waker: Option<Waker>,
}

impl Default for Timer {
    fn default() -> Self {
        Self {
            now: unsafe { Timestamp::from_duration(Duration::ZERO) },
            target: None,
            waker: None,
        }
    }
}

impl Timer {
    pub fn advance_pair(&mut self, other: &mut Self) -> Option<Timestamp> {
        let target = match (self.target, other.target) {
            (Some(a), Some(b)) => Some(a.min(b)),
            (a, None) => a,
            (None, b) => b,
        };

        if let Some(target) = target {
            self.set(target);
            other.set(target);
            Some(self.now)
        } else {
            None
        }
    }

    fn set(&mut self, target: Timestamp) {
        self.now = target;
        self.target = None;
        if let Some(waker) = self.waker.take() {
            waker.wake();
        }
    }
}

impl super::Timer for Timer {
    fn now(&self) -> Timestamp {
        self.now
    }

    fn poll(&mut self, target: Timestamp, cx: &mut Context<'_>) -> Poll<()> {
        if target == self.now {
            Poll::Ready(())
        } else {
            self.target = Some(target);
            self.waker = Some(cx.waker().clone());
            Poll::Pending
        }
    }
}


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

use super::Timestamp;
use core::{
    future::Future,
    pin::Pin,
    task::{Context, Poll},
    time::Duration,
};
use futures::ready;
use tokio::time::{sleep, Instant, Sleep};

#[derive(Debug)]
pub struct Timer {
    start: Instant,
    target: Option<Timestamp>,
    sleep: Pin<Box<Sleep>>,
}

impl Timer {
    fn set_target(&mut self, target: Timestamp) {
        self.target = Some(target);
        let duration = unsafe { target.as_duration() };
        self.sleep.as_mut().reset(self.start + duration);
    }
}

impl Default for Timer {
    fn default() -> Self {
        Self {
            start: Instant::now(),
            target: None,
            sleep: Box::pin(sleep(Duration::from_secs(0))),
        }
    }
}

impl super::Timer for Timer {
    fn now(&self) -> Timestamp {
        let duration = self.start.elapsed();
        unsafe { Timestamp::from_duration(duration) }
    }

    fn poll(&mut self, target: Timestamp, cx: &mut Context) -> Poll<()> {
        if let Some(prev_target) = self.target.as_mut() {
            if *prev_target != target {
                self.set_target(target);
            }
        } else {
            self.set_target(target);
        }

        ready!(self.sleep.as_mut().poll(cx));

        self.target = None;

        Poll::Ready(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::timer::Timer as _;
    use futures_test::task::new_count_waker;

    #[tokio::test(start_paused = true)]
    async fn timer_test() {
        let mut timer = Timer::default();
        let (waker, _count) = new_count_waker();
        let mut cx = Context::from_waker(&waker);

        tokio::time::advance(Duration::from_secs(1)).await;

        let mut now = timer.now();

        let mut times = [now; 5];
        for (idx, time) in times.iter_mut().enumerate() {
            *time += Duration::from_secs(idx as _);
        }

        assert!(timer.poll(now, &mut cx).is_ready());

        for _ in 0..times.len() {
            // poll a bunch of different times in loop to make sure all of the branches are covered
            for time in times.iter().chain(times.iter().rev()).copied() {
                assert_eq!(timer.poll(time, &mut cx).is_ready(), time <= now);
                assert_eq!(timer.poll(time, &mut cx).is_ready(), time <= now);
            }

            tokio::time::advance(Duration::from_secs(1)).await;
            now += Duration::from_secs(1);
        }
    }
}


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

// if the probe is disabled, then the variables won't be used.
#![allow(unused_variables)]

use super::Trace;
use crate::{timer::Timestamp, units::*};
use probe::probe;

#[derive(Clone, Debug, Default)]
pub struct Usdt {
    connection_id: u64,
}

impl Trace for Usdt {
    #[inline(always)]
    fn enter_connection(&mut self, id: u64) {
        self.connection_id = id;
    }

    #[inline(never)]
    fn send(&mut self, _now: Timestamp, stream_id: u64, len: u64) {
        probe!(netbench, netbench__send, self.connection_id, stream_id, len);
    }

    #[inline(never)]
    fn send_finish(&mut self, _now: Timestamp, stream_id: u64) {
        probe!(
            netbench,
            netbench__send__finish,
            self.connection_id,
            stream_id
        );
    }

    #[inline(never)]
    fn receive(&mut self, _now: Timestamp, stream_id: u64, len: u64) {
        probe!(
            netbench,
            netbench__receive,
            self.connection_id,
            stream_id,
            len
        );
    }

    #[inline(never)]
    fn receive_finish(&mut self, _now: Timestamp, stream_id: u64) {
        probe!(
            netbench,
            netbench__receive__finish,
            self.connection_id,
            stream_id
        );
    }

    #[inline(never)]
    fn accept(&mut self, _now: Timestamp, stream_id: u64) {
        probe!(netbench, netbench__accept, self.connection_id, stream_id);
    }

    #[inline(never)]
    fn open(&mut self, _now: Timestamp, stream_id: u64) {
        probe!(netbench, netbench__open, self.connection_id, stream_id);
    }

    #[inline(never)]
    fn trace(&mut self, _now: Timestamp, id: u64) {
        probe!(netbench, netbench__trace, self.connection_id, id);
    }

    #[inline(never)]
    fn profile(&mut self, now: Timestamp, id: u64, time: Duration) {
        let time = time.as_micros() as u64;
        probe!(netbench, netbench__profile, self.connection_id, id, time);
    }

    #[inline(never)]
    fn park(&mut self, _now: Timestamp, id: u64) {
        probe!(netbench, netbench__park, self.connection_id, id);
    }

    #[inline(never)]
    fn unpark(&mut self, _now: Timestamp, id: u64) {
        probe!(netbench, netbench__unpark, self.connection_id, id);
    }

    #[inline(never)]
    fn connect(&mut self, _now: Timestamp, id: u64, time: Duration) {
        let time = time.as_micros() as u64;
        probe!(netbench, netbench__connect, self.connection_id, id, time);
    }
}


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

use crate::units::{Duration, Rate};
use core::{fmt, ops};
use serde::{Deserialize, Serialize};

#[derive(
    Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize,
)]
pub struct Byte(u64);

impl Byte {
    pub const MAX: Self = Self(u64::MAX);
    pub const MIN: Self = Self(u64::MIN);
}

impl ops::Add for Byte {
    type Output = Self;

    fn add(self, rhs: Self) -> Self::Output {
        Self(self.0 + rhs.0)
    }
}

impl ops::AddAssign for Byte {
    fn add_assign(&mut self, rhs: Self) {
        self.0 += rhs.0;
    }
}

impl ops::Sub for Byte {
    type Output = Self;

    fn sub(self, rhs: Self) -> Self::Output {
        Self(self.0 - rhs.0)
    }
}

impl ops::SubAssign for Byte {
    fn sub_assign(&mut self, rhs: Self) {
        self.0 -= rhs.0;
    }
}

impl ops::Mul<u64> for Byte {
    type Output = Self;

    fn mul(self, rhs: u64) -> Self::Output {
        Self(self.0 * rhs)
    }
}

impl ops::Div<u64> for Byte {
    type Output = Self;

    fn div(self, rhs: u64) -> Self::Output {
        Self(self.0 / rhs)
    }
}

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

    fn div(self, period: Duration) -> Self::Output {
        Rate {
            bytes: self,
            period,
        }
    }
}

impl ops::Deref for Byte {
    type Target = u64;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl fmt::Display for Byte {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let opts = humansize::file_size_opts::FileSizeOpts {
            space: false,
            ..humansize::file_size_opts::DECIMAL
        };
        humansize::FileSize::file_size(&self.0, opts)
            .unwrap()
            .fmt(f)
    }
}

impl core::str::FromStr for Byte {
    type Err = crate::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        // try parsing without the suffix
        if let Ok(v) = s.parse() {
            return Ok(Self(v));
        }

        let number_index = s
            .char_indices()
            .find_map(|(idx, c)| {
                if !(c.is_numeric() || c == '.') {
                    Some(idx)
                } else {
                    None
                }
            })
            .unwrap_or(s.len());

        let mut v = Self(s[..number_index].parse()?);

        let mut suffix = s[number_index..].trim();
        let mut is_bits = false;

        if let Some(s) = suffix.strip_suffix('B') {
            suffix = s;
        } else if let Some(s) = suffix.strip_suffix('b') {
            is_bits = true;
            suffix = s;
        }

        v.0 *= *match suffix.trim() {
            "" => 1.bytes(),
            "K" | "k" => 1.kilobytes(),
            "Ki" | "ki" => 1.kibibytes(),
            "M" | "m" => 1.megabytes(),
            "Mi" | "mi" => 1.mebibytes(),
            "G" | "g" => 1.gigabytes(),
            "Gi" | "gi" => 1.gibibytes(),
            "T" | "t" => 1.terabytes(),
            "Ti" | "ti" => 1.tebibytes(),
            _ => return Err(format!("invalid bytes: {s:?}").into()),
        };

        if is_bits {
            // round up to the nearest byte
            if v.0 % 8 != 0 {
                v += 8.bytes();
            }
            v.0 /= 8;
        }

        Ok(v)
    }
}

pub trait ByteExt {
    fn bytes(&self) -> Byte;
    fn kilobytes(&self) -> Byte {
        self.bytes() * 1000
    }
    fn kibibytes(&self) -> Byte {
        self.bytes() * 1024
    }
    fn megabytes(&self) -> Byte {
        self.kilobytes() * 1000
    }
    fn mebibytes(&self) -> Byte {
        self.kibibytes() * 1024
    }
    fn gigabytes(&self) -> Byte {
        self.megabytes() * 1000
    }
    fn gibibytes(&self) -> Byte {
        self.mebibytes() * 1024
    }
    fn terabytes(&self) -> Byte {
        self.gigabytes() * 1000
    }
    fn tebibytes(&self) -> Byte {
        self.gibibytes() * 1024
    }
}

impl ByteExt for u64 {
    fn bytes(&self) -> Byte {
        Byte(*self)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use insta::assert_debug_snapshot;

    #[test]
    fn ext_test() {
        assert_debug_snapshot!([
            42.bytes(),
            42.kilobytes(),
            42.kibibytes(),
            42.megabytes(),
            42.mebibytes(),
            42.gigabytes(),
            42.gibibytes(),
            42.kibibytes() + 42.bytes()
        ]);
    }

    fn p(s: &str) -> crate::Result<Byte> {
        s.parse()
    }

    #[test]
    fn parse_test() {
        assert_debug_snapshot!([
            p("42b"),
            p("42B"),
            p("42Kb"),
            p("42KB"),
            p("42Kib"),
            p("42KiB"),
            p("42Mb"),
            p("42MB"),
            p("42Mib"),
            p("42MiB"),
            p("42Gb"),
            p("42GB"),
            p("42Gib"),
            p("42GiB"),
            p("42Tb"),
            p("42TB"),
            p("42Tib"),
            p("42TiB"),
        ]);
    }
}


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

pub use core::time::Duration;

pub fn parse_duration(s: &str) -> crate::Result<Duration> {
    Ok(humantime::parse_duration(s)?)
}

pub trait DurationExt {
    fn millis(&self) -> Duration;
    fn seconds(&self) -> Duration {
        self.millis() * 1000
    }
    fn minutes(&self) -> Duration {
        self.seconds() * 60
    }
}

impl DurationExt for f32 {
    fn millis(&self) -> Duration {
        Duration::from_secs_f32(*self / 1000.0)
    }
    fn seconds(&self) -> Duration {
        Duration::from_secs_f32(*self)
    }
}

impl DurationExt for u64 {
    fn millis(&self) -> Duration {
        Duration::from_millis(*self)
    }
}

pub(crate) mod duration_format {
    use core::time::Duration;
    use serde::{self, Deserialize, Deserializer, Serializer};

    pub fn serialize<S>(duration: &Duration, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        serializer.serialize_u64(duration.as_millis() as u64)
    }

    pub fn deserialize<'de, D>(deserializer: D) -> Result<Duration, D::Error>
    where
        D: Deserializer<'de>,
    {
        let millis = u64::deserialize(deserializer)?;
        Ok(Duration::from_millis(millis))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use insta::assert_debug_snapshot;

    #[test]
    fn ext_test() {
        assert_debug_snapshot!([
            42.millis(),
            42.seconds(),
            42.minutes(),
            42.minutes() + 42.seconds(),
            4.2.millis(),
            4.2.seconds(),
            4.2.minutes(),
            4.2.minutes() + 4.2.seconds(),
        ]);
    }
}


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

use crate::units::{duration_format, Byte, ByteExt, Duration, DurationExt};
use core::fmt;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, Deserialize, Serialize)]
pub struct Rate {
    pub bytes: Byte,
    #[serde(with = "duration_format", rename = "period_ms")]
    pub period: Duration,
}

impl fmt::Display for Rate {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.period == 1.seconds() {
            return write!(f, "{}ps", self.bytes);
        }

        // force the period to be in seconds
        if f.alternate() {
            let factor = 1.0 / self.period.as_secs_f64();
            let bytes = (*self.bytes as f64 * factor) as u64;
            let bytes = bytes.bytes();
            return write!(f, "{bytes}ps");
        }

        write!(f, "{}/{:?}", self.bytes, self.period)
    }
}

impl core::str::FromStr for Rate {
    type Err = crate::Error;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if let Some(bytes) = s.strip_suffix("ps") {
            return Ok(Self {
                bytes: bytes.parse()?,
                period: 1.seconds(),
            });
        }

        if let Some((bytes, period)) = s.split_once('/') {
            Ok(Self {
                bytes: bytes.trim().parse()?,
                period: *humantime::Duration::from_str(period.trim())?,
            })
        } else {
            Err(format!("invalid rate: {s}").into())
        }
    }
}

#[derive(Debug, Default)]
pub(crate) struct Rates {
    pub send: HashMap<u64, Rate>,
    pub receive: HashMap<u64, Rate>,
}

#[cfg(test)]
mod tests {
    use crate::units::*;
    use insta::assert_debug_snapshot;

    #[test]
    fn ext_test() {
        assert_debug_snapshot!([
            42.bytes() / 42.millis(),
            42.mebibytes() / 42.seconds(),
            42.gigabytes() / 42.minutes()
        ]);
    }

    fn p(s: &str) -> crate::Result<Rate> {
        s.parse()
    }

    #[test]
    fn parse_test() {
        assert_debug_snapshot!([
            p("42bps"),
            p("42Bps"),
            p("42KBps"),
            p("42Kb / 50ms"),
            p("42Mb/5ms"),
        ]);
    }
}


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

use netbench::Result;
use structopt::StructOpt;

mod report;
mod report_tree;

#[derive(StructOpt)]
enum Args {
    Report(report::Report),
    ReportTree(report_tree::ReportTree),
}

fn main() -> Result<()> {
    match Args::from_args() {
        Args::Report(report) => report.run(),
        Args::ReportTree(report) => report.run(),
    }
}