// PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
// Copyright 2025 TimeHexOn & foxhop & russell@unturf
// https://www.permacomputer.com

/*
 * UncloseAI C Client using nghttp2 (HTTP/2)
 * OpenAI-compatible API client with HTTP/2 and streaming support
 * Compatible with vLLM, Ollama, and OpenAI-compatible endpoints
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <netdb.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <nghttp2/nghttp2.h>

#define MAX_ENDPOINTS 100
#define MAX_MODELS 100
#define MAX_URL_LEN 512
#define MAX_MODEL_LEN 256
#define BUFFER_SIZE 8192

// Structure to hold discovered model info
struct ModelInfo {
    char id[MAX_MODEL_LEN];
    char endpoint[MAX_URL_LEN];
    char host[256];
    int port;
    int max_tokens;
};

struct ModelInfo models[MAX_MODELS];
int model_count = 0;

char tts_endpoints[MAX_ENDPOINTS][MAX_URL_LEN];
int tts_count = 0;

// HTTP/2 connection context
struct Connection {
    nghttp2_session *session;
    int sock;
    char *response_data;
    size_t response_size;
    size_t response_capacity;
    int done;
    int streaming;
    char stream_buffer[8192];
    size_t stream_buffer_pos;
};

// Request data provider
struct RequestData {
    const char *data;
    size_t length;
    size_t position;
};

// Parse URL into host, port, path
int parse_url(const char *url, char *host, int *port, char *path) {
    const char *p = url;

    // Skip protocol
    if (strncmp(p, "http://", 7) == 0) {
        p += 7;
        *port = 80;
    } else if (strncmp(p, "https://", 8) == 0) {
        p += 8;
        *port = 443;
    } else {
        return -1;
    }

    // Extract host and optional port
    const char *slash = strchr(p, '/');
    const char *colon = strchr(p, ':');

    if (colon && (!slash || colon < slash)) {
        // Port specified
        size_t host_len = colon - p;
        strncpy(host, p, host_len);
        host[host_len] = '\0';
        *port = atoi(colon + 1);
        p = slash ? slash : (colon + strlen(colon));
    } else if (slash) {
        // No port, use default
        size_t host_len = slash - p;
        strncpy(host, p, host_len);
        host[host_len] = '\0';
        p = slash;
    } else {
        // No path
        strcpy(host, p);
        p = "";
    }

    strcpy(path, p[0] ? p : "/");
    return 0;
}

// Connect to server
int connect_to_server(const char *host, int port) {
    struct addrinfo hints, *res, *rp;
    int sock = -1;

    memset(&hints, 0, sizeof(hints));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;

    char port_str[16];
    snprintf(port_str, sizeof(port_str), "%d", port);

    if (getaddrinfo(host, port_str, &hints, &res) != 0) {
        return -1;
    }

    for (rp = res; rp; rp = rp->ai_next) {
        sock = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if (sock == -1) continue;

        if (connect(sock, rp->ai_addr, rp->ai_addrlen) == 0) {
            break;
        }

        close(sock);
        sock = -1;
    }

    freeaddrinfo(res);
    return sock;
}

// Send callback for nghttp2
ssize_t send_callback(nghttp2_session *session, const uint8_t *data,
                     size_t length, int flags, void *user_data) {
    struct Connection *conn = (struct Connection *)user_data;
    ssize_t rv = send(conn->sock, data, length, 0);
    if (rv < 0) {
        return NGHTTP2_ERR_CALLBACK_FAILURE;
    }
    return rv;
}

// Receive callback for nghttp2
ssize_t recv_callback(nghttp2_session *session, uint8_t *buf, size_t length,
                     int flags, void *user_data) {
    struct Connection *conn = (struct Connection *)user_data;
    ssize_t rv = recv(conn->sock, buf, length, 0);
    if (rv < 0) {
        return NGHTTP2_ERR_CALLBACK_FAILURE;
    }
    if (rv == 0) {
        return NGHTTP2_ERR_EOF;
    }
    return rv;
}

// Extract content from SSE data chunk
void extract_sse_content(const char *data, char *content, size_t max_len) {
    const char *content_marker = "\"content\":\"";
    const char *start = strstr(data, content_marker);
    if (!start) return;

    start += strlen(content_marker);
    const char *end = start;

    // Find closing quote, handling escaped quotes
    while (*end && *end != '"') {
        if (*end == '\\' && *(end + 1)) {
            end += 2;
        } else {
            end++;
        }
    }

    size_t len = end - start;
    if (len > max_len - 1) len = max_len - 1;
    strncpy(content, start, len);
    content[len] = '\0';
}

// Data chunk received callback
int on_data_chunk_recv_callback(nghttp2_session *session, uint8_t flags,
                                int32_t stream_id, const uint8_t *data,
                                size_t len, void *user_data) {
    struct Connection *conn = (struct Connection *)user_data;

    if (conn->streaming) {
        // Process streaming SSE data
        for (size_t i = 0; i < len; i++) {
            if (conn->stream_buffer_pos >= sizeof(conn->stream_buffer) - 1) {
                conn->stream_buffer_pos = 0; // Reset if overflow
            }

            conn->stream_buffer[conn->stream_buffer_pos++] = data[i];

            // Check for line ending
            if (data[i] == '\n' && conn->stream_buffer_pos >= 2 &&
                conn->stream_buffer[conn->stream_buffer_pos - 2] == '\n') {
                conn->stream_buffer[conn->stream_buffer_pos] = '\0';

                // Process SSE line
                if (strncmp(conn->stream_buffer, "data: ", 6) == 0) {
                    const char *json_data = conn->stream_buffer + 6;

                    // Check for [DONE]
                    if (strncmp(json_data, "[DONE]", 6) == 0) {
                        conn->stream_buffer_pos = 0;
                        break;
                    }

                    // Extract and print content
                    char content[1024];
                    extract_sse_content(json_data, content, sizeof(content));

                    if (strlen(content) > 0) {
                        printf("%s", content);
                        fflush(stdout);
                    }
                }

                conn->stream_buffer_pos = 0;
            }
        }
    } else {
        // Non-streaming: accumulate response
        if (conn->response_size + len > conn->response_capacity) {
            size_t new_capacity = (conn->response_capacity + len) * 2;
            char *new_data = realloc(conn->response_data, new_capacity);
            if (!new_data) return 0;
            conn->response_data = new_data;
            conn->response_capacity = new_capacity;
        }

        memcpy(conn->response_data + conn->response_size, data, len);
        conn->response_size += len;
    }

    return 0;
}

// Stream close callback
int on_stream_close_callback(nghttp2_session *session, int32_t stream_id,
                             uint32_t error_code, void *user_data) {
    struct Connection *conn = (struct Connection *)user_data;
    conn->done = 1;
    return 0;
}

// Data provider read callback
ssize_t data_provider_read_callback(nghttp2_session *session, int32_t stream_id,
                                   uint8_t *buf, size_t length,
                                   uint32_t *data_flags,
                                   nghttp2_data_source *source,
                                   void *user_data) {
    struct RequestData *req_data = (struct RequestData *)source->ptr;

    size_t remaining = req_data->length - req_data->position;
    size_t to_copy = (remaining < length) ? remaining : length;

    memcpy(buf, req_data->data + req_data->position, to_copy);
    req_data->position += to_copy;

    if (req_data->position >= req_data->length) {
        *data_flags |= NGHTTP2_DATA_FLAG_EOF;
    }

    return to_copy;
}

// Setup HTTP/2 headers
void setup_http2_headers(nghttp2_nv *hdrs, const char *method, const char *path, const char *host) {
    hdrs[0].name = (uint8_t *)":method";
    hdrs[0].value = (uint8_t *)method;
    hdrs[0].namelen = strlen(":method");
    hdrs[0].valuelen = strlen(method);
    hdrs[0].flags = NGHTTP2_NV_FLAG_NONE;

    hdrs[1].name = (uint8_t *)":path";
    hdrs[1].value = (uint8_t *)path;
    hdrs[1].namelen = strlen(":path");
    hdrs[1].valuelen = strlen(path);
    hdrs[1].flags = NGHTTP2_NV_FLAG_NONE;

    hdrs[2].name = (uint8_t *)":scheme";
    hdrs[2].value = (uint8_t *)"http";
    hdrs[2].namelen = strlen(":scheme");
    hdrs[2].valuelen = 4;
    hdrs[2].flags = NGHTTP2_NV_FLAG_NONE;

    hdrs[3].name = (uint8_t *)":authority";
    hdrs[3].value = (uint8_t *)host;
    hdrs[3].namelen = strlen(":authority");
    hdrs[3].valuelen = strlen(host);
    hdrs[3].flags = NGHTTP2_NV_FLAG_NONE;

    hdrs[4].name = (uint8_t *)"content-type";
    hdrs[4].value = (uint8_t *)"application/json";
    hdrs[4].namelen = strlen("content-type");
    hdrs[4].valuelen = strlen("application/json");
    hdrs[4].flags = NGHTTP2_NV_FLAG_NONE;
}

// Setup nghttp2 session with callbacks
void setup_nghttp2_session(struct Connection *conn, nghttp2_session_callbacks *callbacks) {
    nghttp2_session_callbacks_new(&callbacks);
    nghttp2_session_callbacks_set_send_callback(callbacks, send_callback);
    nghttp2_session_callbacks_set_recv_callback(callbacks, recv_callback);
    nghttp2_session_callbacks_set_on_data_chunk_recv_callback(callbacks, on_data_chunk_recv_callback);
    nghttp2_session_callbacks_set_on_stream_close_callback(callbacks, on_stream_close_callback);

    nghttp2_session_client_new(&conn->session, callbacks, conn);
    nghttp2_session_callbacks_del(callbacks);

    nghttp2_submit_settings(conn->session, NGHTTP2_FLAG_NONE, NULL, 0);
}

// Make HTTP/2 request
int http2_request(const char *host, int port, const char *path,
                 const char *method, const char *body,
                 char **response, size_t *response_size) {
    int sock = connect_to_server(host, port);
    if (sock < 0) return -1;

    // Initialize connection
    struct Connection conn = {0};
    conn.sock = sock;
    conn.response_data = malloc(BUFFER_SIZE);
    conn.response_capacity = BUFFER_SIZE;
    conn.response_size = 0;
    conn.done = 0;
    conn.streaming = 0;
    conn.stream_buffer_pos = 0;

    // Setup nghttp2 session and callbacks
    nghttp2_session_callbacks *callbacks;
    setup_nghttp2_session(&conn, callbacks);

    // Prepare headers
    nghttp2_nv hdrs[5];
    setup_http2_headers(hdrs, method, path, host);

    // Submit request
    if (body && strlen(body) > 0) {
        struct RequestData req_data = {body, strlen(body), 0};
        nghttp2_data_provider data_prd;
        data_prd.source.ptr = &req_data;
        data_prd.read_callback = data_provider_read_callback;
        nghttp2_submit_request(conn.session, NULL, hdrs, 5, &data_prd, NULL);
    } else {
        nghttp2_submit_request(conn.session, NULL, hdrs, 4, NULL, NULL);
    }

    // Send and receive
    while (!conn.done) {
        if (nghttp2_session_send(conn.session) != 0) break;
        if (nghttp2_session_recv(conn.session) != 0) break;
    }

    // Cleanup
    *response = conn.response_data;
    *response_size = conn.response_size;

    nghttp2_session_del(conn.session);
    close(sock);

    return 0;
}

// Make streaming HTTP/2 request
int http2_stream_request(const char *host, int port, const char *path,
                         const char *method, const char *body) {
    int sock = connect_to_server(host, port);
    if (sock < 0) return -1;

    // Initialize connection for streaming
    struct Connection conn = {0};
    conn.sock = sock;
    conn.response_data = malloc(BUFFER_SIZE);
    conn.response_capacity = BUFFER_SIZE;
    conn.response_size = 0;
    conn.done = 0;
    conn.streaming = 1;  // Enable streaming mode
    conn.stream_buffer_pos = 0;

    // Setup nghttp2 session and callbacks
    nghttp2_session_callbacks *callbacks;
    setup_nghttp2_session(&conn, callbacks);

    // Prepare headers
    nghttp2_nv hdrs[5];
    setup_http2_headers(hdrs, method, path, host);

    // Submit request
    if (body && strlen(body) > 0) {
        struct RequestData req_data = {body, strlen(body), 0};
        nghttp2_data_provider data_prd;
        data_prd.source.ptr = &req_data;
        data_prd.read_callback = data_provider_read_callback;
        nghttp2_submit_request(conn.session, NULL, hdrs, 5, &data_prd, NULL);
    } else {
        nghttp2_submit_request(conn.session, NULL, hdrs, 4, NULL, NULL);
    }

    // Send and receive (streaming prints directly in callback)
    while (!conn.done) {
        if (nghttp2_session_send(conn.session) != 0) break;
        if (nghttp2_session_recv(conn.session) != 0) break;
    }

    // Cleanup
    free(conn.response_data);
    nghttp2_session_del(conn.session);
    close(sock);

    return 0;
}

// Simple JSON string extractor
void extract_model_ids(const char *json, const char *endpoint) {
    const char *search = json;
    const char *id_marker = "\"id\":\"";

    while ((search = strstr(search, id_marker)) != NULL && model_count < MAX_MODELS) {
        search += strlen(id_marker);
        const char *end = strchr(search, '"');
        if (end) {
            size_t len = end - search;
            if (len < MAX_MODEL_LEN) {
                strncpy(models[model_count].id, search, len);
                models[model_count].id[len] = '\0';

                // Filter out modelperm-* entries
                if (strncmp(models[model_count].id, "modelperm-", 10) == 0 ||
                    strncmp(models[model_count].id, "chatcmpl-", 9) == 0) {
                    search = end + 1;
                    continue;
                }

                strncpy(models[model_count].endpoint, endpoint, MAX_URL_LEN - 1);

                char host[256];
                int port;
                char path[512];
                if (parse_url(endpoint, host, &port, path) == 0) {
                    strncpy(models[model_count].host, host, sizeof(models[model_count].host) - 1);
                    models[model_count].port = port;
                }

                models[model_count].max_tokens = 8192;
                model_count++;
            }
        }
        search = end + 1;
    }
}

// Discover models from endpoint
void discover_models_from_endpoint(const char *endpoint) {
    printf("Discovering from: %s\n", endpoint);

    char host[256];
    int port;
    char base_path[512];

    if (parse_url(endpoint, host, &port, base_path) != 0) return;

    char path[512];
    snprintf(path, sizeof(path), "%s%s", base_path,
             base_path[strlen(base_path)-1] == '/' ? "models" : "/models");

    char *response = NULL;
    size_t response_size = 0;

    if (http2_request(host, port, path, "GET", NULL, &response, &response_size) == 0) {
        if (response) {
            extract_model_ids(response, endpoint);
            free(response);
        }
    }
}

// Discover all models from environment
void discover_all_models() {
    printf("=== Model Discovery ===\n");

    for (int i = 1; i < 10000; i++) {
        char var_name[32];
        snprintf(var_name, sizeof(var_name), "MODEL_ENDPOINT_%d", i);
        char *endpoint = getenv(var_name);
        if (!endpoint) break;
        discover_models_from_endpoint(endpoint);
    }

    for (int i = 1; i < 10000; i++) {
        char var_name[32];
        snprintf(var_name, sizeof(var_name), "TTS_ENDPOINT_%d", i);
        char *endpoint = getenv(var_name);
        if (!endpoint) break;
        strncpy(tts_endpoints[tts_count++], endpoint, MAX_URL_LEN - 1);
    }

    printf("\nDiscovered %d model(s)\n\n", model_count);
}

int main(void) {
    printf("=== UncloseAI C Client (nghttp2 HTTP/2) ===\n\n");

    discover_all_models();

    if (model_count == 0) {
        printf("ERROR: No models discovered. Set environment variables:\n");
        printf("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.\n");
        return 1;
    }

    for (int i = 0; i < model_count; i++) {
        printf("  - %s (max_tokens: %d)\n", models[i].id, models[i].max_tokens);
    }
    printf("\n");

    // Non-streaming chat (Hermes)
    printf("=== Non-Streaming Chat (Hermes) ===\n");
    printf("Model: %s\n", models[0].id);

    char json[2048];
    snprintf(json, sizeof(json),
        "{\"model\":\"%s\","
        "\"messages\":[{\"role\":\"user\",\"content\":\"Give a Python Fizzbuzz solution in one line of code?\"}],"
        "\"stream\":false,"
        "\"max_tokens\":150,"
        "\"temperature\":0.5}",
        models[0].id);

    char *response = NULL;
    size_t response_size = 0;

    char path[512];
    snprintf(path, sizeof(path), "/v1/chat/completions");

    if (http2_request(models[0].host, models[0].port, path, "POST", json,
                     &response, &response_size) == 0) {
        printf("Response received (%zu bytes)\n", response_size);
        printf("(Full response requires JSON parsing library)\n\n");
        free(response);
    } else {
        printf("Request failed\n\n");
    }

    // Streaming chat (Hermes)
    printf("=== Streaming Chat (Hermes) ===\n");
    printf("Model: %s\n", models[0].id);
    printf("Response: ");

    char stream_json_hermes[2048];
    snprintf(stream_json_hermes, sizeof(stream_json_hermes),
        "{\"model\":\"%s\","
        "\"messages\":[{\"role\":\"user\",\"content\":\"Give a Python Fizzbuzz solution in one line of code?\"}],"
        "\"stream\":true,"
        "\"max_tokens\":500,"
        "\"temperature\":0.7}",
        models[0].id);

    if (http2_stream_request(models[0].host, models[0].port, path, "POST", stream_json_hermes) != 0) {
        printf("\nStreaming request failed");
    }
    printf("\n\n");

    // Streaming chat (Qwen Coder 3) - if second model available
    if (model_count >= 2) {
        printf("=== Streaming Chat (Qwen Coder 3) ===\n");
        printf("Model: %s\n", models[1].id);
        printf("Response: ");

        char stream_json_qwen[2048];
        snprintf(stream_json_qwen, sizeof(stream_json_qwen),
            "{\"model\":\"%s\","
            "\"messages\":[{\"role\":\"user\",\"content\":\"Write a hello world program in C\"}],"
            "\"stream\":true,"
            "\"max_tokens\":500,"
            "\"temperature\":0.7}",
            models[1].id);

        if (http2_stream_request(models[1].host, models[1].port, path, "POST", stream_json_qwen) != 0) {
            printf("\nStreaming request failed");
        }
        printf("\n\n");
    }

    // TTS
    if (tts_count > 0) {
        printf("=== TTS Speech Generation ===\n");
        printf("Model: tts-1\n");

        char host[256];
        int port;
        char base_path[512];
        if (parse_url(tts_endpoints[0], host, &port, base_path) == 0) {
            const char *tts_json = "{"
                "\"model\":\"tts-1\","
                "\"voice\":\"alloy\","
                "\"input\":\"Hello from UncloseAI C client with nghttp2!\""
                "}";

            char *tts_response = NULL;
            size_t tts_size = 0;

            if (http2_request(host, port, "/audio/speech", "POST", tts_json,
                            &tts_response, &tts_size) == 0) {
                FILE *fp = fopen("/tmp/speech.mp3", "wb");
                if (fp) {
                    fwrite(tts_response, 1, tts_size, fp);
                    fclose(fp);
                    printf("Audio saved to /tmp/speech.mp3\n");
                } else {
                    printf("TTS failed\n");
                }
                free(tts_response);
            } else {
                printf("TTS failed\n");
            }
        }
    }

    printf("\n=== Examples Complete ===\n");
    return 0;
}
