#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <openssl/err.h>
#include <openssl/ssl.h>
#include <poll.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/stat.h>
#include <sys/un.h>
#include <termios.h>
#include <time.h>
#include <unistd.h>

#include "auth.h"
#include "config.h"
#include "db.h"
#include "log.h"
#include "protocol.h"
#include "sessions.h"
#include "util.h"
#include "version.h"

#define MAX_CONNS 64
#define READ_CHUNK 65536

struct conn {
    int fd;
    SSL *ssl;
    int handshaking;
    struct buf in;
    struct buf out;
    size_t out_sent;
    int closing;
};

static volatile sig_atomic_t g_stop = 0;
static volatile sig_atomic_t g_reload = 0;
static size_t g_line_limit = 1024 * 1024;
static SSL_CTX *g_tls_ctx = NULL;
static time_t g_tls_cert_mtime = 0;
static time_t g_tls_key_mtime = 0;
static time_t g_tls_next_try = 0;

static void on_signal(int sig)
{
    (void)sig;
    g_stop = 1;
}

static void on_reload(int sig)
{
    (void)sig;
    g_reload = 1;
}

static int mkdir_p(const char *path, mode_t mode)
{
    char tmp[4096];
    if (!path || strlen(path) >= sizeof tmp)
        return -1;
    strcpy(tmp, path);
    for (char *p = tmp + 1; *p; p++) {
        if (*p == '/') {
            *p = '\0';
            if (mkdir(tmp, mode) != 0 && errno != EEXIST)
                return -1;
            *p = '/';
        }
    }
    if (mkdir(tmp, mode) != 0 && errno != EEXIST)
        return -1;
    return 0;
}

static void dirname_of(const char *path, char *out, size_t n)
{
    snprintf(out, n, "%s", path);
    char *slash = strrchr(out, '/');
    if (slash)
        *slash = '\0';
    else
        snprintf(out, n, ".");
}

static int read_password(const char *prompt, char *buf, size_t n)
{
    if (!isatty(STDIN_FILENO)) {
        const char *env = getenv("BOKFD_PASSWORD");
        if (!env)
            return -1;
        snprintf(buf, n, "%s", env);
        return 0;
    }
    struct termios old, noecho;
    fprintf(stderr, "%s", prompt);
    if (tcgetattr(STDIN_FILENO, &old) != 0)
        return -1;
    noecho = old;
    noecho.c_lflag &= ~(tcflag_t)ECHO;
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &noecho);
    char *r = fgets(buf, (int)n, stdin);
    tcsetattr(STDIN_FILENO, TCSAFLUSH, &old);
    fprintf(stderr, "\n");
    if (!r)
        return -1;
    util_str_trim(buf);
    return 0;
}

static int cmd_init(const char *db_path, const char *username)
{
    char dir[4096];
    dirname_of(db_path, dir, sizeof dir);
    if (mkdir_p(dir, 0700) != 0) {
        fprintf(stderr, "bokfd: cannot create %s\n", dir);
        return 1;
    }
    sqlite3 *db = NULL;
    char *err = NULL;
    if (db_open(db_path, &db, &err) != 0) {
        fprintf(stderr, "bokfd: %s\n", err ? err : "cannot open database");
        free(err);
        return 1;
    }
    int64_t users = db_count(db, "SELECT count(*) FROM users");
    if (users > 0) {
        fprintf(stderr, "bokfd: database already initialized (%lld users)\n",
                (long long)users);
        sqlite3_close(db);
        return 1;
    }
    char pw[256];
    if (read_password("Password for admin: ", pw, sizeof pw) != 0 || !pw[0]) {
        fprintf(stderr, "bokfd: no password provided\n");
        sqlite3_close(db);
        return 1;
    }
    if (isatty(STDIN_FILENO)) {
        char pw2[256];
        if (read_password("Repeat password: ", pw2, sizeof pw2) != 0 ||
            strcmp(pw, pw2) != 0) {
            fprintf(stderr, "bokfd: passwords do not match\n");
            sqlite3_close(db);
            return 1;
        }
    }
    int64_t uid = 0;
    if (db_create_user(db, username, username, pw, 1, &uid, &err) != 0) {
        fprintf(stderr, "bokfd: %s\n", err ? err : "cannot create user");
        free(err);
        sqlite3_close(db);
        return 1;
    }
    printf("initialized %s\nadmin user: %s (id %lld)\n", db_path, username,
           (long long)uid);
    sqlite3_close(db);
    return 0;
}

static int unix_listen(const char *path)
{
    struct sockaddr_un sa;
    memset(&sa, 0, sizeof sa);
    sa.sun_family = AF_UNIX;
    if (strlen(path) >= sizeof sa.sun_path)
        return -1;
    snprintf(sa.sun_path, sizeof sa.sun_path, "%s", path);
    unlink(path);
    int fd = socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (fd < 0)
        return -1;
    if (bind(fd, (struct sockaddr *)&sa, sizeof sa) != 0) {
        close(fd);
        return -1;
    }
    chmod(path, 0660);
    if (listen(fd, 64) != 0) {
        close(fd);
        return -1;
    }
    return fd;
}

static int tcp_listen(const char *addrport)
{
    char host[256] = "127.0.0.1";
    char port[16] = "8787";
    const char *colon = strrchr(addrport, ':');
    if (colon) {
        size_t hl = (size_t)(colon - addrport);
        if (hl < sizeof host) {
            memcpy(host, addrport, hl);
            host[hl] = '\0';
        }
        snprintf(port, sizeof port, "%s", colon + 1);
    } else {
        snprintf(port, sizeof port, "%s", addrport);
    }
    struct addrinfo hints, *res = NULL;
    memset(&hints, 0, sizeof hints);
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    hints.ai_flags = AI_PASSIVE;
    if (getaddrinfo(host, port, &hints, &res) != 0)
        return -1;
    int fd = -1;
    for (struct addrinfo *ai = res; ai; ai = ai->ai_next) {
        fd = socket(ai->ai_family, ai->ai_socktype | SOCK_NONBLOCK | SOCK_CLOEXEC,
                    ai->ai_protocol);
        if (fd < 0)
            continue;
        int one = 1;
        setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof one);
        if (bind(fd, ai->ai_addr, ai->ai_addrlen) == 0 && listen(fd, 64) == 0)
            break;
        close(fd);
        fd = -1;
    }
    freeaddrinfo(res);
    return fd;
}

static void log_tls_error(const char *what)
{
    unsigned long e = ERR_get_error();
    char buf[256] = "unknown error";
    if (e)
        ERR_error_string_n(e, buf, sizeof buf);
    log_error("%s: %s", what, buf);
    ERR_clear_error();
}

static int tls_load_certs(void)
{
    SSL_CTX *ctx = SSL_CTX_new(TLS_server_method());
    if (!ctx) {
        log_tls_error("TLS context");
        return -1;
    }
    SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
    SSL_CTX_set_options(ctx, SSL_OP_NO_COMPRESSION |
                                 SSL_OP_CIPHER_SERVER_PREFERENCE |
                                 SSL_OP_NO_RENEGOTIATION);
    if (SSL_CTX_use_certificate_chain_file(ctx, g_cfg.tls_cert) != 1) {
        log_tls_error(g_cfg.tls_cert);
        SSL_CTX_free(ctx);
        return -1;
    }
    if (SSL_CTX_use_PrivateKey_file(ctx, g_cfg.tls_key, SSL_FILETYPE_PEM) !=
        1) {
        log_tls_error(g_cfg.tls_key);
        SSL_CTX_free(ctx);
        return -1;
    }
    if (SSL_CTX_check_private_key(ctx) != 1) {
        log_error("TLS key does not match certificate %s", g_cfg.tls_cert);
        SSL_CTX_free(ctx);
        return -1;
    }
    SSL_CTX *old = g_tls_ctx;
    g_tls_ctx = ctx;
    if (old)
        SSL_CTX_free(old);
    struct stat st;
    if (stat(g_cfg.tls_cert, &st) == 0)
        g_tls_cert_mtime = st.st_mtime;
    if (stat(g_cfg.tls_key, &st) == 0)
        g_tls_key_mtime = st.st_mtime;
    return 0;
}

static void tls_reload_check(void)
{
    struct stat st;
    time_t cert = 0, key = 0;
    if (stat(g_cfg.tls_cert, &st) == 0)
        cert = st.st_mtime;
    if (stat(g_cfg.tls_key, &st) == 0)
        key = st.st_mtime;
    if (cert == 0 || key == 0 ||
        (cert == g_tls_cert_mtime && key == g_tls_key_mtime))
        return;
    time_t now = time(NULL);
    if (now < g_tls_next_try)
        return;
    g_tls_next_try = now + 60;
    if (tls_load_certs() == 0)
        log_info("TLS certificate reloaded");
    else
        log_error("TLS reload failed; keeping previous certificate");
}

static void conn_handshake(struct conn *c)
{
    int r = SSL_accept(c->ssl);
    if (r == 1) {
        c->handshaking = 0;
        return;
    }
    int e = SSL_get_error(c->ssl, r);
    if (e == SSL_ERROR_WANT_READ || e == SSL_ERROR_WANT_WRITE)
        return;
    log_warn("TLS handshake failed from fd %d", c->fd);
    ERR_clear_error();
    c->closing = 1;
}

static void accept_conns(int lfd, struct conn *conns, size_t *nconns, int tls)
{
    for (;;) {
        int fd = accept4(lfd, NULL, NULL, SOCK_NONBLOCK | SOCK_CLOEXEC);
        if (fd < 0)
            return;
        if (*nconns >= MAX_CONNS) {
            close(fd);
            continue;
        }
        struct conn *c = &conns[(*nconns)++];
        memset(c, 0, sizeof *c);
        c->fd = fd;
        buf_init(&c->in);
        buf_init(&c->out);
        if (tls) {
            c->ssl = SSL_new(g_tls_ctx);
            if (!c->ssl) {
                log_tls_error("TLS connection");
                close(fd);
                buf_free(&c->in);
                buf_free(&c->out);
                (*nconns)--;
                continue;
            }
            SSL_set_fd(c->ssl, fd);
            SSL_set_accept_state(c->ssl);
            c->handshaking = 1;
        }
    }
}

static void conn_read(struct conn *c, sqlite3 *db)
{
    if (c->ssl && c->handshaking) {
        conn_handshake(c);
        if (c->closing || c->handshaking)
            return;
    }
    char tmp[READ_CHUNK];
    ssize_t r;
    if (c->ssl) {
        r = SSL_read(c->ssl, tmp, sizeof tmp);
        if (r <= 0) {
            int e = SSL_get_error(c->ssl, (int)r);
            if (e == SSL_ERROR_WANT_READ || e == SSL_ERROR_WANT_WRITE)
                return;
            ERR_clear_error();
            c->closing = 1;
            return;
        }
    } else {
        r = read(c->fd, tmp, sizeof tmp);
        if (r == 0) {
            c->closing = 1;
            return;
        }
        if (r < 0) {
            if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR)
                return;
            c->closing = 1;
            return;
        }
    }
    buf_append(&c->in, tmp, (size_t)r);
    for (;;) {
        unsigned char *nl = memchr(c->in.p, '\n', c->in.len);
        if (!nl)
            break;
        size_t linelen = (size_t)(nl - c->in.p);
        char *resp = protocol_handle_line(db, (const char *)c->in.p, linelen);
        buf_append(&c->out, resp, strlen(resp));
        buf_append(&c->out, "\n", 1);
        free(resp);
        size_t consumed = linelen + 1;
        memmove(c->in.p, c->in.p + consumed, c->in.len - consumed);
        c->in.len -= consumed;
    }
    if (c->in.len > g_line_limit) {
        log_warn("dropping oversized request from fd %d", c->fd);
        c->closing = 1;
    }
}

static void conn_flush(struct conn *c)
{
    if (c->ssl && c->handshaking) {
        conn_handshake(c);
        if (c->closing || c->handshaking)
            return;
    }
    while (c->out_sent < c->out.len) {
        ssize_t w;
        if (c->ssl) {
            w = SSL_write(c->ssl, c->out.p + c->out_sent,
                          (int)(c->out.len - c->out_sent));
            if (w <= 0) {
                int e = SSL_get_error(c->ssl, (int)w);
                if (e == SSL_ERROR_WANT_READ || e == SSL_ERROR_WANT_WRITE)
                    return;
                ERR_clear_error();
                c->closing = 1;
                return;
            }
        } else {
            w = write(c->fd, c->out.p + c->out_sent,
                      c->out.len - c->out_sent);
            if (w > 0) {
                c->out_sent += (size_t)w;
                continue;
            }
            if (w < 0 && (errno == EAGAIN || errno == EWOULDBLOCK))
                return;
            if (w < 0 && errno == EINTR)
                continue;
            c->closing = 1;
            return;
        }
        c->out_sent += (size_t)w;
    }
    c->out.len = 0;
    c->out_sent = 0;
}

static void usage(FILE *f)
{
    fprintf(f,
            "usage: bokfd [options]\n"
            "       bokfd init [--db PATH] [--user NAME]\n"
            "       bokfd --version\n"
            "\n"
            "options:\n"
            "  --config FILE     config file (key = value)\n"
            "  --db PATH         SQLite database (default /var/lib/bokfd/bokfd.db)\n"
            "  --socket PATH     unix socket (default /run/bokfd/bokfd.sock)\n"
            "  --backup-dir DIR  snapshot destination\n"
            "  --export-dir DIR  SIE export destination\n"
            "  --log-level L     error|warn|info|debug\n"
            "\n"
            "config/env: tls (host:port, BOKFD_TLS), tls_cert (BOKFD_TLS_CERT),\n"
            "  tls_key (BOKFD_TLS_KEY); TLS is off unless tls is set\n");
}

static const char *opt_value(const char *arg, const char *name, int *i,
                             int argc, char **argv)
{
    size_t n = strlen(name);
    if (strncmp(arg, name, n) != 0)
        return NULL;
    if (arg[n] == '=')
        return arg + n + 1;
    if (arg[n] == '\0' && *i + 1 < argc)
        return argv[++(*i)];
    return NULL;
}

static void parse_args(int argc, char **argv, const char **config_file,
                       const char **init_user, int *init_mode)
{
    for (int i = 1; i < argc; i++) {
        const char *a = argv[i];
        const char *v;
        if ((v = opt_value(a, "--config", &i, argc, argv)))
            *config_file = v;
        else if ((v = opt_value(a, "--db", &i, argc, argv)))
            g_cfg.db_path = xstrdup(v);
        else if ((v = opt_value(a, "--socket", &i, argc, argv)))
            g_cfg.socket_path = xstrdup(v);
        else if ((v = opt_value(a, "--backup-dir", &i, argc, argv)))
            g_cfg.backup_dir = xstrdup(v);
        else if ((v = opt_value(a, "--export-dir", &i, argc, argv)))
            g_cfg.export_dir = xstrdup(v);
        else if ((v = opt_value(a, "--log-level", &i, argc, argv)))
            g_cfg.log_level = log_level_from_name(v);
        else if ((v = opt_value(a, "--user", &i, argc, argv)))
            *init_user = v;
        else if (strcmp(a, "init") == 0)
            *init_mode = 1;
        else if (strcmp(a, "--version") == 0) {
            printf("bokfd %s (protocol v%d)\n", BOKF_VERSION,
                   BOKF_PROTOCOL_VERSION);
            exit(0);
        } else {
            fprintf(stderr, "bokfd: unknown argument %s\n", a);
            usage(stderr);
            exit(2);
        }
    }
}

int main(int argc, char **argv)
{
    config_defaults();
    const char *config_file = NULL;
    const char *init_user = "admin";
    int init_mode = 0;
    parse_args(argc, argv, &config_file, &init_user, &init_mode);

    if (config_file) {
        char *err = NULL;
        if (config_load_file(config_file, &err) != 0) {
            fprintf(stderr, "bokfd: %s\n", err ? err : "config error");
            free(err);
            return 1;
        }
    }
    config_apply_env();
    log_set_level(g_cfg.log_level);

    if (init_mode)
        return cmd_init(g_cfg.db_path, init_user);

    sqlite3 *db = NULL;
    char *err = NULL;
    if (db_open(g_cfg.db_path, &db, &err) != 0) {
        log_error("%s", err ? err : "cannot open database");
        free(err);
        return 1;
    }
    if (db_count(db, "SELECT count(*) FROM users") <= 0)
        log_warn("no users in database; run 'bokfd init' first");

    mkdir_p(g_cfg.backup_dir, 0700);
    mkdir_p(g_cfg.export_dir, 0700);

    char dir[4096];
    dirname_of(g_cfg.socket_path, dir, sizeof dir);
    if (mkdir_p(dir, 0755) != 0) {
        log_error("cannot create socket directory %s", dir);
        sqlite3_close(db);
        return 1;
    }
    int lfd = unix_listen(g_cfg.socket_path);
    if (lfd < 0) {
        log_error("cannot listen on %s: %s", g_cfg.socket_path,
                  strerror(errno));
        sqlite3_close(db);
        return 1;
    }
    log_info("bokfd %s listening on %s", BOKF_VERSION, g_cfg.socket_path);

    int tfd = -1;
    if (g_cfg.tcp_enabled) {
        tfd = tcp_listen(g_cfg.tcp_addr ? g_cfg.tcp_addr : "127.0.0.1:8787");
        if (tfd < 0)
            log_error("cannot listen on tcp %s",
                      g_cfg.tcp_addr ? g_cfg.tcp_addr : "127.0.0.1:8787");
        else
            log_info("bokfd listening on tcp %s",
                     g_cfg.tcp_addr ? g_cfg.tcp_addr : "127.0.0.1:8787");
    }

    int tlsfd = -1;
    if (g_cfg.tls_enabled) {
        if (tls_load_certs() != 0) {
            log_error("cannot load TLS certificate %s / key %s",
                      g_cfg.tls_cert, g_cfg.tls_key);
            close(lfd);
            if (tfd >= 0)
                close(tfd);
            sqlite3_close(db);
            config_free();
            return 1;
        }
        tlsfd = tcp_listen(g_cfg.tls_addr ? g_cfg.tls_addr : "127.0.0.1:8788");
        if (tlsfd < 0) {
            log_error("cannot listen on tls %s",
                      g_cfg.tls_addr ? g_cfg.tls_addr : "127.0.0.1:8788");
            close(lfd);
            if (tfd >= 0)
                close(tfd);
            sqlite3_close(db);
            config_free();
            return 1;
        }
        log_info("bokfd listening on tls %s", g_cfg.tls_addr);
    }

    signal(SIGINT, on_signal);
    signal(SIGTERM, on_signal);
    signal(SIGHUP, on_reload);
    signal(SIGPIPE, SIG_IGN);
    sessions_init(g_cfg.session_ttl);

    /* base64 expands by 4/3; the JSON envelope needs some room too */
    g_line_limit = (size_t)g_cfg.max_line_bytes;
    size_t att_limit =
        (size_t)g_cfg.max_attachment_bytes / 3 * 4 + 64 * 1024;
    if (att_limit > g_line_limit)
        g_line_limit = att_limit;

    struct conn conns[MAX_CONNS];
    size_t nconns = 0;
    struct pollfd pfds[MAX_CONNS + 3];
    struct conn *map[MAX_CONNS + 3];
    int lfds[3];
    int ltls[3];

    while (!g_stop && !g_reload) {
        int n = 0;
        int nlisteners = 0;
        lfds[nlisteners] = lfd;
        ltls[nlisteners] = 0;
        nlisteners++;
        if (tfd >= 0) {
            lfds[nlisteners] = tfd;
            ltls[nlisteners] = 0;
            nlisteners++;
        }
        if (tlsfd >= 0) {
            lfds[nlisteners] = tlsfd;
            ltls[nlisteners] = 1;
            nlisteners++;
        }
        for (int i = 0; i < nlisteners; i++) {
            pfds[n].fd = lfds[i];
            pfds[n].events = POLLIN;
            map[n] = NULL;
            n++;
        }
        for (size_t i = 0; i < nconns; i++) {
            pfds[n].fd = conns[i].fd;
            pfds[n].events = POLLIN |
                             ((conns[i].out.len > conns[i].out_sent ||
                               conns[i].handshaking)
                                  ? POLLOUT
                                  : 0);
            map[n] = &conns[i];
            n++;
        }
        int pr = poll(pfds, (nfds_t)n, 1000);
        if (pr < 0) {
            if (errno == EINTR)
                continue;
            log_error("poll: %s", strerror(errno));
            break;
        }
        for (int i = 0; i < nlisteners; i++)
            if (pfds[i].revents & POLLIN)
                accept_conns(lfds[i], conns, &nconns, ltls[i]);
        for (int i = nlisteners; i < n; i++) {
            struct conn *c = map[i];
            if (c->fd < 0)
                continue;
            if (pfds[i].revents & (POLLIN | POLLHUP | POLLERR))
                conn_read(c, db);
            if (pfds[i].revents & POLLOUT)
                conn_flush(c);
            if (c->closing && c->out.len == c->out_sent) {
                if (c->ssl) {
                    SSL_free(c->ssl);
                    c->ssl = NULL;
                }
                close(c->fd);
                buf_free(&c->in);
                buf_free(&c->out);
                c->fd = -1;
            }
        }
        if (tlsfd >= 0)
            tls_reload_check();
        size_t keep = 0;
        for (size_t i = 0; i < nconns; i++) {
            if (conns[i].fd >= 0)
                conns[keep++] = conns[i];
        }
        nconns = keep;
    }

    if (g_reload)
        log_info("reloading on SIGHUP");
    else
        log_info("shutting down");
    for (size_t i = 0; i < nconns; i++) {
        if (conns[i].ssl)
            SSL_free(conns[i].ssl);
        close(conns[i].fd);
        buf_free(&conns[i].in);
        buf_free(&conns[i].out);
    }
    close(lfd);
    if (tfd >= 0)
        close(tfd);
    if (tlsfd >= 0)
        close(tlsfd);
    if (g_tls_ctx) {
        SSL_CTX_free(g_tls_ctx);
        g_tls_ctx = NULL;
    }
    unlink(g_cfg.socket_path);
    sessions_free_all();
    sqlite3_close(db);
    config_free();
    if (g_reload) {
        execvp(argv[0], argv);
        log_error("re-exec %s failed: %s", argv[0], strerror(errno));
        return 1;
    }
    return 0;
}