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

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].state = JB_CLIENT_STATE_WAIT_AUTH;
            s->fds[i].fd = fd;
            s->fds[i].events = POLLIN;
            return i;
        }
    }
    return -1;
}

void jb_server_client_remove(struct jb_server *s, int i)
{
    if (s->clients[i].fd >= 0) {
        close(s->clients[i].fd);
    }
    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;
}

int jb_server_client_process_message(struct jb_server *s, int i, uint8_t *data, size_t data_len)
{
    wire_context_t ctx = {
        .byte_cursor = 0,
        .data = data,
        .data_len = data_len,
    };

    wire_message_t msg = {0};

    TRY_NONNEGATIVE(int, wire_parse_message(&ctx, &msg), -1);

    printf("length: %d, serial: %d, header fields length in bytes: %d\n", msg.body_length, msg.serial, msg.header_fields_length);

    switch (msg.type) {
        case DBUS_MESSAGE_METHOD_CALL: {
            bool for_dbus = false;
            char *member;

            // TODO: maybe the fields in the msg.fields array should be indexed based on their type?
            for (int i = 0; i < msg.header_fields_count; i++) {
                if (msg.fields[i].type == DBUS_HEADER_FIELD_DESTINATION) {
                    if (strcmp(msg.fields[i].t.str, "org.freedesktop.DBus") == 0) {
                        for_dbus = true;
                    }
                } else if (msg.fields[i].type == DBUS_HEADER_FIELD_MEMBER) {
                    member = msg.fields[i].t.str;
                }
            }

            if (for_dbus) {
                if (strcmp(member, "Hello") == 0) {
                    printf("jb_server_client_process_message: message was for dbus, got hello\n");
                    uint8_t reply_data[512];
                    memset(reply_data, 0, 512);
                    wire_context_t reply_ctx = {
                        .byte_cursor = 0,
                        .data = reply_data,
                        .data_len = 511,
                    };

                    uint32_t *body_length;
                    TRY_NONNEGATIVE(int, wire_compose_reply(&reply_ctx, &msg, "s", &body_length), -1);
                    uint32_t body_start = reply_ctx.byte_cursor;
                    TRY_NONNULL(char*, wire_set_string(&reply_ctx, ":1.0"), -1);
                    *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;
                    }
                    printf("jb_server_client_process_message: processed hello\n");
                } else {
                    printf("jb_server_client_process_message: message was for dbus, however got unimplemented member '%s'\n", member);
                }
            } else {
                printf("jb_server_client_process_message: skipping message not for dbus for now\n");
            }
        } break;

        default: {
            printf("jb_server_client_process_message: unimplemented message type: %d\n", msg.type);
        } break;
    }

    return 0;
}

void jb_server_free(struct jb_server *s)
{
    if (!s) return;
    if (s->sock_fd) close(s->sock_fd);
    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].state = JB_CLIENT_STATE_NONE;
        s->fds[i].fd = -1;
        s->fds[i].events = 0;
        s->fds[i].revents = 0;
    }

    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 1234deadbeef\r\n";
    static const int data_buffer_len = 256;

    TRY_NONNEGATIVE(int, poll(s->fds, s->fd_num, -1), -1);

    for (int i = 0; i < s->fd_num; i++) {
        int fd = s->fds[i].fd;
        if (s->fds[i].revents & POLLIN) {
            // file descriptor ready for reading

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

                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(), disconnect the client
                // TODO: should we actually do this?
                close(fd);
                jb_server_client_remove(s, i);
                return -1;
            }
            if (bytes == 0) {
                // client disconnected
                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
                    if (*data != '\0' || *(data + 1) == '\0') {
                        _client_die("expected initial auth message to begin with a nul byte");
                    }

                    char *auth_string = data + 1;

                    if (strcmp(auth_string, "AUTH EXTERNAL 31303030\r\n") != 0)
                        _client_die("bad auth");
                    send(fd, auth_ok, sizeof(auth_ok) - 1, 0);
                    c->state = JB_CLIENT_STATE_WAIT_BEGIN;
                } 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: {
                    jb_server_client_process_message(s, i, (uint8_t*)data, data_buffer_len);
                } break;

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

    return 0;
}