#include "naut/peer.h"
#include "test.h"
#include <string.h>

int main(void) {
    uint8_t ih[20], pid[20], ih2[20], pid2[20];
    for (int i = 0; i < 20; i++) { ih[i] = (uint8_t)(i + 1); pid[i] = (uint8_t)(100 + i); }

    /* handshake round-trip, with the BEP-10 extension reserved bit set */
    uint8_t hs[NAUT_HANDSHAKE_LEN];
    uint64_t reserved = 0x0000000000100000ULL;   /* ext protocol bit */
    naut_peer_handshake_build(hs, ih, pid, reserved);
    CHECK_EQ(hs[0], 19);
    CHECK(memcmp(hs + 1, "BitTorrent protocol", 19) == 0);
    uint64_t got_res = 0;
    CHECK(naut_peer_handshake_parse(hs, ih2, pid2, &got_res));
    CHECK(memcmp(ih, ih2, 20) == 0 && memcmp(pid, pid2, 20) == 0);
    CHECK_EQ((long long)got_res, (long long)reserved);
    hs[3] = 'X';   /* corrupt protocol string */
    CHECK(!naut_peer_handshake_parse(hs, ih2, pid2, &got_res));

    naut_msg m;
    uint8_t b[64];

    /* simple messages */
    size_t n = naut_peer_msg_simple(b, NAUT_MSG_INTERESTED);
    CHECK_EQ(n, 5);
    CHECK_EQ(naut_peer_msg_parse(b, n, &m), 5);
    CHECK_EQ(m.type, NAUT_MSG_INTERESTED);

    /* keep-alive */
    n = naut_peer_keepalive(b);
    CHECK_EQ(naut_peer_msg_parse(b, n, &m), 4);
    CHECK_EQ(m.type, NAUT_MSG_KEEPALIVE);

    /* have */
    n = naut_peer_msg_have(b, 0x01020304);
    CHECK_EQ(naut_peer_msg_parse(b, n, &m), 9);
    CHECK(m.type == NAUT_MSG_HAVE && m.index == 0x01020304);

    /* request */
    n = naut_peer_msg_request(b, 7, 16384, 16384);
    CHECK_EQ(naut_peer_msg_parse(b, n, &m), 17);
    CHECK(m.type == NAUT_MSG_REQUEST && m.index == 7 && m.begin == 16384 && m.length == 16384);

    /* piece: header + block, payload aliases input */
    {
        uint8_t blk[16384];
        for (int i = 0; i < 16384; i++) blk[i] = (uint8_t)(i & 0xff);
        uint8_t frame[13 + 16384];
        naut_peer_msg_piece_header(frame, 3, 32768, 16384);
        memcpy(frame + 13, blk, 16384);
        int c = naut_peer_msg_parse(frame, sizeof frame, &m);
        CHECK_EQ(c, (int)sizeof frame);
        CHECK(m.type == NAUT_MSG_PIECE && m.index == 3 && m.begin == 32768);
        CHECK_EQ(m.payload_len, 16384);
        CHECK(m.payload == frame + 13 && memcmp(m.payload, blk, 16384) == 0);
    }

    /* streaming: partial frame returns 0 (need more), completes when filled */
    n = naut_peer_msg_have(b, 42);
    CHECK_EQ(naut_peer_msg_parse(b, 3, &m), 0);   /* < 4 length bytes */
    CHECK_EQ(naut_peer_msg_parse(b, 7, &m), 0);   /* length known, body short */
    CHECK_EQ(naut_peer_msg_parse(b, 9, &m), 9);

    /* bitfield */
    uint8_t bf[4] = { 0xff, 0x80, 0x00, 0x01 };
    n = naut_peer_msg_bitfield(b, bf, 4);
    CHECK_EQ(naut_peer_msg_parse(b, n, &m), (int)n);
    CHECK(m.type == NAUT_MSG_BITFIELD && m.payload_len == 4 && memcmp(m.payload, bf, 4) == 0);

    /* malformed: oversized length prefix rejected */
    uint8_t bad[4] = { 0xff, 0xff, 0xff, 0xff };
    CHECK(naut_peer_msg_parse(bad, 4, &m) < 0);
    /* malformed: HAVE whose body length (5) != required 4 */
    uint8_t badhave[10] = { 0,0,0,6, NAUT_MSG_HAVE, 0,0,0,1, 0 };
    CHECK(naut_peer_msg_parse(badhave, 10, &m) < 0);

    TEST_MAIN_END();
}