#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
//
#include "../src/common.c"
#include "../src/security.c"
#include "../src/sha256.c"

#define MAX_BUFFER 1 << 18
#define MAX_BODY   1 << 22

typedef enum {
    MP_CLIENT_SUCCESS = 0,
    MP_CLIENT_ERROR_NETWORK_FAIL,
    MP_CLIENT_ERROR_INVALID_ADDRESS,
    MP_CLIENT_ERROR_CONNECTION_FAILED,
    MP_CLIENT_ERROR_SEND_FAILED,
    MP_CLIENT_ERROR_RECV_FAILED,
    MP_CLIENT_ERROR_SERVER_BUSY, // 418 teapot
    MP_CLIENT_ERROR_INVALID_RESPONSE,
    MP_CLIENT_ERROR_DESERIALIZE_FAILED,
    MP_CLIENT_ERROR_ALLOC_FAILED,
    MP_CLIENT_ERROR_FILE_NOT_FOUND,
    MP_CLIENT_ERROR_FILE_READ_FAILED,
    MP_CLIENT_ERROR_BAD_ARGUMENT,
} mp_client_error_t;

struct data_packet_list {
    struct data_packet *packets; // caller must free this
    uint16_t            count;
    uint64_t            new_offset;
};

static ssize_t __write_all(int fd, const void *buf, size_t len) {
    size_t      off = 0;
    const char *p   = buf;
    while (off < len) {
        ssize_t r = send(fd, p + off, len - off, MSG_NOSIGNAL);
        if (r < 0) {
            if (errno == EINTR || errno == EAGAIN || errno == EWOULDBLOCK) continue;
            return -1;
        }
        off += (size_t) r;
    }
    return (ssize_t) off;
}

static int __read_exact(int fd, void *buf, size_t count) {
    size_t off = 0;
    char  *p   = buf;
    while (off < count) {
        ssize_t r = recv(fd, p + off, count - off, 0);
        if (r < 0) {
            if (errno == EINTR) continue;
            return -1;
        } else if (r == 0) {
            return -1;
        }
        off += (size_t) r;
    }
    return 0;
}

static mp_client_error_t __read_until_headers(int fd, char **out_buf, size_t *out_len) {
    size_t cap = 4096;
    size_t len = 0;
    char  *buf = malloc(cap);
    if (!buf) return MP_CLIENT_ERROR_ALLOC_FAILED;

    while (1) {
        if (cap - len < 1024) {
            size_t ncap = cap * 2;
            char  *n    = realloc(buf, ncap);
            if (!n) {
                free(buf);
                return MP_CLIENT_ERROR_ALLOC_FAILED;
            }
            buf = n;
            cap = ncap;
        }
        ssize_t r = recv(fd, buf + len, cap - len, 0);
        if (r < 0) {
            if (errno == EINTR) continue;
            free(buf);
            return MP_CLIENT_ERROR_RECV_FAILED;
        } else if (r == 0) {
            break;
        }
        len += (size_t) r;
        if (len >= 4 && memmem(buf, len, "\r\n\r\n", 4) != NULL) { break; }
        if (len > MAX_BUFFER) {
            free(buf);
            return MP_CLIENT_ERROR_INVALID_RESPONSE; // Header too large
        }
    }
    *out_buf = buf;
    *out_len = len;
    return MP_CLIENT_SUCCESS;
}

static ssize_t __parse_content_length(const char *hdr, size_t hdr_len) {
    const char *p   = hdr;
    const char *end = hdr + hdr_len;
    while (p < end) {
        const char *line_end = memchr(p, '\n', (size_t) (end - p));
        if (!line_end) line_end = end;
        if ((size_t) (line_end - p) >= 15 && strncasecmp(p, "Content-Length:", 15) == 0) {
            const char *num = p + 15;
            while (num < line_end && (*num == ' ' || *num == '\t')) num++;
            return strtol(num, NULL, 10);
        }
        p = line_end + 1;
    }
    return -1;
}

static mp_client_error_t mp_client_request(
    const char *ip,
    int         port,
    const byte *data,
    size_t      data_len,
    byte      **response,
    size_t     *resp_len) {
    if (data_len < 8) data_len = 8;

    *response             = NULL;
    *resp_len             = 0;
    mp_client_error_t err = MP_CLIENT_SUCCESS;

    bool is_ipv6 = strchr(ip, ':') != NULL;
    int  domain  = is_ipv6 ? AF_INET6 : AF_INET;

    int sock = socket(domain, SOCK_STREAM, 0);
    if (sock < 0) return MP_CLIENT_ERROR_NETWORK_FAIL;

    struct sockaddr_storage addr = { 0 };
    addr.ss_family               = domain;
    if (domain == AF_INET) {
        struct sockaddr_in *addr4 = (struct sockaddr_in *) &addr;
        addr4->sin_port           = htons(port);
        if (inet_pton(domain, ip, &addr4->sin_addr) != 1) {
            err = MP_CLIENT_ERROR_INVALID_ADDRESS;
            goto cleanup_sock;
        }
    } else {
        struct sockaddr_in6 *addr6 = (struct sockaddr_in6 *) &addr;
        addr6->sin6_port           = htons(port);
        if (inet_pton(domain, ip, &addr6->sin6_addr) != 1) {
            err = MP_CLIENT_ERROR_INVALID_ADDRESS;
            goto cleanup_sock;
        }
    }

    struct timeval timeout = { 5, 0 };
    setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
    setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));

    if (connect(
            sock,
            (struct sockaddr *) &addr,
            (domain == AF_INET) ? sizeof(struct sockaddr_in) : sizeof(struct sockaddr_in6))
        < 0) {
        err = MP_CLIENT_ERROR_CONNECTION_FAILED;
        goto cleanup_sock;
    }

    char req_hdr[ 256 ];
    int  hdr_n = snprintf(
        req_hdr,
        sizeof(req_hdr),
        "POST / HTTP/1.1\r\nHost: %s%s%s\r\nContent-Length: %zu\r\nConnection: close\r\n\r\n",
        is_ipv6 ? "[" : "",
        ip,
        is_ipv6 ? "]" : "",
        data_len);

    if (__write_all(sock, req_hdr, (size_t) hdr_n) < 0
        || (data_len > 0 && __write_all(sock, data, data_len) < 0)) {
        err = MP_CLIENT_ERROR_SEND_FAILED;
        goto cleanup_sock;
    }

    char  *hdr_buf = NULL;
    size_t hdr_len = 0;
    err            = __read_until_headers(sock, &hdr_buf, &hdr_len);
    if (err != MP_CLIENT_SUCCESS) { goto cleanup_sock; }

    if (memmem(hdr_buf, hdr_len, "HTTP/1.1 418", 12) != NULL) {
        err = MP_CLIENT_ERROR_SERVER_BUSY;
        free(hdr_buf);
        goto cleanup_sock;
    }

    char  *sep      = memmem(hdr_buf, hdr_len, "\r\n\r\n", 4);
    byte  *body_ptr = NULL;
    size_t body_len = 0;

    if (sep) {
        size_t  headers_bytes = (size_t) (sep - hdr_buf) + 4;
        ssize_t cl            = __parse_content_length(hdr_buf, headers_bytes);
        size_t  already       = hdr_len - headers_bytes;
        if (cl > 0) {
            if ((size_t) cl > MAX_BODY) cl = MAX_BODY;
            body_ptr = malloc((size_t) cl);
            if (!body_ptr) {
                err = MP_CLIENT_ERROR_ALLOC_FAILED;
                free(hdr_buf);
                goto cleanup_sock;
            }
            size_t to_copy = already > (size_t) cl ? (size_t) cl : already;
            memcpy(body_ptr, sep + 4, to_copy);
            body_len = to_copy;

            if ((size_t) body_len < (size_t) cl) {
                if (__read_exact(sock, body_ptr + body_len, (size_t) cl - body_len) < 0) {
                    free(body_ptr);
                    err = MP_CLIENT_ERROR_RECV_FAILED;
                    free(hdr_buf);
                    goto cleanup_sock;
                }
                body_len = (size_t) cl;
            }
        }
    }

    free(hdr_buf);
    *response = body_ptr;
    *resp_len = body_len;

cleanup_sock:
    close(sock);
    return err;
}

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

    mp_client_error_t ret = mp_client_request(ip, port, data, sizeof(data), &response, &resp_len);
    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }

    if (resp_len < 1 || response[ 0 ] != audit_response) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }

    bool success;
    *out_info = s2relay_info(resp_len - 1, response + 1, &success);
    free(response);

    return success ? MP_CLIENT_SUCCESS : MP_CLIENT_ERROR_DESERIALIZE_FAILED;
}

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

    mp_client_error_t ret = mp_client_request(ip, port, data, sizeof(data), &response, &resp_len);
    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }

    if (resp_len < 1 || response[ 0 ] != neighbors_info) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }

    size_t num_neighbors = (resp_len - 1) / RELAY_NEIGHBOR_WIRE_LEN;
    if (num_neighbors == 0) {
        free(response);
        return MP_CLIENT_SUCCESS;
    }

    *out_neighbors = malloc(num_neighbors * sizeof(struct relay_neighbor));
    if (!*out_neighbors) {
        free(response);
        return MP_CLIENT_ERROR_ALLOC_FAILED;
    }

    byte actual_count = 0;
    for (size_t i = 0; i < num_neighbors; ++i) {
        bool success;
        (*out_neighbors)[ i ] = s2relay_neighbor(
            RELAY_NEIGHBOR_WIRE_LEN, response + 1 + (i * RELAY_NEIGHBOR_WIRE_LEN), &success);
        if (success) { actual_count++; }
    }
    *out_count = actual_count;

    free(response);
    return MP_CLIENT_SUCCESS;
}

mp_client_error_t
mp_push_packet(const char *ip, int port, struct data_packet packet, bool *out_success) {
    const size_t send_len  = 1 + PACKET_SIZE;
    byte        *send_data = malloc(send_len);
    if (!send_data) return MP_CLIENT_ERROR_ALLOC_FAILED;

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

    byte  *response = NULL;
    size_t resp_len = 0;
    *out_success    = false;

    mp_client_error_t ret = mp_client_request(ip, port, send_data, send_len, &response, &resp_len);
    free(send_data);

    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }

    if (resp_len == 2 && response[ 0 ] == push_packet_result) {
        *out_success = response[ 1 ];
    } else {
        ret = MP_CLIENT_ERROR_INVALID_RESPONSE;
    }

    free(response);
    return ret;
}

mp_client_error_t
mp_get_packet(const char *ip, int port, const hash packet_hash, struct data_packet *out_packet) {
    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;

    mp_client_error_t ret = mp_client_request(ip, port, send_data, send_len, &response, &resp_len);
    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }

    if (resp_len < 1 || response[ 0 ] != get_packet_result) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }

    bool success;
    *out_packet = s2data_packet(resp_len - 1, response + 1, &success);
    free(response);

    return success ? MP_CLIENT_SUCCESS : MP_CLIENT_ERROR_DESERIALIZE_FAILED;
}

mp_client_error_t mp_get_firehose_packets(
    const char *ip, int port, uint64_t offset, uint16_t count, struct data_packet_list *out_list) {
    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;
    out_list->packets    = NULL;
    out_list->count      = 0;
    out_list->new_offset = offset;

    mp_client_error_t ret = mp_client_request(ip, port, send_data, send_len, &response, &resp_len);
    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }

    if (resp_len < 9 || response[ 0 ] != firehose_reply) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }

    out_list->new_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);
        return MP_CLIENT_SUCCESS;
    }

    out_list->packets = malloc(available_packets * sizeof(struct data_packet));
    if (!out_list->packets) {
        free(response);
        return MP_CLIENT_ERROR_ALLOC_FAILED;
    }

    out_list->count = 0;
    for (size_t i = 0; i < available_packets; i++) {
        bool               success;
        struct data_packet packet
            = s2data_packet(PACKET_SIZE, response + 9 + (i * PACKET_SIZE), &success);
        if (success) { out_list->packets[ out_list->count++ ] = packet; }
    }

    free(response);
    return MP_CLIENT_SUCCESS;
}

struct data_packet mp_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));

    size_t copy_len = input_len > 65536 ? 65536 : input_len;
    if (input && copy_len > 0) memcpy(packet.contents.data, input, copy_len);

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

    packet.contents.created_time_ms = get_current_time_ms();

    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;

    // Create signature
    size_t msg_len = CONTENTS_SIZE - sizeof(signature);
    byte  *msg_buf = malloc(msg_len);
    if (msg_buf) {
        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);
        free(msg_buf);
    }

    return packet;
}

mp_client_error_t mp_create_packet_from_file(
    const char         *filename,
    byte                recipient_count,
    hash               *recipients,
    keypair             author_keys,
    struct data_packet *out_packet) {
    FILE *file = fopen(filename, "rb");
    if (!file) return MP_CLIENT_ERROR_FILE_NOT_FOUND;

    fseek(file, 0, SEEK_END);
    long file_size = ftell(file);
    fseek(file, 0, SEEK_SET);

    if (file_size < 0) {
        fclose(file);
        return MP_CLIENT_ERROR_FILE_READ_FAILED;
    }

    byte *file_buf = calloc(1, 65536);
    if (!file_buf) {
        fclose(file);
        return MP_CLIENT_ERROR_ALLOC_FAILED;
    }

    size_t to_read = file_size > 65536 ? 65536 : file_size;
    if (fread(file_buf, 1, to_read, file) != to_read) {
        free(file_buf);
        fclose(file);
        return MP_CLIENT_ERROR_FILE_READ_FAILED;
    }
    fclose(file);

    *out_packet
        = mp_create_packet_from_bytes(file_buf, to_read, recipient_count, recipients, author_keys);
    free(file_buf);
    return MP_CLIENT_SUCCESS;
}

struct data_packet mp_run_pow(struct data_packet packet, byte leading_zeros) {
    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) {
        write_u32_le(hash_input_buffer, packet.nonce);
        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++;
    }

    return packet;
}

keypair mp_generate_keypair(void) { return generate_keypair(); }

mp_client_error_t mp_query_author_packets(
    const char              *ip,
    int                      port,
    keypair                  author_keys,
    uint16_t                 count,
    uint64_t                 offset,
    struct data_packet_list *out_list) {
    byte msg[ 32 + 32 + 2 + 8 ];
    memcpy(msg, author_keys.encryption_pub, 32);
    memcpy(msg + 32, author_keys.signing_pub, 32);
    write_u16_le(msg + 64, count);
    write_u64_le(msg + 66, offset);
    byte sig[ 64 ];
    if (crypto_sign_detached(sig, NULL, msg, sizeof(msg), author_keys.signing_priv) != 0) {
        return MP_CLIENT_ERROR_BAD_ARGUMENT;
    }
    const size_t send_len  = 1 + 32 + 32 + 64 + 2 + 8;
    byte        *send_data = malloc(send_len);
    if (!send_data) return MP_CLIENT_ERROR_ALLOC_FAILED;
    size_t off         = 0;
    send_data[ off++ ] = query_packets;
    memcpy(send_data + off, author_keys.encryption_pub, 32);
    off += 32;
    memcpy(send_data + off, author_keys.signing_pub, 32);
    off += 32;
    memcpy(send_data + off, sig, 64);
    off += 64;
    write_u16_le(send_data + off, count);
    off += 2;
    write_u64_le(send_data + off, offset);
    off += 8;
    byte             *response = NULL;
    size_t            resp_len = 0;
    mp_client_error_t ret = mp_client_request(ip, port, send_data, send_len, &response, &resp_len);
    free(send_data);
    if (ret != MP_CLIENT_SUCCESS) {
        free(response);
        return ret;
    }
    if (resp_len < 11 || response[ 0 ] != queried_author_packets) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }
    uint16_t resp_count   = read_u16_le(response + 1);
    uint64_t new_offset   = read_u64_le(response + 3);
    size_t   expected_len = 11 + (size_t) resp_count * PACKET_SIZE;
    if (resp_len != expected_len) {
        free(response);
        return MP_CLIENT_ERROR_INVALID_RESPONSE;
    }
    out_list->packets = malloc((size_t) resp_count * sizeof(struct data_packet));
    if (!out_list->packets) {
        free(response);
        return MP_CLIENT_ERROR_ALLOC_FAILED;
    }
    out_list->count      = 0;
    out_list->new_offset = new_offset;
    for (uint16_t i = 0; i < resp_count; ++i) {
        bool               success;
        struct data_packet packet
            = s2data_packet(PACKET_SIZE, response + 11 + (size_t) i * PACKET_SIZE, &success);
        if (success) { out_list->packets[ out_list->count++ ] = packet; }
    }
    free(response);
    return MP_CLIENT_SUCCESS;
}