#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <string.h>
#include <signal.h>
#include <pthread.h>
#include <assert.h>

#include "server.h"
#include "socket.h"
#include "utils.h"
#include "log.h"

typedef struct client_thread {
    pthread_t self;
    psocket client_sock;
    char host_buf[NI_MAXHOST], service_buf[NI_MAXSERV];
    const server_ctx *server_ctx;
} client_thread;

typedef struct server_event {
    const server_ctx **srv_ctx;
    const size_t siz;
} server_event;

typedef struct client_event {
    const psocket *fwd_sock;
    const client_thread *client_args;
} client_event;

static forward_state
fwd_state_string(const forward_state c_state, const client_thread *args,
                 const psocket *fwd);
static int
server_mainloop(event_ctx *ev_ctx, const server_ctx *srv_ctx[], size_t siz)
    __attribute__((noreturn));
static int server_accept_client(event_ctx *ev_ctx, int fd, void *user_data);
static void *
client_mainloop(void *arg);
static int client_io(event_ctx *ev_ctx, int src_fd, void *user_data);


void server_init_ctx(server_ctx **ctx, forward_ctx *fwd_ctx)
{
    assert(ctx && fwd_ctx);
    if (!*ctx)
        *ctx = (server_ctx *) malloc(sizeof(**ctx));
    assert(*ctx);

    memset(*ctx, 0, sizeof(**ctx));
    (*ctx)->fwd_ctx = fwd_ctx;
}

int server_setup(server_ctx *ctx,
                 const char *listen_addr, const char *listen_port)
{
    int s;
    struct addrinfo *srv_addr = NULL;

    assert(ctx);
    assert(listen_addr || listen_port);

    D2("Try to listen on %s:%s",
       (listen_addr ? listen_addr : "*"), listen_port);
    s = socket_init_in(listen_addr, listen_port, &srv_addr);
    if (s) {
        E_GAIERR(s, "Could not initialise server socket");
        return 1;
    }
    if (socket_bind_in(&ctx->sock, &srv_addr)) {
        E_STRERR("Could not bind server socket to %s:%s",
            listen_addr, listen_port);
        return 1;
    }
    if (socket_listen_in(&ctx->sock)) {
        E_STRERR("Could not listen on server socket on %s:%s",
            listen_addr, listen_port);
        return 1;
    }

    return 0;
}

int server_validate_ctx(const server_ctx *ctx)
{
    assert(ctx && ctx->fwd_ctx);
    assert(ctx->sock.fd >= 0 && ctx->sock.addr_len > 0);

    return 0;
}

int server_setup_event(server_ctx *srv_ctx[], size_t siz, event_ctx **ev_ctx)
{
    int s;

    assert(srv_ctx && ev_ctx);
    assert(siz > 0 && siz < POTD_MAXFD);

    event_init(ev_ctx);
    if (event_setup(*ev_ctx))
        return 1;

    for (size_t i = 0; i < siz; ++i) {
        if (event_add_sock(*ev_ctx, &srv_ctx[i]->sock)) {
            return 1;
        }

        s = socket_addrtostr_in(&srv_ctx[i]->sock,
                                srv_ctx[i]->host_buf, srv_ctx[i]->service_buf);
        if (s) {
            E_GAIERR(s, "Convert socket address to string");
            return 1;
        }
        N("Redirector service listening on %s:%s",
          srv_ctx[i]->host_buf, srv_ctx[i]->service_buf);
    }

    return 0;
}

pid_t server_daemonize(event_ctx *ev_ctx, server_ctx *srv_ctx[], size_t siz)
{
    pid_t p;
    int s;
    size_t i;

    assert(ev_ctx && srv_ctx);
    assert(siz > 0 && siz < POTD_MAXFD); 

    for (i = 0; i < siz; ++i) {
        assert(srv_ctx[i]);
        s = socket_addrtostr_in(&srv_ctx[i]->sock,
            srv_ctx[i]->host_buf,
            srv_ctx[i]->service_buf);
        if (s) {
            E_GAIERR(s, "Could not initialise server daemon socket");
            return 1;
        }
    }

    p = fork();
    switch (p) {
        case -1:
            W_STRERR("%s", "Server daemonize");
            return -1;
        case 0:
            N("%s", "Server daemon mainloop");
            server_mainloop(ev_ctx, (const server_ctx **) srv_ctx, siz);
            break;
    }
    D2("Server daemon pid: %d", p);

    return p;
}

static forward_state
fwd_state_string(const forward_state c_state, const client_thread *args,
                 const psocket *fwd)
{
    switch (c_state) {
        case CON_OK:
            break;
        case CON_IN_ERROR:
            N("Lost connection to %s:%s: %d",
                args->host_buf, args->service_buf,
                args->client_sock.fd);
            break;
        case CON_IN_TERMINATED:
            N("Connection terminated: %s:%s: %d",
                args->host_buf, args->service_buf,
                args->client_sock.fd);
            break;
        case CON_OUT_ERROR:
            N("Lost forward connection to %s:%s: %d",
                args->server_ctx->fwd_ctx->host_buf,
                args->server_ctx->fwd_ctx->service_buf,
                fwd->fd);
            break;
        case CON_OUT_TERMINATED:
            N("Forward connection terminated: %s:%s: %d",
                args->server_ctx->fwd_ctx->host_buf,
                args->server_ctx->fwd_ctx->service_buf,
                fwd->fd);
            break;
    }

    return c_state;
}

static int server_mainloop(event_ctx *ev_ctx, const server_ctx *srv_ctx[], size_t siz)
{
    int rc;
    server_event ev_srv = { srv_ctx, siz };

    set_procname("[potd] server");
    assert( set_child_sighandler() == 0 );

    rc = event_loop(ev_ctx, server_accept_client, &ev_srv);
    event_free(&ev_ctx);

    exit(rc);
}

static int server_accept_client(event_ctx *ev_ctx, int fd, void *user_data)
{
    size_t i;
    int s;
    server_event *ev_srv = (server_event *) user_data;
    client_thread *args;
    const server_ctx *srv_ctx;

    (void) ev_ctx;
    assert(ev_srv);

    for (i = 0; i < ev_srv->siz; ++i) {
        srv_ctx = ev_srv->srv_ctx[i];
        if (srv_ctx->sock.fd == fd) {
            args = (client_thread *) calloc(1, sizeof(*args));
            assert(args);

            if (socket_accept_in(&srv_ctx->sock,
                    &args->client_sock))
            {
                E_STRERR("Could not accept client connection on fd %d",
                    srv_ctx->sock.fd);
                goto error;
            }

            args->server_ctx = srv_ctx;
            s = socket_addrtostr_in(&args->client_sock,
                                    args->host_buf, args->service_buf);
            if (s) {
                E_GAIERR(s, "Convert socket address to string");
                goto error;
            }
            N2("New connection from %s:%s to %s:%s: %d",
                args->host_buf, args->service_buf,
                srv_ctx->host_buf, srv_ctx->service_buf,
                args->client_sock.fd);

            if (pthread_create(&args->self, NULL,
                               client_mainloop, args))
            {
                E_STRERR("Thread creation for %s:%s on fd %d",
                    args->host_buf, args->service_buf,
                    args->client_sock.fd);
                goto error;
            }

            return 1;
error:
            socket_close(&args->client_sock);
            free(args);
            return 0;
        }
    }

    return 0;
}

static void *
client_mainloop(void *arg)
{
    client_thread *args;
    client_event ev_cli;
    int s;
    event_ctx *ev_ctx = NULL;
    psocket fwd;

    assert(arg);
    args = (client_thread *) arg;
    pthread_detach(args->self);

    event_init(&ev_ctx);
    if (event_setup(ev_ctx)) {
        E_STRERR("Client event context creation for server fd %d",
            args->server_ctx->sock.fd);
        goto finish;
    }

    if (fwd_connect_sock(args->server_ctx->fwd_ctx, &fwd)) {
        E_STRERR("Forward connection to %s:%s server fd %d",
            args->server_ctx->fwd_ctx->host_buf,
            args->server_ctx->fwd_ctx->service_buf,
            args->server_ctx->sock.fd);
        goto finish;
    }
    N("Forwarding connection to %s:%s forward fd %d",
        args->server_ctx->fwd_ctx->host_buf,
        args->server_ctx->fwd_ctx->service_buf, fwd.fd);

    if (event_add_sock(ev_ctx, &fwd)) {
        E_STRERR("Forward event context add to %s:%s forward fd %d",
            args->server_ctx->fwd_ctx->host_buf,
            args->server_ctx->fwd_ctx->service_buf, fwd.fd);
        goto finish;
    }

    /*
     * We got the client socket from our main thread, so fd flags like
     * O_NONBLOCK are not inherited!
     */
    s = socket_nonblock(&args->client_sock);
    if (s) {
        E_STRERR("Socket non blocking mode to %s:%s forward fd %d",
            args->server_ctx->fwd_ctx->host_buf,
            args->server_ctx->fwd_ctx->service_buf, fwd.fd);
        goto finish;
    }
    if (event_add_sock(ev_ctx, &args->client_sock)) {
        E_STRERR("Forward event context add to %s:%s forward fd %d",
	        args->server_ctx->fwd_ctx->host_buf,
            args->server_ctx->fwd_ctx->service_buf, fwd.fd);
        goto finish;
    }

    ev_cli.client_args = args;
    ev_cli.fwd_sock = &fwd;
    event_loop(ev_ctx, client_io, &ev_cli);

finish:
    event_free(&ev_ctx);
    socket_close(&fwd);
    socket_close(&args->client_sock);
    free(args);
    return NULL;
}

static int
client_io(event_ctx *ev_ctx, int src_fd, void *user_data)
{
    int dest_fd;
    client_event *ev_cli = (client_event *) user_data;
    const psocket *client_sock = &ev_cli->client_args->client_sock;
    forward_state fwd_state;

    if (src_fd == ev_cli->fwd_sock->fd) {
        dest_fd = client_sock->fd;
    } else if (src_fd == client_sock->fd) {
        dest_fd = ev_cli->fwd_sock->fd;
    } else return 0;

    fwd_state = event_forward_connection(ev_ctx, dest_fd);

    switch (fwd_state_string(fwd_state, ev_cli->client_args,
                             ev_cli->fwd_sock))
    {
        case CON_IN_TERMINATED:
        case CON_OUT_TERMINATED:
            ev_ctx->active = 0;
        case CON_OK:
            return 1;
        case CON_IN_ERROR:
        case CON_OUT_ERROR:
            ev_ctx->active = 0;
            return 0;
    }

    return 1;
}