#define _POSIX_C_SOURCE 200809L
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <poll.h>
#include <fcntl.h>
#include <unistd.h>
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>

#define PAYLOAD 1024
#define HEADER 7
#define DATA 1
#define ACK 2
#define END 3
#define IDLE_TIMEOUT 10000
#define LINGER 6000

static int64_t
now_ms(void)
{
    struct timespec ts;
    if (clock_gettime(CLOCK_MONOTONIC, &ts) < 0) return -1;
    return (int64_t)ts.tv_sec * 1000 + ts.tv_nsec / 1000000;
}

static int
wait_packet(int fd, int64_t deadline)
{
    struct pollfd pfd = {fd, POLLIN, 0};
    for (;;) {
        int64_t now = now_ms();
        int result;
        if (now < 0) return -1;
        if (now >= deadline) return 0;
        result = poll(&pfd, 1, (int)(deadline - now));
        if (result < 0 && errno == EINTR) continue;
        if (result <= 0) return result;
        if (pfd.revents & POLLIN) return 1;
        errno = EIO;
        return -1;
    }
}

static int
write_all(int fd, const unsigned char *buf, size_t len)
{
    while (len > 0) {
        ssize_t n = write(fd, buf, len);
        if (n < 0 && errno == EINTR) continue;
        if (n <= 0) return -1;
        buf += n;
        len -= (size_t)n;
    }
    return 0;
}

static int
send_ack(int fd, const struct sockaddr_in *peer, uint32_t seq)
{
    unsigned char packet[HEADER] = {ACK, 0, 0, 0, 0, 0, 0};
    uint32_t wire_seq = htonl(seq);
    ssize_t n;
    memcpy(packet + 1, &wire_seq, sizeof(wire_seq));
    do {
        n = sendto(fd, packet, sizeof(packet), 0,
                   (const struct sockaddr *)peer, sizeof(*peer));
    } while (n < 0 && errno == EINTR);
    if (n < 0) return -1;
    if (n != HEADER) { errno = EIO; return -1; }
    return 0;
}

int
main(int argc, char **argv)
{
    struct sockaddr_in address = {0}, peer = {0};
    uint32_t expected = 0;
    int fd, output = -1, have_peer = 0, finished = 0, status = 1;
    int64_t deadline;
    char *end;
    long port;

    if (argc != 3) {
        fprintf(stderr, "usage: %s port output\n", argv[0]);
        return 1;
    }
    errno = 0;
    port = strtol(argv[1], &end, 10);
    if (errno || !*argv[1] || *end || port < 1 || port > 65535) {
        fprintf(stderr, "invalid port: %s\n", argv[1]);
        return 1;
    }
    fd = socket(AF_INET, SOCK_DGRAM, 0);
    if (fd < 0) { perror("socket"); return 1; }
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = htonl(INADDR_ANY);
    address.sin_port = htons((unsigned short)port);
    if (bind(fd, (struct sockaddr *)&address, sizeof(address)) < 0) {
        perror("bind");
        close(fd);
        return 1;
    }
    output = open(argv[2], O_WRONLY | O_CREAT | O_TRUNC, 0666);
    if (output < 0) { perror(argv[2]); close(fd); return 1; }
    deadline = now_ms();
    if (deadline < 0) { perror("clock_gettime"); goto done; }
    deadline += IDLE_TIMEOUT;
    for (;;) {
        unsigned char packet[HEADER + PAYLOAD + 1];
        struct sockaddr_in source;
        socklen_t source_len = sizeof(source);
        uint32_t seq;
        uint16_t len;
        ssize_t n;
        int result = wait_packet(fd, deadline);
        if (result < 0) { perror("receive"); break; }
        if (result == 0) {
            if (finished) status = 0;
            else fprintf(stderr, "transfer timed out\n");
            break;
        }
        do {
            n = recvfrom(fd, packet, sizeof(packet), 0,
                         (struct sockaddr *)&source, &source_len);
        } while (n < 0 && errno == EINTR);
        if (n < 0) { perror("recvfrom"); break; }
        if (n < HEADER || n > HEADER + PAYLOAD ||
            source_len != sizeof(source) || source.sin_family != AF_INET)
            continue;
        memcpy(&seq, packet + 1, sizeof(seq));
        memcpy(&len, packet + 5, sizeof(len));
        seq = ntohl(seq);
        len = ntohs(len);
        if (n != HEADER + len ||
            !((packet[0] == DATA && len > 0 && len <= PAYLOAD && seq < UINT32_MAX) ||
              (packet[0] == END && len == 0)))
            continue;
        if (have_peer && (source.sin_addr.s_addr != peer.sin_addr.s_addr ||
                          source.sin_port != peer.sin_port))
            continue;
        if (!have_peer) {
            if (seq != 0) continue;
            peer = source;
            have_peer = 1;
        }
        if (finished) {
            if (packet[0] != END || seq != expected) continue;
        } else if (packet[0] == DATA) {
            if (seq == expected) {
                if (write_all(output, packet + HEADER, len) < 0) {
                    perror("write");
                    break;
                }
                ++expected;
            } else if (expected == 0 || seq != expected - 1) {
                continue;
            }
        } else {
            if (seq != expected) continue;
            result = close(output);
            output = -1;
            if (result < 0) { perror("close output"); break; }
            finished = 1;
            /* Keep the socket briefly to acknowledge a retransmitted END. */
            deadline = now_ms();
            if (deadline < 0) { perror("clock_gettime"); break; }
            deadline += LINGER;
        }
        if (send_ack(fd, &peer, seq) < 0) { perror("sendto"); break; }
        if (!finished) {
            deadline = now_ms();
            if (deadline < 0) { perror("clock_gettime"); break; }
            deadline += IDLE_TIMEOUT;
        }
    }
done:
    if (output >= 0 && close(output) < 0) { perror("close output"); status = 1; }
    close(fd);
    return status;
}