#include <arpa/inet.h>
#include <fcntl.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <unistd.h>
//
#include "common.c"
#include "firehose.c"
#include "omen.c"
#include "packet_storage.c"
#include "security.c"
#include "sha256.c"

struct data_packet_list {
    struct data_packet* packets;
    uint16_t            count;
    uint64_t            new_offset;
};

struct data_packet create_packet_from_bytes(
    const byte* input,
    size_t      input_len,
    byte        recipient_count,
    hash*       recipients,
    keypair     author_keys) {
    struct data_packet packet;
    memset(&packet, 0, sizeof(packet));

    if (input_len > 65536) input_len = 65536;
    if (input_len && input) memcpy(packet.contents.data, input, input_len);
    if (input_len < 65536) memset(packet.contents.data + input_len, 0, 65536 - input_len);

    if (recipient_count > MAX_RECIPIENTS) recipient_count = MAX_RECIPIENTS;
    packet.contents.recipient_count = recipient_count;
    if (recipients) memcpy(packet.contents.recipients, recipients, recipient_count * 32);

    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    packet.contents.created_time_ms = (uint64_t) ts.tv_sec * 1000 + (uint64_t) ts.tv_nsec / 1000000;

    memcpy(packet.contents.author, author_keys.encryption_pub, sizeof(public_key));
    memcpy(packet.contents.author_signing_pub, author_keys.signing_pub, sizeof(public_key));
    packet.contents.encrypted = false;

    size_t msg_len = CONTENTS_SIZE - sizeof(signature);
    byte   msg_buf[ msg_len ];
    msg_buf[ 0 ] = packet.contents.encrypted ? 1 : 0;
    write_u64_le(msg_buf + 1, packet.contents.created_time_ms);
    memcpy(msg_buf + 9, packet.contents.author, 32);
    memcpy(msg_buf + 41, packet.contents.author_signing_pub, 32);
    msg_buf[ 73 ] = packet.contents.recipient_count;
    memcpy(msg_buf + 74, packet.contents.recipients, 32 * MAX_RECIPIENTS);
    memcpy(msg_buf + 74 + 32 * MAX_RECIPIENTS, packet.contents.data, 65536);

    crypto_sign_detached(
        packet.contents.signature, NULL, msg_buf, msg_len, author_keys.signing_priv);

    return packet;
}

struct data_packet
create_packet(const char* filename, keypair author, byte recipient_count, hash* recipients) {
    byte file_buf[ 65536 ];
    memset(file_buf, 0, sizeof(file_buf));

    long  file_size = 0;
    FILE* file      = fopen(filename, "rb");
    if (!file) {
        perror("Failed to open file");
        exit(1);
    }
    if (fseek(file, 0, SEEK_END) == 0) {
        file_size = ftell(file);
        fseek(file, 0, SEEK_SET);
    }
    if (file_size < 0) file_size = 0;
    if (file_size > 65536) file_size = 65536;
    if (file_size > 0) fread(file_buf, 1, (size_t) file_size, file);
    fclose(file);

    return create_packet_from_bytes(
        file_buf, (size_t) file_size, recipient_count, recipients, author);
}

struct data_packet run_pow(struct data_packet packet, byte leading_zeros) {
    printf("running Proof-of-Work with %i leading zeros...\n", leading_zeros);
    uint64_t start_time = get_current_time_ms();

    const size_t hash_input_len = 4 + CONTENTS_SIZE;
    byte         hash_input_buffer[ hash_input_len ];

    byte* contents_buffer = hash_input_buffer + 4;
    contents_buffer[ 0 ]  = packet.contents.encrypted ? 1 : 0;
    write_u64_le(contents_buffer + 1, packet.contents.created_time_ms);
    memcpy(contents_buffer + 9, packet.contents.author, 32);
    memcpy(contents_buffer + 41, packet.contents.author_signing_pub, 32);
    contents_buffer[ 73 ] = packet.contents.recipient_count;
    memcpy(contents_buffer + 74, packet.contents.recipients, 32 * MAX_RECIPIENTS);
    memcpy(contents_buffer + 74 + 32 * MAX_RECIPIENTS, packet.contents.data, 65536);
    memcpy(contents_buffer + 74 + 32 * MAX_RECIPIENTS + 65536, packet.contents.signature, 64);

    packet.nonce = 0;
    while (true) {
        memcpy(hash_input_buffer, &packet.nonce, 4);
        struct sha256 calculated_hash_struct = sha256_hash(hash_input_buffer, hash_input_len);
        memcpy(packet.hash, calculated_hash_struct.hash, 32);
        if (check_leading_zeros(packet.hash, leading_zeros)) { break; }
        packet.nonce++;
    }

    uint64_t end_time = get_current_time_ms();
    printf("PoW completed in %lu ms. nonce: %u\n", end_time - start_time, packet.nonce);
    return packet;
}

struct relay_neighbor* get_server_neighbors(const char* ip, int port, byte count) {
    byte   data[ 8 ] = { get_neighbors, count, 0, 0, 0, 0, 0, 0 };
    byte*  response  = NULL;
    size_t resp_len  = 0;

    int ret = study(ip, port, data, sizeof(data), &response, &resp_len);
    if (ret != 0) {
        printf("Error: study returned %i\n", ret);
        if (response) free(response);
        return NULL;
    }

    if (resp_len < 1) {
        printf("Error: Invalid get_neighbors response\n");
        if (response) free(response);
        return NULL;
    }

    if (response[ 0 ] != neighbors_info) {
        printf("Error: Server returned an error for get_neighbors request.\n");
        if (response) free(response);
        return NULL;
    }

    size_t expected_wire = 1 + (size_t) RELAY_NEIGHBOR_WIRE_LEN * count;
    if (resp_len < expected_wire) {
        printf(
            "Error: get_neighbors response too short (expected %zu, got %zu)\n",
            expected_wire,
            resp_len);
        if (response) free(response);
        return NULL;
    }

    struct relay_neighbor* neighbors = malloc(count * sizeof(struct relay_neighbor));
    if (!neighbors) {
        perror("Failed to allocate memory for neighbors");
        if (response) free(response);
        return NULL;
    }

    for (byte i = 0; i < count; ++i) {
        bool success;
        neighbors[ i ] = s2relay_neighbor(
            RELAY_NEIGHBOR_WIRE_LEN, response + 1 + (i * RELAY_NEIGHBOR_WIRE_LEN), &success);
        if (!success) { printf("Error: Failed to deserialize neighbor %i\n", i); }
    }

    free(response);
    return neighbors;
}

bool push_packet_to_server(const char* ip, int port, struct data_packet packet) {
    const size_t send_len = 1 + PACKET_SIZE;
    byte         send_data[ send_len ];

    send_data[ 0 ] = push_packet;
    data_packet2s(packet, send_data + 1);

    byte*  response = NULL;
    size_t resp_len = 0;
    bool   success  = false;

    int ret = study(ip, port, send_data, send_len, &response, &resp_len);
    if (ret != 0) {
        printf("error: study returned %i\n", ret);
        goto cleanup;
    }

    if (resp_len == 2 && response[ 0 ] == push_packet_result) {
        success = response[ 1 ];
    } else {
        printf("error: invalid push response\n");
    }

cleanup:
    if (response) free(response);
    return success;
}

struct data_packet get_packet_from_server(const char* ip, int port, hash packet_hash) {
    const size_t send_len = 1 + 32;
    byte         send_data[ send_len ];
    send_data[ 0 ] = get_packet;
    memcpy(send_data + 1, packet_hash, 32);

    byte*  response = NULL;
    size_t resp_len = 0;
    bool   success  = false;

    int ret = study(ip, port, send_data, send_len, &response, &resp_len);
    if (ret != 0) {
        printf("Error: study returned %i\n", ret);
        if (response) free(response);
        return (struct data_packet) { 0 };
    }

    if (resp_len < 1 || response[ 0 ] != get_packet_result) {
        printf("Error: Invalid get_packet response\n");
        if (response) free(response);
        return (struct data_packet) { 0 };
    }

    struct data_packet packet = s2data_packet(resp_len - 1, response + 1, &success);
    if (!success) { printf("Error: Failed to deserialize data packet\n"); }

    free(response);
    return packet;
}

struct data_packet_list
get_firehose_packets(const char* ip, int port, uint64_t offset, uint16_t count) {
    const size_t send_len = 1 + 1 + 8;
    byte         send_data[ send_len ];
    send_data[ 0 ] = firehose;
    firehose_request2s(
        (struct firehose_request) { .count = (uint8_t) count, .offset = offset }, send_data + 1);

    byte*  response = NULL;
    size_t resp_len = 0;
    bool   success  = false;

    int ret = study(ip, port, send_data, send_len, &response, &resp_len);
    if (ret != 0) {
        printf("Error: study returned %i\n", ret);
        if (response) free(response);
        return (struct data_packet_list) { 0 };
    }

    if (resp_len < 9 || response[ 0 ] != firehose_reply) {
        printf("Error: Invalid firehose response\n");
        if (response) free(response);
        return (struct data_packet_list) { 0 };
    }

    uint64_t                new_offset = read_u64_le(response + 1);
    struct data_packet_list list       = { .count = 0, .new_offset = new_offset, .packets = NULL };

    size_t available_bytes   = resp_len - 9;
    size_t available_packets = available_bytes / PACKET_SIZE;
    if (available_packets == 0) {
        free(response);
        return list;
    }

    if (available_packets > SIZE_MAX / sizeof(struct data_packet)) {
        perror("allocation overflow");
        free(response);
        return (struct data_packet_list) { 0 };
    }

    list.packets = malloc(available_packets * sizeof(struct data_packet));
    if (!list.packets) {
        perror("Failed to allocate memory for packet list");
        free(response);
        return (struct data_packet_list) { 0 };
    }

    size_t current_offset = 9;
    while (current_offset + PACKET_SIZE <= resp_len && list.count < available_packets) {
        struct data_packet packet = s2data_packet(PACKET_SIZE, response + current_offset, &success);
        if (success) {
            list.packets[ list.count ] = packet;
            list.count++;
            current_offset += PACKET_SIZE;
        } else {
            printf("Error: Failed to deserialize a packet from the firehose stream\n");
            break;
        }
    }

    free(response);
    return list;
}

void get_firehose_stream(const char* ip, int port, uint64_t offset, uint16_t count) {
    size_t send_len = 1 + 1 + 8;
    byte   send_data[ send_len ];
    send_data[ 0 ] = (byte) firehose;

    byte*  response = NULL;
    size_t resp_len = 0;
    bool   success  = false;

    uint64_t last_offset = offset;
    while (true) {
        firehose_request2s(
            (struct firehose_request) { .count = (uint8_t) count, .offset = last_offset },
            send_data + 1);
        int ret = study(ip, port, send_data, send_len, &response, &resp_len);
        if (ret == -1) {
            printf("study returned -1\n");
            return;
        }

        if (ret != 0 || resp_len < 9 || response[ 0 ] != firehose_reply) {
            free(response);
            sleep(5);
            continue;
        }

        uint64_t server_offset     = read_u64_le(response + 1);
        size_t   available_bytes   = resp_len - 9;
        size_t   available_packets = available_bytes / PACKET_SIZE;

        if (available_packets == 0) {
            free(response);
            sleep(5);
            continue;
        }

        size_t current_offset = 9;
        for (size_t i = 0; i < available_packets; i++) {
            struct data_packet packet
                = s2data_packet(PACKET_SIZE, response + current_offset, &success);
            if (!success) break;
            printf("streamed " hash_fmt "\n", hash_ref(packet.hash));
            current_offset += PACKET_SIZE;
        }

        printf("offset: %li\n", last_offset);
        last_offset += available_packets;
        free(response);
        // sleep(5);
    }
}

int main() {
    struct relay_info info = run_server_audit("127.0.0.1", 8040);

    printf("Server Audit Complete:\n");
    printf("  Version: v%i.%i\n", info.major_version, info.minor_version);
    printf("  MOTD: %s\n", info.motd);
    printf("  Leading Zeros Required: %i\n", info.capabilities.leading_zeros);
    printf("  Retention: %i hours\n", info.capabilities.retention_hours);
    printf("  Neighbors: %u\n", info.statistics.neighbors);
    printf("  Held Packets: %lu\n", info.statistics.held_packets);

    keypair author = generate_keypair();

    const char input[] = "this is a test message from the moonpaste client.";

    struct data_packet new_packet
        = create_packet_from_bytes((byte*) input, sizeof(input), 0, NULL, author);
    new_packet = run_pow(new_packet, info.capabilities.leading_zeros);
    printf("hash: " hash_fmt "\n", hash_ref(new_packet.hash));

    printf("pushing new packet to server...\n");
    if (push_packet_to_server("127.0.0.1", 8040, new_packet)) {
        printf("packet pushed successfully!\n");
    } else {
        printf("failed to push packet.\n");
    }

    printf(
        "attempting to retrieve the packet with hash " hash_fmt "...\n", hash_ref(new_packet.hash));
    struct data_packet retrieved_packet
        = get_packet_from_server("127.0.0.1", 8040, new_packet.hash);
    if (retrieved_packet.nonce != 0) {
        printf("packet retrieved successfully! Contents:\n");
        for (int i = 0; i < 65536; ++i) {
            printf("%c", retrieved_packet.contents.data[ i ]);
            if (!retrieved_packet.contents.data[ i ]) break;
        }
        printf("\n");
    } else {
        printf("failed to retrieve packet.\n");
    }

    printf("requesting neighbor list...\n");
    byte                   requested_neighbors = 5;
    struct relay_neighbor* neighbors = get_server_neighbors("127.0.0.1", 8040, requested_neighbors);
    if (neighbors) {
        printf("received %i neighbors.\n", (int) requested_neighbors);
        for (int i = 0; i < requested_neighbors; ++i) {
            printf("Neighbor %i: ", i + 1);
            if (neighbors[ i ].has_ipv4) {
                printf(
                    "%d.%d.%d.%d\n",
                    neighbors[ i ].v4_addr[ 0 ],
                    neighbors[ i ].v4_addr[ 1 ],
                    neighbors[ i ].v4_addr[ 2 ],
                    neighbors[ i ].v4_addr[ 3 ]);
            } else if (neighbors[ i ].has_ipv6) {
                char ip_string[ 46 ];
                int  offset = 0;
                for (int j = 0; j < 8; ++j) {
                    uint16_t word = (neighbors[ i ].v6_addr[ j * 2 ] << 8)
                                    | neighbors[ i ].v6_addr[ j * 2 + 1 ];
                    offset += snprintf(
                        ip_string + offset,
                        sizeof(ip_string) - offset,
                        "%x%s",
                        word,
                        (j < 7) ? ":" : "");
                }
                printf("%s\n", ip_string);
            }
        }
        free(neighbors);
    } else {
        printf("Failed to get neighbors.\n");
    }

    printf("\nstarting firehose stream. Press Ctrl+C to exit.\n");
    get_firehose_stream("127.0.0.1", 8040, 0, 10);

    return 0;
}

#define hash_fmt                                                                                   \
    "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"                             \
    "%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x%02x"

#define hash_ref(_hash)                                                                            \
    _hash[ 0 ], _hash[ 1 ], _hash[ 2 ], _hash[ 3 ], _hash[ 4 ], _hash[ 5 ], _hash[ 6 ],            \
        _hash[ 7 ], _hash[ 8 ], _hash[ 9 ], _hash[ 10 ], _hash[ 11 ], _hash[ 12 ], _hash[ 13 ],    \
        _hash[ 14 ], _hash[ 15 ], _hash[ 16 ], _hash[ 17 ], _hash[ 18 ], _hash[ 19 ], _hash[ 20 ], \
        _hash[ 21 ], _hash[ 22 ], _hash[ 23 ], _hash[ 24 ], _hash[ 25 ], _hash[ 26 ], _hash[ 27 ], \
        _hash[ 28 ], _hash[ 29 ], _hash[ 30 ], _hash[ 31 ]