#include "server_udp_hub.h"

#include <unistd.h>

typedef struct udp_peer_entry {
    struct udp_peer_entry *next;
    char peer_id[OMNI_MAX_PEER_ID];
    struct sockaddr_storage addr;
    socklen_t addr_len;
} udp_peer_entry_t;

struct udp_hub {
    udp_conn_t *conn;
    pthread_rwlock_t lock;
    udp_peer_entry_t *peers;
};

static int udp_addr_equal(const struct sockaddr_storage *a, socklen_t a_len, const struct sockaddr_storage *b, socklen_t b_len) {
    return a_len == b_len && memcmp(a, b, a_len) == 0;
}

static udp_peer_entry_t *udp_hub_find_by_id(udp_hub_t *hub, const char *peer_id) {
    udp_peer_entry_t *entry;
    for (entry = hub->peers; entry != NULL; entry = entry->next) {
        if (strcmp(entry->peer_id, peer_id) == 0) {
            return entry;
        }
    }
    return NULL;
}

static udp_peer_entry_t *udp_hub_find_by_addr(udp_hub_t *hub, const struct sockaddr_storage *addr, socklen_t addr_len) {
    udp_peer_entry_t *entry;
    for (entry = hub->peers; entry != NULL; entry = entry->next) {
        if (udp_addr_equal(&entry->addr, entry->addr_len, addr, addr_len)) {
            return entry;
        }
    }
    return NULL;
}

static int udp_hub_send_error(udp_hub_t *hub, const struct sockaddr_storage *addr, socklen_t addr_len, const char *to, const char *message) {
    message_t msg;
    protocol_message_init(&msg);
    msg.type = MSG_TYPE_ERROR;
    msg.id = 0;
    snprintf(msg.from, sizeof(msg.from), "%s", SERVER_PEER_ID);
    snprintf(msg.to, sizeof(msg.to), "%s", to == NULL || to[0] == '\0' ? "unknown" : to);
    msg.body = (uint8_t *) omni_strdup(message);
    msg.body_len = msg.body == NULL ? 0 : strlen((const char *) msg.body);
    if (msg.body == NULL) {
        return -1;
    }
    if (udp_conn_send_to(hub->conn, &msg, (const struct sockaddr *) addr, addr_len) != 0) {
        protocol_message_clear(&msg);
        return -1;
    }
    protocol_message_clear(&msg);
    return 0;
}

udp_hub_t *udp_hub_open(const char *listen_addr, latency_logger_t *logger, tx_timestamp_debug_logger_t *debug_logger, int enable_timestamping) {
    udp_hub_t *hub = (udp_hub_t *) calloc(1, sizeof(*hub));
    if (hub == NULL) {
        return NULL;
    }
    hub->conn = udp_conn_bind(listen_addr, NULL, enable_timestamping, logger, OMNI_NODE_ROLE_SERVER, "hub", debug_logger);
    if (hub->conn == NULL) {
        free(hub);
        return NULL;
    }
    pthread_rwlock_init(&hub->lock, NULL);
    return hub;
}

int udp_hub_serve(udp_hub_t *hub) {
    message_t msg;
    struct sockaddr_storage addr;
    socklen_t addr_len;
    udp_peer_entry_t *sender;
    udp_peer_entry_t *target;
    udp_peer_entry_t *entry;

    if (hub == NULL) {
        errno = EINVAL;
        return -1;
    }

    protocol_message_init(&msg);
    for (;;) {
        protocol_message_clear(&msg);
        if (udp_conn_receive(hub->conn, &msg, &addr, &addr_len) != 0) {
            return -1;
        }

        if (msg.type == MSG_TYPE_REGISTER) {
            pthread_rwlock_wrlock(&hub->lock);
            entry = udp_hub_find_by_id(hub, msg.from);
            if (entry == NULL) {
                entry = (udp_peer_entry_t *) calloc(1, sizeof(*entry));
                if (entry == NULL) {
                    pthread_rwlock_unlock(&hub->lock);
                    protocol_message_clear(&msg);
                    return -1;
                }
                snprintf(entry->peer_id, sizeof(entry->peer_id), "%s", msg.from);
                entry->next = hub->peers;
                hub->peers = entry;
            }
            memcpy(&entry->addr, &addr, sizeof(addr));
            entry->addr_len = addr_len;
            pthread_rwlock_unlock(&hub->lock);
            continue;
        }
        if (msg.type != MSG_TYPE_TEXT && msg.type != MSG_TYPE_FILE && msg.type != MSG_TYPE_BINARY) {
            if (msg.type == MSG_TYPE_ERROR) {
                udp_hub_send_error(hub, &addr, addr_len, msg.from, "peers cannot send error messages");
            } else {
                char *error_text = omni_strdup_printf("unsupported message type: %s", protocol_message_type_name(msg.type));
                if (error_text != NULL) {
                    udp_hub_send_error(hub, &addr, addr_len, msg.from, error_text);
                    free(error_text);
                }
            }
            continue;
        }

        pthread_rwlock_rdlock(&hub->lock);
        sender = udp_hub_find_by_addr(hub, &addr, addr_len);
        if (sender == NULL) {
            pthread_rwlock_unlock(&hub->lock);
            udp_hub_send_error(hub, &addr, addr_len, msg.from, "not registered; send register first");
            continue;
        }
        snprintf(msg.from, sizeof(msg.from), "%s", sender->peer_id);
        target = udp_hub_find_by_id(hub, msg.to);
        if (target == NULL) {
            char *error_text;
            pthread_rwlock_unlock(&hub->lock);
            error_text = omni_strdup_printf("unknown target: %s", msg.to);
            if (error_text != NULL) {
                udp_hub_send_error(hub, &addr, addr_len, sender->peer_id, error_text);
                free(error_text);
            }
            continue;
        }
        if (udp_conn_send_to(hub->conn, &msg, (const struct sockaddr *) &target->addr, target->addr_len) != 0) {
            char *error_text;
            pthread_rwlock_unlock(&hub->lock);
            error_text = omni_strdup_printf("failed to forward to %s", msg.to);
            if (error_text != NULL) {
                udp_hub_send_error(hub, &addr, addr_len, sender->peer_id, error_text);
                free(error_text);
            }
            continue;
        }
        pthread_rwlock_unlock(&hub->lock);
    }
}

int udp_hub_close(udp_hub_t *hub) {
    if (hub == NULL) {
        return 0;
    }
    return udp_conn_close(hub->conn);
}

void udp_hub_free(udp_hub_t *hub) {
    udp_peer_entry_t *entry;
    udp_peer_entry_t *next;
    if (hub == NULL) {
        return;
    }
    udp_conn_free(hub->conn);
    for (entry = hub->peers; entry != NULL; entry = next) {
        next = entry->next;
        free(entry);
    }
    pthread_rwlock_destroy(&hub->lock);
    free(hub);
}