#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <arpa/inet.h>
#include <netinet/ip.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <time.h>
#include <unistd.h>
#include <errno.h>

#define SESSION_TIMEOUT 120
#define ENTRIES_SIZE 16

struct relay_entry {
    int relay_fd;
    struct sockaddr_in client_addr;
    time_t last_used;
};

int server_fd;

int epoll_fd;

int key;

struct sockaddr_in dest_addr;

struct relay_entry relay_entries[ENTRIES_SIZE];

void recv_server_msg() {
    unsigned char buf[65536];
    struct sockaddr_in client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    ssize_t msg_len = recvfrom(server_fd, buf, sizeof(buf), 0, (struct sockaddr*) &client_addr, &client_addr_len);

    struct relay_entry *entry = NULL;
    for (int i = 0; i < ENTRIES_SIZE; i ++) {
        if (relay_entries[i].relay_fd != -1
                && client_addr.sin_addr.s_addr == relay_entries[i].client_addr.sin_addr.s_addr
                && client_addr.sin_port == relay_entries[i].client_addr.sin_port) {
            entry = &relay_entries[i];
            break;
        }
    }
    if (!entry) {
        int relay_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
        if (relay_fd == -1) {
            perror("socket");
            exit(1);
        }

        int i;
        for (i = 0; i < ENTRIES_SIZE; i ++) {
            if (relay_entries[i].relay_fd == -1) {
                break;
            }
        }
        if (i == ENTRIES_SIZE) {
            printf("%s:%d no more space.\n", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
            return;
        }

        entry = &relay_entries[i];
        memcpy(&entry->client_addr, &client_addr, sizeof(client_addr));
        entry->relay_fd = relay_fd;

        struct epoll_event event = {
            .events = EPOLLIN,
            .data = { .ptr = entry }
        };
        epoll_ctl(epoll_fd, EPOLL_CTL_ADD, relay_fd, &event);

        printf("%s:%d connected.\n", inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
    }
    for (int i = 0; i < msg_len; i ++) {
        buf[i] ^= key;
    }
    sendto(entry->relay_fd, buf, msg_len, 0, (struct sockaddr *) &dest_addr, sizeof(dest_addr));
    entry->last_used = time(NULL);
}

void recv_relay_msg(struct relay_entry *entry) {
    unsigned char buf[65536];
    ssize_t msg_len = recv(entry->relay_fd, buf, sizeof(buf), 0);
    for (int i = 0; i < msg_len; i ++) {
        buf[i] ^= key;
    }
    sendto(server_fd, buf, msg_len, 0, (struct sockaddr *) &entry->client_addr, sizeof(entry->client_addr));
    entry->last_used = time(NULL);
}

void cleanup() {
    static int last_clean = 0;
    int cur_time = time(NULL);
    if (cur_time - last_clean >= SESSION_TIMEOUT) {
        for (int i = 0; i < ENTRIES_SIZE; i ++) {
            struct relay_entry *entry = &relay_entries[i];
            if (entry->relay_fd != -1 && cur_time - entry->last_used >= SESSION_TIMEOUT) {
                printf("%s:%d destroyed.\n", inet_ntoa(entry->client_addr.sin_addr), ntohs(entry->client_addr.sin_port));
                close(entry->relay_fd);
                entry->relay_fd = -1;
            }
        }
        last_clean = cur_time;
    }
}

int main(int argc, char *argv[]) {
	if (argc < 6) {
		puts("usage: udp-proxy <listen-address> <listen-port> <remote-address> <remote-port> <key>");
		exit(1);
	}

	key = atoi(argv[5]);

    for (int i = 0; i < ENTRIES_SIZE; i ++) {
        relay_entries[i].relay_fd = -1;
    }

    server_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (server_fd == -1) {
        perror("socket");
        exit(1);
    }

    struct sockaddr_in server_addr = { .sin_family = AF_INET };
    inet_aton(argv[1], &server_addr.sin_addr);
    server_addr.sin_port = htons(atoi(argv[2]));

    dest_addr.sin_family = AF_INET;
    inet_aton(argv[3], &dest_addr.sin_addr);
    dest_addr.sin_port = htons(atoi(argv[4]));

    int res = bind(server_fd, (struct sockaddr*) &server_addr, sizeof(server_addr));
    if (res == -1) {
        perror("bind");
        exit(1);
    }

    epoll_fd = epoll_create1(0);
    if (epoll_fd == -1) {
        perror("epoll_create1");
        exit(1);
    }

    struct epoll_event event_filter = { .events = EPOLLIN };
    epoll_ctl(epoll_fd, EPOLL_CTL_ADD, server_fd, &event_filter);

    struct epoll_event event;
    while (1) {
        res = epoll_wait(epoll_fd, &event, 1, SESSION_TIMEOUT * 1000);
        if (res == -1 && errno != EINTR) {
            perror("epoll_wait");
            exit(1);
        }
        if (res > 0) {
            if (event.data.ptr == NULL) {
                recv_server_msg();
            } else {
                recv_relay_msg(event.data.ptr);
            }
        }
        cleanup();
    }
}