#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <pthread.h>
#include <netdb.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <signal.h>
#include <time.h>
#include <stdatomic.h>
#include <ctype.h>

static const char *g_host;
static const char *g_port;
static atomic_long g_conn_total;
static atomic_long g_conn_active;
static int         g_verbose = 0;

static void ts(char *buf, size_t n) {
    struct timespec tp;
    clock_gettime(CLOCK_REALTIME, &tp);
    struct tm tm;
    localtime_r(&tp.tv_sec, &tm);
    strftime(buf, n, "%Y-%m-%d %H:%M:%S", &tm);
    size_t len = strlen(buf);
    snprintf(buf + len, n - len, ".%03ld", tp.tv_nsec / 1000000);
}

#define LOG(fmt, ...) do { \
    char _ts[40]; ts(_ts, sizeof(_ts)); \
    printf("[%s] " fmt "\n", _ts, ##__VA_ARGS__); \
    fflush(stdout); \
} while(0)

#define ERR(fmt, ...) do { \
    char _ts[40]; ts(_ts, sizeof(_ts)); \
    fprintf(stderr, "[%s] ERROR " fmt "\n", _ts, ##__VA_ARGS__); \
} while(0)

static void hexdump(const char *label, const unsigned char *data, ssize_t len) {
    char _ts[40]; ts(_ts, sizeof(_ts));
    printf("[%s] %s  (%zd bytes)\n", _ts, label, len);
    for (ssize_t i = 0; i < len; i += 16) {
        printf("  %06zx  ", i);
        for (int j = 0; j < 16; j++) {
            if (i + j < len) printf("%02x ", data[i + j]); else printf("   ");
            if (j == 7) printf(" ");
        }
        printf(" |");
        for (int j = 0; j < 16 && i + j < len; j++)
            putchar(isprint(data[i + j]) ? data[i + j] : '.');
        printf("|\n");
    }
    fflush(stdout);
}

struct pipe_args {
    int       src;
    int       dst;
    char      label[64];
    long     *bytes_out;
};

static void *pipe_data(void *arg) {
    struct pipe_args *p = arg;
    unsigned char buf[65536];
    ssize_t n;
    long bytes = 0;

    LOG("  pipe start  [%s]", p->label);

    while ((n = read(p->src, buf, sizeof(buf))) > 0) {
        if (g_verbose) hexdump(p->label, buf, n);
        ssize_t sent = 0;
        while (sent < n) {
            ssize_t w = write(p->dst, buf + sent, n - sent);
            if (w <= 0) goto done;
            sent += w;
        }
        bytes += n;
    }
done:
    *p->bytes_out = bytes;
    LOG("  pipe end    [%s] %ld bytes", p->label, bytes);
    shutdown(p->dst, SHUT_WR);
    free(p);
    return NULL;
}

static void *handle(void *arg) {
    int client = *(int *)arg;
    free(arg);

    long id     = atomic_fetch_add(&g_conn_total,  1) + 1;
    long active = atomic_fetch_add(&g_conn_active, 1) + 1;

    struct sockaddr_storage peer;
    socklen_t plen = sizeof(peer);
    char paddr[INET6_ADDRSTRLEN] = "?";
    int  pport = 0;
    if (getpeername(client, (struct sockaddr *)&peer, &plen) == 0) {
        if (peer.ss_family == AF_INET) {
            struct sockaddr_in *s = (struct sockaddr_in *)&peer;
            inet_ntop(AF_INET, &s->sin_addr, paddr, sizeof(paddr));
            pport = ntohs(s->sin_port);
        } else {
            struct sockaddr_in6 *s = (struct sockaddr_in6 *)&peer;
            inet_ntop(AF_INET6, &s->sin6_addr, paddr, sizeof(paddr));
            pport = ntohs(s->sin6_port);
        }
    }

    LOG("[#%ld] connect  %s:%d  (active=%ld)", id, paddr, pport, active);

    struct addrinfo hints = {0}, *res;
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    int rc = getaddrinfo(g_host, g_port, &hints, &res);
    if (rc != 0) {
        ERR("[#%ld] getaddrinfo %s:%s: %s", id, g_host, g_port, gai_strerror(rc));
        close(client);
        atomic_fetch_sub(&g_conn_active, 1);
        return NULL;
    }

    int remote = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (remote < 0) {
        ERR("[#%ld] socket: %s", id, strerror(errno));
        freeaddrinfo(res); close(client);
        atomic_fetch_sub(&g_conn_active, 1);
        return NULL;
    }
    if (connect(remote, res->ai_addr, res->ai_addrlen) < 0) {
        ERR("[#%ld] connect %s:%s: %s", id, g_host, g_port, strerror(errno));
        freeaddrinfo(res); close(remote); close(client);
        atomic_fetch_sub(&g_conn_active, 1);
        return NULL;
    }
    freeaddrinfo(res);
    LOG("[#%ld] tunnel   %s:%d <-> %s:%s", id, paddr, pport, g_host, g_port);

    long bytes_c2r = 0, bytes_r2c = 0;

    struct pipe_args *c2r = malloc(sizeof(*c2r));
    struct pipe_args *r2c = malloc(sizeof(*r2c));
    c2r->src = client;  c2r->dst = remote;
    r2c->src = remote;  r2c->dst = client;
    snprintf(c2r->label, sizeof(c2r->label), "#%ld cli->remote", id);
    snprintf(r2c->label, sizeof(r2c->label), "#%ld remote->cli", id);
    c2r->bytes_out = &bytes_c2r;
    r2c->bytes_out = &bytes_r2c;

    pthread_t t;
    pthread_create(&t, NULL, pipe_data, c2r);
    pthread_detach(t);
    pipe_data(r2c);

    close(remote);
    close(client);

    long still_active = atomic_fetch_sub(&g_conn_active, 1) - 1;
    LOG("[#%ld] closed   %s:%d  sent=%ld recv=%ld  (active=%ld)",
        id, paddr, pport, bytes_c2r, bytes_r2c, still_active);

    return NULL;
}

/* ================================================================
 * TCP: listen/accept loop
 * ================================================================ */
static void run_tcp(const char *lport) {
    struct addrinfo hints = {0}, *res;
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags    = AI_PASSIVE;
    if (getaddrinfo("0.0.0.0", lport, &hints, &res) != 0) {
        ERR("getaddrinfo tcp :%s: %s", lport, strerror(errno)); return;
    }
    int srv = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (srv < 0) { perror("socket"); return; }
    int opt = 1;
    setsockopt(srv, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    if (bind(srv, res->ai_addr, res->ai_addrlen) < 0) { perror("bind"); return; }
    freeaddrinfo(res);
    if (listen(srv, 128) < 0) { perror("listen"); return; }

    LOG("[*] listening TCP 0.0.0.0:%s  ->  %s:%s%s",
        lport, g_host, g_port, g_verbose ? "  [verbose]" : "");

    for (;;) {
        int *client = malloc(sizeof(int));
        *client = accept(srv, NULL, NULL);
        if (*client < 0) {
            free(client);
            ERR("accept: %s", strerror(errno));
            continue;
        }
        pthread_t t;
        pthread_create(&t, NULL, handle, client);
        pthread_detach(t);
    }
}

/* ================================================================
 * UDP: session per source address, transparent forwarding
 * ================================================================ */
struct udp_sess {
    struct sockaddr_storage src;
    socklen_t               srclen;
    int                     fwd_fd;     /* connected UDP socket to target */
    int                     listen_fd;  /* shared listen socket */
    atomic_int              alive;
    struct udp_sess        *next;
};

static struct udp_sess  *g_udp_head = NULL;
static pthread_mutex_t   g_udp_lock = PTHREAD_MUTEX_INITIALIZER;

static void *udp_backflow(void *arg) {
    struct udp_sess *s = arg;

    struct timeval tv = {1, 0};
    setsockopt(s->fwd_fd, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));

    unsigned char buf[65536];
    while (atomic_load(&s->alive)) {
        ssize_t n = recvfrom(s->fwd_fd, buf, sizeof(buf), 0, NULL, NULL);
        if (n < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
            if (errno == ECONNREFUSED || errno == EHOSTUNREACH || errno == ENETUNREACH) continue;
            break;
        }
        if (g_verbose) hexdump("UDP<-fwd", buf, n);
        sendto(s->listen_fd, buf, n, 0,
               (struct sockaddr *)&s->src, s->srclen);
    }

    atomic_store(&s->alive, 0);

    pthread_mutex_lock(&g_udp_lock);
    struct udp_sess **pp = &g_udp_head;
    while (*pp && *pp != s) pp = &(*pp)->next;
    if (*pp) *pp = s->next;
    pthread_mutex_unlock(&g_udp_lock);

    long still_active = atomic_fetch_sub(&g_conn_active, 1) - 1;
    LOG("  UDP session closed  (active=%ld)", still_active);

    close(s->fwd_fd);
    free(s);
    return NULL;
}

static struct udp_sess *udp_get_sess(int listen_fd,
    struct sockaddr_storage *src, socklen_t srclen) {

    pthread_mutex_lock(&g_udp_lock);
    struct udp_sess *s;
    for (s = g_udp_head; s; s = s->next)
        if (s->srclen == srclen && memcmp(&s->src, src, srclen) == 0
                && atomic_load(&s->alive)) {
            pthread_mutex_unlock(&g_udp_lock);
            return s;
        }
    pthread_mutex_unlock(&g_udp_lock);

    /* DNS re-resolve on every new session so IP changes take effect */
    struct addrinfo hints = {0}, *res;
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    int rc = getaddrinfo(g_host, g_port, &hints, &res);
    if (rc != 0) {
        ERR("getaddrinfo %s:%s: %s", g_host, g_port, gai_strerror(rc));
        return NULL;
    }
    int fwd_fd = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (fwd_fd < 0 || connect(fwd_fd, res->ai_addr, res->ai_addrlen) < 0) {
        ERR("udp connect %s:%s: %s", g_host, g_port, strerror(errno));
        if (fwd_fd >= 0) close(fwd_fd);
        freeaddrinfo(res); return NULL;
    }
    freeaddrinfo(res);

    s = calloc(1, sizeof(*s));
    memcpy(&s->src, src, srclen);
    s->srclen    = srclen;
    s->fwd_fd    = fwd_fd;
    s->listen_fd = listen_fd;
    atomic_store(&s->alive, 1);

    /* double-checked insert */
    pthread_mutex_lock(&g_udp_lock);
    struct udp_sess *existing;
    for (existing = g_udp_head; existing; existing = existing->next)
        if (existing->srclen == srclen && memcmp(&existing->src, src, srclen) == 0
                && atomic_load(&existing->alive)) {
            pthread_mutex_unlock(&g_udp_lock);
            close(fwd_fd); free(s);
            return existing;
        }
    s->next    = g_udp_head;
    g_udp_head = s;
    pthread_mutex_unlock(&g_udp_lock);

    char ip[INET6_ADDRSTRLEN]; int port = 0;
    if (src->ss_family == AF_INET) {
        inet_ntop(AF_INET,  &((struct sockaddr_in  *)src)->sin_addr,  ip, sizeof(ip));
        port = ntohs(((struct sockaddr_in  *)src)->sin_port);
    } else {
        inet_ntop(AF_INET6, &((struct sockaddr_in6 *)src)->sin6_addr, ip, sizeof(ip));
        port = ntohs(((struct sockaddr_in6 *)src)->sin6_port);
    }
    long id = atomic_fetch_add(&g_conn_total, 1) + 1;
    atomic_fetch_add(&g_conn_active, 1);
    LOG("[#%ld] UDP session  %s:%d  ->  %s:%s", id, ip, port, g_host, g_port);

    pthread_t t;
    pthread_create(&t, NULL, udp_backflow, s);
    pthread_detach(t);
    return s;
}

static void run_udp(const char *lport) {
    struct addrinfo hints = {0}, *res;
    hints.ai_family   = AF_UNSPEC;
    hints.ai_socktype = SOCK_DGRAM;
    hints.ai_flags    = AI_PASSIVE;
    if (getaddrinfo("0.0.0.0", lport, &hints, &res) != 0) {
        ERR("getaddrinfo udp :%s: %s", lport, strerror(errno)); return;
    }
    int fd = socket(res->ai_family, res->ai_socktype, res->ai_protocol);
    if (fd < 0) { perror("udp socket"); return; }
    int opt = 1;
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    if (bind(fd, res->ai_addr, res->ai_addrlen) < 0) { perror("udp bind"); return; }
    freeaddrinfo(res);

    LOG("[*] listening UDP 0.0.0.0:%s  ->  %s:%s%s",
        lport, g_host, g_port, g_verbose ? "  [verbose]" : "");

    unsigned char buf[65536];
    for (;;) {
        struct sockaddr_storage src;
        socklen_t srclen = sizeof(src);
        ssize_t n = recvfrom(fd, buf, sizeof(buf), 0,
                             (struct sockaddr *)&src, &srclen);
        if (n < 0) { ERR("recvfrom: %s", strerror(errno)); continue; }

        struct udp_sess *sess = udp_get_sess(fd, &src, srclen);
        if (!sess) continue;

        if (g_verbose) hexdump("UDP->fwd", buf, n);
        send(sess->fwd_fd, buf, n, 0);
    }
}

static void usage(const char *prog) {
    fprintf(stderr,
        "Usage:\n"
        "  %s --tcp --listen <port> --host <host> --port <port> [--verbose]\n"
        "  %s --udp --listen <port> --host <host> --port <port> [--verbose]\n"
        "  --tcp      TCP listen -> TCP forward\n"
        "  --udp      UDP listen -> UDP forward (session per source address)\n"
        "  --verbose  hex+ascii dump of every payload\n",
        prog, prog);
    exit(1);
}

int main(int argc, char *argv[]) {
    const char *listen_addr = NULL;
    int mode_tcp = 0, mode_udp = 0;
    g_host = NULL;
    g_port = NULL;

    for (int i = 1; i < argc; i++) {
        if      (!strcmp(argv[i], "--tcp"))                        mode_tcp    = 1;
        else if (!strcmp(argv[i], "--udp"))                        mode_udp    = 1;
        else if (!strcmp(argv[i], "--listen") && i+1 < argc)      listen_addr = argv[++i];
        else if (!strcmp(argv[i], "--host")   && i+1 < argc)      g_host      = argv[++i];
        else if (!strcmp(argv[i], "--port")   && i+1 < argc)      g_port      = argv[++i];
        else if (!strcmp(argv[i], "--verbose"))                    g_verbose   = 1;
        else usage(argv[0]);
    }
    if ((mode_tcp + mode_udp) != 1 || !listen_addr || !g_host || !g_port)
        usage(argv[0]);

    signal(SIGPIPE, SIG_IGN);

    if (mode_tcp) run_tcp(listen_addr);
    else          run_udp(listen_addr);
    return 0;
}