// Replace stub with actual ONNX-backed implementation
#include "model_provider.h"
#include <nlohmann/json.hpp>
#include <spdlog/spdlog.h>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstdlib>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <functional>
#include <future>
#include <optional>
#include <thread>
#include "ort_runtime_loader.h"
#include <yams/daemon/components/InternalEventBus.h>
#include <yams/daemon/components/TuneAdvisor.h>
#include <yams/daemon/resource/onnx_colbert_session.h>
#include <yams/daemon/resource/onnx_model_pool.h>

namespace {

namespace fs = std::filesystem;

// Global storage for plugin config JSON passed from daemon at init time
static std::string g_plugin_config_json;

static std::vector<int32_t> readColbertSkiplist(const std::string& modelsRoot,
                                                const std::string& modelName) {
    try {
        fs::path file = fs::path(modelsRoot) / modelName / "skiplist.json";
        if (!fs::exists(file))
            return {};
        std::ifstream in(file);
        if (!in.good())
            return {};
        nlohmann::json j;
        in >> j;
        if (!j.is_array())
            return {};
        std::vector<int32_t> out;
        out.reserve(j.size());
        for (const auto& item : j) {
            if (item.is_number_integer()) {
                out.push_back(item.get<int32_t>());
            }
        }
        return out;
    } catch (...) {
        return {};
    }
}

static bool isColbertModelName(const std::string& name) {
    if (name.empty()) {
        return false;
    }
    std::string lower = name;
    for (auto& c : lower) {
        c = static_cast<char>(std::tolower(c));
    }
    return lower.find("colbert") != std::string::npos;
}

struct ProviderCtx;

extern "C" yams_model_provider_v1* yams_onnx_get_model_provider();

static bool envTruthy(const char* value) {
    if (!value || !*value) {
        return false;
    }
    std::string v(value);
    for (auto& c : v) {
        c = static_cast<char>(std::tolower(c));
    }
    return v != "0" && v != "false" && v != "off" && v != "no";
}

static std::chrono::milliseconds modelAcquireTimeout(bool batchCall) {
    // Batch calls can hold the model session for longer on large chunks/sub-batches.
    // Keep a longer default to avoid turning transient contention into hard failures.
    std::chrono::milliseconds timeout =
        batchCall ? std::chrono::milliseconds(120000) : std::chrono::milliseconds(30000);

    if (const char* global = std::getenv("YAMS_ONNX_MODEL_ACQUIRE_TIMEOUT_MS")) {
        try {
            const auto parsed = std::stoll(global);
            if (parsed > 0) {
                timeout = std::chrono::milliseconds(parsed);
            }
        } catch (...) {
        }
    }

    if (batchCall) {
        if (const char* batch = std::getenv("YAMS_ONNX_MODEL_ACQUIRE_TIMEOUT_BATCH_MS")) {
            try {
                const auto parsed = std::stoll(batch);
                if (parsed > 0) {
                    timeout = std::chrono::milliseconds(parsed);
                }
            } catch (...) {
            }
        }
    }

    if (envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"))) {
        timeout = std::min(timeout, std::chrono::milliseconds(1000));
    }

    constexpr auto kMinTimeout = std::chrono::milliseconds(1000);
    constexpr auto kMaxTimeout = std::chrono::milliseconds(600000);
    timeout = std::clamp(timeout, kMinTimeout, kMaxTimeout);
    return timeout;
}

// Provide a C-callable function for threading control (not supported yet)
static yams_status_t onnx_set_threading(void* /*self*/, const char* /*model_id*/, int /*intra*/,
                                        int /*inter*/);

// ---------------------------------------------------------------------------
// BatchCoalescer — accumulates concurrent embedding requests into larger
// batches before running inference.  Opt-in via YAMS_ONNX_BATCH_COALESCE=1.
// ---------------------------------------------------------------------------
class BatchCoalescer {
public:
    struct Config {
        bool enabled = false;
        int windowMs = 5;
        size_t maxBatch = 64;
    };

    explicit BatchCoalescer(Config cfg) : cfg_(cfg) {}

    bool isEnabled() const { return cfg_.enabled; }

    using InferFn = std::function<yams::Result<std::vector<std::vector<float>>>(
        const std::string& /*modelId*/, const std::vector<std::string>& /*texts*/)>;

    // Submit texts for embedding.  Blocks until the coalesced batch completes
    // and returns only this caller's slice of the result.
    yams::Result<std::vector<std::vector<float>>> submit(const std::string& modelId,
                                                         const std::vector<std::string>& texts,
                                                         const InferFn& inferFn) {
        if (!cfg_.enabled || texts.empty()) {
            return inferFn(modelId, texts);
        }

        std::unique_lock<std::mutex> lk(mu_);

        auto& batch = pending_[modelId];
        if (!batch) {
            batch = std::make_unique<PendingBatch>();
            batch->modelId = modelId;
            batch->deadline =
                std::chrono::steady_clock::now() + std::chrono::milliseconds(cfg_.windowMs);
        }

        // Record this request's slice within the coalesced batch.
        const size_t startIdx = batch->allTexts.size();
        const size_t count = texts.size();
        for (auto& t : texts)
            batch->allTexts.push_back(std::move(t));

        auto future = batch->addRequest(startIdx, count);

        // Should we dispatch now?
        const bool sizeTriggered = batch->allTexts.size() >= cfg_.maxBatch;
        const bool timeTriggered = std::chrono::steady_clock::now() >= batch->deadline;

        if (sizeTriggered || timeTriggered) {
            // This thread becomes the dispatcher.
            auto owned = std::move(batch);
            pending_.erase(modelId);
            lk.unlock();
            dispatch(std::move(owned), inferFn);
        } else {
            // Wait for either the deadline or the dispatcher to signal.
            cv_.wait_until(lk, batch->deadline, [&] {
                // If our batch was consumed (pending_ no longer holds it), we're done.
                auto it = pending_.find(modelId);
                return it == pending_.end() || it->second.get() != batch.get();
            });

            // If we woke up because of timeout and the batch is still ours, dispatch it.
            auto it = pending_.find(modelId);
            if (it != pending_.end() && it->second.get() == batch.get()) {
                auto owned = std::move(it->second);
                pending_.erase(it);
                lk.unlock();
                dispatch(std::move(owned), inferFn);
            } else {
                lk.unlock();
            }
        }

        return future.get();
    }

    // Drain all pending batches with errors (call on shutdown).
    void drain() {
        std::lock_guard<std::mutex> lk(mu_);
        for (auto& [id, batch] : pending_) {
            if (batch) {
                batch->failAll(yams::Error{yams::ErrorCode::InternalError, "plugin shutting down"});
            }
        }
        pending_.clear();
        cv_.notify_all();
    }

private:
    struct PendingBatch {
        struct Request {
            size_t startIdx;
            size_t count;
            std::promise<yams::Result<std::vector<std::vector<float>>>> promise;
        };

        std::string modelId;
        std::vector<std::string> allTexts;
        std::vector<Request> requests;
        std::chrono::steady_clock::time_point deadline;

        std::future<yams::Result<std::vector<std::vector<float>>>> addRequest(size_t startIdx,
                                                                              size_t count) {
            requests.emplace_back();
            auto& req = requests.back();
            req.startIdx = startIdx;
            req.count = count;
            return req.promise.get_future();
        }

        void failAll(const yams::Error& err) {
            for (auto& req : requests) {
                try {
                    req.promise.set_value(yams::Result<std::vector<std::vector<float>>>(err));
                } catch (...) {
                }
            }
        }
    };

    void dispatch(std::unique_ptr<PendingBatch> batch, const InferFn& inferFn) {
        auto result = inferFn(batch->modelId, batch->allTexts);
        for (auto& req : batch->requests) {
            try {
                if (!result) {
                    req.promise.set_value(
                        yams::Result<std::vector<std::vector<float>>>(result.error()));
                } else {
                    auto& mat = result.value();
                    if (req.startIdx + req.count <= mat.size()) {
                        std::vector<std::vector<float>> slice(
                            mat.begin() + static_cast<ptrdiff_t>(req.startIdx),
                            mat.begin() + static_cast<ptrdiff_t>(req.startIdx + req.count));
                        req.promise.set_value(std::move(slice));
                    } else {
                        req.promise.set_value(yams::Error{yams::ErrorCode::InternalError,
                                                          "batch coalescer: output size mismatch"});
                    }
                }
            } catch (...) {
            }
        }
        cv_.notify_all();
    }

    Config cfg_;
    std::mutex mu_;
    std::condition_variable cv_;
    std::unordered_map<std::string, std::unique_ptr<PendingBatch>> pending_;
};

struct ProviderCtx {
    enum class State : uint8_t { Unloaded, Loading, Ready, Failed };

    // Track model load failures for backoff and event signaling
    struct FailureInfo {
        std::chrono::steady_clock::time_point lastAttempt;
        std::chrono::steady_clock::time_point lastFailure;
        uint32_t consecutiveFailures = 0;
        std::string lastError;
        bool eventFired = false; // Whether we've fired ModelLoadFailedEvent
    };

    std::mutex mu;
    yams_model_load_progress_cb progress_cb = nullptr;
    void* progress_user = nullptr;
    std::unique_ptr<yams::daemon::OnnxModelPool> pool;
    std::unique_ptr<yams::daemon::OnnxColbertSession> colbert;   // ColBERT session
    std::unordered_map<std::string, State> model_states;         // per-model FSM
    std::unordered_map<std::string, FailureInfo> model_failures; // failure tracking
    bool ready = false;
    bool disabled = false;
    std::atomic<bool> shutdownRequested{false};
    bool gpuEnabled = false; // Tracks whether GPU acceleration is configured
    std::string actualExecutionProvider{"cpu"};
    std::string last_error;
    std::string colbertModelPath; // Path to ColBERT model
    std::size_t configuredMaxLoadedModels = 0;
    std::size_t configuredHotPoolSize = 0;
    yams::onnx_util::OrtRuntimeInfo runtimeInfo;
    std::unique_ptr<BatchCoalescer> coalescer;

    // Check if model is in cooldown period after failures
    bool isInCooldown(const std::string& modelId) const {
        auto it = model_failures.find(modelId);
        if (it == model_failures.end() || it->second.consecutiveFailures == 0) {
            return false;
        }
        // Exponential backoff: 2^min(failures, 6) seconds (max ~64 seconds)
        auto backoffSeconds =
            std::chrono::seconds(1 << std::min(it->second.consecutiveFailures, 6u));
        auto elapsed = std::chrono::steady_clock::now() - it->second.lastFailure;
        return elapsed < backoffSeconds;
    }

    // Record a model load failure
    void recordFailure(const std::string& modelId, const std::string& error) {
        auto& info = model_failures[modelId];
        auto now = std::chrono::steady_clock::now();
        info.lastAttempt = now;
        info.lastFailure = now;
        info.consecutiveFailures++;
        info.lastError = error;
    }

    // Clear failure tracking on successful load
    void clearFailure(const std::string& modelId) {
        auto it = model_failures.find(modelId);
        if (it != model_failures.end()) {
            it->second.consecutiveFailures = 0;
            it->second.lastError.clear();
            it->second.eventFired = false;
        }
    }

    ProviderCtx() {
        // Allow disabling via env for diagnostics or platform constraints
        if (const char* d = std::getenv("YAMS_ONNX_PLUGIN_DISABLE"); d && *d) {
            disabled = true;
            last_error = "disabled_by_env";
            return;
        }

        runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
        if (!runtimeInfo.available) {
            last_error = runtimeInfo.reason.empty() ? "runtime_not_available" : runtimeInfo.reason;
            spdlog::warn("[ONNX-Plugin] ONNX Runtime unavailable: {}",
                         runtimeInfo.errorMessage.empty() ? last_error : runtimeInfo.errorMessage);
            return;
        }

        yams::daemon::ModelPoolConfig cfg;
        // Defaults: prefer lazy loading to avoid blocking startup
        // Models will be loaded on first use, not during plugin initialization
        cfg.lazyLoading = true;
        cfg.enableGPU = true; // runtime detection in appendGpuExecutionProvider()
        const unsigned hwThreads = std::max(1u, std::thread::hardware_concurrency());
        cfg.numThreads = static_cast<int>(std::clamp<unsigned>(hwThreads / 2, 1u, 4u));
        cfg.loadWorkerThreads = std::max<size_t>(1u, static_cast<size_t>(cfg.numThreads / 2));
        cfg.asyncLoading = false;

        // Variables for config (both file-based and JSON-based)
        std::string dataDir;
        bool keepModelHot = true;
        std::string preferredModel;
        std::vector<std::string> preloadList;

        // Read embeddings settings from config.toml (preferred_model, model_path, keep_model_hot)
        try {
            namespace fs = std::filesystem;
            fs::path cfgPath;
            if (const char* explicitConfig = std::getenv("YAMS_CONFIG");
                explicitConfig && *explicitConfig) {
                cfgPath = explicitConfig;
            } else if (const char* xdg = std::getenv("XDG_CONFIG_HOME")) {
                cfgPath = fs::path(xdg) / "yams" / "config.toml";
            } else if (const char* home = std::getenv("HOME")) {
                cfgPath = fs::path(home) / ".config" / "yams" / "config.toml";
            }

            // First, check for YAMS_STORAGE environment variable
            if (const char* storage = std::getenv("YAMS_STORAGE")) {
                dataDir = storage;
            }

            spdlog::debug("[ONNX-Plugin] Config path: {}", cfgPath.string());
            if (!cfgPath.empty() && fs::exists(cfgPath)) {
                spdlog::debug("[ONNX-Plugin] Config file exists, parsing...");
                std::ifstream file(cfgPath);
                std::string line;
                std::string section;
                auto trim = [](std::string& s) {
                    if (s.empty())
                        return;
                    s.erase(0, s.find_first_not_of(" \t"));
                    auto pos = s.find_last_not_of(" \t");
                    if (pos != std::string::npos)
                        s.erase(pos + 1);
                };
                auto apply_num_threads = [&](const std::string& value) {
                    try {
                        int v = std::stoi(value);
                        if (v > 0) {
                            cfg.numThreads = v;
                        }
                    } catch (...) {
                    }
                };
                auto apply_max_loaded_models = [&](const std::string& value) {
                    try {
                        std::size_t v = static_cast<std::size_t>(std::stoull(value));
                        if (v >= 1 && v <= 64) {
                            cfg.maxLoadedModels = v;
                        }
                    } catch (...) {
                    }
                };
                auto apply_hot_pool_size = [&](const std::string& value) {
                    try {
                        std::size_t v = static_cast<std::size_t>(std::stoull(value));
                        if (v >= 1 && v <= 64) {
                            cfg.hotPoolSize = v;
                        }
                    } catch (...) {
                    }
                };
                auto apply_async_loading = [&](const std::string& value) {
                    std::string v;
                    v.reserve(value.size());
                    for (unsigned char c : value)
                        v.push_back(static_cast<char>(std::tolower(c)));
                    if (v == "1" || v == "true" || v == "yes" || v == "on")
                        cfg.asyncLoading = true;
                    if (v == "0" || v == "false" || v == "no" || v == "off")
                        cfg.asyncLoading = false;
                };
                while (std::getline(file, line)) {
                    if (line.empty() || line[0] == '#')
                        continue;
                    if (line[0] == '[') {
                        auto end = line.find(']');
                        section = (end != std::string::npos) ? line.substr(1, end - 1) : "";
                        spdlog::debug("[ONNX-Plugin] Entering section: [{}]", section);
                        continue;
                    }
                    auto eq = line.find('=');
                    if (eq == std::string::npos)
                        continue;
                    std::string key = line.substr(0, eq);
                    std::string value = line.substr(eq + 1);
                    trim(key);
                    trim(value);
                    auto hash = value.find('#');
                    if (hash != std::string::npos) {
                        value = value.substr(0, hash);
                        trim(value);
                    }
                    if (value.size() >= 2 && value.front() == '"' && value.back() == '"')
                        value = value.substr(1, value.size() - 2);

                    // Read data directory from [storage].base_path if not set by env
                    if (section == "storage" && key == "base_path") {
                        spdlog::debug("[ONNX-Plugin] Found [storage].base_path = '{}'", value);
                        if (dataDir.empty()) {
                            dataDir = value;
                        }
                    }

                    if (section == "embeddings") {
                        if (key == "preferred_model" && preferredModel.empty())
                            preferredModel = std::move(value);
                        else if (key == "num_threads")
                            apply_num_threads(value);
                        else if (key == "keep_model_hot") {
                            std::string v;
                            v.reserve(value.size());
                            for (unsigned char c : value)
                                v.push_back(static_cast<char>(std::tolower(c)));
                            keepModelHot = !(v == "false" || v == "0" || v == "no" || v == "off");
                        }
                    }
                    if (section == "daemon.models") {
                        if (key == "num_threads")
                            apply_num_threads(value);
                        else if (key == "max_loaded_models")
                            apply_max_loaded_models(value);
                        else if (key == "hot_pool_size")
                            apply_hot_pool_size(value);
                    }
                    // New: preload list under [plugins.onnx]
                    if (section == "plugins.onnx" && key == "preload") {
                        // parse CSV style list: preload = "modelA,modelB"
                        std::string s = value;
                        for (char& ch : s) {
                            if (ch == '[' || ch == ']')
                                ch = ' ';
                        }
                        trim(s);
                        // allow ["a","b"] or a,b
                        if (!s.empty() && s.front() == '"' && s.back() == '"')
                            s = s.substr(1, s.size() - 2);
                        size_t start = 0;
                        while (start < s.size()) {
                            auto comma = s.find(',', start);
                            std::string item =
                                s.substr(start, comma == std::string::npos ? std::string::npos
                                                                           : (comma - start));
                            trim(item);
                            if (!item.empty() && item.front() == '"' && item.back() == '"')
                                item = item.substr(1, item.size() - 2);
                            if (!item.empty())
                                preloadList.push_back(item);
                            if (comma == std::string::npos)
                                break;
                            start = comma + 1;
                        }
                    }
                    if (section == "plugins.onnx" && key == "max_loaded_models")
                        apply_max_loaded_models(value);
                    if (section == "plugins.onnx" && key == "hot_pool_size")
                        apply_hot_pool_size(value);
                    if (section == "plugins.onnx" && key == "async_load")
                        apply_async_loading(value);
                    // New: explicit models table entries [plugins.onnx.models.NAME]
                    if (section.rfind("plugins.onnx.models.", 0) == 0 && key == "task") {
                        // Section name encodes model_id; mark for preload as hot
                        std::string model_id =
                            section.substr(std::string("plugins.onnx.models.").size());
                        if (!model_id.empty())
                            preloadList.push_back(model_id);
                    }
                }
            }

            // Set modelsRoot to $dataDir/models if dataDir is available
            spdlog::debug("[ONNX-Plugin] After config parsing: dataDir='{}'", dataDir);
            if (!dataDir.empty()) {
                // Expand ~ if present
                if (dataDir[0] == '~') {
                    if (const char* home = std::getenv("HOME"))
                        dataDir = std::string(home) + dataDir.substr(1);
                }
                cfg.modelsRoot = dataDir + "/models";
                spdlog::info("[ONNX-Plugin] Using models directory: {}", cfg.modelsRoot);
            }
            // Build preload list if any; eager load behavior is controlled later.
            if (!preloadList.empty()) {
                cfg.preloadModels = preloadList;
            } else if (keepModelHot && !preferredModel.empty()) {
                cfg.preloadModels = {preferredModel};
            }
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
            // Force no-preload in test/mock environments to avoid background threads
            if (std::getenv("YAMS_SKIP_MODEL_LOADING") || std::getenv("YAMS_TEST_MODE") ||
                std::getenv("YAMS_USE_MOCK_PROVIDER")) {
                cfg.preloadModels.clear();
                cfg.lazyLoading = true;
            }
#endif
        } catch (const std::exception&) {
            // ignore config parse errors
        }

        // Env overrides for async loading
        auto parse_bool_env = [](const char* name) -> std::optional<bool> {
            if (const char* s = std::getenv(name)) {
                std::string v(s);
                for (auto& c : v)
                    c = static_cast<char>(std::tolower(c));
                if (v == "1" || v == "true" || v == "yes" || v == "on")
                    return true;
                if (v == "0" || v == "false" || v == "no" || v == "off")
                    return false;
            }
            return std::nullopt;
        };
        if (auto val = parse_bool_env("YAMS_ONNX_ASYNC_LOAD")) {
            cfg.asyncLoading = *val;
        }
        if (parse_bool_env("YAMS_ONNX_FORCE_CPU").value_or(false)) {
            cfg.enableGPU = false;
        }
        const bool eagerPreload = parse_bool_env("YAMS_ONNX_EAGER_PRELOAD").value_or(false);
        cfg.lazyLoading = !eagerPreload;
        if (const char* w = std::getenv("YAMS_ONNX_LOAD_WORKERS")) {
            try {
                size_t threads = static_cast<size_t>(std::stoul(w));
                if (threads > 0 && threads <= 32) {
                    cfg.loadWorkerThreads = threads;
                }
            } catch (...) {
            }
        }

        // Parse JSON config from plugin init (overrides file-based config)
        if (!g_plugin_config_json.empty()) {
            try {
                auto j = nlohmann::json::parse(g_plugin_config_json, nullptr, false);
                if (!j.is_discarded() && j.is_object()) {
                    spdlog::info("[ONNX-Plugin] Applying JSON config: {}", g_plugin_config_json);
                    // preferred_model
                    if (j.contains("preferred_model") && j["preferred_model"].is_string()) {
                        preferredModel = j["preferred_model"].get<std::string>();
                        spdlog::info("[ONNX-Plugin] JSON config: preferred_model='{}'",
                                     preferredModel);
                    }
                    // preload (CSV string or single model)
                    if (j.contains("preload") && j["preload"].is_string()) {
                        std::string preloadStr = j["preload"].get<std::string>();
                        preloadList.clear();
                        // Parse CSV: "modelA,modelB" or single "modelA"
                        size_t start = 0;
                        while (start < preloadStr.size()) {
                            auto comma = preloadStr.find(',', start);
                            std::string item = preloadStr.substr(start, comma == std::string::npos
                                                                            ? std::string::npos
                                                                            : (comma - start));
                            // Trim whitespace
                            auto trim = [](std::string& s) {
                                if (s.empty())
                                    return;
                                s.erase(0, s.find_first_not_of(" \t"));
                                auto pos = s.find_last_not_of(" \t");
                                if (pos != std::string::npos)
                                    s.erase(pos + 1);
                            };
                            trim(item);
                            if (!item.empty())
                                preloadList.push_back(item);
                            if (comma == std::string::npos)
                                break;
                            start = comma + 1;
                        }
                        spdlog::info("[ONNX-Plugin] JSON config: preload list has {} models",
                                     preloadList.size());
                    }
                    // keep_model_hot
                    if (j.contains("keep_model_hot") && j["keep_model_hot"].is_boolean()) {
                        keepModelHot = j["keep_model_hot"].get<bool>();
                        spdlog::info("[ONNX-Plugin] JSON config: keep_model_hot={}", keepModelHot);
                    }
                    // num_threads
                    if (j.contains("num_threads") && j["num_threads"].is_number_integer()) {
                        int v = j["num_threads"].get<int>();
                        if (v > 0) {
                            cfg.numThreads = v;
                            spdlog::info("[ONNX-Plugin] JSON config: num_threads={}",
                                         cfg.numThreads);
                        }
                    }
                    if (j.contains("max_loaded_models") &&
                        j["max_loaded_models"].is_number_integer()) {
                        auto v = static_cast<std::size_t>(j["max_loaded_models"].get<int>());
                        if (v >= 1 && v <= 64) {
                            cfg.maxLoadedModels = v;
                            spdlog::info("[ONNX-Plugin] JSON config: max_loaded_models={}",
                                         cfg.maxLoadedModels);
                        }
                    }
                    if (j.contains("hot_pool_size") && j["hot_pool_size"].is_number_integer()) {
                        auto v = static_cast<std::size_t>(j["hot_pool_size"].get<int>());
                        if (v >= 1 && v <= 64) {
                            cfg.hotPoolSize = v;
                            spdlog::info("[ONNX-Plugin] JSON config: hot_pool_size={}",
                                         cfg.hotPoolSize);
                        }
                    }
                    if (j.contains("async_load") && j["async_load"].is_boolean()) {
                        cfg.asyncLoading = j["async_load"].get<bool>();
                        spdlog::info("[ONNX-Plugin] JSON config: async_load={}", cfg.asyncLoading);
                    }
                    if (j.contains("load_worker_threads") &&
                        j["load_worker_threads"].is_number_integer()) {
                        auto v = static_cast<std::size_t>(j["load_worker_threads"].get<int>());
                        if (v >= 1 && v <= 32) {
                            cfg.loadWorkerThreads = v;
                            spdlog::info("[ONNX-Plugin] JSON config: load_worker_threads={}",
                                         cfg.loadWorkerThreads);
                        }
                    }
                    // models_root - override the models directory
                    if (j.contains("models_root") && j["models_root"].is_string()) {
                        cfg.modelsRoot = j["models_root"].get<std::string>();
                        spdlog::info("[ONNX-Plugin] JSON config: models_root='{}'", cfg.modelsRoot);
                    }
                }
            } catch (const std::exception& e) {
                spdlog::warn("[ONNX-Plugin] Failed to parse JSON config: {}", e.what());
            }
            // Re-apply preload settings from JSON config
            if (!preloadList.empty()) {
                cfg.preloadModels = preloadList;
            } else if (keepModelHot && !preferredModel.empty()) {
                cfg.preloadModels = {preferredModel};
            }
        }

        cfg.maxLoadedModels = std::max<std::size_t>(1, cfg.maxLoadedModels);
        cfg.hotPoolSize = std::clamp(cfg.hotPoolSize, std::size_t{1}, cfg.maxLoadedModels);

        if (!cfg.preloadModels.empty()) {
            const std::size_t preloadCap = std::min(cfg.hotPoolSize, cfg.maxLoadedModels);
            if (cfg.preloadModels.size() > preloadCap) {
                spdlog::warn("[ONNX-Plugin] Trimming preload list from {} to {} (hot_pool_size={}, "
                             "max_loaded_models={})",
                             cfg.preloadModels.size(), preloadCap, cfg.hotPoolSize,
                             cfg.maxLoadedModels);
                cfg.preloadModels.resize(preloadCap);
            }
            if (cfg.lazyLoading) {
                spdlog::info("[ONNX-Plugin] preload list configured ({} models), lazy loading "
                             "enabled; models will load on first use",
                             cfg.preloadModels.size());
            }
        }

        configuredMaxLoadedModels = cfg.maxLoadedModels;
        configuredHotPoolSize = cfg.hotPoolSize;
        gpuEnabled = cfg.enableGPU;
        spdlog::info("[ONNX-Plugin] Creating OnnxModelPool with modelsRoot={}, gpuEnabled={}, "
                     "async_load={}",
                     cfg.modelsRoot, gpuEnabled, cfg.asyncLoading);
        pool = std::make_unique<yams::daemon::OnnxModelPool>(cfg);
        spdlog::info("[ONNX-Plugin] OnnxModelPool created, calling initialize()...");
        try {
            auto res = pool->initialize();
            if (res) {
                ready = true;
                last_error.clear();
                spdlog::info("[ONNX-Plugin] Pool initialized successfully, ready=true");

                // Best-effort: derive actual EP from a preloaded model (if any)
                try {
                    if (!cfg.preloadModels.empty()) {
                        auto h = pool->acquireModel(cfg.preloadModels.front(),
                                                    std::chrono::milliseconds(500));
                        if (h) {
                            actualExecutionProvider = h.value()->getExecutionProvider();
                        }
                    }
                } catch (...) {
                }
            } else {
                ready = false;
                last_error = res.error().message;
                spdlog::warn("[ONNX-Plugin] Pool initialize failed: {}", last_error);
            }
        } catch (const std::exception& e) {
            spdlog::error("[ONNX-Plugin] Pool initialize exception: {}", e.what());
            ready = false;
            last_error = e.what();
        }

        std::string colbertModelName = preferredModel;
        if (colbertModelName.empty() && !preloadList.empty()) {
            colbertModelName = preloadList.front();
        }
        if (isColbertModelName(colbertModelName)) {
            colbertModelPath = cfg.modelsRoot + "/" + colbertModelName + "/model.onnx";
            yams::daemon::ColbertConfig colbertCfg{};
            colbertCfg.model_path = colbertModelPath;
            colbertCfg.model_name = colbertModelName;
            colbertCfg.enable_gpu = cfg.enableGPU;
            colbertCfg.num_threads = static_cast<int>(cfg.numThreads);
            colbertCfg.max_sequence_length = 512;
            colbertCfg.token_dim = 48;
            colbertCfg.query_marker_id = 50368;
            colbertCfg.doc_marker_id = 50369;
            colbertCfg.skiplist = readColbertSkiplist(cfg.modelsRoot, colbertModelName);
            colbert = std::make_unique<yams::daemon::OnnxColbertSession>(
                colbertModelPath, colbertModelName, colbertCfg);
        }
        // Initialize batch coalescer (off by default; enable with YAMS_ONNX_BATCH_COALESCE=1).
        // Only useful for workloads with many concurrent small requests (1-4 texts).
        // Callers that already send well-sized batches should NOT enable this — it
        // adds latency (coalesce window) and serializes concurrent inferences.
        {
            BatchCoalescer::Config coalesceCfg;
            coalesceCfg.enabled = envTruthy(std::getenv("YAMS_ONNX_BATCH_COALESCE"));
            if (const char* w = std::getenv("YAMS_ONNX_COALESCE_WINDOW_MS")) {
                try {
                    coalesceCfg.windowMs = std::stoi(w);
                } catch (...) {
                }
            }
            if (const char* m = std::getenv("YAMS_ONNX_COALESCE_MAX_BATCH")) {
                try {
                    coalesceCfg.maxBatch = static_cast<size_t>(std::stoull(m));
                } catch (...) {
                }
            }
            coalescer = std::make_unique<BatchCoalescer>(coalesceCfg);
            if (coalesceCfg.enabled) {
                spdlog::info("[ONNX-Plugin] Batch coalescing enabled: window={}ms max_batch={}",
                             coalesceCfg.windowMs, coalesceCfg.maxBatch);
            }
        }

        spdlog::info("[ONNX-Plugin] ProviderCtx init complete: ready={} pool={}", ready,
                     pool ? "valid" : "null");
    }
};

static yams_status_t onnx_set_threading(void* self, const char* model_id, int intra, int inter) {
    if (!self || !model_id || !*model_id) {
        return YAMS_ERR_INVALID_ARG;
    }
    if (intra == 0 || inter == 0) {
        return YAMS_ERR_INVALID_ARG;
    }
    if ((intra < -1 || intra > 64) || (inter < -1 || inter > 64)) {
        return YAMS_ERR_INVALID_ARG;
    }

    auto* c = static_cast<ProviderCtx*>(self);
    if (c->disabled) {
        return YAMS_ERR_UNSUPPORTED;
    }
    if (!c->pool) {
        return YAMS_ERR_INTERNAL;
    }

    auto r = c->pool->setModelThreading(model_id, intra, inter, true);
    if (!r) {
        spdlog::warn("[ONNX Plugin] set_threading failed model='{}' intra={} inter={}: {}",
                     model_id, intra, inter, r.error().message);
        return YAMS_ERR_INTERNAL;
    }

    spdlog::info("[ONNX Plugin] set_threading applied model='{}' intra={} inter={}", model_id,
                 intra, inter);
    return YAMS_OK;
}

static bool shutdownInProgress(const ProviderCtx* c) {
    if (c && c->shutdownRequested.load(std::memory_order_acquire)) {
        return true;
    }
    return envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"));
}

static ProviderCtx* refreshCtxIfStaleShutdown(ProviderCtx* c) {
    if (!c) {
        return nullptr;
    }

    const bool marker = envTruthy(std::getenv("YAMS_ONNX_SHUTDOWN_IN_PROGRESS"));
    if (marker || !c->shutdownRequested.load(std::memory_order_acquire)) {
        return c;
    }

    // Stale in-process lifecycle: provider self pointer still references a shutdown context
    // from a prior daemon instance. Force provider revival via canonical accessor.
    auto* provider = yams_onnx_get_model_provider();
    if (provider && provider->self) {
        spdlog::info("[ONNX Plugin] Revived stale provider context after prior shutdown");
        return static_cast<ProviderCtx*>(provider->self);
    }

    return c;
}

// Helpers for progress emission
static void emit_progress(ProviderCtx* ctx, const char* model, int phase, const char* msg,
                          uint64_t cur = 0, uint64_t tot = 0) {
    if (!ctx || !ctx->progress_cb)
        return;
    try {
        ctx->progress_cb(ctx->progress_user, model, phase, cur, tot, msg);
    } catch (...) {
    }
}

struct ProviderSingleton {
    ProviderCtx ctx;
    yams_model_provider_v1 vtable;

    std::atomic<bool> shutdownCalled{false};
    std::atomic<bool> explicitShutdownCalled{false}; // Set by yams_onnx_shutdown_provider()
    std::mutex lifecycleMu;

    void ensureActive() {
        std::lock_guard<std::mutex> lk(lifecycleMu);
        if (!shutdownCalled.load(std::memory_order_acquire)) {
            return;
        }

        spdlog::info("[ONNX Plugin] Reviving provider after explicit shutdown");

        try {
            ctx.~ProviderCtx();
        } catch (...) {
        }
        new (&ctx) ProviderCtx();
        vtable.self = &ctx;

        shutdownCalled.store(false, std::memory_order_release);
        explicitShutdownCalled.store(false, std::memory_order_release);
    }

    void shutdown(bool isExplicit = false) noexcept {
        std::lock_guard<std::mutex> lk(lifecycleMu);
        if (shutdownCalled.exchange(true, std::memory_order_acq_rel)) {
            return;
        }

        ctx.shutdownRequested.store(true, std::memory_order_release);
        ctx.ready = false;

        // Drain any pending coalesced batches so waiting threads don't hang.
        if (ctx.coalescer) {
            ctx.coalescer->drain();
        }

        // If this is called from destructor during static destruction (not explicit),
        // skip cleanup to avoid crashes from corrupted global state.
        // At process exit, resources will be freed anyway.
        if (!isExplicit && !explicitShutdownCalled.load()) {
            return; // Skip cleanup during static destruction
        }

        if (ctx.pool) {
            try {
                ctx.pool->shutdown();
            } catch (const std::exception& e) {
                try {
                    spdlog::warn("[ONNX] pool shutdown exception: {}", e.what());
                } catch (...) {
                }
            } catch (...) {
                try {
                    spdlog::warn("[ONNX] pool shutdown unknown exception");
                } catch (...) {
                }
            }
            try {
                ctx.pool.reset();
            } catch (const std::exception& e) {
                try {
                    spdlog::warn("[ONNX] pool reset exception: {}", e.what());
                } catch (...) {
                }
            } catch (...) {
                try {
                    spdlog::warn("[ONNX] pool reset unknown exception");
                } catch (...) {
                }
            }
        }
    }

    ~ProviderSingleton() noexcept {
        shutdown(false); // Called from destructor, not explicit
    }

    ProviderSingleton() {
        vtable.abi_version = 4; // v1.4 (added evict_under_pressure for memory pressure)
        vtable.self = &ctx;

        // Use daemon's default logger - no separate file needed
        spdlog::debug("[ONNX Plugin] Initialized (v1.2) - using daemon logger");

        vtable.set_progress_callback = [](void* self, yams_model_load_progress_cb cb,
                                          void* user) -> yams_status_t {
            if (!self)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            std::lock_guard<std::mutex> lk(c->mu);
            c->progress_cb = cb;
            c->progress_user = user;
            return YAMS_OK;
        };

        vtable.load_model = [](void* self, const char* model_id, const char* /*model_path*/,
                               const char* options_json) -> yams_status_t {
            if (!self || !model_id)
                return YAMS_ERR_INVALID_ARG;
            // Wrap entire lambda body in try/catch to prevent exceptions from crossing ABI
            // boundary. On Windows, ONNX Runtime can throw std::system_error("resource deadlock
            // would occur") when thread pool resources are exhausted (PBI-1c1).
            try {
                auto* c = static_cast<ProviderCtx*>(self);
                if (c->disabled)
                    return YAMS_ERR_UNSUPPORTED;
                if (shutdownInProgress(c)) {
                    spdlog::info("[ONNX Plugin] load_model skipped during shutdown for '{}'",
                                 model_id);
                    return YAMS_ERR_INTERNAL;
                }
                if (!c->ready || !c->pool)
                    return YAMS_ERR_INTERNAL;
                // Parse per-load options (hf.revision, offline)
                try {
                    yams::daemon::OnnxModelPool::ResolutionHints hints;
                    if (options_json && std::strlen(options_json) > 0) {
                        auto j = nlohmann::json::parse(options_json, nullptr, false);
                        if (!j.is_discarded()) {
                            if (j.contains("hf") && j["hf"].is_object()) {
                                auto& hf = j["hf"];
                                if (hf.contains("revision") && hf["revision"].is_string())
                                    hints.hfRevision = hf["revision"].get<std::string>();
                                if (hf.contains("offline") && hf["offline"].is_boolean())
                                    hints.offlineOnly = hf["offline"].get<bool>();
                            }
                            if (j.contains("offline") && j["offline"].is_boolean())
                                hints.offlineOnly = j["offline"].get<bool>();
                        }
                    }
                    if (!hints.hfRevision.empty() || hints.offlineOnly) {
                        c->pool->setResolutionHints(model_id, hints);
                    }
                } catch (...) {
                }
                std::optional<bool> hotModel;
                std::string modelIdStr(model_id);
                if (options_json && *options_json) {
                    try {
                        auto j = nlohmann::json::parse(options_json);
                        if (j.contains("hot") && j["hot"].is_boolean()) {
                            hotModel = j["hot"].get<bool>();
                        }
                    } catch (...) {
                    }
                }

                // FSM: if already loading/ready, return OK
                {
                    std::lock_guard<std::mutex> lk(c->mu);
                    auto& st = c->model_states[modelIdStr];
                    if (st == ProviderCtx::State::Ready) {
                        if (hotModel.has_value()) {
                            (void)c->pool->setModelHot(modelIdStr, hotModel.value());
                        }
                        spdlog::debug("[ONNX Plugin] load_model: '{}' already ready", model_id);
                        return YAMS_OK;
                    }
                    if (st == ProviderCtx::State::Loading) {
                        spdlog::debug("[ONNX Plugin] load_model: '{}' already loading", model_id);
                        return YAMS_OK;
                    }
                    // Check if in cooldown from previous failures
                    if (c->isInCooldown(modelIdStr)) {
                        auto& info = c->model_failures[modelIdStr];
                        spdlog::warn(
                            "[ONNX Plugin] load_model: '{}' in cooldown ({} failures), skipping",
                            model_id, info.consecutiveFailures);
                        return YAMS_ERR_INTERNAL;
                    }
                    st = ProviderCtx::State::Loading;
                }
                spdlog::info("[ONNX Plugin] load_model: starting load for '{}'", model_id);
                emit_progress(c, model_id, YAMS_MODEL_PHASE_PROBE, "probe");

                // Simply call loadModel synchronously - it returns immediately since the pool uses
                // lazy loading
                auto r = c->pool->loadModel(modelIdStr);
                if (r) {
                    {
                        std::lock_guard<std::mutex> lk(c->mu);
                        c->model_states[modelIdStr] = ProviderCtx::State::Ready;
                        c->clearFailure(modelIdStr);
                    }
                    spdlog::info("[ONNX Plugin] load_model: '{}' loaded successfully", model_id);
                    if (hotModel.has_value()) {
                        (void)c->pool->setModelHot(modelIdStr, hotModel.value());
                    }
                    emit_progress(c, model_id, YAMS_MODEL_PHASE_READY, "ready");
                    auto q =
                        yams::daemon::InternalEventBus::instance()
                            .get_or_create_channel<yams::daemon::InternalEventBus::ModelReadyEvent>(
                                "model.events", 256);
                    (void)q->try_push({modelIdStr});
                } else {
                    std::string errMsg = r.error().message;
                    {
                        std::lock_guard<std::mutex> lk(c->mu);
                        c->model_states[modelIdStr] = ProviderCtx::State::Failed;
                        c->recordFailure(modelIdStr, errMsg);
                    }
                    spdlog::error("[ONNX Plugin] load_model: '{}' failed: {}", model_id, errMsg);
                    emit_progress(c, model_id, YAMS_MODEL_PHASE_UNKNOWN, errMsg.c_str());
                    auto q = yams::daemon::InternalEventBus::instance()
                                 .get_or_create_channel<
                                     yams::daemon::InternalEventBus::ModelLoadFailedEvent>(
                                     "model.events", 256);
                    (void)q->try_push({std::string(model_id), r.error().message});
                }
                return YAMS_OK;
            } catch (const std::system_error& e) {
                fprintf(stderr, "[ONNX Plugin] load_model std::system_error: %s\n", e.what());
                spdlog::error("[ONNX Plugin] load_model std::system_error: {}", e.what());
                return YAMS_ERR_INTERNAL;
            } catch (const std::exception& e) {
                fprintf(stderr, "[ONNX Plugin] load_model exception: %s\n", e.what());
                spdlog::error("[ONNX Plugin] load_model exception: {}", e.what());
                return YAMS_ERR_INTERNAL;
            } catch (...) {
                fprintf(stderr, "[ONNX Plugin] load_model unknown exception\n");
                spdlog::error("[ONNX Plugin] load_model unknown exception");
                return YAMS_ERR_INTERNAL;
            }
        };

        vtable.unload_model = [](void* self, const char* model_id) -> yams_status_t {
            if (!self || !model_id)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            if (c->disabled)
                return YAMS_ERR_UNSUPPORTED;
            if (!c->ready || !c->pool)
                return YAMS_ERR_INTERNAL;
            auto r = c->pool->unloadModel(model_id);
            return r ? YAMS_OK : YAMS_ERR_NOT_FOUND;
        };

        vtable.is_model_loaded = [](void* self, const char* model_id,
                                    bool* out_loaded) -> yams_status_t {
            if (!self || !model_id || !out_loaded)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            if (c->disabled) {
                *out_loaded = false;
                return YAMS_ERR_UNSUPPORTED;
            }
            if (!c->ready || !c->pool) {
                *out_loaded = false;
                return YAMS_ERR_INTERNAL;
            }
            bool loaded = c->pool->isModelLoaded(model_id);
            // Consider FSM Ready as loaded (pool may be lazy)
            {
                std::lock_guard<std::mutex> lk(c->mu);
                auto it = c->model_states.find(model_id);
                if (!loaded && it != c->model_states.end() &&
                    it->second == ProviderCtx::State::Ready)
                    loaded = true;
            }
            *out_loaded = loaded;
            return YAMS_OK;
        };

        vtable.get_loaded_models = [](void* self, const char*** out_ids,
                                      size_t* out_count) -> yams_status_t {
            if (!self || !out_ids || !out_count)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            if (c->disabled)
                return YAMS_ERR_UNSUPPORTED;
            if (!c->ready || !c->pool) {
                *out_ids = nullptr;
                *out_count = 0;
                return YAMS_ERR_INTERNAL;
            }
            auto v = c->pool->getLoadedModels();
            const size_t n = v.size();
            const char** ids = nullptr;
            if (n > 0) {
                ids = static_cast<const char**>(std::malloc(sizeof(char*) * n));
                if (!ids)
                    return YAMS_ERR_INTERNAL;
                size_t i = 0;
                for (const auto& s : v) {
                    char* dup = static_cast<char*>(std::malloc(s.size() + 1));
                    if (!dup) {
                        for (size_t j = 0; j < i; ++j)
                            std::free(const_cast<char*>(ids[j]));
                        std::free(ids);
                        return YAMS_ERR_INTERNAL;
                    }
                    std::memcpy(dup, s.c_str(), s.size() + 1);
                    ids[i++] = dup;
                }
            }
            *out_ids = ids;
            *out_count = n;
            return YAMS_OK;
        };

        vtable.free_model_list = [](void* /*self*/, const char** ids, size_t count) {
            if (!ids)
                return;
            for (size_t i = 0; i < count; ++i)
                std::free(const_cast<char*>(ids[i]));
            std::free(const_cast<char**>(ids));
        };

        vtable.generate_embedding = [](void* self, const char* model_id, const uint8_t* input,
                                       size_t input_len, float** out_vec,
                                       size_t* out_dim) -> yams_status_t {
            if (!self || !model_id || !input || !out_vec || !out_dim)
                return YAMS_ERR_INVALID_ARG;
            // Wrap entire lambda body in try/catch to prevent exceptions from crossing ABI
            // boundary. On Windows, ONNX Runtime can throw std::system_error("resource deadlock
            // would occur") when thread pool resources are exhausted (PBI-1c1).
            try {
                auto* c = refreshCtxIfStaleShutdown(static_cast<ProviderCtx*>(self));
                if (c->disabled) {
                    spdlog::warn("[ONNX Plugin] generate_embedding: plugin disabled");
                    return YAMS_ERR_UNSUPPORTED;
                }
                if (shutdownInProgress(c)) {
                    spdlog::info("[ONNX Plugin] generate_embedding: shutdown in progress");
                    *out_vec = nullptr;
                    *out_dim = 0;
                    return YAMS_ERR_INTERNAL;
                }
                if (!c->ready) {
                    spdlog::warn("[ONNX Plugin] generate_embedding: plugin not ready");
                    return YAMS_ERR_INTERNAL;
                }
                if (!c->pool) {
                    spdlog::warn("[ONNX Plugin] generate_embedding: pool is null");
                    return YAMS_ERR_INTERNAL;
                }
                std::string modelIdStr(model_id);

                // Check if model is in cooldown after previous failures
                {
                    std::lock_guard<std::mutex> lk(c->mu);
                    if (c->isInCooldown(modelIdStr)) {
                        auto& info = c->model_failures[modelIdStr];
                        spdlog::debug("[ONNX Plugin] Model '{}' in cooldown ({} failures)",
                                      model_id, info.consecutiveFailures);
                        *out_vec = nullptr;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }
                }

                std::string text(reinterpret_cast<const char*>(input), input_len);
                auto h = c->pool->acquireModel(model_id, modelAcquireTimeout(false));
                if (!h) {
                    const auto err = h.error();
                    const bool acquireTimeout = (err.code == yams::ErrorCode::Timeout);
                    std::string errMsg = h.error().message;
                    spdlog::warn("[ONNX Plugin] acquireModel failed for {}: {}", model_id, errMsg);

                    // Resource-pool acquisition timeout is usually transient contention, not a
                    // model load failure. Do not poison cooldown state for this path.
                    if (acquireTimeout) {
                        *out_vec = nullptr;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }

                    // Record failure and fire event if this is first failure
                    {
                        std::lock_guard<std::mutex> lk(c->mu);
                        c->recordFailure(modelIdStr, errMsg);
                        auto& info = c->model_failures[modelIdStr];
                        c->model_states[modelIdStr] = ProviderCtx::State::Failed;

                        if (!info.eventFired) {
                            info.eventFired = true;
                            spdlog::error("[ONNX Plugin] Model '{}' load failed, firing event: {}",
                                          model_id, errMsg);
                            try {
                                auto q =
                                    yams::daemon::InternalEventBus::instance()
                                        .get_or_create_channel<
                                            yams::daemon::InternalEventBus::ModelLoadFailedEvent>(
                                            "model.events", 256);
                                (void)q->try_push({modelIdStr, errMsg});
                            } catch (...) {
                            }
                        }
                    }
                    *out_vec = nullptr;
                    *out_dim = 0;
                    return YAMS_ERR_INTERNAL;
                }

                // Model acquired successfully - clear failure tracking
                {
                    std::lock_guard<std::mutex> lk(c->mu);
                    c->clearFailure(modelIdStr);
                    c->model_states[modelIdStr] = ProviderCtx::State::Ready;
                }

                auto& session = *h.value();
                auto r =
                    const_cast<yams::daemon::OnnxModelSession&>(session).generateEmbedding(text);
                if (!r) {
                    spdlog::warn("[ONNX Plugin] generateEmbedding failed for {}: {}", model_id,
                                 r.error().message);
                    return YAMS_ERR_INTERNAL;
                }
                auto& vec = r.value();
                const size_t n = vec.size();
                float* buf = static_cast<float*>(std::malloc(sizeof(float) * n));
                if (!buf)
                    return YAMS_ERR_INTERNAL;
                std::memcpy(buf, vec.data(), sizeof(float) * n);
                *out_vec = buf;
                *out_dim = n;
                return YAMS_OK;
            } catch (const std::system_error& e) {
                fprintf(stderr, "[ONNX Plugin] generate_embedding std::system_error: %s\n",
                        e.what());
                spdlog::error("[ONNX Plugin] generate_embedding std::system_error: {}", e.what());
                *out_vec = nullptr;
                *out_dim = 0;
                return YAMS_ERR_INTERNAL;
            } catch (const std::exception& e) {
                fprintf(stderr, "[ONNX Plugin] generate_embedding exception: %s\n", e.what());
                spdlog::error("[ONNX Plugin] generate_embedding exception: {}", e.what());
                *out_vec = nullptr;
                *out_dim = 0;
                return YAMS_ERR_INTERNAL;
            } catch (...) {
                fprintf(stderr, "[ONNX Plugin] generate_embedding unknown exception\n");
                spdlog::error("[ONNX Plugin] generate_embedding unknown exception");
                *out_vec = nullptr;
                *out_dim = 0;
                return YAMS_ERR_INTERNAL;
            }
        };

        vtable.free_embedding = [](void* /*self*/, float* vec, size_t /*dim*/) {
            if (vec)
                std::free(vec);
        };

        vtable.generate_embedding_batch = [](void* self, const char* model_id,
                                             const uint8_t* const* inputs, const size_t* input_lens,
                                             size_t batch_size, float** out_vecs, size_t* out_batch,
                                             size_t* out_dim) -> yams_status_t {
            spdlog::debug("[ONNX Plugin] generate_embedding_batch called: model={} batch={}",
                          model_id ? model_id : "(null)", batch_size);

            if (!self || !model_id || !inputs || !input_lens || !out_vecs || !out_batch ||
                !out_dim) {
                spdlog::error("[ONNX Plugin] Invalid arguments");
                return YAMS_ERR_INVALID_ARG;
            }

            auto* c = refreshCtxIfStaleShutdown(static_cast<ProviderCtx*>(self));
            int retries = shutdownInProgress(c) ? 1 : 3;
            while (retries > 0) {
                retries--;
                try {
                    c = refreshCtxIfStaleShutdown(c);
                    const bool timingDebug = []() {
                        if (const char* s = std::getenv("YAMS_ONNX_DEBUG_BATCH_TIMINGS")) {
                            return std::string{s} == "1" || std::string{s} == "true" ||
                                   std::string{s} == "yes";
                        }
                        return false;
                    }();

                    if (c->disabled) {
                        spdlog::warn("[ONNX Plugin] generate_embedding_batch: plugin disabled");
                        return YAMS_ERR_UNSUPPORTED;
                    }
                    if (shutdownInProgress(c)) {
                        spdlog::info(
                            "[ONNX Plugin] generate_embedding_batch: shutdown in progress");
                        *out_vecs = nullptr;
                        *out_batch = 0;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }
                    std::string modelIdStr(model_id);
                    const bool isColbert = isColbertModelName(modelIdStr);
                    if (!c->ready && !isColbert) {
                        spdlog::warn("[ONNX Plugin] generate_embedding_batch: plugin not ready");
                        *out_vecs = nullptr;
                        *out_batch = 0;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }
                    if (!c->pool && !isColbert) {
                        spdlog::warn("[ONNX Plugin] generate_embedding_batch: pool is null");
                        *out_vecs = nullptr;
                        *out_batch = 0;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }

                    // Check if model is in cooldown
                    {
                        std::lock_guard<std::mutex> lk(c->mu);
                        if (c->isInCooldown(modelIdStr)) {
                            auto& info = c->model_failures[modelIdStr];
                            spdlog::warn("[ONNX Plugin] Model '{}' in cooldown ({} failures)",
                                         model_id, info.consecutiveFailures);
                            *out_vecs = nullptr;
                            *out_batch = 0;
                            *out_dim = 0;
                            return YAMS_ERR_INTERNAL;
                        }
                    }

                    std::vector<std::string> texts;
                    texts.reserve(batch_size);
                    for (size_t i = 0; i < batch_size; ++i) {
                        texts.emplace_back(reinterpret_cast<const char*>(inputs[i]), input_lens[i]);
                    }

                    // Batch coalescing: route small non-ColBERT requests through
                    // the coalescer to accumulate concurrent small batches.
                    // Only coalesce small requests (<=4 texts); callers that
                    // already send well-sized batches skip this to avoid
                    // serialization overhead and the coalesce window delay.
                    constexpr size_t kCoalesceMaxInputBatch = 4;
                    if (!isColbert && c->coalescer && c->coalescer->isEnabled() &&
                        batch_size <= kCoalesceMaxInputBatch) {
                        auto inferFn = [c](const std::string& mid,
                                           const std::vector<std::string>& batch)
                            -> yams::Result<std::vector<std::vector<float>>> {
                            auto h2 = c->pool->acquireModel(mid.c_str(), modelAcquireTimeout(true));
                            if (!h2)
                                return h2.error();
                            return const_cast<yams::daemon::OnnxModelSession&>(*h2.value())
                                .generateBatchEmbeddings(batch);
                        };
                        auto r = c->coalescer->submit(modelIdStr, std::move(texts), inferFn);
                        if (!r) {
                            spdlog::error(
                                "[ONNX Plugin] coalesced generateBatchEmbeddings failed: {}",
                                r.error().message);
                            *out_vecs = nullptr;
                            *out_batch = 0;
                            *out_dim = 0;
                            return YAMS_ERR_INTERNAL;
                        }
                        auto& mat = r.value();
                        if (mat.empty()) {
                            *out_vecs = nullptr;
                            *out_batch = 0;
                            *out_dim = 0;
                            return YAMS_OK;
                        }
                        const size_t b = mat.size();
                        const size_t d = mat[0].size();
                        float* buf = static_cast<float*>(std::malloc(sizeof(float) * b * d));
                        if (!buf)
                            return YAMS_ERR_INTERNAL;
                        for (size_t i = 0; i < b; ++i) {
                            std::memcpy(buf + i * d, mat[i].data(), sizeof(float) * d);
                        }
                        *out_vecs = buf;
                        *out_batch = b;
                        *out_dim = d;
                        return YAMS_OK;
                    }

                    spdlog::debug("[ONNX Plugin] acquiring model '{}'...", model_id);
                    const auto tAcquire = std::chrono::steady_clock::now();
                    auto h = c->pool->acquireModel(model_id, modelAcquireTimeout(true));
                    const auto acquireMs = std::chrono::duration_cast<std::chrono::milliseconds>(
                                               std::chrono::steady_clock::now() - tAcquire)
                                               .count();
                    if (timingDebug) {
                        spdlog::info("[ONNX Plugin] batch_timing acquire model='{}' batch={} ms={} "
                                     "retries_left={}",
                                     model_id, batch_size, acquireMs, retries);
                    }
                    if (!h) {
                        const auto err = h.error();
                        const bool acquireTimeout = (err.code == yams::ErrorCode::Timeout);
                        std::string errMsg = h.error().message;
                        spdlog::error("[ONNX Plugin] acquireModel failed for '{}': {}", model_id,
                                      errMsg);

                        // Timeout while waiting for a session is transient contention. Retry
                        // without tripping model cooldown/failure bookkeeping.
                        if (acquireTimeout) {
                            if (retries > 0) {
                                std::this_thread::sleep_for(std::chrono::milliseconds(250));
                                continue;
                            }
                            *out_vecs = nullptr;
                            *out_batch = 0;
                            *out_dim = 0;
                            return YAMS_ERR_INTERNAL;
                        }

                        if (shutdownInProgress(c)) {
                            *out_vecs = nullptr;
                            *out_batch = 0;
                            *out_dim = 0;
                            return YAMS_ERR_INTERNAL;
                        }

                        {
                            std::lock_guard<std::mutex> lk(c->mu);
                            c->recordFailure(modelIdStr, errMsg);
                            auto& info = c->model_failures[modelIdStr];
                            c->model_states[modelIdStr] = ProviderCtx::State::Failed;

                            if (!info.eventFired) {
                                info.eventFired = true;
                                try {
                                    auto q =
                                        yams::daemon::InternalEventBus::instance()
                                            .get_or_create_channel<yams::daemon::InternalEventBus::
                                                                       ModelLoadFailedEvent>(
                                                "model.events", 256);
                                    (void)q->try_push({modelIdStr, errMsg});
                                } catch (...) {
                                }
                            }
                        }
                        *out_vecs = nullptr;
                        *out_batch = 0;
                        *out_dim = 0;
                        return YAMS_ERR_INTERNAL;
                    }

                    {
                        std::lock_guard<std::mutex> lk(c->mu);
                        c->clearFailure(modelIdStr);
                        c->model_states[modelIdStr] = ProviderCtx::State::Ready;
                    }

                    auto& session = *h.value();
                    const auto tInfer = std::chrono::steady_clock::now();
                    yams::Result<std::vector<std::vector<float>>> r;
                    if (isColbert) {
                        if (!c->colbert || !c->colbert->isValid()) {
                            r = yams::Error{yams::ErrorCode::NotSupported,
                                            "ColBERT session not available for embeddings"};
                        } else {
                            r = c->colbert->encodeDocumentEmbeddings(texts);
                        }
                    } else {
                        r = const_cast<yams::daemon::OnnxModelSession&>(session)
                                .generateBatchEmbeddings(texts);
                    }
                    if (!r) {
                        spdlog::error("[ONNX Plugin] generateBatchEmbeddings failed: {}",
                                      r.error().message);
                        return YAMS_ERR_INTERNAL;
                    }
                    const auto inferMs = std::chrono::duration_cast<std::chrono::milliseconds>(
                                             std::chrono::steady_clock::now() - tInfer)
                                             .count();
                    auto& mat = r.value();
                    if (mat.empty()) {
                        *out_vecs = nullptr;
                        *out_batch = 0;
                        *out_dim = 0;
                        return YAMS_OK;
                    }
                    const size_t b = mat.size();
                    const size_t d = mat[0].size();
                    float* buf = static_cast<float*>(std::malloc(sizeof(float) * b * d));
                    if (!buf)
                        return YAMS_ERR_INTERNAL;
                    for (size_t i = 0; i < b; ++i) {
                        std::memcpy(buf + i * d, mat[i].data(), sizeof(float) * d);
                    }
                    *out_vecs = buf;
                    *out_batch = b;
                    *out_dim = d;
                    if (timingDebug) {
                        spdlog::info(
                            "[ONNX Plugin] batch_timing infer model='{}' batch={} out_batch={} "
                            "dim={} acquire_ms={} infer_ms={}",
                            model_id, batch_size, b, d, acquireMs, inferMs);
                    }
                    return YAMS_OK;

                } catch (const std::system_error& e) {
                    spdlog::error("[ONNX Plugin] system_error: {} (code={}). Retries: {}", e.what(),
                                  e.code().value(), retries);
                    if (retries > 0) {
                        std::this_thread::sleep_for(std::chrono::milliseconds(200));
                        continue;
                    }
                    *out_vecs = nullptr;
                    *out_batch = 0;
                    *out_dim = 0;
                    return YAMS_ERR_INTERNAL;
                } catch (const std::exception& e) {
                    spdlog::error("[ONNX Plugin] exception: {}", e.what());
                    *out_vecs = nullptr;
                    *out_batch = 0;
                    *out_dim = 0;
                    return YAMS_ERR_INTERNAL;
                } catch (...) {
                    spdlog::error("[ONNX Plugin] unknown exception");
                    *out_vecs = nullptr;
                    *out_batch = 0;
                    *out_dim = 0;
                    return YAMS_ERR_INTERNAL;
                }
            } // end while

            return YAMS_ERR_INTERNAL;
        };

        vtable.free_embedding_batch = [](void* /*self*/, float* vecs, size_t /*batch*/,
                                         size_t /*dim*/) {
            if (vecs)
                std::free(vecs);
        };

        // v1.2: get_embedding_dim (non-loading path)
        vtable.get_embedding_dim = [](void* self, const char* model_id,
                                      size_t* out_dim) -> yams_status_t {
            if (!self || !model_id || !out_dim)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            if (!c->pool)
                return YAMS_ERR_INTERNAL;

            // 1) If a hot session already exists, read its dim without blocking
            if (c->ready && c->pool->isModelLoaded(model_id)) {
                auto h = c->pool->acquireModel(model_id, std::chrono::milliseconds(0));
                if (h) {
                    *out_dim = h.value()->getEmbeddingDim();
                    if (*out_dim > 0)
                        return YAMS_OK;
                }
            }

            // 2) Probe model directory metadata without loading the model
            namespace fs = std::filesystem;
            size_t dim = 0;
            try {
                auto read_dim = [&](const fs::path& p) -> size_t {
                    try {
                        if (!fs::exists(p))
                            return 0u;
                        std::ifstream in(p);
                        if (!in)
                            return 0u;
                        nlohmann::json j;
                        in >> j;
                        if (j.contains("embedding_dimension"))
                            return j.value("embedding_dimension", 0u);
                        if (j.contains("embedding_dim"))
                            return j.value("embedding_dim", 0u);
                        // Standard models use hidden_size, Nomic models use n_embd
                        if (j.contains("hidden_size"))
                            return j.value("hidden_size", 0u);
                        if (j.contains("n_embd"))
                            return j.value("n_embd", 0u);
                    } catch (...) {
                    }
                    return 0u;
                };

                // Candidate roots: XDG_DATA_HOME/yams/models, HOME/.local/share/yams/models
                std::vector<fs::path> roots;
                if (const char* xdg = std::getenv("XDG_DATA_HOME"))
                    roots.emplace_back(fs::path(xdg) / "yams" / "models");
                if (const char* home = std::getenv("HOME"))
                    roots.emplace_back(fs::path(home) / ".local" / "share" / "yams" / "models");

                for (const auto& r : roots) {
                    fs::path base = r / model_id;
                    for (const auto& fn :
                         {"sentence_bert_config.json", "config.json", "model.onnx.yams.meta.json",
                          "config.json.yams.meta.json",
                          "sentence_bert_config.json.yams.meta.json"}) {
                        dim = read_dim(base / fn);
                        if (dim > 0)
                            break;
                    }
                    if (dim > 0)
                        break;
                }
            } catch (...) {
            }
            if (dim > 0) {
                *out_dim = dim;
                return YAMS_OK;
            }

            // 3) Heuristic from common embedding model names (no load)
            std::string lm(model_id);
            for (auto& ch : lm)
                ch = static_cast<char>(std::tolower(ch));
            if (lm.find("minilm") != std::string::npos)
                dim = 384;
            else if (lm.find("mpnet") != std::string::npos ||
                     lm.find("nomic") != std::string::npos || lm.find("jina") != std::string::npos)
                dim = 768;
            else if (lm.find("bge-m3") != std::string::npos ||
                     (lm.find("bge") != std::string::npos && lm.find("large") != std::string::npos))
                dim = 1024;
            else if (lm.find("bge") != std::string::npos && lm.find("small") != std::string::npos)
                dim = 384;
            else if (lm.find("colbert") != std::string::npos)
                dim = 48;
            else if (lm.find("e5") != std::string::npos && lm.find("large") != std::string::npos)
                dim = 1024;
            else if (lm.find("e5") != std::string::npos && lm.find("small") != std::string::npos)
                dim = 384;
            else if (lm.find("gte") != std::string::npos && lm.find("large") != std::string::npos)
                dim = 1024;
            else if (lm.find("gte") != std::string::npos && lm.find("small") != std::string::npos)
                dim = 384;
            else if (lm.find("text-embedding-3-large") != std::string::npos)
                dim = 3072;
            else if (lm.find("text-embedding-3-small") != std::string::npos)
                dim = 1536;
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
            else if (lm.find("test-model") != std::string::npos &&
                     (std::getenv("YAMS_TEST_MODE") || std::getenv("YAMS_USE_MOCK_PROVIDER")))
                dim = 384;
#endif

            if (dim > 0) {
                *out_dim = dim;
                return YAMS_OK;
            }

            *out_dim = 0;
            return YAMS_ERR_UNSUPPORTED;
        };

        // v1.2: get_runtime_info_json
        vtable.get_runtime_info_json = [](void* self, const char* model_id,
                                          char** out_json) -> yams_status_t {
            if (!self || !model_id || !out_json)
                return YAMS_ERR_INVALID_ARG;
            auto* c = static_cast<ProviderCtx*>(self);
            if (!c->pool)
                return YAMS_ERR_INTERNAL;
            nlohmann::json j;
            j["backend"] = "onnxruntime";
            j["pipeline"] = "raw_ort";
            j["model"] = model_id;
            // Best-effort dimension and runtime hints
            j["graph_optimization"] = "enabled";

            // Best-effort: report ONNX Runtime build's available providers
            auto providers = yams::onnx_util::OrtRuntimeLoader::instance().availableProviders();
            if (!providers.empty()) {
                auto arr = nlohmann::json::array();
                for (const auto& p : providers)
                    arr.push_back(p);
                j["available_providers"] = std::move(arr);
            }

            const auto& runtimeInfo = c->runtimeInfo;
            if (!runtimeInfo.version.empty())
                j["runtime_version"] = runtimeInfo.version;
            if (!runtimeInfo.libraryPath.empty())
                j["runtime_library"] = runtimeInfo.libraryPath;
            if (!runtimeInfo.reason.empty())
                j["runtime_status"] = runtimeInfo.reason;

            // Report the actual execution provider used by a session (not a platform guess)
            std::string ep = "cpu";
            if (c->ready && c->pool) {
                try {
                    auto h = c->pool->acquireModel(model_id, std::chrono::milliseconds(500));
                    if (h) {
                        auto* session = h.value().operator->();
                        if (!session) {
                            return YAMS_ERR_INTERNAL;
                        }
                        ep = session->getExecutionProvider();
                        auto [intraThreads, interThreads] = session->getThreading();
                        j["intra_threads"] = intraThreads;
                        j["inter_threads"] = interThreads;
                        j["learned_batch_cap"] = session->getLearnedBatchLimit();
                        {
                            std::lock_guard<std::mutex> lk(c->mu);
                            c->actualExecutionProvider = ep;
                        }
                    }
                } catch (...) {
                }
            }
            if (ep == "cpu") {
                std::lock_guard<std::mutex> lk(c->mu);
                if (!c->actualExecutionProvider.empty())
                    ep = c->actualExecutionProvider;
            }
            j["execution_provider"] = ep;
            size_t dim = 0;
            if (c->ready && c->pool) {
                auto h = c->pool->acquireModel(model_id, std::chrono::seconds(2));
                if (h) {
                    auto* session = h.value().operator->();
                    if (!session) {
                        return YAMS_ERR_INTERNAL;
                    }
                    dim = session->getEmbeddingDim();
                    try {
                        const auto& info = session->getInfo();
                        if (!info.path.empty())
                            j["path"] = info.path;
                    } catch (...) {
                    }
                }
            }
            j["dim"] = dim;
            if (!j.contains("intra_threads")) {
                j["intra_threads"] =
                    static_cast<int>(std::max(1u, std::thread::hardware_concurrency()));
            }
            if (!j.contains("inter_threads")) {
                j["inter_threads"] = 1;
            }
            // Source hint from path when present
            try {
                if (j.contains("path") && j["path"].is_string()) {
                    std::string p = j["path"].get<std::string>();
                    if (p.find("/.cache/huggingface/") != std::string::npos)
                        j["source"] = "huggingface-cache";
                    else if (p.find(".yams/models/") != std::string::npos)
                        j["source"] = "yams-models";
                    else if (p.rfind("http://", 0) == 0 || p.rfind("https://", 0) == 0)
                        j["source"] = "remote";
                    else
                        j["source"] = "local";
                }
            } catch (...) {
            }
            std::string s;
            try {
                s = j.dump();
            } catch (...) {
                s = "{}";
            }
            char* buf = static_cast<char*>(std::malloc(s.size() + 1));
            if (!buf)
                return YAMS_ERR_INTERNAL;
            std::memcpy(buf, s.c_str(), s.size() + 1);
            *out_json = buf;
            return YAMS_OK;
        };

        vtable.free_string = [](void* /*self*/, char* s) {
            if (s)
                std::free(s);
        };

        // Assign function pointer for set_threading
        vtable.set_threading = &onnx_set_threading;

        // v1.3: Cross-encoder reranking
        vtable.score_documents = [](void* self, const char* /*reranker_model_id*/,
                                    const char* query, const char* const* documents,
                                    size_t doc_count, float** out_scores,
                                    size_t* out_count) -> yams_status_t {
            if (!self || !query || !documents || !out_scores || !out_count)
                return YAMS_ERR_INVALID_ARG;
            if (doc_count == 0) {
                *out_scores = nullptr;
                *out_count = 0;
                return YAMS_OK;
            }

            try {
                auto* c = static_cast<ProviderCtx*>(self);
                if (c->disabled)
                    return YAMS_ERR_UNSUPPORTED;

                if (c->colbert && c->colbert->isValid()) {
                    auto queryEmb = c->colbert->encodeQuery(query);
                    if (!queryEmb) {
                        spdlog::warn("[ONNX Plugin] ColBERT encode query failed: {}",
                                     queryEmb.error().message);
                        return YAMS_ERR_INTERNAL;
                    }

                    std::vector<std::string> docs;
                    docs.reserve(doc_count);
                    for (size_t i = 0; i < doc_count; ++i) {
                        docs.emplace_back(documents[i] ? documents[i] : "");
                    }
                    auto docEmbeddings = c->colbert->encodeDocuments(docs);
                    if (!docEmbeddings) {
                        spdlog::warn("[ONNX Plugin] ColBERT encode documents failed: {}",
                                     docEmbeddings.error().message);
                        return YAMS_ERR_INTERNAL;
                    }

                    const auto& docVecs = docEmbeddings.value();
                    std::vector<float> scores;
                    scores.reserve(docVecs.size());
                    for (const auto& docEmb : docVecs) {
                        scores.push_back(c->colbert->computeMaxSim(queryEmb.value(), docEmb));
                    }

                    float* buf = static_cast<float*>(std::malloc(sizeof(float) * scores.size()));
                    if (!buf)
                        return YAMS_ERR_INTERNAL;
                    std::memcpy(buf, scores.data(), sizeof(float) * scores.size());
                    *out_scores = buf;
                    *out_count = scores.size();
                    return YAMS_OK;
                }

                spdlog::debug("[ONNX Plugin] Cross-encoder reranker removed; ColBERT is the only "
                              "ONNX reranker path");
                return YAMS_ERR_UNSUPPORTED;

            } catch (const std::exception& e) {
                spdlog::error("[ONNX Plugin] score_documents exception: {}", e.what());
                return YAMS_ERR_INTERNAL;
            } catch (...) {
                spdlog::error("[ONNX Plugin] score_documents unknown exception");
                return YAMS_ERR_INTERNAL;
            }
        };

        vtable.free_scores = [](void* /*self*/, float* scores, size_t /*count*/) {
            if (scores)
                std::free(scores);
        };

        // v1.4: Memory pressure eviction
        vtable.evict_under_pressure = [](void* self, double pressure_level, bool allow_hot_eviction,
                                         size_t* out_evicted) -> yams_status_t {
            if (!self || !out_evicted)
                return YAMS_ERR_INVALID_ARG;
            *out_evicted = 0;

            try {
                auto* c = static_cast<ProviderCtx*>(self);
                if (c->disabled)
                    return YAMS_ERR_UNSUPPORTED;
                if (!c->ready || !c->pool)
                    return YAMS_ERR_INTERNAL;

                // Get profile-aware eviction thresholds from TuneAdvisor
                const double warningThresh =
                    yams::daemon::TuneAdvisor::modelEvictWarningThreshold();
                const double criticalThresh =
                    yams::daemon::TuneAdvisor::modelEvictCriticalThreshold();
                const double emergencyThresh =
                    yams::daemon::TuneAdvisor::modelEvictEmergencyThreshold();

                // Determine how many models to evict based on pressure level:
                // - Emergency: evict all non-hot (or all if allow_hot_eviction)
                // - Critical: evict up to 2 models
                // - Warning: evict 1 model
                size_t numToEvict = 0;
                if (pressure_level >= emergencyThresh) {
                    numToEvict = allow_hot_eviction
                                     ? c->pool->getLoadedModelCount()
                                     : std::max<size_t>(1, c->pool->getLoadedModelCount() / 2);
                } else if (pressure_level >= criticalThresh) {
                    numToEvict = 2;
                } else if (pressure_level >= warningThresh) {
                    numToEvict = 1;
                }

                if (numToEvict == 0) {
                    // pressure below warning threshold: do idle maintenance instead
                    c->pool->performMaintenance();
                    return YAMS_OK;
                }

                // Use existing evictLRU method
                size_t countBefore = c->pool->getLoadedModelCount();
                c->pool->evictLRU(numToEvict);
                size_t countAfter = c->pool->getLoadedModelCount();

                size_t evicted = (countBefore > countAfter) ? (countBefore - countAfter) : 0;
                *out_evicted = evicted;

                if (evicted > 0) {
                    spdlog::info("[ONNX Plugin] Evicted {} models under memory pressure ({:.0f}%)",
                                 evicted, pressure_level * 100.0);
                }

                return YAMS_OK;
            } catch (const std::exception& e) {
                spdlog::error("[ONNX Plugin] evict_under_pressure exception: {}", e.what());
                return YAMS_ERR_INTERNAL;
            } catch (...) {
                spdlog::error("[ONNX Plugin] evict_under_pressure unknown exception");
                return YAMS_ERR_INTERNAL;
            }
        };
    }
};

static ProviderSingleton& singleton() {
    static ProviderSingleton s;
    return s;
}

} // namespace

// Set the plugin config JSON (called from yams_plugin_init before provider is accessed)
extern "C" void yams_onnx_set_config_json(const char* json) {
    if (json && *json) {
        g_plugin_config_json = json;
        spdlog::info("[ONNX Plugin] Config JSON set: {}", g_plugin_config_json);
    }
}

extern "C" void yams_onnx_shutdown_provider() {
    singleton().explicitShutdownCalled.store(true);
    singleton().shutdown(true);
}

extern "C" yams_model_provider_v1* yams_onnx_get_model_provider() {
    auto& s = singleton();
    s.ensureActive();
    return &s.vtable;
}

// Provide a small JSON health snapshot for the ABI host
extern "C" const char* yams_onnx_get_health_json_cstr() {
    static thread_local std::string json;
    auto& c = singleton().ctx;
    // Build a minimal JSON reflecting readiness/disabled and last error
    nlohmann::json j;
    if (c.disabled) {
        j["status"] = "disabled";
        j["reason"] = c.last_error.empty() ? "disabled_by_env" : c.last_error;
    } else if (!c.runtimeInfo.available) {
        j["status"] = "unavailable";
        j["reason"] = c.last_error.empty() ? "runtime_not_available" : c.last_error;
    } else if (!c.pool) {
        j["status"] = "unavailable";
        j["reason"] = "no_pool";
    } else if (!c.ready) {
        j["status"] = "degraded";
        j["reason"] = c.last_error.empty() ? "init_failed" : c.last_error;
    } else {
        j["status"] = "ok";
    }

    // Execution provider + available providers (best-effort)
    {
        std::string ep = "cpu";
        {
            std::lock_guard<std::mutex> lk(c.mu);
            ep = c.actualExecutionProvider;
        }

        if (c.pool) {
            try {
                // Try a quick acquire of any ready model to get the real EP.
                std::string candidateModelId;
                {
                    std::lock_guard<std::mutex> lk(c.mu);
                    for (const auto& [modelId, state] : c.model_states) {
                        if (state == ProviderCtx::State::Ready) {
                            candidateModelId = modelId;
                            break;
                        }
                    }
                }

                if (!candidateModelId.empty()) {
                    auto h = c.pool->acquireModel(candidateModelId, std::chrono::milliseconds(50));
                    if (h) {
                        ep = h.value()->getExecutionProvider();
                        std::lock_guard<std::mutex> lk(c.mu);
                        c.actualExecutionProvider = ep;
                    }
                }
            } catch (...) {
                // Keep cached EP on any failure.
            }
        }

        j["execution_provider"] = ep;

        auto providers = yams::onnx_util::OrtRuntimeLoader::instance().availableProviders();
        if (!providers.empty()) {
            j["available_providers"] = providers;
        }
    }

    if (!c.runtimeInfo.version.empty()) {
        j["runtime_version"] = c.runtimeInfo.version;
    }
    if (!c.runtimeInfo.libraryPath.empty()) {
        j["runtime_library"] = c.runtimeInfo.libraryPath;
    }
    if (!c.runtimeInfo.reason.empty()) {
        j["runtime_status"] = c.runtimeInfo.reason;
    }
    if (!c.runtimeInfo.errorMessage.empty()) {
        j["runtime_error"] = c.runtimeInfo.errorMessage;
    }

    // Add model states for diagnostics
    {
        std::lock_guard<std::mutex> lk(c.mu);
        if (!c.model_states.empty() || !c.model_failures.empty()) {
            nlohmann::json models = nlohmann::json::object();
            for (const auto& [modelId, state] : c.model_states) {
                nlohmann::json m;
                switch (state) {
                    case ProviderCtx::State::Unloaded:
                        m["state"] = "unloaded";
                        break;
                    case ProviderCtx::State::Loading:
                        m["state"] = "loading";
                        break;
                    case ProviderCtx::State::Ready:
                        m["state"] = "ready";
                        break;
                    case ProviderCtx::State::Failed:
                        m["state"] = "failed";
                        break;
                }
                // Add failure info if present
                auto fit = c.model_failures.find(modelId);
                if (fit != c.model_failures.end() && fit->second.consecutiveFailures > 0) {
                    m["failures"] = fit->second.consecutiveFailures;
                    m["last_error"] = fit->second.lastError;
                    m["in_cooldown"] = c.isInCooldown(modelId);
                }
                models[modelId] = m;
            }
            j["models"] = std::move(models);
        }
    }

    try {
        json = j.dump();
    } catch (...) {
        json = "{\"status\":\"unknown\"}";
    }
    return json.c_str();
}

#ifdef YAMS_TESTING
extern "C" void yams_onnx_test_get_pool_config(std::size_t* max_loaded, std::size_t* hot_pool) {
    auto& c = singleton().ctx;
    if (max_loaded)
        *max_loaded = c.configuredMaxLoadedModels;
    if (hot_pool)
        *hot_pool = c.configuredHotPoolSize;
}
#endif