#include "test_server.h"

#include <arpa/inet.h>
#include <errno.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <pthread.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <unistd.h>

#define TS_MAX_ROUTES 16
#define TS_MAX_REPLIES 8
#define TS_MAX_CALLS 2048
#define TS_REQUEST_MAX 65536

typedef struct {
        long status;
        char *body;
        char *headers;
} ts_stored_t;

typedef struct {
        char *path;
        ts_stored_t replies[TS_MAX_REPLIES];
        size_t count;
        size_t served;
} ts_route_t;

struct test_server {
        int listen_fd;
        int port;
        char base_url[64];
        pthread_t thread;
        pthread_mutex_t lock;
        bool stopping;
        ts_route_t routes[TS_MAX_ROUTES];
        size_t route_count;
        ts_route_t fallback;
        char *calls[TS_MAX_CALLS];
        size_t call_count;
};

static char *dup_or_empty(const char *text)
{
        size_t len = (text != NULL) ? strlen(text) : 0;
        char *copy = malloc(len + 1);

        if (copy == NULL)
                return NULL;
        if (len > 0)
                memcpy(copy, text, len);
        copy[len] = '\0';
        return copy;
}

static void route_clear(ts_route_t *route)
{
        size_t i;

        for (i = 0; i < route->count; i++) {
                free(route->replies[i].body);
                free(route->replies[i].headers);
        }
        route->count = 0;
        route->served = 0;
}

static void route_fill(ts_route_t *route, const ts_reply_t *replies,
                       size_t count)
{
        size_t i;

        route_clear(route);
        if (count > TS_MAX_REPLIES)
                count = TS_MAX_REPLIES;
        for (i = 0; i < count; i++) {
                route->replies[i].status = replies[i].status;
                route->replies[i].body = dup_or_empty(replies[i].body);
                route->replies[i].headers = dup_or_empty(replies[i].headers);
        }
        route->count = count;
}

static ts_route_t *route_for(test_server_t *server, const char *path,
                             bool create)
{
        size_t i;

        for (i = 0; i < server->route_count; i++)
                if (strcmp(server->routes[i].path, path) == 0)
                        return &server->routes[i];
        if (!create || server->route_count >= TS_MAX_ROUTES)
                return NULL;
        server->routes[server->route_count].path = dup_or_empty(path);
        return &server->routes[server->route_count++];
}

static const char *reason_phrase(long status)
{
        switch (status) {
        case 200:
                return "OK";
        case 400:
                return "Bad Request";
        case 404:
                return "Not Found";
        case 429:
                return "Too Many Requests";
        case 500:
                return "Internal Server Error";
        case 502:
                return "Bad Gateway";
        case 503:
                return "Service Unavailable";
        case 504:
                return "Gateway Timeout";
        default:
                return "Status";
        }
}

static bool read_request_path(int fd, char *out, size_t out_size)
{
        char buffer[TS_REQUEST_MAX];
        size_t used = 0;
        const char *start;
        const char *space;
        size_t len;

        for (;;) {
                ssize_t got;
                struct pollfd pfd;

                pfd.fd = fd;
                pfd.events = POLLIN;
                pfd.revents = 0;
                if (poll(&pfd, 1, 5000) <= 0)
                        return false;
                got = read(fd, buffer + used, sizeof(buffer) - used - 1);
                if (got <= 0)
                        return false;
                used += (size_t)got;
                buffer[used] = '\0';
                if (strstr(buffer, "\r\n\r\n") != NULL)
                        break;
                if (used + 1 >= sizeof(buffer))
                        return false;
        }

        start = strchr(buffer, ' ');
        if (start == NULL)
                return false;
        start++;
        space = strchr(start, ' ');
        if (space == NULL)
                return false;
        len = (size_t)(space - start);
        if (len + 1 > out_size)
                return false;
        memcpy(out, start, len);
        out[len] = '\0';
        return true;
}

static void write_all(int fd, const char *data, size_t len)
{
        size_t sent = 0;

        while (sent < len) {
                ssize_t wrote = write(fd, data + sent, len - sent);
                if (wrote <= 0) {
                        if (wrote < 0 && errno == EINTR)
                                continue;
                        return;
                }
                sent += (size_t)wrote;
        }
}

static void serve(test_server_t *server, int fd)
{
        char path[TS_REQUEST_MAX];
        char head[1024];
        ts_stored_t reply;
        ts_route_t *route;
        size_t index;
        int head_len;

        if (!read_request_path(fd, path, sizeof(path)))
                return;

        pthread_mutex_lock(&server->lock);
        if (server->call_count < TS_MAX_CALLS)
                server->calls[server->call_count++] = dup_or_empty(path);

        route = route_for(server, path, false);
        if (route == NULL || route->count == 0)
                route = &server->fallback;

        if (route->count == 0) {
                head_len = snprintf(head, sizeof(head),
                                    "HTTP/1.1 404 Not Found\r\n"
                                    "Content-Length: 0\r\n"
                                    "Connection: close\r\n\r\n");
                pthread_mutex_unlock(&server->lock);
                write_all(fd, head, (size_t)head_len);
                return;
        }

        index = route->served;
        if (index >= route->count)
                index = route->count - 1;
        route->served++;

        reply.status = route->replies[index].status;
        reply.body = dup_or_empty(route->replies[index].body);
        reply.headers = dup_or_empty(route->replies[index].headers);
        pthread_mutex_unlock(&server->lock);

        head_len = snprintf(head, sizeof(head),
                            "HTTP/1.1 %ld %s\r\n"
                            "Content-Type: application/octet-stream\r\n"
                            "Content-Length: %zu\r\n"
                            "Connection: close\r\n"
                            "%s%s"
                            "\r\n",
                            reply.status, reason_phrase(reply.status),
                            reply.body != NULL ? strlen(reply.body) : 0,
                            (reply.headers != NULL && reply.headers[0] != '\0')
                                    ? reply.headers
                                    : "",
                            (reply.headers != NULL && reply.headers[0] != '\0')
                                    ? "\r\n"
                                    : "");
        if (head_len > 0)
                write_all(fd, head, (size_t)head_len);
        if (reply.body != NULL)
                write_all(fd, reply.body, strlen(reply.body));
        free(reply.body);
        free(reply.headers);
}

static void *accept_loop(void *argument)
{
        test_server_t *server = argument;

        for (;;) {
                struct pollfd pfd;
                int fd;

                pthread_mutex_lock(&server->lock);
                if (server->stopping) {
                        pthread_mutex_unlock(&server->lock);
                        return NULL;
                }
                pthread_mutex_unlock(&server->lock);

                pfd.fd = server->listen_fd;
                pfd.events = POLLIN;
                pfd.revents = 0;
                if (poll(&pfd, 1, 20) <= 0)
                        continue;

                fd = accept(server->listen_fd, NULL, NULL);
                if (fd < 0)
                        continue;
                serve(server, fd);
                shutdown(fd, SHUT_WR);
                close(fd);
        }
}

test_server_t *ts_start(void)
{
        struct sockaddr_in address;
        socklen_t length = sizeof(address);
        test_server_t *server = calloc(1, sizeof(*server));
        int reuse = 1;

        if (server == NULL)
                return NULL;
        pthread_mutex_init(&server->lock, NULL);

        server->listen_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (server->listen_fd < 0) {
                free(server);
                return NULL;
        }
        setsockopt(server->listen_fd, SOL_SOCKET, SO_REUSEADDR, &reuse,
                   sizeof(reuse));

        memset(&address, 0, sizeof(address));
        address.sin_family = AF_INET;
        address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
        address.sin_port = 0;

        if (bind(server->listen_fd, (struct sockaddr *)&address,
                 sizeof(address)) != 0 ||
            listen(server->listen_fd, 32) != 0 ||
            getsockname(server->listen_fd, (struct sockaddr *)&address,
                        &length) != 0) {
                close(server->listen_fd);
                free(server);
                return NULL;
        }

        server->port = ntohs(address.sin_port);
        snprintf(server->base_url, sizeof(server->base_url),
                 "http://127.0.0.1:%d", server->port);

        if (pthread_create(&server->thread, NULL, accept_loop, server) != 0) {
                close(server->listen_fd);
                free(server);
                return NULL;
        }
        return server;
}

void ts_stop(test_server_t *server)
{
        size_t i;

        if (server == NULL)
                return;
        pthread_mutex_lock(&server->lock);
        server->stopping = true;
        pthread_mutex_unlock(&server->lock);
        pthread_join(server->thread, NULL);
        close(server->listen_fd);

        for (i = 0; i < server->route_count; i++) {
                route_clear(&server->routes[i]);
                free(server->routes[i].path);
        }
        route_clear(&server->fallback);
        for (i = 0; i < server->call_count; i++)
                free(server->calls[i]);
        pthread_mutex_destroy(&server->lock);
        free(server);
}

const char *ts_base_url(const test_server_t *server)
{
        return server->base_url;
}

void ts_respond(test_server_t *server, const char *path, long status,
                const char *body, const char *headers)
{
        ts_reply_t reply;

        reply.status = status;
        reply.body = body;
        reply.headers = headers;
        ts_respond_seq(server, path, &reply, 1);
}

void ts_respond_seq(test_server_t *server, const char *path,
                    const ts_reply_t *replies, size_t count)
{
        ts_route_t *route;

        pthread_mutex_lock(&server->lock);
        route = route_for(server, path, true);
        if (route != NULL)
                route_fill(route, replies, count);
        pthread_mutex_unlock(&server->lock);
}

void ts_respond_default(test_server_t *server, const ts_reply_t *replies,
                        size_t count)
{
        pthread_mutex_lock(&server->lock);
        route_fill(&server->fallback, replies, count);
        pthread_mutex_unlock(&server->lock);
}

size_t ts_count(test_server_t *server, const char *prefix)
{
        size_t total = 0;
        size_t len = strlen(prefix);
        size_t i;

        pthread_mutex_lock(&server->lock);
        for (i = 0; i < server->call_count; i++)
                if (server->calls[i] != NULL &&
                    strncmp(server->calls[i], prefix, len) == 0)
                        total++;
        pthread_mutex_unlock(&server->lock);
        return total;
}

size_t ts_total(test_server_t *server)
{
        size_t total;

        pthread_mutex_lock(&server->lock);
        total = server->call_count;
        pthread_mutex_unlock(&server->lock);
        return total;
}

const char *ts_path(test_server_t *server, size_t index)
{
        const char *path = NULL;

        pthread_mutex_lock(&server->lock);
        if (index < server->call_count)
                path = server->calls[index];
        pthread_mutex_unlock(&server->lock);
        return path;
}