#include "ollama.hpp"
#include <sstream>
#include <cstring>

static size_t curl_wbuf(void* ptr, size_t sz, size_t n, void* ud) {
    std::string* s = static_cast<std::string*>(ud);
    s->append(static_cast<char*>(ptr), sz * n);
    return sz * n;
}

OllamaClient::OllamaClient(const std::string& base_url) : base_url_(base_url), curl_(curl_easy_init()) {}

OllamaClient::~OllamaClient() {
    if (curl_) curl_easy_cleanup(curl_);
}

std::string OllamaClient::request(const std::string& path, const Json::Value& body) {
    Json::Value b = body;
    b["stream"] = false;
    Json::FastWriter w;
    std::string post = w.write(b);
    std::lock_guard<std::mutex> lk(curl_mtx_);
    std::string resp;
    curl_easy_setopt(curl_, CURLOPT_URL, (base_url_ + path).c_str());
    curl_easy_setopt(curl_, CURLOPT_WRITEFUNCTION, curl_wbuf);
    curl_easy_setopt(curl_, CURLOPT_WRITEDATA, &resp);
    curl_easy_setopt(curl_, CURLOPT_TIMEOUT, 300L);
    struct curl_slist* hlist = curl_slist_append(nullptr, "Content-Type: application/json");
    curl_easy_setopt(curl_, CURLOPT_HTTPHEADER, hlist);
    curl_easy_setopt(curl_, CURLOPT_POSTFIELDS, post.c_str());
    curl_easy_setopt(curl_, CURLOPT_POST, 1L);
    curl_easy_setopt(curl_, CURLOPT_SSL_VERIFYPEER, 0L);
    CURLcode r = curl_easy_perform(curl_);
    curl_slist_free_all(hlist);
    curl_easy_setopt(curl_, CURLOPT_HTTPHEADER, nullptr);
    curl_easy_setopt(curl_, CURLOPT_POST, 0L);
    if (r != CURLE_OK) {
        LOG(INFO, "ollama: curl failed");
        return "";
    }
    return resp;
}

std::string OllamaClient::generate(const std::string& model, const std::string& prompt,
                                    const std::string& system,
                                    const std::map<std::string, Json::Value>& options) {
    Json::Value b;
    b["model"] = model;
    b["prompt"] = prompt;
    if (!system.empty()) b["system"] = system;
    if (!options.empty()) {
        Json::Value opts;
        for (const auto& kv : options) opts[kv.first] = kv.second;
        b["options"] = opts;
    }
    std::string r = request("/api/generate", b);
    Json::Reader rd;
    Json::Value root;
    if (rd.parse(r, root) && root.isMember("response")) return root["response"].asString();
    return "";
}

OllamaResponse OllamaClient::chat(const std::string& model, const std::vector<OllamaMsg>& messages,
                                   const std::string& system,
                                   const std::map<std::string, Json::Value>& options) {
    Json::Value b;
    b["model"] = model;
    for (const auto& m : messages) {
        Json::Value msg;
        msg["role"] = m.role;
        msg["content"] = m.content;
        b["messages"].append(msg);
    }
    if (!system.empty()) b["system"] = system;
    if (!options.empty()) {
        Json::Value opts;
        for (const auto& kv : options) opts[kv.first] = kv.second;
        b["options"] = opts;
    }
    std::string r = request("/api/chat", b);
    Json::Reader rd;
    Json::Value root;
    OllamaResponse resp;
    if (rd.parse(r, root)) {
        if (root.isMember("message") && root["message"].isMember("content"))
            resp.response_text = root["message"]["content"].asString();
        if (root.isMember("done")) resp.done = root["done"].asBool();
        if (root.isMember("eval_count")) resp.eval_count = root["eval_count"].asUInt64();
        if (root.isMember("total_duration")) resp.total_duration = root["total_duration"].asUInt64();
    }
    return resp;
}

std::pair<std::string, std::string> OllamaClient::analyze_support_ticket(const std::string& ticket_text) {
    std::string prompt = "Categorize this support ticket as billing, technical, product, or spam. "
                          "Respond with ONLY the category name, then a newline, then a brief suggested response.\n\n" + ticket_text;
    std::string r = generate("mistral", prompt, "You are a support ticket classifier.");
    size_t nl = r.find('\n');
    std::string cat = nl != std::string::npos ? r.substr(0, nl) : r;
    std::string sug = nl != std::string::npos ? r.substr(nl + 1) : "";
    if (cat.empty()) cat = "unknown";
    return {cat, sug};
}

int OllamaClient::detect_fraud(const std::string& payment_data_json) {
    std::string prompt = "Rate fraud risk 0-100 for this payment data. Respond with ONLY a number.\n" + payment_data_json;
    std::string r = generate("mistral", prompt, "You are a fraud detection scorer.");
    try {
        return std::stoi(r);
    } catch (...) {
        return 50;
    }
}

std::string OllamaClient::generate_waf_rule(const std::string& attack_pattern) {
    std::string prompt = "Generate an nftables rule to block this attack pattern. Respond with ONLY the rule.\n" + attack_pattern;
    return generate("mistral", prompt, "You are a network security expert.");
}
