#include "r1_voice/llm.hh"

#include <curl/curl.h>

#include <iostream>
#include <sstream>
#include <stdexcept>

namespace r1_voice {

namespace {

size_t writeCallback(
    void* contents,
    size_t size,
    size_t nmemb,
    void* user)
{
    const size_t total = size * nmemb;

    auto* output =
        static_cast<std::string*>(user);

    output->append(
        static_cast<char*>(contents),
        total
    );

    return total;
}

} // namespace

LLM::LLM(LlmConfig config)
    : config_(std::move(config))
{
    curl_global_init(CURL_GLOBAL_DEFAULT);
}

std::string LLM::escapeJson(
    const std::string& text) const
{
    std::string result;
    result.reserve(text.size() + 16);

    for (const char c : text) {
        switch (c) {
        case '"':
            result += "\\\"";
            break;

        case '\\':
            result += "\\\\";
            break;

        case '\n':
            result += "\\n";
            break;

        case '\r':
            result += "\\r";
            break;

        case '\t':
            result += "\\t";
            break;

        default:
            result += c;
            break;
        }
    }

    return result;
}

std::string LLM::httpPost(
    const std::string& body) const
{
    CURL* curl = curl_easy_init();

    if (!curl) {
        throw std::runtime_error(
            "curl_easy_init() failed"
        );
    }

    std::string response;

    curl_slist* headers = nullptr;

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

    curl_easy_setopt(
        curl,
        CURLOPT_URL,
        config_.url.c_str()
    );

    curl_easy_setopt(
        curl,
        CURLOPT_HTTPHEADER,
        headers
    );

    curl_easy_setopt(
        curl,
        CURLOPT_POST,
        1L
    );

    curl_easy_setopt(
        curl,
        CURLOPT_POSTFIELDS,
        body.c_str()
    );

    curl_easy_setopt(
        curl,
        CURLOPT_WRITEFUNCTION,
        writeCallback
    );

    curl_easy_setopt(
        curl,
        CURLOPT_WRITEDATA,
        &response
    );

    curl_easy_setopt(
        curl,
        CURLOPT_CONNECTTIMEOUT,
        5L
    );

    curl_easy_setopt(
        curl,
        CURLOPT_TIMEOUT,
        120L
    );

    const CURLcode ret =
        curl_easy_perform(curl);

    long status = 0;

    curl_easy_getinfo(
        curl,
        CURLINFO_RESPONSE_CODE,
        &status
    );

    curl_slist_free_all(headers);
    curl_easy_cleanup(curl);

    if (ret != CURLE_OK) {
        throw std::runtime_error(
            std::string("HTTP request failed: ") +
            curl_easy_strerror(ret)
        );
    }

    if (status < 200 || status >= 300) {
        throw std::runtime_error(
            "LLM HTTP status " +
            std::to_string(status) +
            "\nResponse:\n" +
            response
        );
    }

    return response;
}

std::string LLM::parseContent(
    const std::string& json) const
{
    const std::string key = "\"content\"";

    const size_t key_pos =
        json.find(key);

    if (key_pos == std::string::npos)
        return {};

    const size_t colon =
        json.find(
            ':',
            key_pos + key.size()
        );

    if (colon == std::string::npos)
        return {};

    const size_t quote =
        json.find(
            '"',
            colon + 1
        );

    if (quote == std::string::npos)
        return {};

    std::string result;
    result.reserve(256);

    bool escaped = false;

    for (size_t i = quote + 1;
         i < json.size();
         ++i) {

        const char c = json[i];

        if (escaped) {
            switch (c) {
            case 'n':
                result += '\n';
                break;

            case 'r':
                result += '\r';
                break;

            case 't':
                result += '\t';
                break;

            case '"':
                result += '"';
                break;

            case '\\':
                result += '\\';
                break;

            case '/':
                result += '/';
                break;

            default:
                result += c;
                break;
            }

            escaped = false;
            continue;
        }

        if (c == '\\') {
            escaped = true;
            continue;
        }

        if (c == '"')
            break;

        result += c;
    }

    return result;
}

static std::string removeThinkBlocks(
    std::string text)
{
    for (;;) {
        const size_t begin =
            text.find("<think>");

        if (begin == std::string::npos)
            break;

        const size_t end =
            text.find(
                "</think>",
                begin
            );

        if (end == std::string::npos) {
            text.erase(begin);
            break;
        }

        const size_t len =
            end +
            std::string("</think>").size() -
            begin;

        text.erase(begin, len);
    }

    return text;
}

std::string LLM::chat(
    const std::string& user,
    const std::string& transient_context)
{
    history_.push_back({
        "user",
        user
    });

    std::ostringstream body;

    body
        << '{'
        << "\"model\":\""
        << escapeJson(config_.model)
        << "\","

        << "\"messages\":["

        << '{'
        << "\"role\":\"system\","
        << "\"content\":\""
        << escapeJson(
               config_.system_prompt
           )
        << "\""
        << '}';

    /*
     * Temporary context such as current YOLO detections.
     *
     * It is sent as another system message but is NOT
     * stored in history.
     */
    if (!transient_context.empty()) {

        body
            << ",{"
            << "\"role\":\"system\","
            << "\"content\":\""
            << escapeJson(
                   transient_context
               )
            << "\""
            << '}';
    }

    const size_t begin =
        history_.size() > 12
            ? history_.size() - 12
            : 0;

    for (size_t i = begin;
         i < history_.size();
         ++i) {

        body
            << ",{"
            << "\"role\":\""
            << escapeJson(
                   history_[i].role
               )
            << "\","
            << "\"content\":\""
            << escapeJson(
                   history_[i].content
               )
            << "\""
            << '}';
    }

    body
        << "],"

        << "\"max_tokens\":"
        << config_.max_tokens
        << ','

        << "\"temperature\":"
        << config_.temperature
        << ','

        << "\"chat_template_kwargs\":{"
        << "\"enable_thinking\":false"
        << '}'

        << '}';

    std::string response;

    try {
        response =
            httpPost(
                body.str()
            );
    }
    catch (...) {

        /*
         * Don't leave a failed user request in
         * conversation history.
         */
        if (!history_.empty() &&
            history_.back().role == "user" &&
            history_.back().content == user) {

            history_.pop_back();
        }

        throw;
    }

    std::string answer =
        parseContent(response);

    answer =
        removeThinkBlocks(answer);

    if (answer.empty()) {

        std::cerr
            << "LLM returned empty content:\n"
            << response
            << '\n';

        return {};
    }

    history_.push_back({
        "assistant",
        answer
    });

    return answer;
}

} // namespace r1_voice
