#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>
#include <string.h>
#include <inttypes.h>
#include "try.h"
#include "server.h"
#include "wire.h"

char *string_dup(const char *s)
{
    size_t len = strlen(s) + 1;
    char *p = malloc(len);
    return p ? memcpy(p, s, len) : NULL;
}


// https://en.wikipedia.org/wiki/Fowler-Noll-Vo_hash_function#FNV-1a_hash
uint64_t hashmap_hash(const char *bytes, size_t bytes_n, size_t map_len)
{
    uint64_t hash = 0xcbf29ce484222325;

    for (size_t i = 0; i < bytes_n; i++) {
        hash *= 0x100000001b3;
        hash ^= bytes[i];
    }

    return (hash % map_len);
}

int jb_server_client_add(struct jb_server *s, int fd)
{
    for (int i = 0; i < JB_MAX_CLIENTS; i++) {
        if (s->clients[i].fd < 0) {
            s->clients[i].fd = fd;
            s->clients[i].owned_name_index = -1;
            s->clients[i].unique_name_index = -1;
            s->clients[i].state = JB_CLIENT_STATE_WAIT_AUTH;
            s->fds[i].fd = fd;
            s->fds[i].events = POLLIN;
            return i;
        }
    }
    return -1;
}

int jb_server_name_find(struct jb_server *s, char *name)
{
    int bucket = hashmap_hash(name, strlen(name), JB_MAX_NAMES);

    for (int i = bucket; i < bucket + 12 && i < JB_MAX_NAMES; i++) {
        if (s->names[i].name && strcmp(s->names[i].name, name) == 0) {
            return i;
        }
    }

    return -1;
}

struct jb_client *jb_server_name_find_client(struct jb_server *s, char *name)
{
    int name_index = jb_server_name_find(s, name);
    if (name_index < 0) {
        return NULL;
    }
    
    if (s->names[name_index].client_index < 0) {
        return NULL;
    }

    if (s->clients[s->names[name_index].client_index].state != JB_CLIENT_STATE_READY) {
        return NULL;
    }

    return &s->clients[s->names[name_index].client_index];
}

int jb_server_name_add(struct jb_server *s, char *name, int client_index)
{
    if (jb_server_name_find(s, name) >= 0) {
        return -1;
    }

    int bucket = hashmap_hash(name, strlen(name), JB_MAX_NAMES);

    for (int i = bucket; i < bucket + 12 && i < JB_MAX_NAMES; i++) {
        if (s->names[i].client_index == -1) {
            s->names[i].client_index = client_index;
            s->names[i].name = name;
            s->names_count++;
            return i;
        }
    }

    return -1;
}

int jb_server_client_assign_unique_name(struct jb_server *s, int i)
{
    struct jb_client *c = &s->clients[i];
    if (c->unique_name_index != -1) {
        return -1;
    }
    uint32_t id = 0;
    FILE *urandom_file = fopen("/dev/urandom", "rb");
    if (!urandom_file) {
        return -1;
    }
    if (fread(&id, 1, sizeof(uint32_t), urandom_file) != sizeof(uint32_t)) {
        return -1;
    }
    char *name = malloc(sizeof(char) * 16);
    if (!name) {
        return -1;
    }

    if (snprintf(name, 16, ":%"PRIu32, id) < 0) {
        return -1;
    }

    int name_index = jb_server_name_add(s, name, i);
    if (name_index < 0) {
        free(name);
        return -1;
    }
    c->unique_name_index = name_index;

    return name_index;
}

int jb_server_client_assign_own_name(struct jb_server *s, int i, char *name)
{
    struct jb_client *c = &s->clients[i];
    if (c->owned_name_index != -1 || !name || *name == ':' || *name == '\0') {
        return -1;
    }

    int name_index = jb_server_name_add(s, name, i);
    if (name_index < 0) {
        return -1;
    }
    c->owned_name_index = name_index;

    return 0;
}

void jb_server_name_remove(struct jb_server *s, int i)
{
    if (i >= 0) {
        if (s->names[i].name) {
            free(s->names[i].name);
            s->names[i].name = NULL;
        }
        s->names[i].client_index = -1;
        s->names_count--;
    }
}

void jb_server_client_remove(struct jb_server *s, int i)
{
    if (s->clients[i].fd >= 0) {
        close(s->clients[i].fd);
    }
    jb_server_name_remove(s, s->clients[i].unique_name_index);
    jb_server_name_remove(s, s->clients[i].owned_name_index);
    s->clients[i].unique_name_index = -1;
    s->clients[i].owned_name_index = -1;
    s->clients[i].fd = -1;
    s->clients[i].state = JB_CLIENT_STATE_NONE;
    s->fds[i].fd = -1;
    s->fds[i].events = 0;
    s->fds[i].revents = 0;
}

void jb_server_client_error(struct jb_server *s, int i, const char *msg)
{
    printf("jb_server_client_error: %s\n", msg);
    send(s->clients[i].fd, msg, strlen(msg), 0);
    jb_server_client_remove(s, i);
}

ssize_t jb_server_client_recv(struct jb_server *s, int i, void *buf, size_t n)
{
    ssize_t status = recv(s->fds[i].fd, buf, n, 0);
    if (status <= 0) {
        jb_server_client_remove(s, i);
    }
    return status;
}

#define _reply_begin(M_sig) \
    do { \
        TRYST(wire_compose_reply(&reply_ctx, &msg, (M_sig), &body_length)); \
        body_start = reply_ctx.byte_cursor; \
    } while(0); \

#define _reply_end() \
    do { \
        *body_length = reply_ctx.byte_cursor - body_start; \
        if (send(s->fds[i].fd, reply_data, reply_ctx.byte_cursor, 0) != reply_ctx.byte_cursor) { \
            return -1; \
        } \
    } while(0); \

#define _reply_error(message) \
    do { \
        TRYST(wire_compose_error(&reply_ctx, &msg, (message))); \
        if (send(s->fds[i].fd, reply_data, reply_ctx.byte_cursor, 0) != reply_ctx.byte_cursor) { \
            return -1; \
        } \
    } while(0) \

int jb_server_client_process_message(struct jb_server *s, int i, uint8_t *data, size_t data_len)
{
    struct jb_client *client = &s->clients[i];

    wire_message_t msg = {0};
    wire_context_t ctx = {
        .byte_cursor = 0,
        .data = data,
        .data_len = data_len,
    };

    TRYST(wire_parse_message(&ctx, &msg));

    wire_message_field_t *destination_field = &msg.fields[DBUS_HEADER_FIELD_DESTINATION];
    wire_message_field_t *member_field = &msg.fields[DBUS_HEADER_FIELD_MEMBER];
    bool should_unicast = false;
    uint32_t *body_length;
    uint32_t body_start;

    uint8_t reply_data[4096];
    memset(reply_data, 0, 4096);
    wire_context_t reply_ctx = {
        .byte_cursor = 0,
        .data = reply_data,
        .data_len = 4096,
    };

    switch (msg.type) {
        case DBUS_MESSAGE_METHOD_CALL: {
            if (!destination_field->present || !member_field->present) {
                return -1;
            }

            char *member = member_field->t.str;

            if (strcmp(destination_field->t.str, "org.freedesktop.DBus") != 0) {
                // not for dbus.
                should_unicast = true;
                break;
            }

            if (strcmp(member, "Hello") == 0) {
                int unique_name_index = jb_server_client_assign_unique_name(s, i);
                if (unique_name_index < 0) {
                    return -1;
                }

                _reply_begin("s") {
                    TRYPTR(wire_set_string(&reply_ctx, s->names[unique_name_index].name));
                } _reply_end()

                printf("assigned unique name '%s' to connection %d\n", s->names[unique_name_index].name, i);
            } else if (strcmp(member, "RequestName") == 0) {
                char *name = TRYPTR(wire_get_string(&ctx));
                int name_len = strlen(name);
                if (name_len < 1 || name_len > 256) {
                    return -1;
                }

                char *name_str = string_dup(name);
                if (!name_str) {
                    return -1;
                }

                int status_code = DBUS_REQUEST_NAME_REPLY_PRIMARY_OWNER;

                if (jb_server_client_assign_own_name(s, i, name_str) < 0) {
                    free(name_str);
                    // TODO: report the actual error
                    status_code = DBUS_REQUEST_NAME_REPLY_EXISTS;
                }

                _reply_begin("u") {
                    TRYPTR(wire_set_u32(&reply_ctx, status_code));
                } _reply_end()

                if (status_code == DBUS_REQUEST_NAME_REPLY_PRIMARY_OWNER) {
                    printf("client '%s' (index=%d) now owns name '%s'\n", s->names[client->unique_name_index].name, i, name_str);
                }
            } else if (strcmp(member, "GetNameOwner") == 0) {
                char *name = TRYPTR(wire_get_string(&ctx));
                int name_len = strlen(name);
                if (name_len < 1 || name_len > 256) {
                    return -1;
                }

                char *name_str = TRYPTR(string_dup(name));

                struct jb_client *target = jb_server_name_find_client(s, name_str);
                if (!target || target->unique_name_index < 0) {
                    _reply_error("org.freedesktop.DBus.Error.NameHasNoOwner");
                    return 0;
                }

                _reply_begin("s") {
                    TRYPTR(wire_set_string(&reply_ctx, s->names[target->unique_name_index].name));
                } _reply_end()
            } else if (strcmp(member, "NameHasOwner") == 0) {
                char *name = TRYPTR(wire_get_string(&ctx));
                int name_len = strlen(name);
                if (name_len < 1 || name_len > 256) {
                    return -1;
                }

                char *name_str = TRYPTR(string_dup(name));

                struct jb_client *target = jb_server_name_find_client(s, name_str);

                _reply_begin("b") {
                    TRYPTR(wire_set_u32(&reply_ctx, target ? 1 : 0));
                } _reply_end()
            } else if (strcmp(member, "ListNames") == 0) {
                _reply_begin("as") {
                    TRYPTR(wire_set_u32(&reply_ctx, s->names_count));

                    int left = s->names_count;
                    for (int i = 0; i < JB_MAX_NAMES; i++) {
                        if (s->names[i].name && s->names[i].client_index >= 0) {
                            left--;
                            TRYPTR(wire_set_string(&reply_ctx, s->names[i].name));
                        }
                        if (left <= 0) {
                            break;
                        }
                    }
                } _reply_end()
            } else if (strcmp(member, "ListActivatableNames") == 0) {
                // TODO: stub (always returns empty array)
                printf("STUB: ListActivatableNames\n");

                _reply_begin("as") {
                    TRYPTR(wire_set_u32(&reply_ctx, 0));
                    // empty arrays still need to align
                    TRYPTR(wire_write_align(&reply_ctx, 4));
                } _reply_end()
            } else if (strcmp(member, "StartServiceByName") == 0) {
                // TODO: stub (does nothing and always returns success)

                char *name = TRYPTR(wire_get_string(&ctx));
                int name_len = strlen(name);
                if (name_len < 1 || name_len > 256) {
                    return -1;
                }

                printf("STUB: StartServiceByName: %s\n", name);

                _reply_begin("u") {
                    TRYPTR(wire_set_u32(&reply_ctx, 1));
                } _reply_end()
                return 0;
            } else {
                _reply_error("org.freedesktop.DBus.Error.UnknownMethod");
                return 0;
            }
        } break;

        case DBUS_MESSAGE_METHOD_RETURN: {
            if (!destination_field->present) {
                return -1;
            }

            should_unicast = true;
        } break;

        default: {
            _reply_error("xyz.hippoz.jitterbug.NotImplemented");
            return 0;
        } break;
    }

    if (should_unicast) {
        if (client->unique_name_index < 0 || !destination_field->present) {
            return -1;
        }
        struct jb_client *target = jb_server_name_find_client(s, destination_field->t.str);
        if (!target) {
            _reply_error("org.freedesktop.DBus.Error.NameHasNoOwner");
            return 0;
        }
        TRYST(wire_compose_unicast_reply(&reply_ctx, &ctx, &msg, s->names[client->unique_name_index].name));
        TRYST(send(target->fd, reply_data, reply_ctx.byte_cursor, 0));
    }

    return 0;
}

void jb_server_free(struct jb_server *s)
{
    if (!s) return;
    if (s->sock_fd) close(s->sock_fd);
    // names are allocated on the heap
    for (int i = 0; i < JB_MAX_NAMES; i++) {
        if (s->names[i].name) {
            free(s->names[i].name);
        }
    }
    free(s);
}

struct jb_server *jb_server_create(const char *socket_path)
{
    struct jb_server *s = malloc(sizeof(struct jb_server));
    if (s == NULL) {
        return NULL;
    }

    s->sock_fd = socket(AF_UNIX, SOCK_STREAM, 0);
    if (s->sock_fd == -1) {
        free(s);
        return NULL;
    }

    struct sockaddr_un name;
    memset(&name, 0, sizeof(name));
    name.sun_family = AF_UNIX;
    strncpy(name.sun_path, socket_path, sizeof(name.sun_path) - 1);

    unlink(socket_path);

    if (bind(s->sock_fd, (const struct sockaddr *)&name, sizeof(name)) == -1) {
        close(s->sock_fd);
        free(s);
        return NULL;
    }
    if (listen(s->sock_fd, JB_BACKLOG) == -1) {
        close(s->sock_fd);
        free(s);
        return NULL;
    }

    for (int i = 0; i < JB_MAX_CLIENTS; i++) {
        s->clients[i].fd = -1;
        s->clients[i].owned_name_index = -1;
        s->clients[i].unique_name_index = -1;
        s->clients[i].state = JB_CLIENT_STATE_NONE;
        s->fds[i].fd = -1;
        s->fds[i].events = 0;
        s->fds[i].revents = 0;
    }

    for (int i = 0; i < JB_MAX_NAMES; i++) {
        s->names[i].client_index = -1;
        s->names[i].name = NULL;
    }

    s->fds[JB_MAX_CLIENTS].fd = s->sock_fd;
    s->fds[JB_MAX_CLIENTS].events = POLLIN;
    s->fd_num = JB_MAX_CLIENTS + 1;

    return s;
}

#define _client_die(m) ({jb_server_client_error(s,i,m);continue;})
int jb_server_turn(struct jb_server *s)
{
    static const char agree_unix_fd[] = "AGREE_UNIX_FD\r\n";
    static const char auth_ok[] = "OK aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa\r\n";
    static const char auth_list[] = "REJECTED EXTERNAL\r\n";
    static const char auth_data[] = "DATA\r\n";
    static const int data_buffer_len = 4096;

    TRYST(poll(s->fds, s->fd_num, -1));

    for (int i = 0; i < s->fd_num; i++) {
        int fd = s->fds[i].fd;
        if (s->fds[i].revents & POLLNVAL) {
            fprintf(stderr, "jb_server_turn: error: got POLLNVAL for fds[%d]. This is considered a bug in the server implementation.\n", i);
            return -1;
        }
        if (s->fds[i].revents & POLLERR) {
            if (fd == s->sock_fd) {
                fprintf(stderr, "jb_server_turn: error: got POLLERR for sock_fd\n");
                return -1;
            }
            jb_server_client_remove(s, i);
            continue;
        }
        if (s->fds[i].revents & POLLHUP) {
            if (fd == s->sock_fd) {
                fprintf(stderr, "jb_server_turn: error: got POLLHUP for sock_fd\n");
                return -1;
            }
            jb_server_client_remove(s, i);
            continue;
        }
        if (s->fds[i].revents & POLLIN) {
            // file descriptor ready for reading

            if (fd == s->sock_fd) {
                // new connection
                int fd = TRYST(accept(s->sock_fd, NULL, NULL));

                if (jb_server_client_add(s, fd) == -1) {
                    close(fd);
                }
                continue;
            }

            char data[data_buffer_len];
            memset(data, 0, sizeof(data));
            ssize_t bytes = recv(fd, data, sizeof(data) - 1, 0);
            if (bytes <= 0) {
                // error during recv() OR client disconnected, disconnect the client
                // TODO: should we actually do this?
                jb_server_client_remove(s, i);
                continue;
            }

            struct jb_client *c = &s->clients[i];
            switch (c->state) {
                case JB_CLIENT_STATE_WAIT_AUTH: {
                    // The D-Bus authentication protocol is a simple plain text protocol.
                    // Before the flow of messages can begin, the two applications must authenticate.
                    // SPEC: https://dbus.freedesktop.org/doc/dbus-specification.html#auth-protocol

                    // Immediately after connecting, clients must send a null byte
                    // SPEC: https://dbus.freedesktop.org/doc/dbus-specification.html#auth-nul-byte
                    char *auth_string = data;
                    if (*auth_string == '\0') {
                        auth_string++;
                    }

                    // TODO: the code below is hacky and does not reflect the specification
                    if (strcmp(auth_string, "AUTH\r\n") == 0) {
                        send(fd, auth_list, sizeof(auth_list) - 1, 0);
                    } else if (strcmp(auth_string, "AUTH EXTERNAL 31303030\r\n") == 0 || strcmp(auth_string, "DATA\r\n") == 0) {
                        send(fd, auth_ok, sizeof(auth_ok) - 1, 0);
                        c->state = JB_CLIENT_STATE_WAIT_BEGIN;
                    } else if (strcmp(auth_string, "AUTH EXTERNAL\r\n") == 0) {
                        send(fd, auth_data, sizeof(auth_data) - 1, 0);
                    }
                } break;
                case JB_CLIENT_STATE_WAIT_BEGIN: {
                    // Right now, we're expecting the client to either immediately begin the connection,
                    // or to negotiate UNIX file descriptor passing.
                    if (strncmp(data, "BEGIN\r\n", 7) == 0) {
                        c->state = JB_CLIENT_STATE_READY;
                        // At this point, a D-Bus connection has been established.
                        // The first octet after the \r\n of the BEGIN command is the first octet of the D-Bus communication.
                        // SPEC: https://dbus.freedesktop.org/doc/dbus-specification.html#auth-command-begin
                        char *first_message_begin = data + 7; /* 7 = length of "BEGIN\r\n" */
                        if (*first_message_begin == 'l' || *first_message_begin == 'B') {
                            // This looks like a D-Bus message. Let's process it!
                            if (jb_server_client_process_message(s, i, (uint8_t *)first_message_begin, data_buffer_len - 7) < 0) {
                                _client_die("failed to process message");
                            }
                        }
                    } else if (strcmp(data, "NEGOTIATE_UNIX_FD\r\n") == 0) {
                        // SPEC: https://dbus.freedesktop.org/doc/dbus-specification.html#auth-command-negotiate-unix-fd
                        send(fd, agree_unix_fd, sizeof(agree_unix_fd) - 1, 0);
                    } else {
                        _client_die("bad auth response (expected BEGIN or NEGOTIATE_UNIX_FD)");
                    }
                } break;

                case JB_CLIENT_STATE_READY: {
                    if (jb_server_client_process_message(s, i, (uint8_t*)data, data_buffer_len) < 0) {
                        _client_die("failed to process message");
                    }
                } break;

                case JB_CLIENT_STATE_NONE: {} /* through */
                default: {
                    _client_die("bad state");
                } break;
            }
        }
    }

    return 0;
}