#include "naut/bencode.h"
#include "naut/dht.h"
#include "test.h"

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

static void check_get_peers_query(void) {
    uint8_t tx[2] = { 0x12, 0x34 };
    uint8_t id[20], hash[20];
    for (size_t i = 0; i < 20; i++) {
        id[i] = (uint8_t)i;
        hash[i] = (uint8_t)(0x80 + i);
    }

    uint8_t *query = NULL;
    size_t query_len = 0;
    CHECK(naut_dht_build_get_peers(tx, sizeof tx, id, hash,
                                   &query, &query_len) == NAUT_OK);

    naut_bc_doc *doc = NULL;
    CHECK(naut_bc_parse(query, query_len, &doc) == NAUT_OK);
    const naut_bc *root = naut_bc_root(doc);
    CHECK(naut_bc_str_eq(naut_bc_dict_get(root, "y"), "q"));
    CHECK(naut_bc_str_eq(naut_bc_dict_get(root, "q"), "get_peers"));

    const naut_bc *args = naut_bc_dict_get(root, "a");
    const uint8_t *p = NULL;
    size_t n = 0;
    CHECK(naut_bc_get_str(naut_bc_dict_get(args, "id"), &p, &n));
    CHECK(n == 20 && memcmp(p, id, 20) == 0);
    CHECK(naut_bc_get_str(naut_bc_dict_get(args, "info_hash"), &p, &n));
    CHECK(n == 20 && memcmp(p, hash, 20) == 0);

    naut_bc_free(doc);
    free(query);
}

static void check_response(void) {
    uint8_t packet[256];
    size_t len = 0;
    const char *prefix = "d1:rd2:id20:";
    memcpy(packet + len, prefix, strlen(prefix));
    len += strlen(prefix);
    for (size_t i = 0; i < 20; i++) packet[len++] = (uint8_t)(0x20 + i);

    const char *nodes = "5:nodes26:";
    memcpy(packet + len, nodes, strlen(nodes));
    len += strlen(nodes);
    for (size_t i = 0; i < 20; i++) packet[len++] = (uint8_t)(0x40 + i);
    packet[len++] = 192; packet[len++] = 0; packet[len++] = 2; packet[len++] = 9;
    packet[len++] = 0x1a; packet[len++] = 0xe1;

    const char *suffix = "5:token3:abc6:valuesl6:";
    memcpy(packet + len, suffix, strlen(suffix));
    len += strlen(suffix);
    packet[len++] = 203; packet[len++] = 0; packet[len++] = 113; packet[len++] = 7;
    packet[len++] = 0xc8; packet[len++] = 0xd5;
    const char *tail = "ee1:t2:aa1:y1:re";
    memcpy(packet + len, tail, strlen(tail));
    len += strlen(tail);

    naut_dht_response response;
    CHECK(naut_dht_parse_response(packet, len, &response) == NAUT_OK);
    CHECK(response.type == NAUT_DHT_RESPONSE);
    CHECK(response.transaction_len == 2 &&
          memcmp(response.transaction, "aa", 2) == 0);
    CHECK(response.has_id && response.id[0] == 0x20);
    CHECK(response.token_len == 3 &&
          memcmp(response.token, "abc", 3) == 0);
    CHECK(response.num_nodes == 1);
    CHECK(response.nodes[0].ip[0] == 192 &&
          response.nodes[0].port == 6881);
    CHECK(response.num_peers == 1);
    CHECK(response.peers[0].ip[0] == 203 &&
          response.peers[0].port == 51413);
    naut_dht_response_free(&response);
}

int main(void) {
    check_get_peers_query();
    check_response();

    const char error[] = "d1:eli203e12:Server errore1:t2:zz1:y1:ee";
    naut_dht_response response;
    CHECK(naut_dht_parse_response((const uint8_t *)error, sizeof error - 1,
                                  &response) == NAUT_OK);
    CHECK(response.type == NAUT_DHT_ERROR && response.error_code == 203);
    naut_dht_response_free(&response);

    const char malformed[] =
        "d1:rd2:id20:abcdefghijklmnopqrst5:nodes1:xe1:t1:a1:y1:re";
    CHECK(naut_dht_parse_response((const uint8_t *)malformed,
                                  sizeof malformed - 1,
                                  &response) == NAUT_ERR_PROTO);

    TEST_MAIN_END();
}