#include "naut/extension.h"
#include "naut/peer.h"
#include "test.h"

#include <stdlib.h>
#include <string.h>

static const uint8_t *ext_payload(const uint8_t *frame, size_t len,
                                  uint8_t *ext_id, size_t *payload_len) {
    naut_msg msg;
    CHECK(naut_peer_msg_parse(frame, len, &msg) == (int)len);
    CHECK(msg.type == NAUT_MSG_EXTENDED && msg.payload_len >= 1);
    *ext_id = msg.payload[0];
    *payload_len = msg.payload_len - 1;
    return msg.payload + 1;
}

int main(void) {
    uint8_t *frame = NULL;
    size_t frame_len = 0, payload_len = 0;
    uint8_t ext_id = 0;

    CHECK(naut_ext_build_handshake(NAUT_EXT_UT_METADATA, NAUT_EXT_UT_PEX,
                                   31235, 6881, &frame, &frame_len) == NAUT_OK);
    const uint8_t *payload = ext_payload(frame, frame_len, &ext_id, &payload_len);
    CHECK_EQ(ext_id, 0);
    naut_ext_handshake hs;
    CHECK(naut_ext_parse_handshake(payload, payload_len, &hs) == NAUT_OK);
    CHECK_EQ(hs.ut_metadata, NAUT_EXT_UT_METADATA);
    CHECK_EQ(hs.ut_pex, NAUT_EXT_UT_PEX);
    CHECK_EQ(hs.metadata_size, 31235);
    CHECK_EQ(hs.port, 6881);
    CHECK_EQ(hs.reqq, 256);
    free(frame);

    CHECK(naut_metadata_build(3, NAUT_METADATA_REQUEST, 7, 0, NULL, 0,
                              &frame, &frame_len) == NAUT_OK);
    payload = ext_payload(frame, frame_len, &ext_id, &payload_len);
    CHECK_EQ(ext_id, 3);
    naut_metadata_msg mm;
    CHECK(naut_metadata_parse(payload, payload_len, &mm) == NAUT_OK);
    CHECK(mm.type == NAUT_METADATA_REQUEST && mm.piece == 7 && mm.data_len == 0);
    free(frame);

    uint8_t metadata[100];
    for (size_t i = 0; i < sizeof metadata; i++) metadata[i] = (uint8_t)i;
    CHECK(naut_metadata_build(3, NAUT_METADATA_DATA, 1,
                              NAUT_METADATA_BLOCK + sizeof metadata,
                              metadata, sizeof metadata,
                              &frame, &frame_len) == NAUT_OK);
    payload = ext_payload(frame, frame_len, &ext_id, &payload_len);
    CHECK(naut_metadata_parse(payload, payload_len, &mm) == NAUT_OK);
    CHECK(mm.type == NAUT_METADATA_DATA && mm.piece == 1);
    CHECK_EQ(mm.total_size, NAUT_METADATA_BLOCK + sizeof metadata);
    CHECK_EQ(mm.data_len, sizeof metadata);
    CHECK(memcmp(mm.data, metadata, sizeof metadata) == 0);
    free(frame);

    naut_peer_addr added[2] = {
        { { 1, 2, 3, 4 }, 6881 },
        { { 10, 0, 0, 2 }, 51413 },
    };
    naut_peer_addr dropped[1] = { { { 192, 0, 2, 7 }, 7000 } };
    uint8_t flags[2] = { 0x01, 0x12 };
    naut_pex_msg pex = {
        .added = added, .added_flags = flags, .num_added = 2,
        .dropped = dropped, .num_dropped = 1,
    };
    CHECK(naut_pex_build(4, &pex, &frame, &frame_len) == NAUT_OK);
    payload = ext_payload(frame, frame_len, &ext_id, &payload_len);
    CHECK_EQ(ext_id, 4);
    naut_pex_msg parsed;
    CHECK(naut_pex_parse(payload, payload_len, &parsed) == NAUT_OK);
    CHECK_EQ(parsed.num_added, 2);
    CHECK_EQ(parsed.num_dropped, 1);
    CHECK(parsed.added[0].port == 6881 && parsed.added[1].port == 51413);
    CHECK(parsed.added_flags[0] == 0x01 && parsed.added_flags[1] == 0x12);
    CHECK(parsed.dropped[0].ip[0] == 192 && parsed.dropped[0].port == 7000);
    naut_pex_free(&parsed);
    free(frame);

    const uint8_t malformed[] = "d5:added5:abcdee";
    CHECK(naut_pex_parse(malformed, sizeof malformed - 1, &parsed) ==
          NAUT_ERR_PROTO);

    TEST_MAIN_END();
}