#include "sender_harness.hh"
#include "wrapping_integers.hh"

#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <string>

using namespace std;

int main() {
    try {
        auto rd = get_random_generator();

        {
            TCPConfig cfg;
            WrappingInt32 isn(rd());
            cfg.fixed_isn = isn;

            TCPSenderTestHarness test{"Three short writes", cfg};
            test.execute(ExpectSegment{}.with_no_flags().with_syn(true).with_payload_size(0).with_seqno(isn));
            test.execute(AckReceived{WrappingInt32{isn + 1}});
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
            test.execute(WriteBytes{"ab"});
            test.execute(ExpectSegment{}.with_data("ab").with_seqno(isn + 1));
            test.execute(WriteBytes{"cd"});
            test.execute(ExpectSegment{}.with_data("cd").with_seqno(isn + 3));
            test.execute(WriteBytes{"abcd"});
            test.execute(ExpectSegment{}.with_data("abcd").with_seqno(isn + 5));
            test.execute(ExpectSeqno{WrappingInt32{isn + 9}});
            test.execute(ExpectBytesInFlight{8});
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
        }

        {
            TCPConfig cfg;
            WrappingInt32 isn(rd());
            cfg.fixed_isn = isn;

            TCPSenderTestHarness test{"Many short writes, continuous acks", cfg};
            test.execute(ExpectSegment{}.with_no_flags().with_syn(true).with_payload_size(0).with_seqno(isn));
            test.execute(AckReceived{WrappingInt32{isn + 1}});
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
            uint32_t max_block_size = 10;
            uint32_t n_rounds = 10000;
            size_t bytes_sent = 0;
            for (uint32_t i = 0; i < n_rounds; ++i) {
                string data;
                uint32_t block_size = uniform_int_distribution<uint32_t>{1, max_block_size}(rd);
                for (uint8_t j = 0; j < block_size; ++j) {
                    uint8_t c = 'a' + ((i + j) % 26);
                    data.push_back(c);
                }
                test.execute(ExpectSeqno{WrappingInt32{isn + uint32_t(bytes_sent) + 1}});
                test.execute(WriteBytes(string(data)));
                bytes_sent += block_size;
                test.execute(ExpectBytesInFlight{block_size});
                test.execute(
                    ExpectSegment{}.with_seqno(isn + 1 + uint32_t(bytes_sent - block_size)).with_data(move(data)));
                test.execute(ExpectNoSegment{});
                test.execute(AckReceived{WrappingInt32{isn + 1 + uint32_t(bytes_sent)}});
            }
        }

        {
            TCPConfig cfg;
            WrappingInt32 isn(rd());
            cfg.fixed_isn = isn;

            TCPSenderTestHarness test{"Many short writes, ack at end", cfg};
            test.execute(ExpectSegment{}.with_no_flags().with_syn(true).with_payload_size(0).with_seqno(isn));
            test.execute(AckReceived{WrappingInt32{isn + 1}}.with_win(65000));
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
            uint32_t max_block_size = 10;
            uint32_t n_rounds = 1000;
            size_t bytes_sent = 0;
            for (uint32_t i = 0; i < n_rounds; ++i) {
                string data;
                uint32_t block_size = uniform_int_distribution<uint32_t>{1, max_block_size}(rd);
                for (uint8_t j = 0; j < block_size; ++j) {
                    uint8_t c = 'a' + ((i + j) % 26);
                    data.push_back(c);
                }
                test.execute(ExpectSeqno{WrappingInt32{isn + uint32_t(bytes_sent) + 1}});
                test.execute(WriteBytes(string(data)));
                bytes_sent += block_size;
                test.execute(ExpectBytesInFlight{bytes_sent});
                test.execute(
                    ExpectSegment{}.with_seqno(isn + 1 + uint32_t(bytes_sent - block_size)).with_data(move(data)));
                test.execute(ExpectNoSegment{});
            }
            test.execute(ExpectBytesInFlight{bytes_sent});
            test.execute(AckReceived{WrappingInt32{isn + 1 + uint32_t(bytes_sent)}});
            test.execute(ExpectBytesInFlight{0});
        }

        {
            TCPConfig cfg;
            WrappingInt32 isn(rd());
            cfg.fixed_isn = isn;

            TCPSenderTestHarness test{"Window filling", cfg};
            test.execute(ExpectSegment{}.with_no_flags().with_syn(true).with_payload_size(0).with_seqno(isn));
            test.execute(AckReceived{WrappingInt32{isn + 1}}.with_win(3));
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
            test.execute(WriteBytes("01234567"));
            test.execute(ExpectBytesInFlight{3});
            test.execute(ExpectSegment{}.with_data("012"));
            test.execute(ExpectNoSegment{});
            test.execute(ExpectSeqno{WrappingInt32{isn + 1 + 3}});
            test.execute(AckReceived{WrappingInt32{isn + 1 + 3}}.with_win(3));
            test.execute(ExpectBytesInFlight{3});
            test.execute(ExpectSegment{}.with_data("345"));
            test.execute(ExpectNoSegment{});
            test.execute(ExpectSeqno{WrappingInt32{isn + 1 + 6}});
            test.execute(AckReceived{WrappingInt32{isn + 1 + 6}}.with_win(3));
            test.execute(ExpectBytesInFlight{2});
            test.execute(ExpectSegment{}.with_data("67"));
            test.execute(ExpectNoSegment{});
            test.execute(ExpectSeqno{WrappingInt32{isn + 1 + 8}});
            test.execute(AckReceived{WrappingInt32{isn + 1 + 8}}.with_win(3));
            test.execute(ExpectBytesInFlight{0});
            test.execute(ExpectNoSegment{});
        }

        {
            TCPConfig cfg;
            WrappingInt32 isn(rd());
            cfg.fixed_isn = isn;

            TCPSenderTestHarness test{"Immediate writes respect the window", cfg};
            test.execute(ExpectSegment{}.with_no_flags().with_syn(true).with_payload_size(0).with_seqno(isn));
            test.execute(AckReceived{WrappingInt32{isn + 1}}.with_win(3));
            test.execute(ExpectState{TCPSenderStateSummary::SYN_ACKED});
            test.execute(WriteBytes("01"));
            test.execute(ExpectBytesInFlight{2});
            test.execute(ExpectSegment{}.with_data("01"));
            test.execute(ExpectNoSegment{});
            test.execute(ExpectSeqno{WrappingInt32{isn + 1 + 2}});
            test.execute(WriteBytes("23"));
            test.execute(ExpectBytesInFlight{3});
            test.execute(ExpectSegment{}.with_data("2"));
            test.execute(ExpectNoSegment{});
            test.execute(ExpectSeqno{WrappingInt32{isn + 1 + 3}});
        }

    } catch (const exception &e) {
        cerr << e.what() << endl;
        return 1;
    }

    return EXIT_SUCCESS;
}