#include "gem/chat.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

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

static struct gem_sampler_config defaults(void)
{
    return (struct gem_sampler_config){
        .temperature = 0.8f, .top_k = 40, .top_p = 0.95f,
        .repeat_penalty = 1.0f, .repeat_last_n = 64, .seed = 1,
    };
}

static int basic_prompt(void)
{
    const char body[] =
        "{\"model\":\"gemma-4-31b\",\"stream\":true,"
        "\"messages\":[{\"role\":\"system\",\"content\":\"  Be useful.  \"},"
        "{\"role\":\"user\",\"content\":\"Hello\\nworld\"}],"
        "\"stream_options\":{\"include_usage\":true},\"temperature\":0} ";
    struct gem_chat_request request;
    struct gem_sampler_config config = defaults();
    char error[256];
    CHECK(gem_chat_request_parse(&request, body, sizeof(body) - 1, 128, &config,
                                 error, sizeof(error)) == 0);
    CHECK(strcmp(request.model, "gemma-4-31b") == 0);
    CHECK(strcmp(request.prompt,
                 "<|turn>system\nBe useful.<turn|>\n"
                 "<|turn>user\nHello\nworld<turn|>\n<|turn>model\n") == 0);
    CHECK(request.stream && request.include_usage && request.max_tokens == 128);
    CHECK(request.sampler.temperature == 0.0f);
    gem_chat_request_destroy(&request);
    return 0;
}

static int tool_prompt(void)
{
    const char body[] =
        "{\"model\":\"gem\",\"stream\":true,\"reasoning_effort\":\"high\","
        "\"messages\":[{\"role\":\"user\",\"content\":\"Weather?\"}],"
        "\"tools\":[{\"type\":\"function\",\"function\":{"
        "\"name\":\"weather\",\"description\":\"Gets weather\","
        "\"parameters\":{\"type\":\"object\",\"properties\":{"
        "\"city\":{\"type\":\"string\",\"description\":\"City\"}},"
        "\"required\":[\"city\"]}}}]}";
    struct gem_chat_request request;
    struct gem_sampler_config config = defaults();
    char error[256];
    CHECK(gem_chat_request_parse(&request, body, sizeof(body) - 1, 64, &config,
                                 error, sizeof(error)) == 0);
    CHECK(strstr(request.prompt, "<|turn>system\n<|think|>\n") == request.prompt);
    CHECK(strstr(request.prompt, "<|tool>declaration:weather{") != NULL);
    CHECK(strstr(request.prompt, "description:<|\"|>Gets weather<|\"|>") != NULL);
    CHECK(strstr(request.prompt, "type:<|\"|>OBJECT<|\"|>") != NULL);
    CHECK(strstr(request.prompt, "<|turn>user\nWeather?<turn|>\n<|turn>model\n") != NULL);
    gem_chat_request_destroy(&request);
    return 0;
}

static int tool_round_trip_prompt(void)
{
    const char body[] =
        "{\"model\":\"gem\",\"messages\":["
        "{\"role\":\"user\",\"content\":\"List files\"},"
        "{\"role\":\"assistant\",\"content\":null,\"tool_calls\":[{"
        "\"id\":\"call_7\",\"type\":\"function\",\"function\":{"
        "\"name\":\"shell\",\"arguments\":\"{\\\"path\\\":\\\"/tmp\\\"}\"}}]},"
        "{\"role\":\"tool\",\"tool_call_id\":\"call_7\",\"content\":\"a\\nb\"}]}";
    struct gem_chat_request request;
    struct gem_sampler_config config = defaults();
    char error[256];
    CHECK(gem_chat_request_parse(&request, body, sizeof(body) - 1, 64, &config,
                                 error, sizeof(error)) == 0);
    CHECK(strstr(request.prompt,
                 "<|tool_call>call:shell{path:<|\"|>/tmp<|\"|>}<tool_call|>") != NULL);
    CHECK(strstr(request.prompt,
                 "<|tool_response>response:shell{value:<|\"|>a\nb<|\"|>}<tool_response|>") != NULL);
    CHECK(request.prompt[strlen(request.prompt) - 1] == '>');
    gem_chat_request_destroy(&request);
    return 0;
}

static int tool_continuation_prompt(void)
{
    const char body[] =
        "{\"model\":\"gem\",\"messages\":["
        "{\"role\":\"user\",\"content\":\"Where?\"},"
        "{\"role\":\"assistant\",\"content\":null,\"tool_calls\":[{"
        "\"id\":\"c\",\"function\":{\"name\":\"pwd\",\"arguments\":\"{}\"}}]},"
        "{\"role\":\"tool\",\"tool_call_id\":\"c\",\"content\":\"/tmp\"},"
        "{\"role\":\"assistant\",\"content\":\"You are in /tmp.\"},"
        "{\"role\":\"user\",\"content\":\"Thanks\"}]}";
    struct gem_chat_request request;
    struct gem_sampler_config config = defaults();
    char error[256];
    CHECK(gem_chat_request_parse(&request, body, sizeof(body) - 1, 64, &config,
                                 error, sizeof(error)) == 0);
    const char *response = strstr(request.prompt, "<tool_response|>");
    CHECK(response != NULL);
    CHECK(strstr(response,
                 "<tool_response|>You are in /tmp.<turn|>\n<|turn>user\n"
                 "Thanks<turn|>\n<|turn>model\n") == response);
    gem_chat_request_destroy(&request);
    return 0;
}

static int tool_call_parse(void)
{
    const char input[] =
        "call:write{path:<|\"|>/tmp/a,b<|\"|>,count:2,flags:[true,false],"
        "meta:{note:<|\"|>x{y}<|\"|>}}";
    char *name = NULL;
    char *arguments = NULL;
    char error[256];
    CHECK(gem_chat_tool_call_parse(input, sizeof(input) - 1, &name, &arguments,
                                   error, sizeof(error)) == 0);
    CHECK(strcmp(name, "write") == 0);
    CHECK(strcmp(arguments,
                 "{\"path\":\"/tmp/a,b\",\"count\":2,\"flags\":[true,false],"
                 "\"meta\":{\"note\":\"x{y}\"}}") == 0);
    free(arguments);
    free(name);

    static const char bare[] = "call:bash{command:pwd}";
    name = NULL;
    arguments = NULL;
    CHECK(gem_chat_tool_call_parse(bare, sizeof(bare) - 1, &name, &arguments,
                                   error, sizeof(error)) == 0);
    CHECK(strcmp(name, "bash") == 0);
    CHECK(strcmp(arguments, "{\"command\":\"pwd\"}") == 0);
    free(arguments);
    free(name);
    return 0;
}

static int malformed(void)
{
    struct gem_chat_request request;
    struct gem_sampler_config config = defaults();
    char error[256];
    const char body[] = "{\"model\":\"gem\",\"messages\":[]}";
    CHECK(gem_chat_request_parse(&request, body, sizeof(body) - 1, 64, &config,
                                 error, sizeof(error)) != 0);
    CHECK(strstr(error, "malformed") != NULL);
    return 0;
}

int main(void)
{
    CHECK(basic_prompt() == 0);
    CHECK(tool_prompt() == 0);
    CHECK(tool_round_trip_prompt() == 0);
    CHECK(tool_continuation_prompt() == 0);
    CHECK(tool_call_parse() == 0);
    CHECK(malformed() == 0);
    puts("chat tests passed");
    return 0;
}