#include <doctest/doctest.h>

#include <string>
#include <vector>

#include "packeteer/l7/http.hpp"
#include "packeteer/net/tcp.hpp"
#include "packeteer/net/tcp_reassembly.hpp"

using namespace packeteer::net;

namespace {

Ipv4Address addr(unsigned char a, unsigned char b, unsigned char c, unsigned char d) {
    return Ipv4Address{{a, b, c, d}};
}

std::vector<unsigned char> to_bytes(const std::string& s) {
    return std::vector<unsigned char>(s.begin(), s.end());
}

}  // namespace

TEST_CASE("TcpReassembler ignores payload before SYN is seen") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    auto data = to_bytes("data before syn");
    auto result = r.process_segment(client, 40000, server, 80, 1000, 0, data);
    CHECK_FALSE(result.has_value());
}

TEST_CASE("TcpReassembler joins two in-order segments into one contiguous buffer") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    // SYN: seq 1000, consumes seq 1000 itself, next data starts at 1001.
    auto syn = r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});
    CHECK_FALSE(syn.has_value());

    auto part1 = to_bytes("GET /index.html HTTP/1.1\r\n");
    auto r1 = r.process_segment(client, 40000, server, 80, 1001, kTcpPsh | kTcpAck, part1);
    REQUIRE(r1.has_value());
    CHECK(r1->size() == part1.size());

    auto part2 = to_bytes("Host: example.com\r\n\r\n");
    std::uint32_t seq2 = 1001 + static_cast<std::uint32_t>(part1.size());
    auto r2 = r.process_segment(client, 40000, server, 80, seq2, kTcpPsh | kTcpAck, part2);
    REQUIRE(r2.has_value());

    std::string joined(r2->begin(), r2->end());
    CHECK(joined == "GET /index.html HTTP/1.1\r\nHost: example.com\r\n\r\n");

    auto http = parse_http(*r2);
    REQUIRE(http.has_value());
    CHECK(http->is_request);
    CHECK(http->method_or_version == "GET");
    CHECK(http->target_or_status == "/index.html");
    REQUIRE(http->host.has_value());
    CHECK(*http->host == "example.com");
}

TEST_CASE("TcpReassembler drops an out-of-order segment rather than buffering it") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});

    auto part1 = to_bytes("first ");
    r.process_segment(client, 40000, server, 80, 1001, kTcpAck, part1);

    // Skip ahead instead of continuing at 1001 + part1.size(): out of order.
    auto part3 = to_bytes("third ");
    auto result = r.process_segment(client, 40000, server, 80, 9999, kTcpAck, part3);
    CHECK_FALSE(result.has_value());
}

TEST_CASE("TcpReassembler drops a retransmitted (already-seen) segment") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});

    auto part1 = to_bytes("hello");
    auto r1 = r.process_segment(client, 40000, server, 80, 1001, kTcpAck, part1);
    REQUIRE(r1.has_value());

    // Same seq again: a retransmission, not new data.
    auto retransmit = r.process_segment(client, 40000, server, 80, 1001, kTcpAck, part1);
    CHECK_FALSE(retransmit.has_value());
}

TEST_CASE("TcpReassembler tracks each direction of a flow independently") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});
    r.process_segment(server, 80, client, 40000, 5000, kTcpSyn | kTcpAck, {});

    auto request = to_bytes("GET / HTTP/1.1\r\n\r\n");
    auto req_result = r.process_segment(client, 40000, server, 80, 1001, kTcpPsh | kTcpAck,
                                         request);
    REQUIRE(req_result.has_value());
    CHECK(std::string(req_result->begin(), req_result->end()) == "GET / HTTP/1.1\r\n\r\n");

    auto response = to_bytes("HTTP/1.1 200 OK\r\n\r\n");
    auto resp_result = r.process_segment(server, 80, client, 40000, 5001, kTcpPsh | kTcpAck,
                                          response);
    REQUIRE(resp_result.has_value());
    CHECK(std::string(resp_result->begin(), resp_result->end()) == "HTTP/1.1 200 OK\r\n\r\n");

    // Requesting side's buffer should be untouched by the response.
    CHECK(std::string(req_result->begin(), req_result->end()) == "GET / HTTP/1.1\r\n\r\n");
}

TEST_CASE("TcpReassembler canonicalizes both directions of a connection to the same flow") {
    TcpReassembler r;
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});
    CHECK(r.flow_count() == 1);

    // A segment in the reverse direction of the *same* connection must
    // not create a second flow entry.
    r.process_segment(server, 80, client, 40000, 5000, kTcpSyn | kTcpAck, {});
    CHECK(r.flow_count() == 1);
}

TEST_CASE("TcpReassembler caps buffered bytes per direction and stops growing past the limit") {
    TcpReassembler r(/*max_buffer_per_direction=*/10, /*max_flows=*/16);
    auto client = addr(10, 0, 0, 1);
    auto server = addr(10, 0, 0, 2);

    r.process_segment(client, 40000, server, 80, 1000, kTcpSyn, {});

    auto part1 = to_bytes("12345");  // 5 bytes, fits
    auto r1 = r.process_segment(client, 40000, server, 80, 1001, kTcpAck, part1);
    REQUIRE(r1.has_value());
    CHECK(r1->size() == 5);

    // Next 5 bytes would land exactly at the 10-byte cap.
    auto part2 = to_bytes("67890");
    auto r2 = r.process_segment(client, 40000, server, 80, 1006, kTcpAck, part2);
    REQUIRE(r2.has_value());
    CHECK(r2->size() == 10);

    // A further segment would exceed the cap: sequence tracking still
    // advances (so future in-order segments aren't misjudged), but the
    // buffer itself does not grow past max_buffer_.
    auto part3 = to_bytes("overflow");
    auto r3 = r.process_segment(client, 40000, server, 80, 1011, kTcpAck, part3);
    REQUIRE(r3.has_value());
    CHECK(r3->size() == 10);
}

TEST_CASE("TcpReassembler caps the number of tracked flows") {
    TcpReassembler r(/*max_buffer_per_direction=*/1024, /*max_flows=*/1);
    auto server = addr(10, 0, 0, 2);

    auto client1 = addr(10, 0, 0, 1);
    r.process_segment(client1, 40000, server, 80, 1000, kTcpSyn, {});
    CHECK(r.flow_count() == 1);

    // A second, distinct flow should be refused: table is full.
    auto client2 = addr(10, 0, 0, 3);
    auto data = to_bytes("x");
    auto result = r.process_segment(client2, 40000, server, 80, 2000, kTcpSyn, data);
    CHECK_FALSE(result.has_value());
    CHECK(r.flow_count() == 1);
}