#include "gem/server.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <time.h>
#include <unistd.h>

#define CHECK(condition) do { if (!(condition)) { \
    fprintf(stderr, "check failed at %s:%d: %s\n", __FILE__, __LINE__, #condition); return 1; \
} } while (0)

static int send_all_test(int fd, const char *data, size_t size)
{
    while (size) {
        ssize_t count = write(fd, data, size);
        if (count <= 0) return -1;
        data += count; size -= (size_t)count;
    }
    return 0;
}

static char *exchange(const struct gem_server_config *config, const char *request)
{
    int sockets[2];
    if (socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) != 0) return NULL;
    if (send_all_test(sockets[1], request, strlen(request)) != 0) {
        close(sockets[0]); close(sockets[1]); return NULL;
    }
    (void)shutdown(sockets[1], SHUT_WR);
    char error[256];
    (void)gem_server_handle_connection(sockets[0], config, error, sizeof(error));
    close(sockets[0]);
    size_t size = 0, capacity = 4096;
    char *response = malloc(capacity);
    if (!response) { close(sockets[1]); return NULL; }
    for (;;) {
        if (size + 1024 + 1 > capacity) {
            capacity *= 2;
            void *memory = realloc(response, capacity);
            if (!memory) { free(response); close(sockets[1]); return NULL; }
            response = memory;
        }
        ssize_t count = read(sockets[1], response + size, capacity - size - 1);
        if (count < 0) { free(response); close(sockets[1]); return NULL; }
        if (!count) break;
        size += (size_t)count;
    }
    close(sockets[1]); response[size] = '\0';
    return response;
}

static int fake_generate(void *backend, const struct gem_chat_request *request,
                         gem_stream_emit_fn emit, void *user,
                         size_t *input_tokens, size_t *output_tokens,
                         const char **finish_reason, char *error, size_t error_size)
{
    (void)backend; (void)error; (void)error_size;
    if (!strstr(request->prompt, "<|turn>user\nhello<turn|>")) return -1;
    struct gem_stream_event reasoning = {
        .kind = GEM_STREAM_REASONING, .text = "brief thought", .text_size = 13,
    };
    struct gem_stream_event heartbeat = {.kind = GEM_STREAM_HEARTBEAT};
    struct gem_stream_event text = {
        .kind = GEM_STREAM_TEXT, .text = "hi\n", .text_size = 3,
    };
    struct gem_stream_event tool = {
        .kind = GEM_STREAM_TOOL_CALL, .tool_id = "call_1", .tool_name = "shell",
        .tool_arguments_json = "{\"cmd\":\"pwd\"}",
    };
    if (emit(&heartbeat, user) != 0 || emit(&reasoning, user) != 0 ||
        emit(&text, user) != 0 || emit(&tool, user) != 0)
        return -1;
    *input_tokens = 11; *output_tokens = 7; *finish_reason = "tool_calls";
    return 0;
}

static struct gem_server_config config(void)
{
    return (struct gem_server_config){
        .model_id = "gemma-4-31b", .context_length = 32768, .default_max_tokens = 128,
        .sampler = {.temperature = 0.8f, .top_k = 40, .top_p = 0.95f,
                    .repeat_penalty = 1.0f, .repeat_last_n = 64, .seed = 1},
        .generate = fake_generate,
    };
}

static int endpoints(void)
{
    struct gem_server_config server = config();
    char *response = exchange(&server, "GET /health HTTP/1.1\r\nHost: localhost\r\n\r\n");
    CHECK(response && strstr(response, "HTTP/1.1 200 OK") && strstr(response, "{\"status\":\"ok\"}"));
    free(response);
    response = exchange(&server, "GET /v1/models HTTP/1.1\r\nHost: localhost\r\n\r\n");
    CHECK(response && strstr(response, "\"id\":\"gemma-4-31b\"") &&
          strstr(response, "\"n_ctx\":32768"));
    free(response);
    response = exchange(&server, "GET /props?model=gemma HTTP/1.1\r\nHost: localhost\r\n\r\n");
    CHECK(response && strstr(response, "\"vision\":false"));
    free(response);
    return 0;
}

static int stream(void)
{
    struct gem_server_config server = config();
    const char body[] =
        "{\"model\":\"gemma-4-31b\",\"stream\":true,"
        "\"stream_options\":{\"include_usage\":true},"
        "\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
    char request[1024];
    int size = snprintf(request, sizeof(request),
        "POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
        "Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
        sizeof(body) - 1, body);
    CHECK(size > 0 && (size_t)size < sizeof(request));
    char *response = exchange(&server, request);
    CHECK(response && strstr(response, "Content-Type: text/event-stream"));
    CHECK(strstr(response, "\"role\":\"assistant\""));
    CHECK(strstr(response, ": keepalive\n\n"));
    CHECK(strstr(response, "\"reasoning_content\":\"brief thought\""));
    CHECK(strstr(response, "\"content\":\"hi\\n\""));
    CHECK(strstr(response, "\"tool_calls\":[{\"index\":0,\"id\":\"call_1\""));
    CHECK(strstr(response, "\"arguments\":\"{\\\"cmd\\\":\\\"pwd\\\"}\""));
    CHECK(strstr(response, "\"finish_reason\":\"tool_calls\""));
    CHECK(strstr(response, "\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18"));
    CHECK(strstr(response, "data: [DONE]\n\n"));
    free(response);
    return 0;
}

static int completion(void)
{
    struct gem_server_config server = config();
    const char body[] =
        "{\"model\":\"gemma-4-31b\",\"stream\":false,"
        "\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
    char request[1024];
    int size = snprintf(request, sizeof(request),
        "POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
        "Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
        sizeof(body) - 1, body);
    CHECK(size > 0 && (size_t)size < sizeof(request));
    char *response = exchange(&server, request);
    CHECK(response && strstr(response, "HTTP/1.1 200 OK"));
    CHECK(strstr(response, "Content-Type: application/json"));
    CHECK(strstr(response, "\"object\":\"chat.completion\""));
    CHECK(strstr(response, "\"role\":\"assistant\",\"content\":\"hi\\n\""));
    CHECK(strstr(response, "\"reasoning_content\":\"brief thought\""));
    CHECK(strstr(response, "\"tool_calls\":[{\"id\":\"call_1\""));
    CHECK(strstr(response, "\"arguments\":\"{\\\"cmd\\\":\\\"pwd\\\"}\""));
    CHECK(strstr(response, "\"finish_reason\":\"tool_calls\""));
    CHECK(strstr(response, "\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18"));
    CHECK(!strstr(response, "keepalive"));
    CHECK(!strstr(response, "data: "));
    free(response);
    return 0;
}

static int cancellation_generate(void *backend, const struct gem_chat_request *request,
                                 gem_stream_emit_fn emit, void *user,
                                 size_t *input_tokens, size_t *output_tokens,
                                 const char **finish_reason, char *error, size_t error_size)
{
    (void)request; (void)input_tokens; (void)output_tokens; (void)finish_reason;
    int *cancelled = backend;
    struct gem_stream_event heartbeat = {.kind = GEM_STREAM_HEARTBEAT};
    const struct timespec delay = {.tv_nsec = 1000000};
    for (size_t attempt = 0; attempt < 5000; ++attempt) {
        if (emit(&heartbeat, user) != 0) {
            *cancelled = 1;
            (void)snprintf(error, error_size, "client disconnected");
            return -1;
        }
        (void)nanosleep(&delay, NULL);
    }
    return -1;
}

static int cancellation(void)
{
    int sockets[2];
    CHECK(socketpair(AF_UNIX, SOCK_STREAM, 0, sockets) == 0);
    pid_t reader = fork();
    CHECK(reader >= 0);
    if (reader == 0) {
        close(sockets[0]);
        const char body[] =
            "{\"model\":\"gemma-4-31b\",\"stream\":true,"
            "\"messages\":[{\"role\":\"user\",\"content\":\"hello\"}]}";
        char request[1024];
        int size = snprintf(request, sizeof(request),
            "POST /v1/chat/completions HTTP/1.1\r\nHost: localhost\r\n"
            "Content-Type: application/json\r\nContent-Length: %zu\r\n\r\n%s",
            sizeof(body) - 1, body);
        if (size <= 0 || (size_t)size >= sizeof(request) ||
            send_all_test(sockets[1], request, (size_t)size) != 0)
            _exit(1);
        char response[4096] = {0};
        size_t used = 0;
        while (used + 1 < sizeof(response)) {
            ssize_t count = read(sockets[1], response + used,
                                 sizeof(response) - used - 1);
            if (count <= 0) _exit(1);
            used += (size_t)count;
            response[used] = '\0';
            if (strstr(response, "\"role\":\"assistant\"")) break;
        }
        close(sockets[1]);
        _exit(0);
    }
    close(sockets[1]);
    int cancelled = 0;
    struct gem_server_config server = config();
    server.backend = &cancelled;
    server.generate = cancellation_generate;
    char error[256] = {0};
    CHECK(gem_server_handle_connection(sockets[0], &server, error, sizeof(error)) != 0);
    close(sockets[0]);
    int reader_status = 0;
    CHECK(waitpid(reader, &reader_status, 0) == reader);
    CHECK(WIFEXITED(reader_status) && WEXITSTATUS(reader_status) == 0);
    CHECK(cancelled == 1);
    CHECK(strstr(error, "client disconnected"));
    return 0;
}

static int errors(void)
{
    struct gem_server_config server = config();
    const char request[] = "POST /v1/chat/completions HTTP/1.1\r\nHost: x\r\nContent-Length: 2\r\n\r\n{}";
    char *response = exchange(&server, request);
    CHECK(response && strstr(response, "HTTP/1.1 400 Bad Request"));
    CHECK(strstr(response, "invalid_request_error"));
    free(response);
    return 0;
}

int main(void)
{
    CHECK(endpoints() == 0);
    CHECK(stream() == 0);
    CHECK(completion() == 0);
    CHECK(cancellation() == 0);
    CHECK(errors() == 0);
    puts("server tests passed");
    return 0;
}