/* ring_reduce: un all-reduce de verdad sobre TCP, en anillo.
 *
 *   ring_reduce <rango> <n> <puerto> <ip_siguiente> <bytes> <rondas>
 *
 * Cada nodo tiene un vector de enteros de 64 bits y al final de cada ronda todos deben tener la
 * suma de todos los vectores. En cada uno de los n-1 pasos, cada nodo le pasa al siguiente lo
 * último que recibió del anterior y lo suma a su resultado. Nadie termina una ronda hasta que
 * le llegó el aporte de todos. Eso es lo que queremos medir: un enlace lento atrasa al grupo.
 *
 * Enviar y recibir ocurren a la vez, con poll. Si todos enviaran primero y recibieran después,
 * con vectores grandes los buffers se llenarían y el anillo entero se quedaría trabado.
 */
#include <arpa/inet.h>
#include <errno.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>

#define CONNECT_TRIES 200

static double now_ms(void) {
    struct timespec ts;
    if (clock_gettime(CLOCK_MONOTONIC, &ts) < 0) {
        perror("clock_gettime");
        exit(1);
    }
    return ts.tv_sec * 1000.0 + ts.tv_nsec / 1e6;
}

static long parse_long(const char *text, long low, long high, const char *name) {
    char *end;
    long value = strtol(text, &end, 10);
    if (*text == '\0' || *end != '\0' || value < low || value > high) {
        fprintf(stderr, "%s fuera de rango: %s\n", name, text);
        exit(2);
    }
    return value;
}

/* Cada entero viaja en dos palabras de 32 bits en orden de red. */
static void encode(unsigned char *wire, const uint64_t *values, size_t count) {
    for (size_t i = 0; i < count; i++) {
        uint32_t high = htonl((uint32_t)(values[i] >> 32));
        uint32_t low = htonl((uint32_t)values[i]);
        memcpy(wire + 8 * i, &high, 4);
        memcpy(wire + 8 * i + 4, &low, 4);
    }
}

static void decode(uint64_t *values, const unsigned char *wire, size_t count) {
    for (size_t i = 0; i < count; i++) {
        uint32_t high, low;
        memcpy(&high, wire + 8 * i, 4);
        memcpy(&low, wire + 8 * i + 4, 4);
        values[i] = ((uint64_t)ntohl(high) << 32) | ntohl(low);
    }
}

static int set_nonblocking(int fd) {
    int flags = fcntl(fd, F_GETFL);
    if (flags < 0 || fcntl(fd, F_SETFL, flags | O_NONBLOCK) < 0) {
        perror("fcntl");
        return -1;
    }
    return 0;
}

/* Abre el puerto propio, se conecta al siguiente y acepta al anterior. */
static int open_ring(int port, const char *next_ip, int *to_next, int *from_prev) {
    int listener = socket(AF_INET, SOCK_STREAM, 0);
    if (listener < 0) {
        perror("socket");
        return -1;
    }
    int yes = 1;
    if (setsockopt(listener, SOL_SOCKET, SO_REUSEADDR, &yes, sizeof yes) < 0) {
        perror("setsockopt");
        return -1;
    }
    struct sockaddr_in local = {.sin_family = AF_INET,
                                .sin_port = htons((uint16_t)port),
                                .sin_addr.s_addr = htonl(INADDR_ANY)};
    if (bind(listener, (struct sockaddr *)&local, sizeof local) < 0) {
        perror("bind");
        return -1;
    }
    if (listen(listener, 1) < 0) {
        perror("listen");
        return -1;
    }
    struct sockaddr_in remote = {.sin_family = AF_INET, .sin_port = htons((uint16_t)port)};
    if (inet_pton(AF_INET, next_ip, &remote.sin_addr) != 1) {
        fprintf(stderr, "dirección inválida: %s\n", next_ip);
        return -1;
    }
    /* Los nodos no arrancan a la vez: se reintenta mientras el siguiente no esté escuchando. */
    *to_next = -1;
    for (int attempt = 0; attempt < CONNECT_TRIES; attempt++) {
        int fd = socket(AF_INET, SOCK_STREAM, 0);
        if (fd < 0) {
            perror("socket");
            return -1;
        }
        if (connect(fd, (struct sockaddr *)&remote, sizeof remote) == 0) {
            *to_next = fd;
            break;
        }
        if (errno != ECONNREFUSED && errno != ETIMEDOUT && errno != EHOSTUNREACH && errno != ENETUNREACH) {
            perror("connect");
            close(fd);
            return -1;
        }
        close(fd);
        struct timespec pause = {.tv_sec = 0, .tv_nsec = 100000000};
        if (nanosleep(&pause, NULL) < 0 && errno != EINTR) {
            perror("nanosleep");
            return -1;
        }
    }
    if (*to_next < 0) {
        fprintf(stderr, "no se pudo conectar con %s\n", next_ip);
        return -1;
    }
    *from_prev = accept(listener, NULL, NULL);
    if (*from_prev < 0) {
        perror("accept");
        return -1;
    }
    if (close(listener) < 0) {
        perror("close");
    }
    if (setsockopt(*to_next, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof yes) < 0) {
        perror("setsockopt");
        return -1;
    }
    if (set_nonblocking(*to_next) < 0 || set_nonblocking(*from_prev) < 0) {
        return -1;
    }
    return 0;
}

/* Envía out[] y recibe in[], los dos de size bytes, sin que uno bloquee al otro. */
static int exchange(int to_next, int from_prev, const unsigned char *out, unsigned char *in, size_t size) {
    size_t sent = 0, received = 0;
    while (sent < size || received < size) {
        struct pollfd fds[2] = {{.fd = sent < size ? to_next : -1, .events = POLLOUT},
                                {.fd = received < size ? from_prev : -1, .events = POLLIN}};
        if (poll(fds, 2, -1) < 0) {
            if (errno == EINTR) {
                continue;
            }
            perror("poll");
            return -1;
        }
        if (fds[0].revents & (POLLOUT | POLLERR | POLLHUP)) {
            ssize_t n = send(to_next, out + sent, size - sent, MSG_NOSIGNAL);
            if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
                perror("send");
                return -1;
            }
            if (n > 0) {
                sent += (size_t)n;
            }
        }
        if (fds[1].revents & (POLLIN | POLLERR | POLLHUP)) {
            ssize_t n = recv(from_prev, in + received, size - received, 0);
            if (n < 0 && errno != EAGAIN && errno != EWOULDBLOCK) {
                perror("recv");
                return -1;
            }
            if (n == 0) {
                fprintf(stderr, "el nodo anterior cerró la conexión\n");
                return -1;
            }
            if (n > 0) {
                received += (size_t)n;
            }
        }
    }
    return 0;
}

int main(int argc, char **argv) {
    if (argc != 7) {
        fprintf(stderr, "uso: %s <rango> <n> <puerto> <ip_siguiente> <bytes> <rondas>\n", argv[0]);
        return 2;
    }
    int nodes = (int)parse_long(argv[2], 2, 1024, "n");
    int rank = (int)parse_long(argv[1], 0, nodes - 1, "rango");
    int port = (int)parse_long(argv[3], 1, 65535, "puerto");
    size_t count = (size_t)parse_long(argv[5], 8, 1L << 30, "bytes") / 8;
    int rounds = (int)parse_long(argv[6], 1, 100000, "rondas");
    size_t size = count * 8;

    uint64_t *transit = malloc(size), *result = malloc(size);
    unsigned char *out = malloc(size), *in = malloc(size);
    if (!transit || !result || !out || !in) {
        perror("malloc");
        return 1;
    }
    int to_next, from_prev;
    if (open_ring(port, argv[4], &to_next, &from_prev) < 0) {
        return 1;
    }
    for (int round = 0; round < rounds; round++) {
        double start = now_ms();
        for (size_t i = 0; i < count; i++) {
            transit[i] = (uint64_t)(rank + 1 + round) + i;
            result[i] = transit[i];
        }
        for (int step = 0; step < nodes - 1; step++) {
            encode(out, transit, count);
            if (exchange(to_next, from_prev, out, in, size) < 0) {
                return 1;
            }
            decode(transit, in, count);
            for (size_t i = 0; i < count; i++) {
                result[i] += transit[i];
            }
        }
        /* Se comprueba la suma: si el anillo perdiera o cruzara datos, aquí se notaría. */
        for (size_t i = 0; i < count; i++) {
            uint64_t expected = (uint64_t)nodes * ((uint64_t)round + i) +
                                (uint64_t)nodes * ((uint64_t)nodes + 1) / 2;
            if (result[i] != expected) {
                fprintf(stderr, "suma incorrecta en la ronda %d\n", round);
                return 1;
            }
        }
        printf("ronda %d %.1f ms\n", round, now_ms() - start);
        if (fflush(stdout) == EOF) {
            perror("fflush");
            return 1;
        }
    }
    if (close(to_next) < 0 || close(from_prev) < 0) {
        perror("close");
    }
    free(transit);
    free(result);
    free(out);
    free(in);
    return 0;
}