#include <gtest/gtest.h>
#include <cstring>
#include <string>
#include <algorithm>
#include <stdint.h>
#include "tcp.h"
#include "utils.h"

using namespace std;
using namespace Tins;

class TCPTest : public testing::Test {
public:
    static const uint8_t expected_packet[];
    
    void test_equals(const TCP &tcp1, const TCP &tcp2);
};

const uint8_t TCPTest::expected_packet[] = {'\x7f', 'M', 'O', '\x1d', '\xf1', '\xda', '\xe5', 
'F', '_', '\xae', '\xd1', '#', '\xd0', '\x02', 'q', '\xda', '\x00', '\x00', '\x1f', '\xae', 
'\x02', '\x04', '\x98', '\xfa', '\x08', '\n', 'O', '\xd2', ':', '\xcb', '\x89', '\xfe', 
'\x12', '4', '\x03', '\x03', 'z', '\x04', '\x02', '\x05', '\n', '\x00', '\x01', '\x02', 
'\x03', '\x04', '\x05', '\x06', '\x07', '\x00', '\x00', '\x00'};


TEST_F(TCPTest, DefaultConstructor) {
    TCP tcp;
    EXPECT_EQ(tcp.dport(), 0);
    EXPECT_EQ(tcp.sport(), 0);
    EXPECT_EQ(tcp.pdu_type(), PDU::TCP);
}

TEST_F(TCPTest, CopyConstructor) {
    TCP tcp1(0x6d1f, 0x78f2);
    TCP tcp2(tcp1);
    test_equals(tcp1, tcp2);
}

TEST_F(TCPTest, CopyAssignmentOperator) {
    TCP tcp1(0x6d1f, 0x78f2);
    TCP tcp2 = tcp1;
    test_equals(tcp1, tcp2);
}

TEST_F(TCPTest, NestedCopy) {
    TCP *nested_tcp = new TCP(0x6d1f, 0x78f2);
    TCP tcp1(0x6d1f, 0x78f2);
    tcp1.inner_pdu(nested_tcp);
    TCP tcp2(tcp1);
    test_equals(tcp1, tcp2);
}

TEST_F(TCPTest, CompleteConstructor) {
    TCP tcp(0x6d1f, 0x78f2);
    EXPECT_EQ(tcp.dport(), 0x6d1f);
    EXPECT_EQ(tcp.sport(), 0x78f2);
}

TEST_F(TCPTest, DPort) {
    TCP tcp;
    tcp.dport(0x5fad);
    EXPECT_EQ(tcp.dport(), 0x5fad);
}

TEST_F(TCPTest, SPort) {
    TCP tcp;
    tcp.sport(0x5fad);
    EXPECT_EQ(tcp.sport(), 0x5fad);
}

TEST_F(TCPTest, Seq) {
    TCP tcp;
    tcp.seq(0x5fad65fb);
    EXPECT_EQ(tcp.seq(), 0x5fad65fb);
}

TEST_F(TCPTest, AckSeq) {
    TCP tcp;
    tcp.ack_seq(0x5fad65fb);
    EXPECT_EQ(tcp.ack_seq(), 0x5fad65fb);
}

TEST_F(TCPTest, Window) {
    TCP tcp;
    tcp.window(0x5fad);
    EXPECT_EQ(tcp.window(), 0x5fad);
}

TEST_F(TCPTest, Check) {
    TCP tcp;
    tcp.check(0x5fad);
    EXPECT_EQ(tcp.check(), 0x5fad);
}

TEST_F(TCPTest, UrgPtr) {
    TCP tcp;
    tcp.urg_ptr(0x5fad);
    EXPECT_EQ(tcp.urg_ptr(), 0x5fad);
}

TEST_F(TCPTest, DataOffset) {
    TCP tcp;
    tcp.data_offset(0xe);
    EXPECT_EQ(tcp.data_offset(), 0xe);
}

TEST_F(TCPTest, SetFlag) {
    TCP tcp;
    tcp.set_flag(TCP::SYN, 1);
    tcp.set_flag(TCP::FIN, 1);
    
    EXPECT_EQ(tcp.get_flag(TCP::SYN), 1);
    EXPECT_EQ(tcp.get_flag(TCP::FIN), 1);
    EXPECT_EQ(tcp.get_flag(TCP::RST), 0);
    EXPECT_EQ(tcp.get_flag(TCP::PSH), 0);
    EXPECT_EQ(tcp.get_flag(TCP::ACK), 0);
    EXPECT_EQ(tcp.get_flag(TCP::URG), 0);
    EXPECT_EQ(tcp.get_flag(TCP::ECE), 0);
    EXPECT_EQ(tcp.get_flag(TCP::CWR), 0);
}

TEST_F(TCPTest, MSS) {
    TCP tcp;
    uint16_t mss = 0x456f, found_mss;
    tcp.add_mss_option(mss);
    ASSERT_TRUE(tcp.search_mss_option(&found_mss));
    EXPECT_EQ(mss, found_mss);
}

TEST_F(TCPTest, WindowScale) {
    TCP tcp;
    uint8_t scale = 0x4f, found_scale;
    tcp.add_winscale_option(scale);
    ASSERT_TRUE(tcp.search_winscale_option(&found_scale));
    EXPECT_EQ(scale, found_scale);
}

TEST_F(TCPTest, SackPermitted) {
    TCP tcp;
    tcp.add_sack_permitted_option();
    ASSERT_TRUE(tcp.search_sack_permitted_option());
}

TEST_F(TCPTest, Sack) {
    TCP tcp;
    list<uint32_t> edges, edges_found;
    edges.push_back(0x13);
    edges.push_back(0x63fa1d7a);
    edges.push_back(0xff1c);
    tcp.add_sack_option(edges);
    ASSERT_TRUE(tcp.search_sack_option(&edges_found));
    ASSERT_EQ(edges.size(), edges_found.size());
    while(edges.size()) {
        EXPECT_EQ(edges.front(), edges_found.front());
        edges.pop_front();
        edges_found.pop_front();
    }
}

TEST_F(TCPTest, AlternateChecksum) {
    TCP tcp;
    uint8_t found;
    tcp.add_altchecksum_option(TCP::CHK_16FLETCHER);
    ASSERT_TRUE(tcp.search_altchecksum_option(&found));
    EXPECT_EQ(found, (uint8_t)TCP::CHK_16FLETCHER);
}

TEST_F(TCPTest, Timestamp) {
    TCP tcp;
    uint32_t value = 0x456fa23d, found_value;
    uint32_t reply = 0xfa12d345, found_reply;
    tcp.add_timestamp_option(value, reply);
    ASSERT_TRUE(tcp.search_timestamp_option(&found_value, &found_reply));
    EXPECT_EQ(value, found_value);
    EXPECT_EQ(reply, found_reply);
}

void TCPTest::test_equals(const TCP &tcp1, const TCP &tcp2) {
    EXPECT_EQ(tcp1.dport(), tcp2.dport());
    EXPECT_EQ(tcp2.sport(), tcp2.sport());
    EXPECT_EQ(tcp1.seq(), tcp2.seq());
    EXPECT_EQ(tcp1.ack_seq(), tcp2.ack_seq());
    EXPECT_EQ(tcp1.window(), tcp2.window());
    EXPECT_EQ(tcp1.check(), tcp2.check());
    EXPECT_EQ(tcp1.urg_ptr(), tcp2.urg_ptr());
    EXPECT_EQ(tcp1.data_offset(), tcp2.data_offset());
    EXPECT_EQ((bool)tcp1.inner_pdu(), (bool)tcp2.inner_pdu());
}

// This is not working, but i don't want to fix it right now.
TEST_F(TCPTest, ConstructorFromBuffer) {
    TCP tcp1(expected_packet, sizeof(expected_packet));
    uint32_t value32, ovalue32;
    uint16_t value16;
    uint8_t value8;
    
    EXPECT_EQ(tcp1.dport(), 0x4f1d);
    EXPECT_EQ(tcp1.sport(), 0x7f4d);
    EXPECT_EQ(tcp1.seq(), 0xf1dae546);
    EXPECT_EQ(tcp1.ack_seq(), 0x5faed123);
    EXPECT_EQ(tcp1.window(), 0x71da);
    EXPECT_EQ(tcp1.urg_ptr(), 0x1fae);
    EXPECT_EQ(tcp1.data_offset(), 0xd);
    
    ASSERT_TRUE(tcp1.search_timestamp_option(&value32, &ovalue32));
    EXPECT_EQ(value32, 0x4fd23acb);
    EXPECT_EQ(ovalue32, 0x89fe1234);
    
    EXPECT_TRUE(tcp1.search_sack_permitted_option());
    
    ASSERT_TRUE(tcp1.search_winscale_option(&value8));
    EXPECT_EQ(value8, 0x7a);
    
    ASSERT_TRUE(tcp1.search_mss_option(&value16));
    EXPECT_EQ(value16, 0x98fa);
    
    list<uint32_t> edges;
    ASSERT_TRUE(tcp1.search_sack_option(&edges));
    ASSERT_EQ(edges.size(), 2);
    EXPECT_EQ(edges.front(), 0x00010203); edges.pop_front();
    EXPECT_EQ(edges.front(), 0x04050607); 
    
    uint32_t size;
    uint8_t *buffer = tcp1.serialize(size);
    
    TCP tcp2(buffer, size);
    test_equals(tcp1, tcp2);
    delete[] buffer;
}

TEST_F(TCPTest, Serialize) {
    TCP tcp1(expected_packet, sizeof(expected_packet));
    uint32_t size;
    uint8_t *buffer = tcp1.serialize(size);
    ASSERT_EQ(size, sizeof(expected_packet));
    ASSERT_TRUE(std::equal(buffer, buffer + size, expected_packet));
    delete[] buffer;
}