/*
 * UncloseAI C++ Library using libcurl
 * OpenAI-compatible API client with streaming support
 */

#include <iostream>
#include <string>
#include <vector>
#include <functional>
#include <cstring>
#include <cstdlib>
#include <fstream>
#include <curl/curl.h>

#define MAX_CONTENT_LEN 1024

struct ModelInfo {
    std::string id;
    std::string endpoint;
    int max_tokens;
};

struct MemoryStruct {
    std::string data;
};

// Callback for non-streaming responses
static size_t WriteCallback(void *contents, size_t size, size_t nmemb, void *userp) {
    size_t realsize = size * nmemb;
    MemoryStruct *mem = static_cast<MemoryStruct*>(userp);
    mem->data.append(static_cast<char*>(contents), realsize);
    return realsize;
}

// Streaming context
struct StreamContext {
    std::function<void(const std::string&)> callback;
    std::string buffer;
};

// Extract content from SSE JSON
static void extract_sse_content(const std::string& data, std::string& content) {
    const char *content_marker = "\"content\":\"";
    size_t start = data.find(content_marker);
    if(start == std::string::npos) return;

    start += strlen(content_marker);
    size_t end = start;

    while(end < data.size() && data[end] != '"') {
        if(data[end] == '\\' && end + 1 < data.size()) {
            end += 2;
        } else {
            end++;
        }
    }

    content = data.substr(start, end - start);
}

// Streaming callback
static size_t StreamCallback(void *contents, size_t size, size_t nmemb, void *userp) {
    size_t realsize = size * nmemb;
    StreamContext *ctx = static_cast<StreamContext*>(userp);

    ctx->buffer.append(static_cast<char*>(contents), realsize);

    size_t pos = 0;
    while((pos = ctx->buffer.find("\n\n")) != std::string::npos) {
        std::string line = ctx->buffer.substr(0, pos);
        ctx->buffer.erase(0, pos + 2);

        if(line.substr(0, 6) == "data: ") {
            std::string data = line.substr(6);

            if(data == "[DONE]") break;

            std::string content;
            extract_sse_content(data, content);

            if(!content.empty() && ctx->callback) {
                ctx->callback(content);
            }
        }
    }

    return realsize;
}

class UncloseAI {
private:
    std::vector<ModelInfo> models;
    std::vector<std::string> tts_endpoints;
    int timeout;
    bool debug;

    void discover_endpoints_from_env(const std::string& prefix, std::vector<std::string>& endpoints) {
        for(int i = 1; i < 10000; i++) {
            std::string var_name = prefix + "_" + std::to_string(i);
            const char* endpoint = std::getenv(var_name.c_str());
            if(!endpoint) break;
            endpoints.push_back(endpoint);
        }
    }

    void discover_models(const std::vector<std::string>& endpoints) {
        for(const auto& endpoint : endpoints) {
            if(debug) {
                std::cout << "[DEBUG] Discovering from: " << endpoint << std::endl;
            }

            CURL *curl = curl_easy_init();
            if(curl) {
                MemoryStruct response;
                std::string url = endpoint + "/models";

                curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
                curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, WriteCallback);
                curl_easy_setopt(curl, CURLOPT_WRITEDATA, &response);
                curl_easy_setopt(curl, CURLOPT_TIMEOUT, 10L);

                CURLcode res = curl_easy_perform(curl);
                curl_easy_cleanup(curl);

                if(res == CURLE_OK) {
                    // Simple JSON parsing for model IDs
                    size_t pos = 0;
                    while((pos = response.data.find("\"id\":\"", pos)) != std::string::npos) {
                        pos += 6;
                        size_t end = response.data.find("\"", pos);
                        if(end != std::string::npos) {
                            std::string model_id = response.data.substr(pos, end - pos);

                            if(model_id.substr(0, 10) != "modelperm-") {
                                ModelInfo info;
                                info.id = model_id;
                                info.endpoint = endpoint;
                                info.max_tokens = 8192;
                                models.push_back(info);

                                if(debug) {
                                    std::cout << "[DEBUG] Discovered: " << model_id << std::endl;
                                }
                            }
                        }
                        pos = end + 1;
                    }
                }
            }
        }
    }

public:
    UncloseAI(int timeout = 30, bool debug = false) : timeout(timeout), debug(debug) {
        std::vector<std::string> endpoints;
        std::vector<std::string> tts_eps;

        discover_endpoints_from_env("MODEL_ENDPOINT", endpoints);
        discover_endpoints_from_env("TTS_ENDPOINT", tts_eps);

        if(debug) {
            std::cout << "[DEBUG] Initialized with " << endpoints.size() << " endpoint(s)" << std::endl;
        }

        discover_models(endpoints);
        tts_endpoints = tts_eps;
    }

    const std::vector<ModelInfo>& list_models() const {
        return models;
    }

    int chat(const std::string& prompt, std::string& response, int model_idx = 0) {
        if(model_idx >= static_cast<int>(models.size())) return -1;

        const ModelInfo& model = models[model_idx];
        std::string url = model.endpoint + "/chat/completions";

        std::string json = "{\"model\":\"" + model.id + "\","
                          "\"messages\":[{\"role\":\"user\",\"content\":\"" + prompt + "\"}],"
                          "\"max_tokens\":100}";

        CURL *curl = curl_easy_init();
        if(!curl) return -1;

        MemoryStruct mem;
        struct curl_slist *headers = nullptr;
        headers = curl_slist_append(headers, "Content-Type: application/json");

        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);
        curl_easy_setopt(curl, CURLOPT_POSTFIELDS, json.c_str());
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, WriteCallback);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &mem);
        curl_easy_setopt(curl, CURLOPT_TIMEOUT, (long)timeout);

        CURLcode res = curl_easy_perform(curl);

        curl_slist_free_all(headers);
        curl_easy_cleanup(curl);

        response = mem.data;
        return (res == CURLE_OK) ? 0 : -1;
    }

    int chat_stream(const std::string& prompt, std::function<void(const std::string&)> callback, int model_idx = 0) {
        if(model_idx >= static_cast<int>(models.size())) return -1;

        const ModelInfo& model = models[model_idx];
        std::string url = model.endpoint + "/chat/completions";

        std::string json = "{\"model\":\"" + model.id + "\","
                          "\"messages\":[{\"role\":\"user\",\"content\":\"" + prompt + "\"}],"
                          "\"stream\":true,"
                          "\"max_tokens\":500}";

        CURL *curl = curl_easy_init();
        if(!curl) return -1;

        StreamContext ctx;
        ctx.callback = callback;

        struct curl_slist *headers = nullptr;
        headers = curl_slist_append(headers, "Content-Type: application/json");

        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);
        curl_easy_setopt(curl, CURLOPT_POSTFIELDS, json.c_str());
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, StreamCallback);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &ctx);
        curl_easy_setopt(curl, CURLOPT_TIMEOUT, (long)timeout);

        CURLcode res = curl_easy_perform(curl);

        curl_slist_free_all(headers);
        curl_easy_cleanup(curl);

        return (res == CURLE_OK) ? 0 : -1;
    }

    int tts(const std::string& text, const std::string& voice, const std::string& output_file) {
        if(tts_endpoints.empty()) return -1;

        std::string url = tts_endpoints[0] + "/audio/speech";
        std::string json = "{\"model\":\"tts-1\","
                          "\"voice\":\"" + voice + "\","
                          "\"input\":\"" + text + "\"}";

        CURL *curl = curl_easy_init();
        if(!curl) return -1;

        MemoryStruct mem;
        struct curl_slist *headers = nullptr;
        headers = curl_slist_append(headers, "Content-Type: application/json");

        curl_easy_setopt(curl, CURLOPT_URL, url.c_str());
        curl_easy_setopt(curl, CURLOPT_HTTPHEADER, headers);
        curl_easy_setopt(curl, CURLOPT_POSTFIELDS, json.c_str());
        curl_easy_setopt(curl, CURLOPT_WRITEFUNCTION, WriteCallback);
        curl_easy_setopt(curl, CURLOPT_WRITEDATA, &mem);
        curl_easy_setopt(curl, CURLOPT_TIMEOUT, (long)timeout);

        CURLcode res = curl_easy_perform(curl);

        curl_slist_free_all(headers);
        curl_easy_cleanup(curl);

        if(res == CURLE_OK) {
            std::ofstream file(output_file, std::ios::binary);
            if(file.is_open()) {
                file.write(mem.data.c_str(), mem.data.size());
                file.close();
                return 0;
            }
        }

        return -1;
    }
};

// Demo
int main() {
    std::cout << "=== UncloseAI C++ Client (with Streaming) ===\n\n";

    curl_global_init(CURL_GLOBAL_ALL);

    UncloseAI client(30, true);

    if(client.list_models().empty()) {
        std::cout << "ERROR: No models discovered. Set environment variables:\n";
        std::cout << "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.\n";
        curl_global_cleanup();
        return 1;
    }

    auto models = client.list_models();
    std::cout << "\nDiscovered " << models.size() << " model(s):\n";
    for(const auto& m : models) {
        std::cout << "  - " << m.id << " (max_tokens: " << m.max_tokens << ")\n";
    }
    std::cout << "\n";

    // Non-streaming chat
    std::cout << "=== Non-Streaming Chat ===\n";
    std::string response;
    if(client.chat("Explain quantum computing in one sentence", response) == 0) {
        std::cout << "Response: (" << response.size() << " bytes received)\n\n";
    }

    // Streaming chat
    std::cout << "=== Streaming Chat ===\n";
    int model_idx = (models.size() > 1) ? 1 : 0;
    std::cout << "Model: " << models[model_idx].id << "\n";
    std::cout << "Response: ";

    client.chat_stream("Write a C++ function to check if a number is prime",
                      [](const std::string& content) {
                          std::cout << content << std::flush;
                      }, model_idx);

    std::cout << "\n\n";

    // TTS
    std::cout << "=== TTS Speech Generation ===\n";
    if(client.tts("Hello from UncloseAI C++ client! This demonstrates streaming support.",
                  "alloy", "speech.mp3") == 0) {
        std::cout << "[OK] Speech file created: speech.mp3\n\n";
    } else {
        std::cout << "[ERROR] TTS Error\n\n";
    }

    std::cout << "=== Examples Complete ===\n";

    curl_global_cleanup();
    return 0;
}