#include "client.h"

#include <arpa/inet.h>
#include <errno.h>
#include <netdb.h>
#include <openssl/err.h>
#include <openssl/ssl.h>
#include <openssl/x509v3.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>

#include "util.h"
#include "yyjson.h"

static char g_last_error[256];

const char *client_last_error(void)
{
    return g_last_error[0] ? g_last_error : "connection error";
}

static void set_error(const char *msg)
{
    snprintf(g_last_error, sizeof g_last_error, "%s", msg);
}

static void set_tls_error(const char *what)
{
    unsigned long e = ERR_get_error();
    char buf[200] = "unknown TLS error";
    if (e)
        ERR_error_string_n(e, buf, sizeof buf);
    snprintf(g_last_error, sizeof g_last_error, "%s: %s", what, buf);
    ERR_clear_error();
}

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

static void split_addrport(const char *addrport, char *host, size_t host_sz,
                           char *port, size_t port_sz, const char *def_port)
{
    snprintf(host, host_sz, "127.0.0.1");
    snprintf(port, port_sz, "%s", def_port);
    const char *colon = strrchr(addrport, ':');
    if (colon) {
        size_t hl = (size_t)(colon - addrport);
        if (hl < host_sz) {
            memcpy(host, addrport, hl);
            host[hl] = '\0';
        }
        snprintf(port, port_sz, "%s", colon + 1);
    } else if (*addrport) {
        snprintf(port, port_sz, "%s", addrport);
    }
}

static int connect_tcp(const char *host, const char *port)
{
    struct addrinfo hints, *res = NULL;
    memset(&hints, 0, sizeof hints);
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    if (getaddrinfo(host, port, &hints, &res) != 0) {
        set_error("cannot resolve host");
        return -1;
    }
    int fd = -1;
    for (struct addrinfo *ai = res; ai; ai = ai->ai_next) {
        fd = socket(ai->ai_family, ai->ai_socktype, ai->ai_protocol);
        if (fd < 0)
            continue;
        if (connect(fd, ai->ai_addr, ai->ai_addrlen) == 0)
            break;
        close(fd);
        fd = -1;
    }
    if (fd < 0)
        set_error(strerror(errno));
    freeaddrinfo(res);
    return fd;
}

static int set_verify_host(SSL *ssl, const char *host)
{
    X509_VERIFY_PARAM *param = SSL_get0_param(ssl);
    struct in_addr in4;
    struct in6_addr in6;
    X509_VERIFY_PARAM_set_hostflags(param,
                                    X509_CHECK_FLAG_NO_PARTIAL_WILDCARDS);
    if (inet_pton(AF_INET, host, &in4) == 1 ||
        inet_pton(AF_INET6, host, &in6) == 1)
        return X509_VERIFY_PARAM_set1_ip_asc(param, host);
    return X509_VERIFY_PARAM_set1_host(param, host, 0);
}

static int tls_connect_addr(const char *addrport, struct client_conn *out)
{
    char host[256], port[16];
    split_addrport(addrport, host, sizeof host, port, sizeof port, "8788");
    int fd = connect_tcp(host, port);
    if (fd < 0)
        return -1;

    SSL_CTX *ctx = SSL_CTX_new(TLS_client_method());
    if (!ctx) {
        set_tls_error("TLS context");
        close(fd);
        return -1;
    }
    SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
    SSL_CTX_set_verify(ctx, SSL_VERIFY_PEER, NULL);
    SSL_CTX_set_default_verify_paths(ctx);
    const char *ca = getenv("BOKFD_TLS_CA");
    if (ca && *ca && SSL_CTX_load_verify_locations(ctx, ca, NULL) != 1) {
        set_tls_error("TLS CA file");
        SSL_CTX_free(ctx);
        close(fd);
        return -1;
    }
    SSL *ssl = SSL_new(ctx);
    if (!ssl) {
        set_tls_error("TLS connection");
        SSL_CTX_free(ctx);
        close(fd);
        return -1;
    }
    SSL_set_fd(ssl, fd);
    SSL_set_tlsext_host_name(ssl, host);
    if (set_verify_host(ssl, host) != 1) {
        set_error("invalid TLS host name");
        SSL_free(ssl);
        SSL_CTX_free(ctx);
        close(fd);
        return -1;
    }
    if (SSL_connect(ssl) != 1) {
        long vr = SSL_get_verify_result(ssl);
        if (vr != X509_V_OK)
            snprintf(g_last_error, sizeof g_last_error,
                     "certificate verification failed: %s",
                     X509_verify_cert_error_string(vr));
        else
            set_tls_error("TLS handshake");
        SSL_free(ssl);
        SSL_CTX_free(ctx);
        close(fd);
        return -1;
    }
    out->fd = fd;
    out->ssl = ssl;
    out->ctx = ctx;
    return 0;
}

int client_connect(const char *target, struct client_conn *out)
{
    memset(out, 0, sizeof *out);
    out->fd = -1;
    g_last_error[0] = '\0';

    if (strncmp(target, "tls:", 4) == 0)
        return tls_connect_addr(target + 4, out);
    if (strncmp(target, "tcp:", 4) == 0) {
        char host[256], port[16];
        split_addrport(target + 4, host, sizeof host, port, sizeof port,
                       "8787");
        out->fd = connect_tcp(host, port);
        return out->fd < 0 ? -1 : 0;
    }

    struct sockaddr_un sa;
    memset(&sa, 0, sizeof sa);
    sa.sun_family = AF_UNIX;
    if (strlen(target) >= sizeof sa.sun_path) {
        set_error("socket path too long");
        errno = ENAMETOOLONG;
        return -1;
    }
    snprintf(sa.sun_path, sizeof sa.sun_path, "%s", target);
    int fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (fd < 0) {
        set_error(strerror(errno));
        return -1;
    }
    if (connect(fd, (struct sockaddr *)&sa, sizeof sa) != 0) {
        set_error(strerror(errno));
        close(fd);
        return -1;
    }
    out->fd = fd;
    return 0;
}

void client_close(struct client_conn *c)
{
    if (c->ssl) {
        SSL_shutdown(c->ssl);
        SSL_free(c->ssl);
        c->ssl = NULL;
    }
    if (c->ctx) {
        SSL_CTX_free(c->ctx);
        c->ctx = NULL;
    }
    if (c->fd >= 0) {
        close(c->fd);
        c->fd = -1;
    }
}

int client_send_line(struct client_conn *c, const char *line)
{
    size_t len = strlen(line);
    if (c->ssl) {
        size_t off = 0;
        while (off < len) {
            int n = SSL_write(c->ssl, line + off, (int)(len - off));
            if (n <= 0) {
                int e = SSL_get_error(c->ssl, n);
                if (e == SSL_ERROR_WANT_READ || e == SSL_ERROR_WANT_WRITE)
                    set_error("TLS write would block");
                else
                    set_tls_error("TLS write");
                return -1;
            }
            off += (size_t)n;
        }
        int n = SSL_write(c->ssl, "\n", 1);
        if (n <= 0) {
            set_tls_error("TLS write");
            return -1;
        }
        return 0;
    }
    if (write_all(c->fd, line, len) < 0)
        return -1;
    return write_all(c->fd, "\n", 1) < 0 ? -1 : 0;
}

char *client_read_line(struct client_conn *c)
{
    struct buf b;
    buf_init(&b);
    char chunk[4096];
    for (;;) {
        ssize_t r;
        if (c->ssl) {
            int n = SSL_read(c->ssl, chunk, sizeof chunk);
            if (n > 0) {
                r = n;
            } else {
                int e = SSL_get_error(c->ssl, n);
                if (e == SSL_ERROR_ZERO_RETURN) {
                    r = 0;
                } else {
                    if (e == SSL_ERROR_WANT_READ || e == SSL_ERROR_WANT_WRITE)
                        set_error("TLS read would block");
                    else
                        set_tls_error("TLS read");
                    buf_free(&b);
                    return NULL;
                }
            }
        } else {
            r = read(c->fd, chunk, sizeof chunk);
            if (r < 0) {
                if (errno == EINTR)
                    continue;
                set_error(strerror(errno));
                buf_free(&b);
                return NULL;
            }
        }
        if (r == 0)
            break;
        unsigned char *nl = memchr(chunk, '\n', (size_t)r);
        if (nl) {
            buf_append(&b, chunk, (size_t)(nl - (unsigned char *)chunk));
            break;
        }
        buf_append(&b, chunk, (size_t)r);
    }
    char *out = xmalloc(b.len + 1);
    memcpy(out, b.p ? (char *)b.p : "", b.len);
    out[b.len] = '\0';
    buf_free(&b);
    return out;
}

char *client_make_request(const char *cmd, const char *session, int64_t org,
                          const char *args_json, const char *id)
{
    yyjson_doc *adoc = NULL;
    if (args_json) {
        adoc = yyjson_read(args_json, strlen(args_json), 0);
        if (!adoc || !yyjson_is_obj(yyjson_doc_get_root(adoc))) {
            yyjson_doc_free(adoc);
            return NULL;
        }
    }
    yyjson_mut_doc *d = yyjson_mut_doc_new(NULL);
    yyjson_mut_val *o = yyjson_mut_obj(d);
    yyjson_mut_doc_set_root(d, o);
    yyjson_mut_obj_add_int(d, o, "v", 1);
    yyjson_mut_obj_add_strcpy(d, o, "id", id ? id : "cli");
    yyjson_mut_obj_add_strcpy(d, o, "cmd", cmd);
    if (session)
        yyjson_mut_obj_add_strcpy(d, o, "session", session);
    if (org > 0)
        yyjson_mut_obj_add_int(d, o, "org", org);
    if (adoc) {
        yyjson_mut_val *args = yyjson_val_mut_copy(d, yyjson_doc_get_root(adoc));
        yyjson_mut_obj_add_val(d, o, "args", args);
        yyjson_doc_free(adoc);
    }
    char *s = yyjson_mut_write(d, 0, NULL);
    yyjson_mut_doc_free(d);
    return s;
}

char *client_make_login_args(const char *user, const char *password)
{
    yyjson_mut_doc *d = yyjson_mut_doc_new(NULL);
    yyjson_mut_val *o = yyjson_mut_obj(d);
    yyjson_mut_doc_set_root(d, o);
    yyjson_mut_obj_add_strcpy(d, o, "method", "password");
    yyjson_mut_obj_add_strcpy(d, o, "username", user);
    yyjson_mut_obj_add_strcpy(d, o, "password", password);
    char *s = yyjson_mut_write(d, 0, NULL);
    yyjson_mut_doc_free(d);
    return s;
}

char *client_rpc(struct client_conn *c, const char *cmd, const char *session,
                 int64_t org, const char *args_json)
{
    char *req = client_make_request(cmd, session, org, args_json, "rpc");
    if (!req)
        return NULL;
    int rc = client_send_line(c, req);
    free(req);
    if (rc != 0)
        return NULL;
    return client_read_line(c);
}

static int session_open(struct client_conn *c, const char *args,
                        char **session_out, char **err_out)
{
    *session_out = NULL;
    *err_out = NULL;
    char *resp = client_rpc(c, "session.open", NULL, 0, args);
    if (!resp) {
        *err_out = xstrdup(client_last_error());
        return -1;
    }
    if (!client_ok(resp)) {
        *err_out = resp;
        return -1;
    }
    yyjson_doc *d = yyjson_read(resp, strlen(resp), 0);
    yyjson_val *root = d ? yyjson_doc_get_root(d) : NULL;
    yyjson_val *res = root ? yyjson_obj_get(root, "result") : NULL;
    yyjson_val *s = res ? yyjson_obj_get(res, "session") : NULL;
    if (!s || !yyjson_is_str(s)) {
        *err_out = xstrdup("login response had no session");
        yyjson_doc_free(d);
        free(resp);
        return -1;
    }
    *session_out = xstrdup(yyjson_get_str(s));
    yyjson_doc_free(d);
    return 0;
}

int client_login(struct client_conn *c, const char *user, const char *password,
                 char **session_out, char **err_out)
{
    char *args = client_make_login_args(user, password);
    if (!args) {
        *session_out = NULL;
        *err_out = xstrdup("could not build login request");
        return -1;
    }
    int rc = session_open(c, args, session_out, err_out);
    free(args);
    return rc;
}

int client_token_login(struct client_conn *c, const char *token,
                       char **session_out, char **err_out)
{
    yyjson_mut_doc *d = yyjson_mut_doc_new(NULL);
    if (!d) {
        *session_out = NULL;
        *err_out = xstrdup("out of memory");
        return -1;
    }
    yyjson_mut_val *o = yyjson_mut_obj(d);
    yyjson_mut_doc_set_root(d, o);
    yyjson_mut_obj_add_strcpy(d, o, "method", "token");
    yyjson_mut_obj_add_strcpy(d, o, "token", token);
    char *args = yyjson_mut_write(d, 0, NULL);
    yyjson_mut_doc_free(d);
    if (!args) {
        *session_out = NULL;
        *err_out = xstrdup("could not build token request");
        return -1;
    }
    int rc = session_open(c, args, session_out, err_out);
    free(args);
    return rc;
}

int client_ok(const char *response)
{
    if (!response)
        return 0;
    yyjson_doc *d = yyjson_read(response, strlen(response), 0);
    yyjson_val *ok = d ? yyjson_obj_get(yyjson_doc_get_root(d), "ok") : NULL;
    int result = ok && yyjson_is_bool(ok) && yyjson_get_bool(ok);
    yyjson_doc_free(d);
    return result;
}

int client_session_from(const char *response, char *buf, unsigned long cap)
{
    if (!response)
        return -1;
    yyjson_doc *d = yyjson_read(response, strlen(response), 0);
    yyjson_val *root = d ? yyjson_doc_get_root(d) : NULL;
    yyjson_val *res = root ? yyjson_obj_get(root, "result") : NULL;
    yyjson_val *s = res ? yyjson_obj_get(res, "session") : NULL;
    if (!s || !yyjson_is_str(s)) {
        yyjson_doc_free(d);
        return -1;
    }
    snprintf(buf, cap, "%s", yyjson_get_str(s));
    yyjson_doc_free(d);
    return 0;
}