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

#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>
#include <pthread.h>
#include <stdatomic.h>

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

/* Request queue management */
typedef struct {
    AIRequest *head;
    AIRequest *tail;
    size_t count;
} AIRequestQueue;

static AIRequestQueue ai_queue = {NULL, NULL, 0};
static ExaRequest *exa_queue_head = NULL;
static ExaRequest *exa_queue_tail = NULL;
static size_t exa_queue_count = 0;

static CURLM *curl_multi = NULL;
static AIContext current_context;

/* Thread synchronization - unified mutex hierarchy to prevent deadlocks */
static pthread_mutex_t ai_mutex = PTHREAD_MUTEX_INITIALIZER;
static atomic_int ai_shutting_down = 0;

/* Thread-safe buffer for JSON escaping */
typedef struct {
    char *data;
    size_t size;
    size_t capacity;
} ResponseBuffer;

/* Safe memory reallocation with overflow checking */
static void *safe_realloc(void *ptr, size_t new_size, size_t old_size)
{
    (void)old_size; /* Used for debugging/logging if needed */
    if (new_size == 0) {
        dwn_free(ptr);
        return NULL;
    }
    /* Check for overflow in multiplication scenarios */
    if (new_size > (size_t)-1 / 2) {
        LOG_ERROR("Allocation size overflow: %zu", new_size);
        return NULL;
    }
    return dwn_realloc(ptr, new_size);
}

/* JSON string escaping with proper buffer management */
static char *escape_json_string(const char *input, size_t *out_len)
{
    if (input == NULL) {
        if (out_len) *out_len = 0;
        return NULL;
    }
    
    size_t input_len = strlen(input);
    size_t max_escaped_len = input_len * 2 + 1;  /* Worst case: every char escaped */
    
    char *escaped = dwn_malloc(max_escaped_len);
    if (escaped == NULL) {
        if (out_len) *out_len = 0;
        return NULL;
    }
    
    size_t j = 0;
    for (size_t i = 0; i < input_len && j < max_escaped_len - 1; i++) {
        unsigned char c = (unsigned char)input[i];
        switch (c) {
            case '"':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = '"'; }
                break;
            case '\\':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = '\\'; }
                break;
            case '\n':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = 'n'; }
                break;
            case '\r':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = 'r'; }
                break;
            case '\t':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = 't'; }
                break;
            case '\b':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = 'b'; }
                break;
            case '\f':
                if (j + 2 < max_escaped_len) { escaped[j++] = '\\'; escaped[j++] = 'f'; }
                break;
            default:
                if (c < 0x20) {
                    /* Control characters - use unicode escape if room */
                    if (j + 6 < max_escaped_len) {
                        snprintf(&escaped[j], 7, "\\u%04x", c);
                        j += 6;
                    }
                } else {
                    escaped[j++] = (char)c;
                }
                break;
        }
    }
    escaped[j] = '\0';
    
    if (out_len) *out_len = j;
    return escaped;
}

/* Calculate required buffer size for JSON payload safely */
static size_t calculate_json_buffer_size(const char *model, const char *escaped_prompt)
{
    size_t model_len = (model != NULL) ? strlen(model) : 0;
    size_t prompt_len = (escaped_prompt != NULL) ? strlen(escaped_prompt) : 0;
    
    /* Base JSON structure size + model + prompt + safety margin */
    size_t base_size = 256;  /* {"model":"","messages":[{"role":"user","content":""}]} */
    size_t total = base_size + model_len + prompt_len;
    
    /* Overflow check */
    if (total < base_size || total < model_len || total < prompt_len) {
        return 0;  /* Overflow detected */
    }
    
    return total;
}

/* CURL write callback with dynamic buffer */
static size_t write_callback(void *contents, size_t size, size_t nmemb, void *userp)
{
    size_t realsize = size * nmemb;
    ResponseBuffer *buf = (ResponseBuffer *)userp;
    
    if (buf == NULL) return 0;
    
    /* Overflow check for addition */
    if (realsize > (size_t)-1 - buf->size - 1) {
        LOG_ERROR("Response buffer size overflow");
        return 0;
    }
    
    size_t new_size = buf->size + realsize + 1;
    char *ptr = safe_realloc(buf->data, new_size, buf->capacity);
    if (ptr == NULL) {
        LOG_ERROR("Failed to grow response buffer");
        return 0;
    }
    
    buf->data = ptr;
    memcpy(buf->data + buf->size, contents, realsize);
    buf->size += realsize;
    buf->data[buf->size] = '\0';
    buf->capacity = new_size;
    
    return realsize;
}

/* Remove request from queue - must hold ai_mutex */
static void ai_queue_remove(AIRequest *req)
{
    if (req == NULL) return;
    
    AIRequest **current = &ai_queue.head;
    while (*current != NULL) {
        if (*current == req) {
            *current = req->next;
            if (ai_queue.tail == req) {
                ai_queue.tail = (req->next == NULL) ? 
                    (ai_queue.head != NULL ? ai_queue.head : NULL) : NULL;
                /* Find new tail if needed */
                if (ai_queue.head != NULL && ai_queue.tail == NULL) {
                    AIRequest *t = ai_queue.head;
                    while (t->next != NULL) t = t->next;
                    ai_queue.tail = t;
                }
            }
            ai_queue.count--;
            req->next = NULL;
            break;
        }
        current = &(*current)->next;
    }
}

/* Free AI request resources safely */
static void ai_request_free(AIRequest *req)
{
    if (req == NULL) return;
    
    /* Free headers first */
    if (req->headers != NULL) {
        curl_slist_free_all(req->headers);
        req->headers = NULL;
    }
    
    /* Clean up CURL handle - should already be removed from multi */
    if (req->curl_handle != NULL) {
        curl_easy_cleanup(req->curl_handle);
        req->curl_handle = NULL;
    }
    
    /* Free user data (ResponseBuffer) */
    if (req->user_data != NULL) {
        ResponseBuffer *buf = (ResponseBuffer *)req->user_data;
        if (buf->data != NULL) {
            dwn_free(buf->data);
            buf->data = NULL;
        }
        dwn_free(buf);
        req->user_data = NULL;
    }
    
    /* Free strings */
    if (req->prompt != NULL) {
        dwn_free(req->prompt);
        req->prompt = NULL;
    }
    if (req->response != NULL) {
        dwn_free(req->response);
        req->response = NULL;
    }
    
    dwn_free(req);
}

bool ai_init(void)
{
    if (dwn == NULL || dwn->config == NULL) {
        return false;
    }
    
    atomic_store(&ai_shutting_down, 0);
    
    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);
    
    pthread_mutex_lock(&ai_mutex);
    curl_multi = curl_multi_init();
    if (curl_multi == NULL) {
        pthread_mutex_unlock(&ai_mutex);
        LOG_ERROR("Failed to initialize curl multi handle");
        return false;
    }
    pthread_mutex_unlock(&ai_mutex);
    
    dwn->ai_enabled = true;
    LOG_INFO("AI features enabled");
    
    return true;
}

void ai_cleanup(void)
{
    atomic_store(&ai_shutting_down, 1);
    
    pthread_mutex_lock(&ai_mutex);
    
    /* Cancel and free all pending AI requests */
    AIRequest *req = ai_queue.head;
    while (req != NULL) {
        AIRequest *next = req->next;
        
        /* Remove from curl multi if present */
        if (curl_multi != NULL && req->curl_handle != NULL) {
            curl_multi_remove_handle(curl_multi, req->curl_handle);
        }
        
        ai_request_free(req);
        req = next;
    }
    ai_queue.head = NULL;
    ai_queue.tail = NULL;
    ai_queue.count = 0;
    
    /* Cancel and free all pending Exa requests */
    ExaRequest *exa_req = exa_queue_head;
    while (exa_req != NULL) {
        ExaRequest *next = exa_req->next;
        
        if (curl_multi != NULL && exa_req->curl_handle != NULL) {
            curl_multi_remove_handle(curl_multi, exa_req->curl_handle);
        }
        
        if (exa_req->headers != NULL) {
            curl_slist_free_all(exa_req->headers);
        }
        if (exa_req->curl_handle != NULL) {
            curl_easy_cleanup(exa_req->curl_handle);
        }
        if (exa_req->user_data != NULL) {
            ResponseBuffer *buf = (ResponseBuffer *)exa_req->user_data;
            if (buf->data != NULL) dwn_free(buf->data);
            dwn_free(buf);
        }
        if (exa_req->query != NULL) dwn_free(exa_req->query);
        dwn_free(exa_req);
        
        exa_req = next;
    }
    exa_queue_head = NULL;
    exa_queue_tail = NULL;
    exa_queue_count = 0;
    
    if (curl_multi != NULL) {
        curl_multi_cleanup(curl_multi);
        curl_multi = NULL;
    }
    
    pthread_mutex_unlock(&ai_mutex);
    
    curl_global_cleanup();
}

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

AIRequest *ai_send_request(const char *prompt, void (*callback)(AIRequest *))
{
    if (!ai_is_available() || prompt == NULL) {
        return NULL;
    }
    
    /* Escape the prompt for JSON */
    size_t escaped_len = 0;
    char *escaped_prompt = escape_json_string(prompt, &escaped_len);
    if (escaped_prompt == NULL) {
        LOG_ERROR("Failed to escape prompt");
        return NULL;
    }
    
    /* Calculate safe buffer size including model name */
    size_t json_size = calculate_json_buffer_size(dwn->config->ai_model, escaped_prompt);
    if (json_size == 0) {
        LOG_ERROR("JSON buffer size calculation overflow");
        dwn_free(escaped_prompt);
        return NULL;
    }
    
    char *json_payload = dwn_malloc(json_size);
    if (json_payload == NULL) {
        LOG_ERROR("Failed to allocate JSON payload buffer");
        dwn_free(escaped_prompt);
        return NULL;
    }
    
    /* Build JSON payload safely */
    int written = snprintf(json_payload, json_size,
             "{\"model\":\"%s\",\"messages\":[{\"role\":\"user\",\"content\":\"%s\"}]}",
             dwn->config->ai_model, escaped_prompt);
    
    dwn_free(escaped_prompt);
    
    if (written < 0 || (size_t)written >= json_size) {
        LOG_ERROR("JSON payload truncation detected");
        dwn_free(json_payload);
        return NULL;
    }
    
    /* Allocate request structure */
    AIRequest *req = dwn_calloc(1, sizeof(AIRequest));
    if (req == NULL) {
        dwn_free(json_payload);
        return NULL;
    }
    
    req->prompt = dwn_strdup(prompt);
    req->state = AI_STATE_PENDING;
    req->callback = callback;
    
    /* Allocate response buffer */
    ResponseBuffer *response = dwn_calloc(1, sizeof(ResponseBuffer));
    if (response == NULL) {
        dwn_free(json_payload);
        dwn_free(req->prompt);
        dwn_free(req);
        return NULL;
    }
    req->user_data = response;
    
    /* Create CURL handle */
    CURL *easy = curl_easy_init();
    if (easy == NULL) {
        LOG_ERROR("Failed to initialize CURL easy handle");
        dwn_free(json_payload);
        ai_request_free(req);
        return NULL;
    }
    req->curl_handle = easy;
    
    /* Build headers - store in request for later cleanup */
    struct curl_slist *headers = NULL;
    headers = curl_slist_append(headers, "Content-Type: application/json");
    
    char auth_header[300];
    snprintf(auth_header, sizeof(auth_header), "Authorization: Bearer %s",
             dwn->config->openrouter_api_key);
    headers = curl_slist_append(headers, auth_header);
    
    req->headers = headers;
    
    /* Configure CURL */
    curl_easy_setopt(easy, CURLOPT_URL, OPENROUTER_URL);
    curl_easy_setopt(easy, CURLOPT_HTTPHEADER, headers);
    curl_easy_setopt(easy, CURLOPT_COPYPOSTFIELDS, json_payload);
    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_easy_setopt(easy, CURLOPT_NOSIGNAL, 1L);  /* Prevent SIGALRM issues */
    
    /* Add to queue and multi handle under lock */
    pthread_mutex_lock(&ai_mutex);
    
    if (curl_multi != NULL) {
        CURLMcode mc = curl_multi_add_handle(curl_multi, easy);
        if (mc != CURLM_OK) {
            pthread_mutex_unlock(&ai_mutex);
            LOG_ERROR("Failed to add handle to curl multi: %s", curl_multi_strerror(mc));
            dwn_free(json_payload);
            ai_request_free(req);
            return NULL;
        }
    } else {
        pthread_mutex_unlock(&ai_mutex);
        LOG_ERROR("CURL multi not initialized");
        dwn_free(json_payload);
        ai_request_free(req);
        return NULL;
    }
    
    /* Add to queue tail for FIFO processing */
    req->next = NULL;
    if (ai_queue.tail != NULL) {
        ai_queue.tail->next = req;
    } else {
        ai_queue.head = req;
    }
    ai_queue.tail = req;
    ai_queue.count++;
    
    pthread_mutex_unlock(&ai_mutex);
    
    /* json_payload is copied by CURL, we can free our copy */
    dwn_free(json_payload);
    
    LOG_DEBUG("AI request sent: %.50s...", prompt);
    
    return req;
}

void ai_cancel_request(AIRequest *req)
{
    if (req == NULL) return;
    
    pthread_mutex_lock(&ai_mutex);
    
    /* Remove from curl multi if still active */
    if (curl_multi != NULL && req->curl_handle != NULL) {
        curl_multi_remove_handle(curl_multi, req->curl_handle);
    }
    
    /* Remove from queue */
    ai_queue_remove(req);
    
    pthread_mutex_unlock(&ai_mutex);
    
    /* Mark as cancelled and free */
    req->state = AI_STATE_ERROR;
    ai_request_free(req);
    
    LOG_DEBUG("AI request cancelled");
}

void ai_process_pending(void)
{
    if (atomic_load(&ai_shutting_down)) {
        return;
    }
    
    pthread_mutex_lock(&ai_mutex);
    
    if (curl_multi == NULL) {
        pthread_mutex_unlock(&ai_mutex);
        return;
    }
    
    /* Perform CURL operations */
    int running_handles;
    curl_multi_perform(curl_multi, &running_handles);
    
    /* Process completed transfers */
    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) {
                    /* Parse JSON response */
                    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) && err_msg->valuestring != NULL) {
                                    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));
                }
                
                /* Remove from curl multi */
                curl_multi_remove_handle(curl_multi, easy);
                req->curl_handle = NULL;  /* Mark as removed */
                
                /* Invoke callback */
                if (req->callback != NULL) {
                    /* Unlock for callback to prevent deadlock */
                    pthread_mutex_unlock(&ai_mutex);
                    req->callback(req);
                    pthread_mutex_lock(&ai_mutex);
                }
                
                /* Remove from queue and free */
                ai_queue_remove(req);
                ai_request_free(req);
            } else {
                /* Unknown request - just cleanup */
                curl_multi_remove_handle(curl_multi, easy);
                curl_easy_cleanup(easy);
            }
        }
    }
    
    pthread_mutex_unlock(&ai_mutex);
}

/* Context update and analysis functions */
void ai_update_context(void)
{
    memset(&current_context, 0, sizeof(current_context));
    
    Workspace *ws = workspace_get_current();
    if (ws != NULL && ws->focused != NULL) {
        strncpy(current_context.focused_window, ws->focused->title, 
                sizeof(current_context.focused_window) - 1);
        current_context.focused_window[sizeof(current_context.focused_window) - 1] = '\0';
        strncpy(current_context.focused_class, ws->focused->class,
                sizeof(current_context.focused_class) - 1);
        current_context.focused_class[sizeof(current_context.focused_class) - 1] = '\0';
    }
    
    size_t 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 && (size_t)written < sizeof(current_context.workspace_windows) - offset) {
                offset += (size_t)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;
}

/* Command palette callback */
static void ai_command_response_callback(AIRequest *req)
{
    if (req == NULL) return;
    
    if (req->state == AI_STATE_COMPLETED && req->response != NULL) {
        /* Parse and execute WM commands */
        char *wm_cmd = strstr(req->response, "[WM_CMD:");
        if (wm_cmd != NULL) {
            char *cmd_start = strchr(wm_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 = (size_t)(cmd_end - cmd_start);
                    if (cmd_len > 0 && cmd_len < 4096) {  /* Sanity check */
                        char *cmd_json = dwn_malloc(cmd_len + 1);
                        if (cmd_json != NULL) {
                            memcpy(cmd_json, cmd_start, cmd_len);
                            cmd_json[cmd_len] = '\0';
                            
                            LOG_INFO("AI executing WM command: %s", cmd_json);
                            notification_show("DWN AI", "Executing WM Command", cmd_json, NULL, 2000);
                            
                            extern void api_handle_json_command(const char *json_str);
                            api_handle_json_command(cmd_json);
                            
                            dwn_free(cmd_json);
                        }
                    }
                }
            }
        }
        
        /* Parse and execute shell commands */
        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 = (size_t)(cmd_end - cmd_start);
                    if (cmd_len > 0 && cmd_len < 4096) {  /* Sanity check */
                        char *cmd = dwn_malloc(cmd_len + 1);
                        if (cmd != NULL) {
                            memcpy(cmd, cmd_start, cmd_len);
                            cmd[cmd_len] = '\0';
                            
                            /* Trim trailing whitespace */
                            while (cmd_len > 0 && (cmd[cmd_len - 1] == ' ' || cmd[cmd_len - 1] == '\t')) {
                                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();
    
    /* Build prompt with bounds checking */
    char prompt[4096];
    const char *prompt_template =
        "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"
        "When the user asks for window management tasks (switching workspaces, focusing windows), "
        "respond with the JSON command in this exact format: [WM_CMD: {\"command\": \"...\", ...}]\n"
        "Available WM Commands:\n"
        "- switch_workspace (e.g. [WM_CMD: {\"command\": \"switch_workspace\", \"workspace\": 2}])\n"
        "- focus_client (e.g. [WM_CMD: {\"command\": \"focus_client\", \"window\": <window_id>}])\n"
        "Examples:\n"
        "- User: 'open chrome' -> [RUN: google-chrome]\n"
        "- User: 'switch to workspace 3' -> [WM_CMD: {\"command\": \"switch_workspace\", \"workspace\": 2}]\n"
        "- User: 'go to desktop 1' -> [WM_CMD: {\"command\": \"switch_workspace\", \"workspace\": 0}]\n"
        "For questions or non-command requests, respond briefly (1-2 sentences) without the [RUN:] or [WM_CMD:] format.\n\n"
        "User's current task: %s\n"
        "User's request: %s";
    
    int written = snprintf(prompt, sizeof(prompt), prompt_template, task, input);
    if (written < 0 || (size_t)written >= sizeof(prompt)) {
        LOG_ERROR("Prompt too long, truncating");
        prompt[sizeof(prompt) - 1] = '\0';
    }
    
    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);
}

/* Placeholder functions */
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;
}

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

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

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;
    }
    
    /* Escape query for JSON */
    size_t escaped_len = 0;
    char *escaped = escape_json_string(query, &escaped_len);
    if (escaped == NULL) {
        return NULL;
    }
    
    /* Calculate buffer size */
    size_t json_size = escaped_len + 256;
    if (json_size < escaped_len) {  /* Overflow check */
        dwn_free(escaped);
        return NULL;
    }
    
    char *json_query = dwn_malloc(json_size);
    if (json_query == NULL) {
        dwn_free(escaped);
        return NULL;
    }
    
    int written = snprintf(json_query, json_size,
             "{\"query\":\"%s\",\"type\":\"auto\",\"numResults\":10,\"contents\":{\"text\":true}}",
             escaped);
    dwn_free(escaped);
    
    if (written < 0 || (size_t)written >= json_size) {
        dwn_free(json_query);
        return NULL;
    }
    
    ExaRequest *req = dwn_calloc(1, sizeof(ExaRequest));
    if (req == NULL) {
        dwn_free(json_query);
        return NULL;
    }
    
    req->query = dwn_strdup(query);
    req->state = AI_STATE_PENDING;
    req->callback = callback;
    req->result_count = 0;
    
    ResponseBuffer *response = dwn_calloc(1, sizeof(ResponseBuffer));
    if (response == NULL) {
        dwn_free(json_query);
        dwn_free(req->query);
        dwn_free(req);
        return NULL;
    }
    req->user_data = response;
    
    CURL *easy = curl_easy_init();
    if (easy == NULL) {
        dwn_free(json_query);
        dwn_free(req->query);
        dwn_free(response);
        dwn_free(req);
        return NULL;
    }
    req->curl_handle = easy;
    
    /* Build headers */
    struct curl_slist *headers = NULL;
    headers = curl_slist_append(headers, "Content-Type: application/json");
    
    char api_header[300];
    snprintf(api_header, sizeof(api_header), "x-api-key: %s",
             dwn->config->exa_api_key);
    headers = curl_slist_append(headers, api_header);
    
    req->headers = headers;
    
    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);
    curl_easy_setopt(easy, CURLOPT_NOSIGNAL, 1L);
    
    pthread_mutex_lock(&ai_mutex);
    
    if (curl_multi == NULL) {
        curl_multi = curl_multi_init();
    }
    
    if (curl_multi != NULL) {
        curl_multi_add_handle(curl_multi, easy);
    }
    
    /* Add to queue */
    req->next = NULL;
    if (exa_queue_tail != NULL) {
        exa_queue_tail->next = req;
    } else {
        exa_queue_head = req;
    }
    exa_queue_tail = req;
    exa_queue_count++;
    
    pthread_mutex_unlock(&ai_mutex);
    
    dwn_free(json_query);
    
    LOG_DEBUG("Exa search sent: %s", query);
    
    return req;
}

void exa_process_pending(void)
{
    if (atomic_load(&ai_shutting_down)) {
        return;
    }
    
    pthread_mutex_lock(&ai_mutex);
    
    if (curl_multi == NULL) {
        pthread_mutex_unlock(&ai_mutex);
        return;
    }
    
    if (exa_queue_head == NULL) {
        pthread_mutex_unlock(&ai_mutex);
        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);
            
            if (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));
                }
                
                curl_multi_remove_handle(curl_multi, easy);
                req->curl_handle = NULL;
                
                /* Invoke callback */
                if (req->callback != NULL) {
                    pthread_mutex_unlock(&ai_mutex);
                    req->callback(req);
                    pthread_mutex_lock(&ai_mutex);
                }
                
                /* Remove from queue and free */
                ExaRequest **pp = &exa_queue_head;
                while (*pp != NULL) {
                    if (*pp == req) {
                        *pp = req->next;
                        if (exa_queue_tail == req) {
                            exa_queue_tail = (pp == &exa_queue_head) ? NULL : exa_queue_head;
                            if (exa_queue_head != NULL && exa_queue_tail == NULL) {
                                ExaRequest *t = exa_queue_head;
                                while (t->next != NULL) t = t->next;
                                exa_queue_tail = t;
                            }
                        }
                        exa_queue_count--;
                        break;
                    }
                    pp = &(*pp)->next;
                }
                
                /* Cleanup */
                if (req->headers != NULL) curl_slist_free_all(req->headers);
                if (easy != NULL) curl_easy_cleanup(easy);
                if (req->user_data != NULL) {
                    ResponseBuffer *b = (ResponseBuffer *)req->user_data;
                    if (b->data != NULL) dwn_free(b->data);
                    dwn_free(b);
                }
                if (req->query != NULL) dwn_free(req->query);
                dwn_free(req);
            } else {
                curl_multi_remove_handle(curl_multi, easy);
                curl_easy_cleanup(easy);
            }
        }
    }
    
    pthread_mutex_unlock(&ai_mutex);
}

/* Exa launcher callback and UI */
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 = (size_t)(req->result_count * 300);
    char *choices = dwn_malloc(choices_size);
    if (choices == NULL) {
        notification_show("Exa Search", "Error", "Memory allocation failed", NULL, 3000);
        return;
    }
    
    size_t offset = 0;
    choices[0] = '\0';
    
    for (int i = 0; i < req->result_count; i++) {
        int w = snprintf(choices + offset, choices_size - offset,
                         "%s%s", offset > 0 ? "\n" : "", req->results[i].title);
        if (w > 0 && (size_t)w < choices_size - offset) {
            offset += (size_t)w;
        }
    }
    
    char *escaped_choices = shell_escape(choices);
    size_t cmd_size = (escaped_choices != NULL) ? strlen(escaped_choices) + 64 : 64;
    char *cmd = dwn_malloc(cmd_size);
    
    if (cmd != NULL && escaped_choices != NULL) {
        snprintf(cmd, cmd_size, "echo %s | dmenu -l 10 -p 'Results:'", escaped_choices);
        
        char *selected = spawn_capture(cmd);
        
        if (selected != NULL && selected[0] != '\0') {
            for (int i = 0; i < req->result_count; i++) {
                size_t title_len = strlen(req->results[i].title);
                if (strncmp(selected, req->results[i].title, title_len) == 0) {
                    char *escaped_url = shell_escape(req->results[i].url);
                    if (escaped_url != NULL) {
                        size_t open_cmd_size = strlen(escaped_url) + 32;
                        char *open_cmd = dwn_malloc(open_cmd_size);
                        if (open_cmd != NULL) {
                            snprintf(open_cmd, open_cmd_size, "xdg-open %s &", escaped_url);
                            spawn_async(open_cmd);
                            dwn_free(open_cmd);
                        }
                        dwn_free(escaped_url);
                    }
                    break;
                }
            }
            dwn_free(selected);
        }
    }
    
    dwn_free(cmd);
    dwn_free(escaped_choices);
    dwn_free(choices);
    
    /* Cleanup request */
    if (req->query != NULL) dwn_free(req->query);
    if (req->user_data != NULL) {
        ResponseBuffer *b = (ResponseBuffer *)req->user_data;
        if (b->data != NULL) dwn_free(b->data);
        dwn_free(b);
    }
    if (req->headers != NULL) curl_slist_free_all(req->headers);
    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);
    dwn_free(query);
}