#pragma once

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <time.h>
//
#include "omen.c"

#define likely(A)   __builtin_expect(A, 1)
#define unlikely(A) __builtin_expect(A, 0)

typedef uint8_t byte;

typedef byte ipv4[ 4 ];
typedef byte ipv6[ 16 ];

typedef byte hash[ 32 ];
typedef byte public_key[ 32 ];
typedef byte private_key[ 32 ];
typedef byte signature[ 64 ];

const uint16_t default_port = 8040;

enum client_input {
    // ask server for its information (responds with audit_response)
    audit_server = 0x9a,

    // ask server for a random set of neighbors
    // (input uint8_t neighbor_count, responds with neighbors_info)
    get_neighbors = 0x20,

    // monitor incoming packets
    // (input struct firehose_request, responds with firehose_reply)
    firehose = 0xf1,

    // give a packet to the relay
    // (input struct data_packet, responds with push_packet_result)
    push_packet = 0x88,

    // get a packet from the server
    // (input hash, responds with get_packet_result)
    get_packet = 0x87,

    // register a new neighbor relay to this server
    // (input struct relay_neighbor, responds with register_neighbor_result)
    register_neighbor = 0x73,
};

enum server_reply {
    // pretty self-explanatory.
    error = 0x1,

    // type `struct relay_info`
    audit_response = 0xaf,
    // type `struct relay_neighbor[ neighbor_count ]`
    neighbors_info = 0x43,
    // type `uint64_t new_offset, struct data_packet[ count ]`
    firehose_reply = 0x6c,
    // type `bool result_code`
    push_packet_result = 0x4a,
    // type `struct data_packet`
    get_packet_result = 0x42,
    // type `bool result_code`
    register_neighbor_result = 0x6e,
};

typedef uint8_t get_neighbors_request;

struct firehose_request {
    uint8_t  count;
    uint64_t offset;
};

typedef uint64_t firehose_reply_new_offset;

struct relay_neighbor {
    bool active; // set to 0 to indicate empty quick connect slot

    bool has_ipv4; // whether the neighbor supports ipv4 connection
    bool has_ipv6; // whether the neighbor supports ipv6 connection
    ipv4 v4_addr;  // ipv4 address of the neighbor (if supported)
    ipv6 v6_addr;  // ipv6 address of the neighbor (if supported)
};

struct relay_info {
    // -- version info --
    byte major_version;
    byte minor_version;
    // -- motd --
    byte motd[ 128 ];
    // -- capabilities --
    struct {
        uint16_t retention_hours; // how many hours packets will be retained for
        bool     can_store;       // whether or not the relay is accepting new packets
        byte     leading_zeros;   // relay challenge level
        bool     accepts_plain;   // whether or not the relay accepts unencrypted packets
    } capabilities;
    // -- statistics --
    struct {
        uint16_t              neighbors;          // current neighboring relay count for this relay
        uint64_t              held_packets;       // how many packets the relay is holding
        struct relay_neighbor quick_connect[ 4 ]; // four best neighbors for quick connection
    } statistics;
};

// do not change
#define MAX_RECIPIENTS 32

struct data_packet {
    // nonce + hash for PoW
    uint32_t nonce;
    hash     hash;

    // content inside to be nonce + hashed
    struct {
        bool       encrypted;          // whether the contents is encrypted or plaintext
        uint64_t   created_time_ms;    // unix timestamp for creation time in milliseconds
        public_key author;             // message author
        public_key author_signing_pub; // message signing pubkey
        uint8_t    recipient_count;    // number of recipients. leave as 0 for public.
        public_key recipients[ 32 ];   // up to 32 recipients for packet
        byte       data[ 65536 ];      // packet data
        signature  signature;          // data signature
    } contents;
};

const size_t CONTENTS_SIZE = 1 + 8 + 32 + 32 + 1 + (32 * MAX_RECIPIENTS) + 65536 + 64;
const size_t PACKET_SIZE   = 4 + 32 + (CONTENTS_SIZE);

struct relay_neighbor s2relay_neighbor(size_t len, byte *stream, bool *success) {
    if (len != 1 + 1 + 1 + 4 + 16) {
        *success = false;
        return (struct relay_neighbor) { 0 };
    }

    struct relay_neighbor out;
    out.active   = stream[ 0 ] != 0;
    out.has_ipv4 = stream[ 1 ] != 0;
    out.has_ipv6 = stream[ 2 ] != 0;

    memcpy(out.v4_addr, stream + 3, 4);
    memcpy(out.v6_addr, stream + 7, 16);

    *success = true;
    return out;
}

void relay_neighbor2s(struct relay_neighbor data, byte *stream) {
    stream[ 0 ] = data.active ? 1 : 0;
    stream[ 1 ] = data.has_ipv4 ? 1 : 0;
    stream[ 2 ] = data.has_ipv6 ? 1 : 0;

    memcpy(stream + 3, data.v4_addr, 4);
    memcpy(stream + 7, data.v6_addr, 16);
}

struct relay_info s2relay_info(size_t len, byte *stream, bool *success) {
    if (len != 1 + 1 + 128 + 2 + 1 + 1 + 1 + 2 + 8 + (4 * (1 + 1 + 1 + 4 + 16))) {
        *success = false;
        return (struct relay_info) { 0 };
    }

    struct relay_info out;
    size_t            off = 0;

    out.major_version = stream[ off++ ];
    out.minor_version = stream[ off++ ];

    memcpy(out.motd, stream + off, 128);
    off += 128;

    out.capabilities.retention_hours = stream[ off ] + (stream[ off + 1 ] << 8);
    off += 2;
    out.capabilities.can_store     = stream[ off++ ] != 0;
    out.capabilities.leading_zeros = stream[ off++ ];
    out.capabilities.accepts_plain = stream[ off++ ] != 0;

    out.statistics.neighbors = stream[ off ] + (stream[ off + 1 ] << 8);
    off += 2;
    out.statistics.held_packets
        = ((uint64_t) stream[ off ]) | ((uint64_t) stream[ off + 1 ] << 8)
          | ((uint64_t) stream[ off + 2 ] << 16) | ((uint64_t) stream[ off + 3 ] << 24)
          | ((uint64_t) stream[ off + 4 ] << 32) | ((uint64_t) stream[ off + 5 ] << 40)
          | ((uint64_t) stream[ off + 6 ] << 48) | ((uint64_t) stream[ off + 7 ] << 56);
    off += 8;

    for (int i = 0; i < 4; i++) {
        bool ok;
        out.statistics.quick_connect[ i ] = s2relay_neighbor(1 + 1 + 1 + 4 + 16, stream + off, &ok);
        if (!ok) {
            *success = false;
            return (struct relay_info) { 0 };
        }
        off += 1 + 1 + 1 + 4 + 16;
    }

    *success = true;
    return out;
}

void relay_info2s(struct relay_info data, byte *stream) {
    size_t off = 0;

    stream[ off++ ] = data.major_version;
    stream[ off++ ] = data.minor_version;

    memcpy(stream + off, data.motd, 128);
    off += 128;

    stream[ off++ ] = data.capabilities.retention_hours & 0xFF;
    stream[ off++ ] = (data.capabilities.retention_hours >> 8) & 0xFF;
    stream[ off++ ] = data.capabilities.can_store ? 1 : 0;
    stream[ off++ ] = data.capabilities.leading_zeros;
    stream[ off++ ] = data.capabilities.accepts_plain ? 1 : 0;

    stream[ off++ ] = data.statistics.neighbors & 0xFF;
    stream[ off++ ] = (data.statistics.neighbors >> 8) & 0xFF;

    stream[ off++ ] = data.statistics.held_packets & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 8) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 16) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 24) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 32) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 40) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 48) & 0xFF;
    stream[ off++ ] = (data.statistics.held_packets >> 56) & 0xFF;

    for (int i = 0; i < 4; i++) {
        relay_neighbor2s(data.statistics.quick_connect[ i ], stream + off);
        off += 1 + 1 + 1 + 4 + 16;
    }
}

struct data_packet s2data_packet(size_t len, byte *stream, bool *success) {
    *success               = false;
    struct data_packet out = { 0 };
    if (len != PACKET_SIZE) {
        printf("s2data_packet: length mismatch. expected %zu, got %zu\n", PACKET_SIZE, len);
        return out;
    }

    size_t off = 0;

    out.nonce = stream[ off ] | (stream[ off + 1 ] << 8) | (stream[ off + 2 ] << 16)
                | (stream[ off + 3 ] << 24);
    off += 4;

    memcpy(out.hash, stream + off, 32);
    off += 32;

    out.contents.encrypted = stream[ off++ ] != 0;

    out.contents.created_time_ms
        = ((uint64_t) stream[ off ]) | ((uint64_t) stream[ off + 1 ] << 8)
          | ((uint64_t) stream[ off + 2 ] << 16) | ((uint64_t) stream[ off + 3 ] << 24)
          | ((uint64_t) stream[ off + 4 ] << 32) | ((uint64_t) stream[ off + 5 ] << 40)
          | ((uint64_t) stream[ off + 6 ] << 48) | ((uint64_t) stream[ off + 7 ] << 56);
    off += 8;

    memcpy(out.contents.author, stream + off, 32);
    off += 32;

    memcpy(out.contents.author_signing_pub, stream + off, 32);
    off += 32;

    out.contents.recipient_count = stream[ off++ ];
    memcpy(out.contents.recipients, stream + off, 32 * MAX_RECIPIENTS);
    off += 32 * MAX_RECIPIENTS;

    memcpy(out.contents.data, stream + off, 65536);
    off += 65536;

    memcpy(out.contents.signature, stream + off, 64);
    off += 64;

    *success = true;
    return out;
}

bool data_packet2s(struct data_packet data, byte *stream) {
    if (data.contents.recipient_count > MAX_RECIPIENTS) { return false; }

    size_t off = 0;

    stream[ off++ ] = data.nonce & 0xFF;
    stream[ off++ ] = (data.nonce >> 8) & 0xFF;
    stream[ off++ ] = (data.nonce >> 16) & 0xFF;
    stream[ off++ ] = (data.nonce >> 24) & 0xFF;

    memcpy(stream + off, data.hash, 32);
    off += 32;

    stream[ off++ ] = data.contents.encrypted ? 1 : 0;

    for (int i = 0; i < 8; i++) stream[ off++ ] = (data.contents.created_time_ms >> (i * 8)) & 0xFF;

    memcpy(stream + off, data.contents.author, 32);
    off += 32;

    memcpy(stream + off, data.contents.author_signing_pub, 32);
    off += 32;

    stream[ off++ ] = data.contents.recipient_count;
    memcpy(stream + off, data.contents.recipients, 32 * MAX_RECIPIENTS);
    off += 32 * MAX_RECIPIENTS;

    memcpy(stream + off, data.contents.data, 65536);
    off += 65536;

    memcpy(stream + off, data.contents.signature, 64);
    off += 64;

    if (off != PACKET_SIZE) {
        printf("data_packet2s: size mismatch wrote %zu expected %zu\n", off, PACKET_SIZE);
        return false;
    }

    return true;
}

get_neighbors_request s2get_neighbors_request(size_t len, byte *stream, bool *success) {
    if (len != 1) {
        *success = false;
        return 0;
    }
    *success = true;
    return stream[ 0 ];
}

void get_neighbors_request2s(get_neighbors_request data, byte *stream) { stream[ 0 ] = data; }

struct firehose_request s2firehose_request(size_t len, byte *stream, bool *success) {
    if (len != 1 + 8) {
        *success = false;
        return (struct firehose_request) { 0 };
    }
    struct firehose_request out;
    out.count  = stream[ 0 ];
    out.offset = ((uint64_t) stream[ 1 ]) | ((uint64_t) stream[ 2 ] << 8)
                 | ((uint64_t) stream[ 3 ] << 16) | ((uint64_t) stream[ 4 ] << 24)
                 | ((uint64_t) stream[ 5 ] << 32) | ((uint64_t) stream[ 6 ] << 40)
                 | ((uint64_t) stream[ 7 ] << 48) | ((uint64_t) stream[ 8 ] << 56);
    *success = true;
    return out;
}

void firehose_request2s(struct firehose_request data, byte *stream) {
    stream[ 0 ] = data.count;
    stream[ 1 ] = data.offset & 0xFF;
    stream[ 2 ] = (data.offset >> 8) & 0xFF;
    stream[ 3 ] = (data.offset >> 16) & 0xFF;
    stream[ 4 ] = (data.offset >> 24) & 0xFF;
    stream[ 5 ] = (data.offset >> 32) & 0xFF;
    stream[ 6 ] = (data.offset >> 40) & 0xFF;
    stream[ 7 ] = (data.offset >> 48) & 0xFF;
    stream[ 8 ] = (data.offset >> 56) & 0xFF;
}

void s2hash(size_t len, byte *stream, bool *success, hash out) {
    if (len != 32) {
        *success = false;
        return;
    }
    memcpy(out, stream, 32);
    *success = true;
}

void hash2s(hash data, byte *stream) { memcpy(stream, data, 32); }

bool check_leading_zeros(const hash _hash, byte required_zeros) {
    if (required_zeros == 0) return true;

    int full_bytes     = required_zeros / 8;
    int remaining_bits = required_zeros % 8;

    for (int i = 0; i < full_bytes; i++) {
        if (_hash[ i ] != 0) return false;
    }

    if (remaining_bits > 0) {
        byte mask = (0xFF << (8 - remaining_bits));
        if ((_hash[ full_bytes ] & mask) != 0) return false;
    }

    return true;
}

uint64_t get_current_time_ms() {
    struct timespec ts;
    clock_gettime(CLOCK_REALTIME, &ts);
    return (uint64_t) ts.tv_sec * 1000 + (uint64_t) ts.tv_nsec / 1000000;
}

#define RELAY_NEIGHBOR_WIRE_LEN (1 + 1 + 1 + 4 + 16)

static inline uint64_t read_u64_le(const byte *b) {
    uint64_t v = 0;
    for (int i = 0; i < 8; ++i) v |= (uint64_t) b[ i ] << (8 * i);
    return v;
}

static inline void write_u64_le(byte *b, uint64_t v) {
    for (int i = 0; i < 8; ++i) b[ i ] = (v >> (8 * i)) & 0xFF;
}

static inline uint32_t read_u32_le(const byte *b) {
    uint32_t v = 0;
    for (int i = 0; i < 4; ++i) v |= (uint32_t) b[ i ] << (8 * i);
    return v;
}

static inline void write_u32_le(byte *b, uint32_t v) {
    for (int i = 0; i < 4; ++i) b[ i ] = (v >> (8 * i)) & 0xFF;
}

static inline uint16_t read_u16_le(const byte *b) {
    return (uint16_t) b[ 0 ] | ((uint16_t) b[ 1 ] << 8);
}

static inline void write_u16_le(byte *b, uint16_t v) {
    b[ 0 ] = v & 0xFF;
    b[ 1 ] = (v >> 8) & 0xFF;
}

struct relay_info run_server_audit(const char* ip, int port) {
    byte   data[]   = { audit_server, 0, 0, 0, 0, 0, 0, 0 };
    byte*  response = NULL;
    size_t resp_len = 0;
    bool   success;

    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 (struct relay_info) { 0 };
    }

    struct relay_info info = s2relay_info(resp_len - 1, response + 1, &success);
    if (!success) { printf("error: failed to deserialize relay_info\n"); }

    free(response);
    return info;
}

pthread_mutex_t packet_storage_mutex = PTHREAD_MUTEX_INITIALIZER;