#include "tcp_connection.hh"

#include <chrono>
#include <cstdlib>
#include <iomanip>
#include <iostream>
#include <string>

using namespace std;
using namespace std::chrono;

constexpr size_t len = 100 * 1024 * 1024;

void move_segments(TCPConnection &x, TCPConnection &y, vector<TCPSegment> &segments, const bool reorder) {
    while (not x.segments_out().empty()) {
        segments.emplace_back(move(x.segments_out().front()));
        x.segments_out().pop();
    }
    if (reorder) {
        for (auto it = segments.rbegin(); it != segments.rend(); ++it) {
            y.segment_received(move(*it));
        }
    } else {
        for (auto it = segments.begin(); it != segments.end(); ++it) {
            y.segment_received(move(*it));
        }
    }
    segments.clear();
}

void main_loop(const bool reorder) {
    TCPConfig config;
    TCPConnection x{config}, y{config};

    string string_to_send(len, 'x');
    for (auto &ch : string_to_send) {
        ch = rand();
    }

    Buffer bytes_to_send{string(string_to_send)};
    x.connect();
    y.end_input_stream();

    bool x_closed = false;

    string string_received;
    string_received.reserve(len);

    const auto first_time = high_resolution_clock::now();

    auto loop = [&] {
        // write input into x
        while (bytes_to_send.size() and x.remaining_outbound_capacity()) {
            const auto want = min(x.remaining_outbound_capacity(), bytes_to_send.size());
            const auto written = x.write(string(bytes_to_send.str().substr(0, want)));
            if (want != written) {
                throw runtime_error("want = " + to_string(want) + ", written = " + to_string(written));
            }
            bytes_to_send.remove_prefix(written);
        }

        if (bytes_to_send.size() == 0 and not x_closed) {
            x.end_input_stream();
            x_closed = true;
        }

        // exchange segments between x and y but in reverse order
        vector<TCPSegment> segments;
        move_segments(x, y, segments, reorder);
        move_segments(y, x, segments, false);

        // read output from y
        const auto available_output = y.inbound_stream().buffer_size();
        if (available_output > 0) {
            string_received.append(y.inbound_stream().read(available_output));
        }

        // time passes
        x.tick(1000);
        y.tick(1000);
    };

    while (not y.inbound_stream().eof()) {
        loop();
    }

    if (string_received != string_to_send) {
        throw runtime_error("strings sent vs. received don't match");
    }

    const auto final_time = high_resolution_clock::now();

    const auto duration = duration_cast<nanoseconds>(final_time - first_time).count();

    const auto gigabits_per_second = len * 8.0 / double(duration);

    cout << fixed << setprecision(2);
    cout << "CPU-limited throughput" << (reorder ? " with reordering: " : "                : ") << gigabits_per_second
         << " Gbit/s\n";

    while (x.active() or y.active()) {
        loop();
    }
}

int main() {
    try {
        main_loop(false);
        main_loop(true);
    } catch (const exception &e) {
        cerr << e.what() << "\n";
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}