/*
 *  Copyright (C) 2024  stsp
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 3 of the License, or
 *  (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <unistd.h>
#include <sys/socket.h>
#include <linux/prctl.h>  /* Definition of PR_* constants */
#include <sys/prctl.h>
#include <searpc.h>
#include "priv.h"
#include "sig.h"
#include "utilities.h"
#include "util.h"

typedef struct {
    int fd;
    int in_async;
    void *rpc_priv;
} SockTransport;

static void svc_run(const char *svc_name, int transp_fd, int (*exiting)(void));

/*
 * arg: rpc_client->arg. Normally a socket number
 * fcall_str: the JSON data stream generated by Searpc
 * fcall_len: the length of `fcall_str`
 * ret_len: place to get the length of the returned json data stream
 * Returns: A newly allocated string stores the JSON data stream
 */
static char *transport_callback(void *arg, const char *fcall_str,
        size_t fcall_len, size_t *ret_len)
{
    char buf[4096];
    SockTransport *sock = arg;
    ssize_t sd;

    assert(!sock->in_async);
    sd = send(sock->fd, fcall_str, fcall_len, MSG_DONTWAIT);
    if (sd <= 0)
        return NULL;
    sd = recv(sock->fd, buf, sizeof(buf), 0);
    if (sd <= 0)
        return NULL;
    *ret_len = sd;
    return g_strndup(buf, sd);
}

#if 1
/* work around https://github.com/haiwen/libsearpc/pull/79 */
typedef struct {
    SearpcClient *client;
    AsyncCallback callback;
    const gchar *ret_type;
    GType gtype;                /* to specify the specific gobject type
                                 if ret_type is object or objlist */
    void *cbdata;
} AsyncCallData;
#define _ASYNC_CALL_DATA_SIZEOF sizeof(AsyncCallData)
#endif

static int transport_send(void *arg, char *fcall_str,
                          size_t fcall_len, void *rpc_priv)
{
    SockTransport *sock = arg;
    ssize_t sd;

    assert(!sock->in_async);
    sock->in_async++;
    sd = send(sock->fd, fcall_str, fcall_len, MSG_DONTWAIT);
    if (sd <= 0)
        return -1;
#ifdef _ASYNC_CALL_DATA_SIZEOF
    sock->rpc_priv = malloc(_ASYNC_CALL_DATA_SIZEOF);
    memcpy(sock->rpc_priv, rpc_priv, _ASYNC_CALL_DATA_SIZEOF);
#else
    sock->rpc_priv = rpc_priv;
#endif
    return 0;
}

static int transport_recv(SockTransport *sock)
{
    ssize_t sd;
    char buf[4096];

    assert(sock->in_async == 1);
    sock->in_async--;
    sd = recv(sock->fd, buf, sizeof(buf), 0);
    if (sd <= 0)
        return -1;
    searpc_client_generic_callback(buf, sd, sock->rpc_priv, NULL);
#ifdef _ASYNC_CALL_DATA_SIZEOF
    free(sock->rpc_priv);
#endif
    sock->rpc_priv = NULL;
    return 0;
}

int searpc_async_recv(SearpcClient *clnt)
{
    return transport_recv(clnt->async_arg);
}

static void chld_crash(int sig)
{
    fprintf(stderr, "RPC server crashed\n");
    exit(7);
}

SearpcClient *clnt_init(int *sock_rx, init_cb_t init_cb,
        void *init_arg, int (*svc_ex)(void),
        void (*ex_cb)(void *), const char *svc_name, pid_t *r_pid)
{
    SearpcClient *clnt;
    int socks[2];
    int transp[2];
    SockTransport *sock;
    pid_t pid;
    int err;
    pshared_sem_t svc_sem, svc_sem2;
    struct sigaction act;

    err = socketpair(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0, socks);
    if (err) {
        perror("socketpair()");
        return NULL;
    }
    err = socketpair(AF_UNIX, SOCK_STREAM | SOCK_CLOEXEC, 0, transp);
    if (err) {
        perror("socketpair()");
        goto err0;
    }
    err = pshared_sem_init(&svc_sem, 0);
    if (err)
        goto err1;
    err = pshared_sem_init(&svc_sem2, 0);
    if (err)
        goto err1;
    pid = fork();
    switch (pid) {
        case -1:
            perror("fork()");
            goto err2;
        case 0:
            /* child */
            close(socks[0]);
            close(transp[0]);
            pshared_sem_wait(svc_sem2);
            err = priv_drop();
            if (err)
                _exit(1);
            setsid();
            prctl(PR_SET_PDEATHSIG, SIGQUIT);
            for (int i = STDERR_FILENO + 1; i < sysconf(_SC_OPEN_MAX); i++) {
                if (i == socks[1] || i == transp[1])
                    continue;
                close(i);
            }
            err = init_cb(svc_name, socks[1], init_arg);
            if (err) {
                fprintf(stderr, "%s service failed\n", svc_name);
                _exit(1);
            }
            pshared_sem_post(svc_sem);
            pshared_sem_destroy(&svc_sem);
            pshared_sem_destroy(&svc_sem2);
            svc_run(svc_name, transp[1], svc_ex);
            _exit(1);  // not reached
            break;
    }

    close(socks[1]);
    close(transp[1]);
    sigchld_register_handler(pid, ex_cb, NULL);
    sigchld_set_critical(chld_crash, &act);
    pshared_sem_post(svc_sem2);
    pshared_sem_wait(svc_sem);
    sigchld_unset_critical(&act);
    pshared_sem_destroy(&svc_sem);
    pshared_sem_destroy(&svc_sem2);

    sock = malloc(sizeof(SockTransport));
    memset(sock, 0, sizeof(SockTransport));
    sock->fd = transp[0];
    clnt = searpc_client_new();
    clnt->send = transport_callback;
    clnt->arg = sock;
    clnt->async_send = transport_send;
    clnt->async_arg = sock;
    *sock_rx = socks[0];
    if (r_pid)
        *r_pid = pid;
    return clnt;

err2:
    pshared_sem_destroy(&svc_sem);
err1:
    close(transp[0]);
    close(transp[1]);
err0:
    close(socks[0]);
    close(socks[1]);
    return NULL;
}

static void svc_run(const char *svc_name, int transp_fd, int (*exiting)(void))
{
    while (1) {
        char buf[4096];
        gchar *json;
        gsize len = 0;
        ssize_t rd = recv(transp_fd, buf, sizeof(buf), 0);
        if (rd <= 0)
            _exit(0);
        json = searpc_server_call_function(svc_name, buf, rd, &len);
        if (!len)
            _exit(0);
        rd = send(transp_fd, json, len, 0);
        free(json);
        if (rd <= 0)
            _exit(0);
        if (exiting && exiting())
            _exit(0);
    }
}
