#include "local_llm.h"

#include <chrono>
#include <sstream>
#include <streambuf>

#include "llm/llm.hpp"
#include <hilog/log.h>

#define LOG_TAG "LocalLlm"
#define LOGI(...) OH_LOG_INFO(LOG_APP, __VA_ARGS__)
#define LOGE(...) OH_LOG_ERROR(LOG_APP, __VA_ARGS__)

using MNN::Transformer::ChatMessages;
using MNN::Transformer::Llm;

namespace {

// 计时流：接住 Llm::response 的流式输出，记录首 chunk 时间并累计全文。
class TimingStreamBuf : public std::streambuf {
public:
    std::string text;
    long firstTokenMs = -1;
    std::chrono::steady_clock::time_point start;

    void reset(std::chrono::steady_clock::time_point s) {
        text.clear();
        firstTokenMs = -1;
        start = s;
    }

protected:
    int overflow(int c) override {
        if (c == EOF) {
            return 0;
        }
        markFirst();
        text += static_cast<char>(c);
        return c;
    }

    std::streamsize xsputn(const char *s, std::streamsize n) override {
        if (n > 0) {
            markFirst();
            text.append(s, static_cast<size_t>(n));
        }
        return n;
    }

private:
    void markFirst() {
        if (firstTokenMs < 0) {
            firstTokenMs = std::chrono::duration_cast<std::chrono::milliseconds>(
                std::chrono::steady_clock::now() - start)
                .count();
        }
    }
};

std::string jsonEscape(const std::string &in) {
    std::string out;
    out.reserve(in.size() + 16);
    for (char c : in) {
        switch (c) {
        case '"': out += "\\\""; break;
        case '\\': out += "\\\\"; break;
        case '\n': out += "\\n"; break;
        case '\r': out += "\\r"; break;
        case '\t': out += "\\t"; break;
        default: out += c; break;
        }
    }
    return out;
}

} // namespace

LocalLlm &LocalLlm::instance() {
    static LocalLlm inst;
    return inst;
}

bool LocalLlm::loaded() const {
    return llm_ != nullptr;
}

std::string LocalLlm::load(const std::string &modelDir) {
    std::lock_guard<std::mutex> lock(mutex_);
    unload();
    const std::string configPath = modelDir + "/config.json";
    Llm *llm = Llm::createLLM(configPath);
    if (llm == nullptr) {
        LOGE("createLLM failed: %{public}s", configPath.c_str());
        return "createLLM failed: " + configPath;
    }
    // use_mmap 手机端必开（权重走 mmap 防 OOM）；tmp_path 必须指向沙箱可写目录，
    // 由 ArkTS 侧在调用前 mkdir 好（modelDir/.tmp）。
    const std::string cfg = std::string("{\"tmp_path\":\"") + modelDir + "/.tmp" +
                            "\",\"use_mmap\":true,\"thread_num\":4,"
                            "\"precision\":\"low\",\"memory\":\"low\","
                            "\"reuse_kv\":true,\"max_new_tokens\":256}";
    llm->set_config(cfg);
    if (!llm->load()) {
        LOGE("llm load failed: %{public}s", configPath.c_str());
        Llm::destroy(llm);
        return "llm load failed: " + configPath;
    }
    llm_ = llm;
    modelDir_ = modelDir;
    LOGI("model loaded: %{public}s", configPath.c_str());
    return "";
}

std::string LocalLlm::generate(const std::string &userPrompt, int maxNewTokens) {
    std::lock_guard<std::mutex> lock(mutex_);
    if (llm_ == nullptr) {
        return "{\"error\":\"model not loaded\"}";
    }
    TimingStreamBuf buf;
    std::ostream os(&buf);
    ChatMessages msgs;
    msgs.emplace_back("user", userPrompt);
    const auto start = std::chrono::steady_clock::now();
    buf.reset(start);
    try {
        llm_->response(msgs, &os, nullptr, maxNewTokens);
    } catch (const std::exception &e) {
        LOGE("response exception: %{public}s", e.what());
        return std::string("{\"error\":\"") + jsonEscape(e.what()) + "\"}";
    }
    const long totalMs = std::chrono::duration_cast<std::chrono::milliseconds>(
        std::chrono::steady_clock::now() - start)
        .count();
    LOGI("generate: first=%{public}ldms total=%{public}ldms chars=%{public}zu",
         buf.firstTokenMs, totalMs, buf.text.size());
    std::ostringstream out;
    out << "{\"text\":\"" << jsonEscape(buf.text) << "\",\"first_token_ms\":" << buf.firstTokenMs
        << ",\"total_ms\":" << totalMs << ",\"chars\":" << buf.text.size() << "}";
    return out.str();
}

void LocalLlm::unload() {
    if (llm_ != nullptr) {
        Llm::destroy(llm_);
        llm_ = nullptr;
        modelDir_.clear();
    }
}