#include "tcp_config.hh"
#include "tcp_expectation.hh"
#include "tcp_fsm_test_harness.hh"
#include "tcp_header.hh"
#include "tcp_segment.hh"
#include "test_err_if.hh"
#include "util.hh"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <stdexcept>
#include <string>
#include <vector>

using namespace std;
using State = TCPTestHarness::State;

static constexpr unsigned NREPS = 32;
static constexpr unsigned MIN_SWIN = 2048;
static constexpr unsigned MAX_SWIN = 34816;
static constexpr unsigned MIN_SWIN_MUL = 2;
static constexpr unsigned MAX_SWIN_MUL = 6;

int main() {
    try {
        auto rd = get_random_generator();
        TCPConfig cfg{};
        cfg.send_capacity = MAX_SWIN * MAX_SWIN_MUL;

        // test 1: listen -> established -> check advertised winsize -> check sent bytes before ACK
        for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
            cfg.recv_capacity = 2048 + (rd() % 32768);
            const WrappingInt32 seq_base(rd());
            TCPTestHarness test_1(cfg);

            // connect
            test_1.execute(Listen{});
            test_1.send_syn(seq_base);

            TCPSegment seg = test_1.expect_seg(
                ExpectOneSegment{}.with_ack(true).with_ackno(seq_base + 1).with_win(cfg.recv_capacity),
                "test 1 failed: SYN/ACK invalid");
            auto &seg_hdr = seg.header();

            const WrappingInt32 ack_base = seg_hdr.seqno;

            // ack
            const uint16_t swin = MIN_SWIN + (rd() % (MAX_SWIN - MIN_SWIN));
            test_1.send_ack(seq_base + 1, ack_base + 1, swin);  // adjust send window here

            test_1.execute(ExpectNoSegment{}, "test 1 failed: ACK after acceptable ACK");
            test_1.execute(ExpectState{State::ESTABLISHED});

            // write swin_mul * swin, make sure swin gets sent
            const unsigned swin_mul = MIN_SWIN_MUL + (rd() % (MAX_SWIN_MUL - MIN_SWIN_MUL));
            string d(swin_mul * swin, 0);
            generate(d.begin(), d.end(), [&] { return rd(); });
            test_1.execute(Write{d}.with_bytes_written(swin_mul * swin));
            test_1.execute(Tick(1));

            string d_out(swin_mul * swin, 0);
            size_t bytes_total = 0;
            while (bytes_total < swin_mul * swin) {
                test_1.execute(ExpectSegmentAvailable{}, "test 1 failed: nothing sent after write()");
                size_t bytes_read = 0;
                while (test_1.can_read()) {
                    TCPSegment seg2 = test_1.expect_seg(
                        ExpectSegment{}.with_ack(true).with_ackno(seq_base + 1).with_win(cfg.recv_capacity),
                        "test 1 failed: invalid datagrams carrying write() data");
                    auto &seg2_hdr = seg2.header();
                    bytes_read += seg2.payload().size();
                    size_t seg2_first = seg2_hdr.seqno - ack_base - 1;
                    copy(seg2.payload().str().cbegin(), seg2.payload().str().cend(), d_out.begin() + seg2_first);
                }
                test_err_if(  // correct number of bytes sent
                    bytes_read + TCPConfig::MAX_PAYLOAD_SIZE < swin,
                    "test 1 failed: sender did not fill window");
                test_1.execute(ExpectBytesInFlight{bytes_read}, "test 1 failed: sender wrong bytes_in_flight");

                bytes_total += bytes_read;
                // NOTE that we don't override send window here because cfg should have been updated
                test_1.send_ack(seq_base + 1, ack_base + 1 + bytes_total, swin);
                test_1.execute(Tick(1));
            }
            test_1.execute(ExpectBytesInFlight{0}, "test 1 failed: after acking, bytes still in flight?");
            test_err_if(!equal(d.cbegin(), d.cend(), d_out.cbegin()), "test 1 failed: data mismatch");
        }
    } catch (const exception &e) {
        cerr << e.what() << endl;
        return err_num;
    }

    return EXIT_SUCCESS;
}