#include "li_cil_oc2_common_inet_DefaultSessionLayer.h"

#include <arpa/inet.h>
#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/un.h>
#include <unistd.h>

#define SOCK_PATH_DEFAULT "/var/run/oc2rnet-helper.sock"
#define PROTOCOL_VERSION 1
#define MAX_PAYLOAD 2048

struct PingReq {
    uint8_t version;
    uint8_t reserved;
    uint16_t size;
    uint32_t ip_be;
    uint32_t timeout_ms;
} __attribute__((packed));

struct PingRes {
    int32_t status;
    uint16_t size;
    uint16_t reserved;
} __attribute__((packed));

static int uds_connect(const char *path) {
    int fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        return -1;
    }

    struct sockaddr_un sa;
    memset(&sa, 0, sizeof(sa));
    sa.sun_family = AF_UNIX;
    if (strlcpy(sa.sun_path, path, sizeof(sa.sun_path)) >= sizeof(sa.sun_path)) {
        close(fd);
        errno = ENAMETOOLONG;
        return -1;
    }

    if (connect(fd, (struct sockaddr *)&sa, sizeof(sa)) < 0) {
        close(fd);
        return -1;
    }

    return fd;
}

JNIEXPORT jbyteArray JNICALL
Java_li_cil_oc2_common_inet_DefaultSessionLayer_sendICMP(JNIEnv *env, jobject instance,
                                                         jbyteArray ip, jbyteArray data,
                                                         jint size, jint timeout) {
    (void)instance;

    if (size < 0 || size > MAX_PAYLOAD) {
        return NULL;
    }

    const char *sock_path = getenv("OC2RNET_SOCK");
    if (sock_path == NULL) {
        sock_path = SOCK_PATH_DEFAULT;
    }

    jbyte *addr = (*env)->GetByteArrayElements(env, ip, NULL);
    if (addr == NULL) {
        return NULL;
    }

    jbyte *payload = (*env)->GetByteArrayElements(env, data, NULL);
    if (payload == NULL) {
        (*env)->ReleaseByteArrayElements(env, ip, addr, JNI_ABORT);
        return NULL;
    }

    int fd = uds_connect(sock_path);
    if (fd < 0) {
        (*env)->ReleaseByteArrayElements(env, ip, addr, JNI_ABORT);
        (*env)->ReleaseByteArrayElements(env, data, payload, JNI_ABORT);
        return NULL;
    }

    struct PingReq req;
    memset(&req, 0, sizeof(req));
    req.version = PROTOCOL_VERSION;
    req.size = (uint16_t)size;
    memcpy(&req.ip_be, addr, sizeof(req.ip_be));
    req.timeout_ms = (uint32_t)timeout;

    ssize_t written = write(fd, &req, sizeof(req));
    if (written != (ssize_t)sizeof(req)) {
        goto fail;
    }

    if (size > 0) {
        written = write(fd, payload, (size_t)size);
        if (written != size) {
            goto fail;
        }
    }

    struct PingRes res;
    ssize_t received = read(fd, &res, sizeof(res));
    if (received != (ssize_t)sizeof(res)) {
        goto fail;
    }

    if (res.status != 0 || res.size > (uint16_t)size) {
        goto fail;
    }

    jbyteArray out = (*env)->NewByteArray(env, res.size);
    if (out == NULL) {
        goto fail;
    }

    jbyte *outbuf = malloc(res.size);
    if (outbuf == NULL) {
        (*env)->DeleteLocalRef(env, out);
        goto fail;
    }

    received = read(fd, outbuf, res.size);
    if (received != res.size) {
        free(outbuf);
        (*env)->DeleteLocalRef(env, out);
        goto fail;
    }

    (*env)->SetByteArrayRegion(env, out, 0, res.size, outbuf);
    free(outbuf);

    close(fd);
    (*env)->ReleaseByteArrayElements(env, ip, addr, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, data, payload, JNI_ABORT);
    return out;

fail:
    close(fd);
    (*env)->ReleaseByteArrayElements(env, ip, addr, JNI_ABORT);
    (*env)->ReleaseByteArrayElements(env, data, payload, JNI_ABORT);
    return NULL;
}

#ifdef CLITEST

static ssize_t write_all(int fd, const void *buffer, size_t length) {
    const uint8_t *ptr = buffer;
    size_t total = 0;
    while (total < length) {
        ssize_t written = write(fd, ptr + total, length - total);
        if (written < 0) {
            if (errno == EINTR) {
                continue;
            }
            return -1;
        }
        total += (size_t)written;
    }
    return (ssize_t)total;
}

static ssize_t read_all(int fd, void *buffer, size_t length) {
    uint8_t *ptr = buffer;
    size_t total = 0;
    while (total < length) {
        ssize_t read_bytes = read(fd, ptr + total, length - total);
        if (read_bytes == 0) {
            return (ssize_t)total;
        }
        if (read_bytes < 0) {
            if (errno == EINTR) {
                continue;
            }
            return -1;
        }
        total += (size_t)read_bytes;
    }
    return (ssize_t)total;
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "usage: %s <ipv4> [payload]\n", argv[0]);
        return 1;
    }

    const char *target_ip = argv[1];
    const char *payload_str = (argc >= 3) ? argv[2] : "oc2r";
    size_t payload_len = strlen(payload_str);

    if (payload_len > MAX_PAYLOAD) {
        fprintf(stderr, "payload too large (max %d)\n", MAX_PAYLOAD);
        return 1;
    }

    struct in_addr in_addr_target;
    if (inet_pton(AF_INET, target_ip, &in_addr_target) != 1) {
        perror("inet_pton");
        return 1;
    }

    const char *sock_path = getenv("OC2RNET_SOCK");
    if (sock_path == NULL || *sock_path == '\0') {
        sock_path = SOCK_PATH_DEFAULT;
    }

    int fd = uds_connect(sock_path);
    if (fd < 0) {
        perror("uds_connect");
        return 1;
    }

    struct PingReq req;
    memset(&req, 0, sizeof(req));
    req.version = PROTOCOL_VERSION;
    req.size = (uint16_t)payload_len;
    req.ip_be = in_addr_target.s_addr;
    req.timeout_ms = 1000;

    if (write_all(fd, &req, sizeof(req)) != (ssize_t)sizeof(req)) {
        perror("write request");
        close(fd);
        return 1;
    }

    if (payload_len > 0) {
        if (write_all(fd, payload_str, payload_len) != (ssize_t)payload_len) {
            perror("write payload");
            close(fd);
            return 1;
        }
    }

    struct PingRes res;
    if (read_all(fd, &res, sizeof(res)) != (ssize_t)sizeof(res)) {
        perror("read response");
        close(fd);
        return 1;
    }

    printf("status=%d size=%u\n", res.status, res.size);

    if (res.status == 0 && res.size > 0) {
        uint8_t *buffer = malloc(res.size);
        if (buffer == NULL) {
            fprintf(stderr, "malloc failed\n");
            close(fd);
            return 1;
        }
        if (read_all(fd, buffer, res.size) != res.size) {
            perror("read payload");
            free(buffer);
            close(fd);
            return 1;
        }

        fwrite(buffer, 1, res.size, stdout);
        putchar('\n');
        free(buffer);
    }

    close(fd);
    return 0;
}

#endif /* CLITEST */