#include "peer_udp_client.h"

#include <poll.h>
#include <pthread.h>
#include <string.h>

struct udp_client {
    char id[OMNI_MAX_PEER_ID];
    udp_conn_t *conn;
    latency_logger_t *logger;
    pthread_mutex_t id_mu;
    uint64_t next_message_id;
};

static int client_next_message_id(udp_client_t *client, uint64_t *out_id) {
    pthread_mutex_lock(&client->id_mu);
    *out_id = ++client->next_message_id;
    pthread_mutex_unlock(&client->id_mu);
    return 0;
}

static void udp_client_fill_recv_meta(udp_client_recv_meta_t *meta, const message_t *msg) {
    if (meta == NULL || msg == NULL) {
        return;
    }
    memset(meta, 0, sizeof(*meta));
    meta->type = msg->type;
    meta->id = msg->id;
    meta->body_len = msg->body_len;
    snprintf(meta->from, sizeof(meta->from), "%s", msg->from);
    snprintf(meta->to, sizeof(meta->to), "%s", msg->to);
    snprintf(meta->file_name, sizeof(meta->file_name), "%s", msg->file_name);
}

static int client_persist_message_to_disk(const message_t *msg, const char *inbox_dir, char *out_path, size_t out_path_len) {
    char path[512];
    if (omni_ensure_dir(inbox_dir) != 0) {
        return -1;
    }
    if (msg->type == MSG_TYPE_TEXT) {
        char *body = omni_json_escape_bytes(msg->body, msg->body_len);
        char *from = omni_json_escape(msg->from);
        char *to = omni_json_escape(msg->to);
        char *line;
        if (body == NULL || from == NULL || to == NULL) {
            free(body);
            free(from);
            free(to);
            return -1;
        }
        snprintf(path, sizeof(path), "%s/messages.log", inbox_dir);
        line = omni_strdup_printf("{\"message_type\":\"%s\",\"message_id\":%" PRIu64 ",\"from\":\"%s\",\"to\":\"%s\",\"body\":\"%s\"}\n", protocol_message_type_name(msg->type), msg->id, from, to, body);
        free(body);
        free(from);
        free(to);
        if (line == NULL) {
            return -1;
        }
        if (omni_append_file(path, (const uint8_t *) line, strlen(line)) != 0) {
            free(line);
            return -1;
        }
        free(line);
    } else if (msg->type == MSG_TYPE_FILE) {
        const char *file_name = omni_path_base_name(msg->file_name);
        if (file_name[0] == '\0') {
            file_name = "unnamed";
        }
        snprintf(path, sizeof(path), "%s/%s-%" PRIu64 "-%s", inbox_dir, msg->from, msg->id, file_name);
        if (omni_write_file(path, msg->body, msg->body_len) != 0) {
            return -1;
        }
    } else if (msg->type == MSG_TYPE_BINARY) {
        snprintf(path, sizeof(path), "%s/%s-%" PRIu64 ".bin", inbox_dir, msg->from, msg->id);
        if (omni_write_file(path, msg->body, msg->body_len) != 0) {
            return -1;
        }
    } else {
        errno = EINVAL;
        return -1;
    }
    if (out_path != NULL && out_path_len > 0) {
        snprintf(out_path, out_path_len, "%s", path);
    }
    return 0;
}

udp_client_t *udp_client_dial_with_options(const char *server_addr, const char *peer_id, const char *bind_ip, const char *bind_device, latency_logger_t *logger, tx_timestamp_debug_logger_t *debug_logger, int enable_timestamping) {
    udp_client_t *client;
    message_t register_msg;
    client = (udp_client_t *) calloc(1, sizeof(*client));
    if (client == NULL) {
        return NULL;
    }
    snprintf(client->id, sizeof(client->id), "%s", peer_id);
    pthread_mutex_init(&client->id_mu, NULL);
    client->logger = logger;
    client->conn = udp_conn_dial(server_addr, bind_ip, bind_device, enable_timestamping, logger, OMNI_NODE_ROLE_PEER, peer_id, debug_logger);
    if (client->conn == NULL) {
        udp_client_free(client);
        return NULL;
    }
    protocol_message_init(&register_msg);
    register_msg.type = MSG_TYPE_REGISTER;
    register_msg.id = 0;
    snprintf(register_msg.from, sizeof(register_msg.from), "%s", peer_id);
    snprintf(register_msg.to, sizeof(register_msg.to), "%s", SERVER_PEER_ID);
    if (udp_conn_send(client->conn, &register_msg) != 0) {
        udp_client_free(client);
        return NULL;
    }
    return client;
}

udp_client_t *udp_client_dial(const char *server_addr, const char *peer_id, const char *bind_ip, latency_logger_t *logger, tx_timestamp_debug_logger_t *debug_logger, int enable_timestamping) {
    return udp_client_dial_with_options(server_addr, peer_id, bind_ip, NULL, logger, debug_logger, enable_timestamping);
}

const char *udp_client_id(const udp_client_t *client) {
    return client == NULL ? "" : client->id;
}

int udp_client_send_text(udp_client_t *client, const char *to, const char *text) {
    message_t msg;
    uint64_t id;
    protocol_message_init(&msg);
    client_next_message_id(client, &id);
    msg.type = MSG_TYPE_TEXT;
    msg.id = id;
    snprintf(msg.from, sizeof(msg.from), "%s", client->id);
    snprintf(msg.to, sizeof(msg.to), "%s", to);
    msg.body = (uint8_t *) omni_strdup(text);
    if (msg.body == NULL) {
        return -1;
    }
    msg.body_len = strlen((const char *) msg.body);
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_A_APP_PREP_BEGIN, &msg);
    if (udp_conn_send(client->conn, &msg) != 0) {
        protocol_message_clear(&msg);
        return -1;
    }
    protocol_message_clear(&msg);
    return 0;
}

int udp_client_send_binary(udp_client_t *client, const char *to, const void *data, size_t data_len) {
    message_t msg;
    uint64_t id;

    if (client == NULL || to == NULL || (data == NULL && data_len > 0)) {
        errno = EINVAL;
        return -1;
    }

    protocol_message_init(&msg);
    client_next_message_id(client, &id);
    msg.type = MSG_TYPE_BINARY;
    msg.id = id;
    snprintf(msg.from, sizeof(msg.from), "%s", client->id);
    snprintf(msg.to, sizeof(msg.to), "%s", to);
    if (data_len > 0) {
        msg.body = (uint8_t *) malloc(data_len);
        if (msg.body == NULL) {
            return -1;
        }
        memcpy(msg.body, data, data_len);
    }
    msg.body_len = data_len;
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_A_APP_PREP_BEGIN, &msg);
    if (udp_conn_send(client->conn, &msg) != 0) {
        protocol_message_clear(&msg);
        return -1;
    }
    protocol_message_clear(&msg);
    return 0;
}

int udp_client_send_file_path(udp_client_t *client, const char *to, const char *path) {
    message_t msg;
    uint64_t id;
    uint8_t *body = NULL;
    size_t body_len = 0;
    const char *base_name = strrchr(path, '/');
    if (omni_read_file(path, &body, &body_len) != 0) {
        return -1;
    }
    protocol_message_init(&msg);
    client_next_message_id(client, &id);
    msg.type = MSG_TYPE_FILE;
    msg.id = id;
    snprintf(msg.from, sizeof(msg.from), "%s", client->id);
    snprintf(msg.to, sizeof(msg.to), "%s", to);
    snprintf(msg.file_name, sizeof(msg.file_name), "%s", base_name == NULL ? path : base_name + 1);
    msg.body = body;
    msg.body_len = body_len;
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_A_APP_PREP_BEGIN, &msg);
    if (udp_conn_send(client->conn, &msg) != 0) {
        protocol_message_clear(&msg);
        return -1;
    }
    protocol_message_clear(&msg);
    return 0;
}

int udp_client_receive_timed(udp_client_t *client, message_t *out_msg, int timeout_ms) {
    if (client == NULL || out_msg == NULL) {
        errno = EINVAL;
        return -1;
    }

    if (timeout_ms >= 0) {
        struct pollfd pfd;
        int rc;

        memset(&pfd, 0, sizeof(pfd));
        pfd.fd = udp_conn_fd(client->conn);
        pfd.events = POLLIN | POLLERR | POLLHUP;
        do {
            rc = poll(&pfd, 1, timeout_ms);
        } while (rc < 0 && errno == EINTR);
        if (rc == 0) {
            return 1;
        }
        if (rc < 0) {
            return -1;
        }
        if ((pfd.revents & (POLLERR | POLLHUP | POLLNVAL)) != 0 && (pfd.revents & POLLIN) == 0) {
            errno = ECONNRESET;
            return -1;
        }
    }

    if (udp_conn_receive(client->conn, out_msg, NULL, NULL) != 0) {
        return -1;
    }
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_B_APP_RECV, out_msg);
    return 0;
}

int udp_client_receive(udp_client_t *client, message_t *out_msg) {
    return udp_client_receive_timed(client, out_msg, -1);
}

int udp_client_receive_into(udp_client_t *client, void *buffer, size_t buffer_len, udp_client_recv_meta_t *out_meta, int timeout_ms) {
    message_t msg;
    int rc;

    if (client == NULL || (buffer == NULL && buffer_len > 0) || out_meta == NULL) {
        errno = EINVAL;
        return -1;
    }

    protocol_message_init(&msg);
    rc = udp_client_receive_timed(client, &msg, timeout_ms);
    if (rc != 0) {
        return rc;
    }

    udp_client_fill_recv_meta(out_meta, &msg);
    if (msg.body_len > buffer_len) {
        protocol_message_clear(&msg);
        errno = EMSGSIZE;
        return 2;
    }

    if (msg.body_len > 0) {
        memcpy(buffer, msg.body, msg.body_len);
    }
    protocol_message_clear(&msg);
    return 0;
}

int udp_client_persist_message(udp_client_t *client, const message_t *msg, const char *inbox_dir, char *out_path, size_t out_path_len) {
    if (!latencylog_is_business_message(msg)) {
        errno = EINVAL;
        return -1;
    }
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_B_PERSIST_BEGIN, msg);
    if (client_persist_message_to_disk(msg, inbox_dir, out_path, out_path_len) != 0) {
        return -1;
    }
    latencylog_log_message_event(client->logger, OMNI_NODE_ROLE_PEER, client->id, EVENT_B_PERSIST_END, msg);
    return 0;
}

int udp_client_close(udp_client_t *client) {
    return client == NULL ? 0 : udp_conn_close(client->conn);
}

void udp_client_free(udp_client_t *client) {
    if (client == NULL) {
        return;
    }
    udp_conn_free(client->conn);
    pthread_mutex_destroy(&client->id_mu);
    free(client);
}