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

#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>

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

// in LISTEN, send ACK
static void ack_listen_test(const TCPConfig &cfg,
                            const WrappingInt32 seqno,
                            const WrappingInt32 ackno,
                            const int lineno) {
    try {
        TCPTestHarness test = TCPTestHarness::in_listen(cfg);

        // any ACK should result in a RST
        test.send_ack(seqno, ackno);

        test.execute(ExpectState{State::RESET});
        test.execute(ExpectOneSegment{}.with_rst(true), "test 3 failed: no RST after ACK in LISTEN");
    } catch (const exception &e) {
        throw runtime_error(string(e.what()) + " (ack_listen_test called from line " + to_string(lineno) + ")");
    }
}

// in SYN_SENT, send ACK and maybe RST
static void ack_rst_syn_sent_test(const TCPConfig &cfg,
                                  const WrappingInt32 base_seq,
                                  const WrappingInt32 seqno,
                                  const WrappingInt32 ackno,
                                  const bool send_rst,
                                  const int lineno) {
    try {
        TCPTestHarness test = TCPTestHarness::in_syn_sent(cfg, base_seq);

        // unacceptable ACKs should elicit a RST
        if (send_rst) {
            test.send_rst(seqno, ackno);
        } else {
            test.send_ack(seqno, ackno);
        }

        if (send_rst) {
            test.execute(ExpectNotInState{State::RESET});
            test.execute(ExpectNoSegment{}, "test 5 failed: bad ACK w/ RST should have been ignored in SYN_SENT");
        } else {
            test.execute(ExpectState{State::RESET});
            test.execute(ExpectOneSegment{}.with_rst(true).with_seqno(ackno),
                         "test 5 failed: problem with RST segment for bad ACK");
        }
    } catch (const exception &e) {
        throw runtime_error(string(e.what()) + " (ack_rst_syn_sent_test called from line " + to_string(lineno) + ")");
    }
}

int main() {
    try {
        TCPConfig cfg{};
        const WrappingInt32 base_seq(1 << 31);

        // test #1: in ESTABLISHED, send unacceptable segments and ACKs
        {
            cerr << "Test 1" << endl;
            TCPTestHarness test_1 = TCPTestHarness::in_established(cfg, base_seq - 1, base_seq - 1);

            // acceptable ack---no response
            test_1.send_ack(base_seq, base_seq);

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

            // ack in the past---no response
            test_1.send_ack(base_seq, base_seq - 1);

            test_1.execute(ExpectNoSegment{}, "test 1 failed: ACK after past ACK");

            // ack in the future---should get ACK back
            test_1.send_ack(base_seq, base_seq + 1);

            test_1.execute(ExpectOneSegment{}.with_ack(true).with_ackno(base_seq), "test 1 failed: bad ACK");

            // segment out of the window---should get an ACK
            test_1.send_byte(base_seq - 1, base_seq, 1);

            test_1.execute(ExpectUnassembledBytes{0}, "test 1 failed: seg queued on early seqno");

            test_1.execute(ExpectOneSegment{}.with_ack(true).with_ackno(base_seq), "test 1 failed: bad ACK");

            // segment out of the window---should get an ACK
            test_1.send_byte(base_seq + cfg.recv_capacity, base_seq, 1);

            test_1.execute(ExpectUnassembledBytes{0}, "test 1 failed: seg queued on late seqno");
            test_1.execute(ExpectOneSegment{}.with_ack(true).with_ackno(base_seq),
                           "test 1 failed: bad ACK on late seqno");

            // packet next byte in the window - ack should advance and data should be readable
            test_1.send_byte(base_seq, base_seq, 1);

            test_1.execute(ExpectData{}, "test 1 failed: pkt not processed on next seqno");

            test_1.execute(ExpectOneSegment{}.with_ack(true).with_ackno(base_seq + 1), "test 1 failed: bad ACK");

            // segment not in window with RST set --- should get nothing back
            test_1.send_rst(base_seq);
            test_1.send_rst(base_seq + cfg.recv_capacity + 1);
            test_1.execute(ExpectNoSegment{}, "test 1 failed: got a response to an out-of-window RST");

            test_1.send_rst(base_seq + 1);
            test_1.execute(ExpectState{State::RESET});
        }

        // test #2: in LISTEN, send RSTs
        {
            cerr << "Test 2" << endl;
            TCPTestHarness test_2 = TCPTestHarness::in_listen(cfg);

            // all RSTs should be ignored in LISTEN
            test_2.send_rst(base_seq);
            test_2.send_rst(base_seq - 1);
            test_2.send_rst(base_seq + cfg.recv_capacity);

            test_2.execute(ExpectNoSegment{}, "test 2 failed: RST was not ignored in LISTEN");
        }

        // test 3: ACKs in LISTEN
        cerr << "Test 3" << endl;
        ack_listen_test(cfg, base_seq, base_seq, __LINE__);
        ack_listen_test(cfg, base_seq - 1, base_seq, __LINE__);
        ack_listen_test(cfg, base_seq, base_seq - 1, __LINE__);
        ack_listen_test(cfg, base_seq - 1, base_seq, __LINE__);
        ack_listen_test(cfg, base_seq - 1, base_seq - 1, __LINE__);
        ack_listen_test(cfg, base_seq + cfg.recv_capacity, base_seq, __LINE__);
        ack_listen_test(cfg, base_seq, base_seq + cfg.recv_capacity, __LINE__);
        ack_listen_test(cfg, base_seq + cfg.recv_capacity, base_seq, __LINE__);
        ack_listen_test(cfg, base_seq + cfg.recv_capacity, base_seq + cfg.recv_capacity, __LINE__);

        // test 4: ACK and RST in SYN_SENT
        {
            cerr << "Test 4" << endl;
            TCPTestHarness test_4 = TCPTestHarness::in_syn_sent(cfg, base_seq);

            // data segments with acceptable ACKs should be ignored
            test_4.send_byte(base_seq, base_seq + 1, 1);

            test_4.execute(ExpectState{State::SYN_SENT});

            test_4.execute(ExpectNoSegment{}, "test 4 failed: bytes were not ignored in SYN_SENT");

            // RST segments without ACKs should be ignored
            test_4.send_rst(base_seq);
            test_4.send_rst(base_seq + 1);
            test_4.send_rst(base_seq - 1);

            test_4.execute(ExpectState{State::SYN_SENT});
            test_4.execute(ExpectNoSegment{}, "test 4 failed: RST without ACK was not ignored in SYN_SENT");

            // good ACK with RST should result in a RESET but no RST segment sent
            test_4.send_rst(base_seq, base_seq + 1);

            test_4.execute(ExpectState{State::RESET});
            test_4.execute(ExpectNoSegment{}, "test 4 failed: RST with good ackno should RESET the connection");
        }

        // test 5: ack/rst in SYN_SENT
        cerr << "Test 5" << endl;
        ack_rst_syn_sent_test(cfg, base_seq, base_seq, base_seq, false, __LINE__);
        ack_rst_syn_sent_test(cfg, base_seq, base_seq, base_seq + 2, false, __LINE__);
        ack_rst_syn_sent_test(cfg, base_seq, base_seq, base_seq, true, __LINE__);
        ack_rst_syn_sent_test(cfg, base_seq, base_seq, base_seq + 2, true, __LINE__);
    } catch (const exception &e) {
        cerr << e.what() << endl;
        return EXIT_FAILURE;
    }

    return EXIT_SUCCESS;
}