dwn/src/ai.c

/*
 * DWN - Desktop Window Manager
 * retoor <retoor@molodetz.nl>
 * AI Integration implementation
 */

#include "ai.h"
#include "config.h"
#include "client.h"
#include "workspace.h"
#include "notifications.h"
#include "util.h"
#include "cJSON.h"

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

#define OPENROUTER_URL "https://openrouter.ai/api/v1/chat/completions"

static AIRequest *request_queue = NULL;
static CURLM *curl_multi = NULL;
static AIContext current_context;

typedef struct {
    char *data;
    size_t size;
} ResponseBuffer;


static size_t write_callback(void *contents, size_t size, size_t nmemb, void *userp)
{
    size_t realsize = size * nmemb;
    ResponseBuffer *buf = (ResponseBuffer *)userp;

    char *ptr = realloc(buf->data, buf->size + realsize + 1);
    if (ptr == NULL) {
        return 0;
    }

    buf->data = ptr;
    memcpy(&(buf->data[buf->size]), contents, realsize);
    buf->size += realsize;
    buf->data[buf->size] = '\0';

    return realsize;
}


bool ai_init(void)
{
    if (dwn == NULL || dwn->config == NULL) {
        return false;
    }

    if (dwn->config->openrouter_api_key[0] == '\0') {
        LOG_INFO("AI features disabled (no OPENROUTER_API_KEY)");
        dwn->ai_enabled = false;
        return true;
    }

    curl_global_init(CURL_GLOBAL_DEFAULT);
    curl_multi = curl_multi_init();

    if (curl_multi == NULL) {
        LOG_ERROR("Failed to initialize curl multi handle");
        return false;
    }

    dwn->ai_enabled = true;
    LOG_INFO("AI features enabled");

    return true;
}

void ai_cleanup(void)
{
    while (request_queue != NULL) {
        AIRequest *next = request_queue->next;
        if (request_queue->prompt) free(request_queue->prompt);
        if (request_queue->response) free(request_queue->response);
        free(request_queue);
        request_queue = next;
    }

    if (curl_multi != NULL) {
        curl_multi_cleanup(curl_multi);
        curl_multi = NULL;
    }

    curl_global_cleanup();
}

bool ai_is_available(void)
{
    return dwn != NULL && dwn->ai_enabled;
}


AIRequest *ai_send_request(const char *prompt, void (*callback)(AIRequest *))
{
    if (!ai_is_available() || prompt == NULL) {
        return NULL;
    }

    AIRequest *req = dwn_calloc(1, sizeof(AIRequest));
    req->prompt = dwn_strdup(prompt);
    req->state = AI_STATE_PENDING;
    req->callback = callback;

    char *json_prompt = dwn_malloc(strlen(prompt) * 2 + 256);
    char *escaped_prompt = dwn_malloc(strlen(prompt) * 2 + 1);

    const char *src = prompt;
    char *dst = escaped_prompt;
    while (*src) {
        if (*src == '"' || *src == '\\' || *src == '\n' || *src == '\r' || *src == '\t') {
            *dst++ = '\\';
            if (*src == '\n') *dst++ = 'n';
            else if (*src == '\r') *dst++ = 'r';
            else if (*src == '\t') *dst++ = 't';
            else *dst++ = *src;
        } else {
            *dst++ = *src;
        }
        src++;
    }
    *dst = '\0';

    snprintf(json_prompt, strlen(prompt) * 2 + 256,
             "{\"model\":\"%s\",\"messages\":[{\"role\":\"user\",\"content\":\"%s\"}]}",
             dwn->config->ai_model, escaped_prompt);

    dwn_free(escaped_prompt);

    CURL *easy = curl_easy_init();
    if (easy == NULL) {
        dwn_free(json_prompt);
        dwn_free(req->prompt);
        dwn_free(req);
        return NULL;
    }

    ResponseBuffer *response = dwn_calloc(1, sizeof(ResponseBuffer));

    struct curl_slist *headers = NULL;
    char auth_header[300];
    snprintf(auth_header, sizeof(auth_header), "Authorization: Bearer %s",
             dwn->config->openrouter_api_key);

    headers = curl_slist_append(headers, "Content-Type: application/json");
    headers = curl_slist_append(headers, auth_header);

    curl_easy_setopt(easy, CURLOPT_URL, OPENROUTER_URL);
    curl_easy_setopt(easy, CURLOPT_HTTPHEADER, headers);
    curl_easy_setopt(easy, CURLOPT_POSTFIELDS, json_prompt);
    curl_easy_setopt(easy, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(easy, CURLOPT_WRITEDATA, response);
    curl_easy_setopt(easy, CURLOPT_TIMEOUT, 30L);
    curl_easy_setopt(easy, CURLOPT_PRIVATE, req);
    curl_easy_setopt(easy, CURLOPT_SSL_VERIFYPEER, 1L);
    curl_easy_setopt(easy, CURLOPT_SSL_VERIFYHOST, 2L);

    curl_multi_add_handle(curl_multi, easy);

    req->next = request_queue;
    request_queue = req;

    req->user_data = response;

    LOG_DEBUG("AI request sent: %.50s...", prompt);


    return req;
}

void ai_cancel_request(AIRequest *req)
{
    if (req == NULL) {
        return;
    }

    AIRequest **pp = &request_queue;
    while (*pp != NULL) {
        if (*pp == req) {
            *pp = req->next;
            break;
        }
        pp = &(*pp)->next;
    }

    req->state = AI_STATE_ERROR;

    if (req->prompt) dwn_free(req->prompt);
    if (req->response) dwn_free(req->response);
    if (req->user_data) {
        ResponseBuffer *buf = (ResponseBuffer *)req->user_data;
        if (buf->data) free(buf->data);
        dwn_free(buf);
    }
    dwn_free(req);
}

void ai_process_pending(void)
{
    if (curl_multi == NULL) {
        return;
    }

    int running_handles;
    curl_multi_perform(curl_multi, &running_handles);

    CURLMsg *msg;
    int msgs_left;

    while ((msg = curl_multi_info_read(curl_multi, &msgs_left)) != NULL) {
        if (msg->msg == CURLMSG_DONE) {
            CURL *easy = msg->easy_handle;
            AIRequest *req = NULL;

            curl_easy_getinfo(easy, CURLINFO_PRIVATE, &req);

            if (req != NULL) {
                ResponseBuffer *buf = (ResponseBuffer *)req->user_data;

                if (msg->data.result == CURLE_OK && buf != NULL && buf->data != NULL) {
                    cJSON *root = cJSON_Parse(buf->data);
                    if (root != NULL) {
                        cJSON *choices = cJSON_GetObjectItemCaseSensitive(root, "choices");
                        if (cJSON_IsArray(choices) && cJSON_GetArraySize(choices) > 0) {
                            cJSON *first_choice = cJSON_GetArrayItem(choices, 0);
                            cJSON *message = cJSON_GetObjectItemCaseSensitive(first_choice, "message");
                            if (message != NULL) {
                                cJSON *content = cJSON_GetObjectItemCaseSensitive(message, "content");
                                if (cJSON_IsString(content) && content->valuestring != NULL) {
                                    req->response = dwn_strdup(content->valuestring);
                                    req->state = AI_STATE_COMPLETED;
                                }
                            }
                        }
                        if (req->state != AI_STATE_COMPLETED) {
                            cJSON *error = cJSON_GetObjectItemCaseSensitive(root, "error");
                            if (error != NULL) {
                                cJSON *err_msg = cJSON_GetObjectItemCaseSensitive(error, "message");
                                if (cJSON_IsString(err_msg)) {
                                    req->response = dwn_strdup(err_msg->valuestring);
                                    req->state = AI_STATE_ERROR;
                                    LOG_ERROR("AI API error: %s", err_msg->valuestring);
                                }
                            }
                        }
                        cJSON_Delete(root);
                    }

                    if (req->state != AI_STATE_COMPLETED && req->state != AI_STATE_ERROR) {
                        req->response = dwn_strdup(buf->data);
                        req->state = AI_STATE_COMPLETED;
                        LOG_WARN("Could not parse AI response, returning raw");
                    }
                } else {
                    req->state = AI_STATE_ERROR;
                    LOG_ERROR("AI request failed: %s", curl_easy_strerror(msg->data.result));
                }

                if (req->callback != NULL) {
                    req->callback(req);
                }

                if (buf != NULL) {
                    if (buf->data) free(buf->data);
                    dwn_free(buf);
                }
            }

            curl_multi_remove_handle(curl_multi, easy);
            curl_easy_cleanup(easy);
        }
    }
}


void ai_update_context(void)
{
    memset(&current_context, 0, sizeof(current_context));

    Workspace *ws = workspace_get_current();
    if (ws != NULL && ws->focused != NULL) {
        snprintf(current_context.focused_window, sizeof(current_context.focused_window),
                 "%s", ws->focused->title);
        snprintf(current_context.focused_class, sizeof(current_context.focused_class),
                 "%s", ws->focused->class);
    }

    int offset = 0;
    for (Client *c = dwn->client_list; c != NULL; c = c->next) {
        if (c->workspace == (unsigned int)dwn->current_workspace) {
            int written = snprintf(current_context.workspace_windows + offset,
                                   sizeof(current_context.workspace_windows) - offset,
                                   "%s%s", offset > 0 ? ", " : "", c->title);
            if (written > 0) {
                offset += written;
            }
            current_context.window_count++;
        }
    }
}

const char *ai_analyze_task(void)
{
    const char *class = current_context.focused_class;

    if (strstr(class, "code") || strstr(class, "Code") ||
        strstr(class, "vim") || strstr(class, "emacs")) {
        return "coding";
    }
    if (strstr(class, "firefox") || strstr(class, "chrome") ||
        strstr(class, "Firefox") || strstr(class, "Chrome")) {
        return "browsing";
    }
    if (strstr(class, "slack") || strstr(class, "discord") ||
        strstr(class, "Slack") || strstr(class, "Discord")) {
        return "communication";
    }
    if (strstr(class, "terminal") || strstr(class, "Terminal")) {
        return "terminal";
    }

    return "general";
}

const char *ai_suggest_window(void)
{
    const char *task = ai_analyze_task();

    if (strcmp(task, "coding") == 0) {
        return "Consider opening a terminal for testing";
    }
    if (strcmp(task, "browsing") == 0) {
        return "Documentation or reference material available?";
    }

    return NULL;
}

const char *ai_suggest_app(void)
{
    return NULL;
}


static void ai_command_response_callback(AIRequest *req)
{
    if (req == NULL) {
        return;
    }

    if (req->state == AI_STATE_COMPLETED && req->response != NULL) {
        char *run_cmd = strstr(req->response, "[RUN:");
        if (run_cmd == NULL) {
            run_cmd = strstr(req->response, "[EXEC:");
        }

        if (run_cmd != NULL) {
            char *cmd_start = strchr(run_cmd, ':');
            if (cmd_start != NULL) {
                cmd_start++;
                while (*cmd_start == ' ') cmd_start++;

                char *cmd_end = strchr(cmd_start, ']');
                if (cmd_end != NULL) {
                    size_t cmd_len = cmd_end - cmd_start;
                    char *cmd = dwn_malloc(cmd_len + 1);
                    strncpy(cmd, cmd_start, cmd_len);
                    cmd[cmd_len] = '\0';

                    while (cmd_len > 0 && cmd[cmd_len - 1] == ' ') {
                        cmd[--cmd_len] = '\0';
                    }

                    LOG_INFO("AI executing command: %s", cmd);
                    notification_show("DWN AI", "Running", cmd, NULL, 2000);
                    spawn_async(cmd);
                    dwn_free(cmd);
                }
            }
        } else {
            notification_show("DWN AI", "Response", req->response, NULL, 8000);
        }
    } else {
        notification_show("DWN AI", "Error", "Failed to get AI response", NULL, 3000);
    }

}

void ai_show_command_palette(void)
{
    if (!ai_is_available()) {
        notification_show("DWN", "AI Unavailable",
                          "Set OPENROUTER_API_KEY to enable AI features",
                          NULL, 3000);
        return;
    }

    char *input = NULL;

    if (spawn("command -v dmenu >/dev/null 2>&1") == 0) {
        input = spawn_capture("echo '' | dmenu -p 'Ask AI:'");
    } else if (spawn("command -v rofi >/dev/null 2>&1") == 0) {
        input = spawn_capture("rofi -dmenu -p 'Ask AI:'");
    } else {
        notification_show("DWN AI", "Missing Dependency",
                          "Install dmenu or rofi for AI command palette:\n"
                          "sudo apt install dmenu",
                          NULL, 5000);
        return;
    }

    if (input == NULL || input[0] == '\0') {
        if (input != NULL) {
            dwn_free(input);
        }
        LOG_DEBUG("AI command palette cancelled");
        return;
    }

    LOG_DEBUG("AI command palette input: %s", input);

    notification_show("DWN AI", "Processing...", input, NULL, 2000);

    ai_update_context();
    const char *task = ai_analyze_task();

    char prompt[2048];
    snprintf(prompt, sizeof(prompt),
             "You are an AI assistant integrated into a Linux window manager called DWN. "
             "You can execute shell commands for the user.\n\n"
             "IMPORTANT: When the user asks you to run, open, launch, or start an application, "
             "respond with the command in this exact format: [RUN: command]\n"
             "Examples:\n"
             "- User: 'open chrome' -> [RUN: google-chrome]\n"
             "- User: 'run firefox' -> [RUN: firefox]\n"
             "- User: 'open file manager' -> [RUN: thunar]\n"
             "- User: 'launch terminal' -> [RUN: xfce4-terminal]\n"
             "- User: 'open vs code' -> [RUN: code]\n\n"
             "For questions or non-command requests, respond briefly (1-2 sentences) without the [RUN:] format.\n\n"
             "User's current task: %s\n"
             "User's request: %s",
             task, input);

    dwn_free(input);

    ai_send_request(prompt, ai_command_response_callback);
}

void ai_execute_command(const char *command)
{
    if (!ai_is_available() || command == NULL) {
        return;
    }

    LOG_DEBUG("AI executing command: %s", command);

    char prompt[512];
    snprintf(prompt, sizeof(prompt),
             "User command: %s\nCurrent task: %s\nRespond with a single action to take.",
             command, ai_analyze_task());

    ai_send_request(prompt, NULL);
}


void ai_auto_organize_workspace(void)
{
    LOG_DEBUG("AI auto-organize (placeholder)");
}

void ai_suggest_layout(void)
{
    LOG_DEBUG("AI layout suggestion (placeholder)");
}

void ai_analyze_workflow(void)
{
    LOG_DEBUG("AI workflow analysis (placeholder)");
}


bool ai_should_show_notification(const char *app, const char *summary)
{
    (void)app;
    (void)summary;
    return true;
}

int ai_notification_priority(const char *app, const char *summary)
{
    if (strstr(summary, "urgent") || strstr(summary, "Urgent") ||
        strstr(summary, "error") || strstr(summary, "Error")) {
        return 3;
    }
    if (strstr(app, "slack") || strstr(app, "Slack") ||
        strstr(app, "discord") || strstr(app, "Discord")) {
        return 2;
    }
    return 1;
}


void ai_monitor_performance(void)
{
    LOG_DEBUG("AI performance monitoring (placeholder)");
}

const char *ai_performance_suggestion(void)
{
    return NULL;
}


#define EXA_API_URL "https://api.exa.ai/search"

static ExaRequest *exa_queue = NULL;

bool exa_is_available(void)
{
    return dwn != NULL && dwn->config != NULL &&
           dwn->config->exa_api_key[0] != '\0';
}

static void exa_parse_response(ExaRequest *req, const char *json)
{
    if (req == NULL || json == NULL) {
        return;
    }

    req->result_count = 0;

    cJSON *root = cJSON_Parse(json);
    if (root == NULL) {
        LOG_WARN("Failed to parse Exa response JSON");
        return;
    }

    cJSON *results = cJSON_GetObjectItemCaseSensitive(root, "results");
    if (!cJSON_IsArray(results)) {
        cJSON_Delete(root);
        return;
    }

    int array_size = cJSON_GetArraySize(results);
    for (int i = 0; i < array_size && req->result_count < 10; i++) {
        cJSON *item = cJSON_GetArrayItem(results, i);
        if (!cJSON_IsObject(item)) continue;

        ExaSearchResult *res = &req->results[req->result_count];

        cJSON *title = cJSON_GetObjectItemCaseSensitive(item, "title");
        if (cJSON_IsString(title) && title->valuestring != NULL) {
            strncpy(res->title, title->valuestring, sizeof(res->title) - 1);
            res->title[sizeof(res->title) - 1] = '\0';
        }

        cJSON *url = cJSON_GetObjectItemCaseSensitive(item, "url");
        if (cJSON_IsString(url) && url->valuestring != NULL) {
            strncpy(res->url, url->valuestring, sizeof(res->url) - 1);
            res->url[sizeof(res->url) - 1] = '\0';
        }

        cJSON *text = cJSON_GetObjectItemCaseSensitive(item, "text");
        if (cJSON_IsString(text) && text->valuestring != NULL) {
            strncpy(res->snippet, text->valuestring, sizeof(res->snippet) - 1);
            res->snippet[sizeof(res->snippet) - 1] = '\0';
        }

        req->result_count++;
    }

    cJSON_Delete(root);
    LOG_DEBUG("Parsed %d Exa results", req->result_count);
}

ExaRequest *exa_search(const char *query, void (*callback)(ExaRequest *))
{
    if (!exa_is_available() || query == NULL) {
        return NULL;
    }

    ExaRequest *req = dwn_calloc(1, sizeof(ExaRequest));
    req->query = dwn_strdup(query);
    req->state = AI_STATE_PENDING;
    req->callback = callback;
    req->result_count = 0;

    char *json_query = dwn_malloc(strlen(query) * 2 + 256);
    char *escaped = dwn_malloc(strlen(query) * 2 + 1);

    const char *src = query;
    char *dst = escaped;
    while (*src) {
        if (*src == '"' || *src == '\\') {
            *dst++ = '\\';
        }
        *dst++ = *src++;
    }
    *dst = '\0';

    snprintf(json_query, strlen(query) * 2 + 256,
             "{\"query\":\"%s\",\"type\":\"auto\",\"numResults\":10,\"contents\":{\"text\":true}}",
             escaped);
    dwn_free(escaped);

    CURL *easy = curl_easy_init();
    if (easy == NULL) {
        dwn_free(json_query);
        dwn_free(req->query);
        dwn_free(req);
        return NULL;
    }

    ResponseBuffer *response = dwn_calloc(1, sizeof(ResponseBuffer));

    struct curl_slist *headers = NULL;
    char api_header[300];
    snprintf(api_header, sizeof(api_header), "x-api-key: %s",
             dwn->config->exa_api_key);

    headers = curl_slist_append(headers, "Content-Type: application/json");
    headers = curl_slist_append(headers, api_header);

    curl_easy_setopt(easy, CURLOPT_URL, EXA_API_URL);
    curl_easy_setopt(easy, CURLOPT_HTTPHEADER, headers);
    curl_easy_setopt(easy, CURLOPT_POSTFIELDS, json_query);
    curl_easy_setopt(easy, CURLOPT_WRITEFUNCTION, write_callback);
    curl_easy_setopt(easy, CURLOPT_WRITEDATA, response);
    curl_easy_setopt(easy, CURLOPT_TIMEOUT, 15L);
    curl_easy_setopt(easy, CURLOPT_PRIVATE, req);
    curl_easy_setopt(easy, CURLOPT_SSL_VERIFYPEER, 1L);
    curl_easy_setopt(easy, CURLOPT_SSL_VERIFYHOST, 2L);

    if (curl_multi == NULL) {
        curl_multi = curl_multi_init();
    }
    curl_multi_add_handle(curl_multi, easy);

    req->next = exa_queue;
    exa_queue = req;
    req->user_data = response;

    LOG_DEBUG("Exa search sent: %s", query);

    return req;
}

void exa_process_pending(void)
{
    if (curl_multi == NULL || exa_queue == NULL) {
        return;
    }

    int running_handles;
    curl_multi_perform(curl_multi, &running_handles);

    CURLMsg *msg;
    int msgs_left;

    while ((msg = curl_multi_info_read(curl_multi, &msgs_left)) != NULL) {
        if (msg->msg == CURLMSG_DONE) {
            CURL *easy = msg->easy_handle;
            ExaRequest *req = NULL;

            curl_easy_getinfo(easy, CURLINFO_PRIVATE, &req);

            bool is_exa = false;
            for (ExaRequest *r = exa_queue; r != NULL; r = r->next) {
                if (r == req) {
                    is_exa = true;
                    break;
                }
            }

            if (is_exa && req != NULL) {
                ResponseBuffer *buf = (ResponseBuffer *)req->user_data;

                if (msg->data.result == CURLE_OK && buf != NULL && buf->data != NULL) {
                    exa_parse_response(req, buf->data);
                    req->state = AI_STATE_COMPLETED;
                } else {
                    req->state = AI_STATE_ERROR;
                    LOG_ERROR("Exa request failed: %s", curl_easy_strerror(msg->data.result));
                }

                if (req->callback != NULL) {
                    req->callback(req);
                }

                if (buf != NULL) {
                    if (buf->data) free(buf->data);
                    dwn_free(buf);
                }

                ExaRequest **pp = &exa_queue;
                while (*pp != NULL) {
                    if (*pp == req) {
                        *pp = req->next;
                        break;
                    }
                    pp = &(*pp)->next;
                }
            }

            curl_multi_remove_handle(curl_multi, easy);
            curl_easy_cleanup(easy);
        }
    }
}

static void exa_launcher_callback(ExaRequest *req)
{
    if (req == NULL || req->state != AI_STATE_COMPLETED) {
        notification_show("Exa Search", "Error", "Search failed", NULL, 3000);
        return;
    }

    if (req->result_count == 0) {
        notification_show("Exa Search", "No Results", req->query, NULL, 3000);
        return;
    }

    size_t choices_size = req->result_count * 300;
    char *choices = dwn_malloc(choices_size);
    size_t offset = 0;
    choices[0] = '\0';

    for (int i = 0; i < req->result_count; i++) {
        int written = snprintf(choices + offset, choices_size - offset,
                               "%s%s", offset > 0 ? "\n" : "", req->results[i].title);
        if (written > 0 && (size_t)written < choices_size - offset) {
            offset += written;
        }
    }

    char *escaped_choices = shell_escape(choices);
    char *cmd = dwn_malloc(strlen(escaped_choices) + 64);
    snprintf(cmd, strlen(escaped_choices) + 64, "echo %s | dmenu -l 10 -p 'Results:'", escaped_choices);

    char *selected = spawn_capture(cmd);
    dwn_free(cmd);
    dwn_free(escaped_choices);

    if (selected != NULL && selected[0] != '\0') {
        for (int i = 0; i < req->result_count; i++) {
            if (strncmp(selected, req->results[i].title, strlen(req->results[i].title)) == 0) {
                char *escaped_url = shell_escape(req->results[i].url);
                char *open_cmd = dwn_malloc(strlen(escaped_url) + 32);
                snprintf(open_cmd, strlen(escaped_url) + 32, "xdg-open %s &", escaped_url);
                spawn_async(open_cmd);
                dwn_free(open_cmd);
                dwn_free(escaped_url);
                break;
            }
        }
        dwn_free(selected);
    }

    dwn_free(choices);
    if (req->query) dwn_free(req->query);
    dwn_free(req);
}

void exa_show_app_launcher(void)
{
    if (!exa_is_available()) {
        notification_show("Exa", "Unavailable",
                          "Set EXA_API_KEY in config to enable semantic search",
                          NULL, 3000);
        return;
    }

    char *query = NULL;

    if (spawn("command -v dmenu >/dev/null 2>&1") == 0) {
        query = spawn_capture("echo '' | dmenu -p 'Exa Search:'");
    } else if (spawn("command -v rofi >/dev/null 2>&1") == 0) {
        query = spawn_capture("rofi -dmenu -p 'Exa Search:'");
    } else {
        notification_show("Exa", "Missing Dependency",
                          "Install dmenu or rofi", NULL, 3000);
        return;
    }

    if (query == NULL || query[0] == '\0') {
        if (query != NULL) dwn_free(query);
        return;
    }

    query[strcspn(query, "\n")] = '\0';

    notification_show("Exa", "Searching...", query, NULL, 2000);
    exa_search(query, exa_launcher_callback);
}