#include <yams/common/test_utils.h>
#include <yams/daemon/components/TuneAdvisor.h>
#include <yams/daemon/resource/onnx_model_pool.h>
#include <yams/vector/tokenizer.h>

#include "onnx_gpu_provider.h"
#include <yams/daemon/resource/gpu_info.h>

#include <nlohmann/json.hpp>
#include <spdlog/spdlog.h>
#include <filesystem>
#include <mutex>
#include <sstream>
#include <unordered_set>
#include <boost/asio/post.hpp>
#include <boost/asio/thread_pool.hpp>

#include "../ort_runtime/ort_cxx_api_wrapper.h"
#include "../ort_runtime/ort_runtime_loader.h"
#ifdef _WIN32
#include <fcntl.h>
#include <io.h>
#include <windows.h>
#include <sys/stat.h>
#else
#include <fcntl.h>
#include <unistd.h>
#include <sys/file.h>
#endif

// GPU execution provider headers - conditionally included based on Conan build options
// Note: Prebuilt ONNX Runtime binaries from GitHub releases use the C API for providers.
// The provider-specific headers (cuda_provider_factory.h, etc.) are only available when
// building ONNX Runtime from source. For prebuilt binaries, we use
// OrtSessionOptionsAppendExecutionProvider_* functions from the C API which are declared in
// onnxruntime_c_api.h (included via onnxruntime_cxx_api.h).

// Global ONNX Runtime environment - single instance per process, intentionally leaked.
// Using leaking singleton pattern because ONNX Runtime's LoggingManager destructor
// tries to lock a mutex that may be destroyed during static destruction, causing
// "mutex lock failed: Invalid argument" crashes.
static std::recursive_mutex* g_onnx_init_mutex = new std::recursive_mutex();
static Ort::Env* g_onnx_env = nullptr; // Intentionally leaked to avoid static destruction crash
static bool g_onnx_env_initialized = false;

static Ort::Env& get_global_ort_env() {
    std::lock_guard<std::recursive_mutex> lock(*g_onnx_init_mutex);
    if (!g_onnx_env_initialized) {
        spdlog::info("[ONNX] Initializing global Ort::Env (thread-safe, single instance)");

        OrtThreadingOptions* threading_options = nullptr;
        (void)Ort::GetApi().CreateThreadingOptions(&threading_options);
        if (threading_options) {
            unsigned hw_threads = std::thread::hardware_concurrency();
            if (hw_threads == 0)
                hw_threads = 4;

#ifdef _WIN32
            int intra_threads = 1;
            int inter_threads = 1;
            spdlog::info(
                "[ONNX] Windows: using single-threaded globalOrt::Env (EDEADLK workaround)");
#else
            int intra_threads = static_cast<int>(std::max(1u, hw_threads / 2));
            int inter_threads = 1;
            spdlog::info("[ONNX] Using multi-threaded global Ort::Env (intra={}, inter={})",
                         intra_threads, inter_threads);
#endif
            (void)Ort::GetApi().SetGlobalIntraOpNumThreads(threading_options, intra_threads);
            (void)Ort::GetApi().SetGlobalInterOpNumThreads(threading_options, inter_threads);
            (void)Ort::GetApi().SetGlobalSpinControl(threading_options, 0);

            // Process-lifetime singleton: Ort::Env intentionally never freed — ORT
            // requires the env for the process lifetime.
            g_onnx_env = new Ort::Env(threading_options, ORT_LOGGING_LEVEL_WARNING, "YamsDaemon");
            Ort::GetApi().ReleaseThreadingOptions(threading_options);
            spdlog::info("[ONNX] Global Ort::Env initialized with custom threading options");
        } else {
            // Process-lifetime singleton: see comment above.
            g_onnx_env = new Ort::Env(ORT_LOGGING_LEVEL_WARNING, "YamsDaemon");
            spdlog::info("[ONNX] Global Ort::Env initialized (default threading)");
        }
        g_onnx_env_initialized = true;
    }
    return *g_onnx_env;
}

// Check if we should skip model loading for tests
#ifdef GTEST_API_
#include <yams/common/test_utils.h>
#endif

#include <algorithm>
#include <cmath>
#include <cstdlib>
#include <filesystem>
#include <fstream>
#include <future>
#include <limits>
#include <thread>
#include <unordered_map>

namespace yams::daemon {

namespace fs = std::filesystem;

namespace {

bool envTruthy(const char* value) {
    if (!value || !*value) {
        return false;
    }
    std::string v(value);
    std::transform(v.begin(), v.end(), v.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return !(v == "0" || v == "false" || v == "off" || v == "no");
}

bool allowGpuAutoUnload() {
    return envTruthy(std::getenv("YAMS_ONNX_GPU_AUTO_UNLOAD"));
}

int envIntOr(const char* name, int fallback) {
    if (const char* value = std::getenv(name)) {
        try {
            return std::stoi(value);
        } catch (...) {
        }
    }
    return fallback;
}

std::string lowerCopy(std::string value) {
    std::transform(value.begin(), value.end(), value.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return value;
}

bool supportsCoreMLCpuFallbackModel(const std::string& modelName) {
    const std::string lowerName = lowerCopy(modelName);
    return lowerName.find("embeddinggemma") != std::string::npos ||
           lowerName.find("gemma") != std::string::npos;
}

bool isGpuExecutionProviderName(std::string provider) {
    provider = lowerCopy(std::move(provider));
    return provider.find("cuda") != std::string::npos ||
           provider.find("migraphx") != std::string::npos ||
           provider.find("coreml") != std::string::npos ||
           provider.find("dml") != std::string::npos ||
           provider.find("directml") != std::string::npos;
}

GraphOptimizationLevel resolveGraphOptimizationLevel(bool cpuFallback) {
    GraphOptimizationLevel optLevel = cpuFallback ? GraphOptimizationLevel::ORT_ENABLE_BASIC
                                                  : GraphOptimizationLevel::ORT_ENABLE_EXTENDED;
    if (const char* opt = std::getenv("YAMS_ONNX_OPT_LEVEL")) {
        std::string level = lowerCopy(opt);
        if (level == "all") {
            optLevel = GraphOptimizationLevel::ORT_ENABLE_ALL;
            if (!cpuFallback) {
                spdlog::info("[ONNX] Using ORT_ENABLE_ALL optimization (slow first load)");
            }
        } else if (level == "extended") {
            optLevel = GraphOptimizationLevel::ORT_ENABLE_EXTENDED;
        } else if (level == "basic") {
            optLevel = GraphOptimizationLevel::ORT_ENABLE_BASIC;
        }
    }
    return optLevel;
}

int resolveCpuFallbackIntraThreads(int configuredIntraThreads) {
    if (const char* value = std::getenv("YAMS_ONNX_INTRA_OP_THREADS")) {
        try {
            const int parsed = std::stoi(value);
            if (parsed >= 1 && parsed <= 64) {
                return parsed;
            }
        } catch (...) {
        }
    }
    const int fallback = configuredIntraThreads > 0 ? configuredIntraThreads : 2;
    return std::clamp(fallback, 1, 2);
}

int resolveCpuFallbackInterThreads(int configuredInterThreads) {
    if (const char* value = std::getenv("YAMS_ONNX_INTER_OP_THREADS")) {
        try {
            const int parsed = std::stoi(value);
            if (parsed >= 1 && parsed <= 64) {
                return parsed;
            }
        } catch (...) {
        }
    }
    (void)configuredInterThreads;
    return 1;
}

bool resolveCpuFallbackAllowSpinning() {
    if (const char* value = std::getenv("YAMS_ONNX_ALLOW_SPINNING")) {
        return envTruthy(value);
    }
    return false;
}

// Resolve path for caching ONNX Runtime optimized graphs on disk.
// Precedence: YAMS_DATA_DIR > XDG_DATA_HOME > ~/.local/share > model directory.
fs::path resolveOptimizedModelPath(const std::string& modelPath, GraphOptimizationLevel optLevel) {
    std::string filename = fs::path(modelPath).stem().string() + "." +
                           std::to_string(static_cast<int>(optLevel)) + ".ort";
    if (const char* dir = std::getenv("YAMS_DATA_DIR"); dir && *dir) {
        return fs::path(dir) / "cache" / "onnx-opt" / filename;
    }
    if (const char* xdg = std::getenv("XDG_DATA_HOME"); xdg && *xdg) {
        return fs::path(xdg) / "yams" / "cache" / "onnx-opt" / filename;
    }
    if (const char* home = std::getenv("HOME"); home && *home) {
        return fs::path(home) / ".local" / "share" / "yams" / "cache" / "onnx-opt" / filename;
    }
    return fs::path(modelPath).parent_path() / filename;
}

bool gpuProcessLockEnabled() {
    if (const char* v = std::getenv("YAMS_GPU_PROCESS_LOCK")) {
        return envTruthy(v);
    }
    if (const char* v = std::getenv("YAMS_MIGRAPHX_PROCESS_LOCK")) {
        return envTruthy(v);
    }
    return true;
}

std::filesystem::path resolveGpuProcessLockPath() {
    if (const char* configured = std::getenv("YAMS_GPU_LOCK_PATH"); configured && *configured) {
        std::filesystem::path path(configured);
        if (path.extension() == ".lock") {
            return path;
        }
        auto devEnv = std::getenv("YAMS_GPU_DEVICE_ID");
        int deviceId = devEnv ? std::stoi(devEnv) : 0;
        return path / ("yams-gpu-device-" + std::to_string(deviceId) + ".lock");
    }

    if (const char* configured = std::getenv("YAMS_MIGRAPHX_LOCK_PATH");
        configured && *configured) {
        std::filesystem::path path(configured);
        if (path.extension() == ".lock") {
            return path;
        }
        return path / "migraphx.lock"; // Fallback legacy behavior if needed, or just standard
    }

    int deviceId = envIntOr("YAMS_GPU_DEVICE_ID", -1);
    if (deviceId < 0) {
        deviceId = envIntOr("YAMS_MIGRAPHX_DEVICE_ID", 0);
    }
    return std::filesystem::temp_directory_path() /
           ("yams-gpu-device-" + std::to_string(deviceId) + ".lock");
}

class ScopedProcessFileLock {
public:
    explicit ScopedProcessFileLock(const std::filesystem::path& path) : path_(path) {
        std::error_code ec;
        auto parent = path_.parent_path();
        if (!parent.empty()) {
            std::filesystem::create_directories(parent, ec);
        }

#ifdef _WIN32
        fd_ = _open(path_.string().c_str(), _O_CREAT | _O_RDWR | _O_BINARY, _S_IREAD | _S_IWRITE);
        if (fd_ < 0) {
            return;
        }
        hFile_ = reinterpret_cast<HANDLE>(_get_osfhandle(fd_));
        if (hFile_ == INVALID_HANDLE_VALUE) {
            _close(fd_);
            fd_ = -1;
            return;
        }
        OVERLAPPED overlapped = {0};
        if (!LockFileEx(hFile_, LOCKFILE_EXCLUSIVE_LOCK, 0, 1, 0, &overlapped)) {
            _close(fd_);
            fd_ = -1;
            hFile_ = INVALID_HANDLE_VALUE;
            return;
        }
#else
        fd_ = ::open(path_.c_str(), O_CREAT | O_RDWR, 0644);
        if (fd_ < 0) {
            return;
        }
        if (::flock(fd_, LOCK_EX) != 0) {
            ::close(fd_);
            fd_ = -1;
            return;
        }
#endif
    }

    ScopedProcessFileLock(const ScopedProcessFileLock&) = delete;
    ScopedProcessFileLock& operator=(const ScopedProcessFileLock&) = delete;

    bool locked() const { return fd_ >= 0; }

    ~ScopedProcessFileLock() {
        if (fd_ < 0) {
            return;
        }
#ifdef _WIN32
        if (hFile_ != INVALID_HANDLE_VALUE) {
            OVERLAPPED overlapped = {0};
            (void)UnlockFileEx(hFile_, 0, 1, 0, &overlapped);
        }
        _close(fd_);
#else
        (void)::flock(fd_, LOCK_UN);
        ::close(fd_);
#endif
    }

private:
    std::filesystem::path path_;
    int fd_{-1};
#ifdef _WIN32
    HANDLE hFile_{INVALID_HANDLE_VALUE};
#endif
};

size_t& gpuProcessLockDepth() {
    static thread_local size_t depth = 0;
    return depth;
}

// Thread-local reusable buffers for ONNX tensor inputs.
// Avoids per-inference heap allocation for tokens, masks, and token_type_ids.
struct TensorBufferPool {
    std::vector<int64_t> tokens;
    std::vector<int64_t> masks;
    std::vector<int64_t> token_type_ids;
};

TensorBufferPool& getThreadLocalBuffers() {
    static thread_local TensorBufferPool pool;
    return pool;
}

std::string providerBatchLimitKey(const std::string& modelName, const std::string& provider) {
    std::string providerLower = provider;
    std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
                   [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
    return modelName + "|" + providerLower;
}

std::recursive_mutex& sharedBatchLimitMutex() {
    static auto* mu = new std::recursive_mutex();
    return *mu;
}

std::unordered_map<std::string, size_t>& sharedBatchLimitMap() {
    static auto* caps = new std::unordered_map<std::string, size_t>();
    return *caps;
}

size_t loadSharedBatchLimit(const std::string& modelName, const std::string& provider) {
    const std::string key = providerBatchLimitKey(modelName, provider);
    std::lock_guard<std::recursive_mutex> lock(sharedBatchLimitMutex());
    const auto& caps = sharedBatchLimitMap();
    const auto it = caps.find(key);
    return (it == caps.end()) ? 0 : it->second;
}

size_t mergeSharedBatchLimit(const std::string& modelName, const std::string& provider,
                             size_t observedCap) {
    if (observedCap == 0) {
        return 0;
    }
    const std::string key = providerBatchLimitKey(modelName, provider);
    std::lock_guard<std::recursive_mutex> lock(sharedBatchLimitMutex());
    auto& caps = sharedBatchLimitMap();
    const auto it = caps.find(key);
    const size_t merged = (it == caps.end()) ? observedCap : std::min(it->second, observedCap);
    caps[key] = merged;
    return merged;
}

size_t resolveCpuFallbackBatchCap() {
    constexpr size_t kSafeCpuFallbackBatchCap = 4;
    const size_t configuredCap = TuneAdvisor::getEmbedDocCap();
    if (configuredCap == 0) {
        return kSafeCpuFallbackBatchCap;
    }
    return std::max<size_t>(1u, std::min(configuredCap, kSafeCpuFallbackBatchCap));
}

size_t applyCpuFallbackSafetyCap(const std::string& modelName, const std::string& provider,
                                 std::atomic<size_t>& providerBatchLimit) {
    const size_t safeCap = resolveCpuFallbackBatchCap();
    const size_t mergedCap = mergeSharedBatchLimit(modelName, provider, safeCap);
    providerBatchLimit.store(mergedCap, std::memory_order_relaxed);

    const size_t currentCap = TuneAdvisor::getEmbedDocCap();
    if (currentCap == 0 || currentCap > mergedCap) {
        TuneAdvisor::setEmbedDocCap(mergedCap);
    }

    return mergedCap;
}

} // namespace

// ============================================================================
// Cross-instance CoreML denylist
// ============================================================================
// When a model fails CoreML inference (e.g., mask feature partition bug),
// record it here so other pool instances skip CoreML and go straight to CPU.
// This avoids wasting ~20s per instance on CoreML compilation + failed retries.

static std::mutex g_coreml_denylist_mu;
static std::unordered_set<std::string> g_coreml_denylist;

bool OnnxModelSession::isCoreMLDenylisted(const std::string& modelName) {
    std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
    return g_coreml_denylist.count(modelName) > 0;
}

void OnnxModelSession::denylistCoreML(const std::string& modelName) {
    std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
    g_coreml_denylist.insert(modelName);
}

void OnnxModelSession::clearCoreMLDenylist() {
    std::lock_guard<std::mutex> lk(g_coreml_denylist_mu);
    g_coreml_denylist.clear();
}

// ============================================================================
// OnnxModelSession Implementation
// ============================================================================

class OnnxModelSession::Impl {
public:
    Impl(const std::string& modelPath, const std::string& modelName, const OnnxTextConfig& config)
        : modelPath_(modelPath), modelName_(modelName), config_(config), preprocessor_(config) {
        // Check if dynamic padding is disabled via environment variable
        if (const char* dp = std::getenv("YAMS_ONNX_DYNAMIC_PADDING")) {
            std::string val(dp);
            std::transform(val.begin(), val.end(), val.begin(), ::tolower);
            if (val == "0" || val == "false" || val == "no" || val == "off") {
                dynamicPaddingEnabled_ = false;
            }
        }
#ifdef GTEST_API_
        // Check if we should skip model loading in test mode
        if (yams::test::shouldSkipModelLoading()) {
            test_mode_ = true;
            spdlog::debug("[ONNX] Test mode enabled - skipping actual model loading");
            return;
        }
#endif

#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
        // Lightweight, environment-driven mock mode for tests.
        if (std::getenv("YAMS_USE_MOCK_PROVIDER") || std::getenv("YAMS_SKIP_MODEL_LOADING") ||
            std::getenv("YAMS_TEST_MODE")) {
            test_mode_ = true;
            spdlog::warn("[ONNX] Mock provider mode enabled via env; skipping ONNX init");
            return;
        }
#endif

        const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
        if (!runtimeInfo.available) {
            throw std::runtime_error(runtimeInfo.errorMessage.empty()
                                         ? std::string("ONNX Runtime unavailable")
                                         : "ONNX Runtime unavailable: " + runtimeInfo.errorMessage);
        }

        memoryInfo_ = std::make_unique<Ort::MemoryInfo>(
            Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault));

        // Initialize ONNX Runtime environment
        // Use a single global Ort::Env per process (best practice)
        // The global env is protected by g_onnx_init_mutex during initialization.
        spdlog::debug("[ONNX] Impl constructor: getting global Ort::Env for '{}'", modelName);
        env_ = &get_global_ort_env();
        spdlog::debug("[ONNX] Impl constructor: global Ort::Env obtained for '{}'", modelName);

        // Configure session options
        sessionOptions_ = std::make_unique<Ort::SessionOptions>();
        // Allow env overrides for thread tuning
        auto detect_threads = [](const char* name, int fallback) {
            if (const char* s = std::getenv(name)) {
                try {
                    int v = std::stoi(s);
                    if (v > 0 && v <= 64)
                        return v;
                } catch (...) {
                }
            }
            return fallback;
        };
        auto detect_bool = [](const char* name, bool fallback) {
            if (const char* s = std::getenv(name)) {
                std::string val(s);
                std::transform(val.begin(), val.end(), val.begin(), ::tolower);
                if (val == "0" || val == "false" || val == "no")
                    return false;
                if (val == "1" || val == "true" || val == "yes")
                    return true;
            }
            return fallback;
        };

#ifdef _WIN32
        // Windows-specific: ALWAYS use minimal thread counts to avoid thread pool
        // exhaustion and "resource deadlock would occur" errors. The Windows thread
        // scheduler behaves differently than Linux. Do NOT use config.num_threads
        // here as it defaults to hardware_concurrency() which causes issues.
        int intra_default = 1;               // Single thread for intra-op
        int inter_default = 1;               // Single thread for inter-op
        bool allow_spinning_default = false; // No spinning on Windows
#else
        int intra_default = 4;
        int inter_default = 1;
        bool allow_spinning_default = true;
#endif

        // On Windows, ignore config.num_threads to enforce safe defaults
        // On other platforms, respect the config
#ifdef _WIN32
        int intra = detect_threads("YAMS_ONNX_INTRA_OP_THREADS", intra_default);
#else
        int intra = config.num_threads > 0 ? config.num_threads : intra_default;
        intra = detect_threads("YAMS_ONNX_INTRA_OP_THREADS", intra);
#endif
        int inter = config.inter_op_threads > 0 ? config.inter_op_threads : inter_default;
        inter = detect_threads("YAMS_ONNX_INTER_OP_THREADS", inter);
        configuredIntraThreads_ = intra;
        configuredInterThreads_ = inter;
        sessionOptions_->SetIntraOpNumThreads(intra);
        sessionOptions_->SetInterOpNumThreads(inter);

        // Thread spinning control - spinning threads consume CPU cycles waiting for work
        // which can cause contention issues on Windows with concurrent embedding requests
        bool allow_spinning = detect_bool("YAMS_ONNX_ALLOW_SPINNING", allow_spinning_default);
        sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
                                        allow_spinning ? "1" : "0");
        sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
                                        allow_spinning ? "1" : "0");

        spdlog::debug("[ONNX] SessionOptions threads: intra-op={} inter-op={} spinning={}", intra,
                      inter, allow_spinning);

#ifdef _WIN32
        // On Windows, use sequential execution mode to avoid thread pool issues
        // This ensures ONNX Runtime doesn't spawn additional threads internally
        sessionOptions_->SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
        spdlog::info("[ONNX] Windows: using sequential execution mode");
#endif

        // Graph optimization level: EXTENDED enables operator fusion (MatMul+Add, GELU, etc.)
        // which significantly benefits transformer embedding models.
        // ORT_ENABLE_ALL performs additional expensive transformations (minutes for large models).
        // Override with YAMS_ONNX_OPT_LEVEL=basic|extended|all
        GraphOptimizationLevel optLevel = resolveGraphOptimizationLevel(false);
        sessionOptions_->SetGraphOptimizationLevel(optLevel);

        // Cache the optimized graph to disk so subsequent loads skip optimization passes.
        if (detect_bool("YAMS_ONNX_GRAPH_CACHE", false)) {
            auto optPath = resolveOptimizedModelPath(modelPath, optLevel);
            std::error_code ec;
            fs::create_directories(optPath.parent_path(), ec);
            if (!ec) {
                sessionOptions_->SetOptimizedModelFilePath(optPath.c_str());
                spdlog::info("[ONNX] Optimized graph cache: {}", optPath.string());
            } else {
                spdlog::warn("[ONNX] Could not create graph cache dir {}: {}",
                             optPath.parent_path().string(), ec.message());
            }
        }

        // Enable memory pattern optimization for faster inference
        sessionOptions_->EnableMemPattern();
        // Enable CPU memory arena for better memory reuse
        sessionOptions_->EnableCpuMemArena();

        if (config.enable_gpu) {
            if (OnnxModelSession::isCoreMLDenylisted(modelName)) {
                const int fallbackIntra = resolveCpuFallbackIntraThreads(configuredIntraThreads_);
                const int fallbackInter = resolveCpuFallbackInterThreads(configuredInterThreads_);
                const bool allowSpinning = resolveCpuFallbackAllowSpinning();
                sessionOptions_->SetIntraOpNumThreads(fallbackIntra);
                sessionOptions_->SetInterOpNumThreads(fallbackInter);
                sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
                                                allowSpinning ? "1" : "0");
                sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
                                                allowSpinning ? "1" : "0");
                configuredIntraThreads_ = fallbackIntra;
                configuredInterThreads_ = fallbackInter;
                const size_t fallbackBatchCap =
                    applyCpuFallbackSafetyCap(modelName_, "cpu", providerBatchLimit_);
                spdlog::info("[ONNX] Skipping CoreML for '{}' (previously failed; "
                             "CPU with intra_threads={} inter_threads={} batch_cap={} "
                             "spinning={})",
                             modelName, fallbackIntra, fallbackInter, fallbackBatchCap,
                             allowSpinning ? 1 : 0);
                actualExecutionProvider_ = "cpu";
            } else {
                appendGpuExecutionProvider();
            }
        }
    }

    Result<void> loadModel() {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        if (test_mode_) {
            // In test mode, pretend loading succeeded
            // Derive embeddingDim_/maxSequenceLength_/pooling/normalize from nearby configs when
            // available to better match real model characteristics.
            try {
                parseModelConfigHints();
            } catch (...) {
            }
            // Heuristics based on common model families when config hints are absent
            auto lname = modelName_;
            std::transform(lname.begin(), lname.end(), lname.begin(),
                           [](unsigned char c) { return static_cast<char>(std::tolower(c)); });
            // Dimension defaults
            if (embeddingDim_ == 0) {
                auto has = [&](const char* s) { return lname.find(s) != std::string::npos; };
                // OpenAI TE3 sizes (mock support for external references)
                if (has("text-embedding-3-large")) {
                    embeddingDim_ = 3072;
                } else if (has("text-embedding-3-small")) {
                    embeddingDim_ = 1536;
                } else if (has("e5")) {
                    if (has("large"))
                        embeddingDim_ = 1024;
                    else if (has("small"))
                        embeddingDim_ = 384;
                    else
                        embeddingDim_ = 768; // base/default
                } else if (has("gte")) {
                    if (has("large"))
                        embeddingDim_ = 1024;
                    else if (has("small"))
                        embeddingDim_ = 384;
                    else
                        embeddingDim_ = 768; // base/default
                } else if (has("bge")) {
                    if (has("-m3") || has("m3") || has("large"))
                        embeddingDim_ = 1024;
                    else if (has("small"))
                        embeddingDim_ = 384;
                    else
                        embeddingDim_ = 768; // base/default
                } else if (has("nomic") || has("mpnet") || has("jina")) {
                    embeddingDim_ = 768;
                } else if (has("minilm")) {
                    embeddingDim_ = 384;
                } else if (has("colbert")) {
                    embeddingDim_ = 48;
                }
            }
            // Sequence length defaults
            if (maxSequenceLength_ == 0 || maxSequenceLength_ == 512) {
                auto has = [&](const char* s) { return lname.find(s) != std::string::npos; };
                if (has("nomic") || has("bge-m3") || has("text-embedding-3-")) {
                    maxSequenceLength_ = 8192;
                } else if (has("minilm")) {
                    maxSequenceLength_ = 256;
                } else if (has("mpnet") || has("bge") || has("e5") || has("gte") ||
                           has("colbert")) {
                    maxSequenceLength_ = 512;
                }
            }
            // Respect explicit config cap
            if (config_.max_sequence_length > 0 &&
                maxSequenceLength_ > config_.max_sequence_length) {
                maxSequenceLength_ = config_.max_sequence_length;
            }
            isLoaded_ = true;
            return Result<void>();
        }

        auto providerLower = actualExecutionProvider_;
        std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
                       [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

        std::unique_lock<std::recursive_mutex> gpuLoadLock;
        std::unique_ptr<ScopedProcessFileLock> gpuProcessLoadLock;

        bool usesGpuProvider = false;
        if (providerLower.find("migraphx") != std::string::npos ||
            providerLower.find("cuda") != std::string::npos ||
            providerLower.find("dml") != std::string::npos ||
            providerLower.find("directml") != std::string::npos) {
            usesGpuProvider = true;
        }

        if (usesGpuProvider) {
            // Serialize GPU load paths to avoid non-deterministic GPU faults during startup.
            static std::recursive_mutex* g_gpu_load_mutex = new std::recursive_mutex();
            gpuLoadLock = std::unique_lock<std::recursive_mutex>(*g_gpu_load_mutex);
            spdlog::debug("[ONNX] GPU load gate acquired for '{}'", modelName_);

            if (gpuProcessLockEnabled()) {
                const auto processLockPath = resolveGpuProcessLockPath();
                gpuProcessLoadLock = std::make_unique<ScopedProcessFileLock>(processLockPath);
                if (gpuProcessLoadLock->locked()) {
                    spdlog::debug("[ONNX] GPU process load lock acquired: {}",
                                  processLockPath.string());
                } else {
                    spdlog::warn("[ONNX] Failed to acquire GPU process load lock: {}",
                                 processLockPath.string());
                }
            }
        }

        try {
            spdlog::debug("[ONNX] Creating Ort::Session for model '{}' at {}", modelName_.c_str(),
                          modelPath_.c_str());

            int retries = 3;
            while (retries-- > 0) {
                try {
                    // NOTE: Removed g_onnx_init_mutex lock here - it may have been conflicting
                    // with ONNX Runtime's internal locking. We rely on:
                    // 1. Single-flight pattern in loadModel() to prevent concurrent model loads
                    // 2. ONNX Runtime configured for single-threaded operation on Windows
                    spdlog::debug("[ONNX] Configuring session options for '{}'", modelName_);

                    // Clone session options; on Windows force single-thread to avoid deadlocks
                    auto options = sessionOptions_->Clone();
#ifdef _WIN32
                    // Windows: force single-threaded to avoid "resource deadlock would occur"
                    options.SetIntraOpNumThreads(1);
                    options.SetInterOpNumThreads(1);
#endif
                    // Non-Windows: use configured threads from sessionOptions_ (default intra=4)
                    applyCoreMLFreeDimensionOverrides(options);

                    // Create session directly - no async wrapper needed for local file operations
                    spdlog::debug("[ONNX] Creating Ort::Session for '{}' at '{}'", modelName_,
                                  modelPath_);
                    session_ = std::make_unique<Ort::Session>(
                        *env_, std::filesystem::path(modelPath_).c_str(), options);
                    spdlog::info("[ONNX] Ort::Session created successfully for '{}'", modelName_);

                    // Success!
                    break;
                } catch (const std::system_error& e) {
                    spdlog::warn("[ONNX] loadModel: system_error '{}' (code={}). Retries left: {}",
                                 e.what(), e.code().value(), retries);
                    if (retries == 0)
                        throw;
                    std::this_thread::sleep_for(std::chrono::milliseconds(200));
                } catch (const Ort::Exception& e) {
                    spdlog::warn("[ONNX] loadModel: Ort::Exception '{}'. Retries left: {}",
                                 e.what(), retries);
                    if (retries == 0)
                        throw;
                    std::this_thread::sleep_for(std::chrono::milliseconds(200));
                }
            }

            // Get input/output information
            size_t numInputs = session_->GetInputCount();
            size_t numOutputs = session_->GetOutputCount();

            if (numInputs < 2) {
                return Error{ErrorCode::InvalidData,
                             "Model must have at least 2 inputs (input_ids, attention_mask)"};
            }

            // Get input/output names
            Ort::AllocatorWithDefaultOptions allocator;
            for (size_t i = 0; i < numInputs; ++i) {
                auto inputName = session_->GetInputNameAllocated(i, allocator);
                inputNames_.push_back(inputName.get());
                // Inspect first input shape to detect model sequence length when available
                if (i == 0) {
                    try {
                        auto inInfo = session_->GetInputTypeInfo(i);
                        auto inTensor = inInfo.GetTensorTypeAndShapeInfo();
                        auto inShape = inTensor.GetShape();
                        // Expect shape like [batch, seq_len]; seq_len > 0 means fixed
                        if (inShape.size() >= 2) {
                            auto seq = inShape[1];
                            if (seq > 0) {
                                maxSequenceLength_ = static_cast<size_t>(seq);
                                hasFixedInputShape_ = true;
                            }
                        }
                    } catch (const std::exception& e) {
                        spdlog::debug("[ONNX] Failed to read input shape: {}", e.what());
                    }
                }
            }

            for (size_t i = 0; i < numOutputs; ++i) {
                auto outputName = session_->GetOutputNameAllocated(i, allocator);
                outputNames_.push_back(outputName.get());
            }

            // Get output shape to determine embedding dimension
            if (numOutputs > 0) {
                auto outputInfo = session_->GetOutputTypeInfo(0);
                auto tensorInfo = outputInfo.GetTensorTypeAndShapeInfo();
                auto shape = tensorInfo.GetShape();
                if (shape.size() >= 2 && shape.back() > 0) {
                    embeddingDim_ = static_cast<size_t>(shape.back());
                }
            }

            // Try to parse nearby config (Sentence-Transformers) for pooling/max
            // length/normalization
            parseModelConfigHints();

            // Load tokenizer.json from model directory (same search dirs as config.json)
            try {
                namespace fs = std::filesystem;
                fs::path mp(modelPath_);
                for (const auto& dir : {mp.parent_path(), mp.parent_path().has_parent_path()
                                                              ? mp.parent_path().parent_path()
                                                              : fs::path{}}) {
                    if (dir.empty())
                        continue;
                    fs::path tok = dir / "tokenizer.json";
                    if (fs::exists(tok) && tokenizer_.load(tok.string())) {
                        spdlog::info("[ONNX] Loaded tokenizer from {} (vocab={})", tok.string(),
                                     tokenizer_.vocabSize());
                        break;
                    }
                }
                if (!tokenizer_.isLoaded()) {
                    spdlog::warn(
                        "[ONNX] tokenizer.json not found; using TextPreprocessor fallback");
                }
            } catch (...) {
            }

            // If input shape and config did not give a positive seq_len, fall back to a heuristic.
            if (maxSequenceLength_ == 0) {
                if (modelName_.find("MiniLM") != std::string::npos) {
                    maxSequenceLength_ = 256;
                } else if (modelName_.find("mpnet") != std::string::npos) {
                    maxSequenceLength_ = 512;
                } else if (modelName_.find("nomic") != std::string::npos) {
                    // Nomic BERT supports very long contexts (e.g., 8192). Default to 2048 unless
                    // user config overrides to avoid excessive memory on defaults.
                    maxSequenceLength_ = 2048;
                } else if (modelName_.find("embeddinggemma") != std::string::npos ||
                           modelName_.find("gemma") != std::string::npos) {
                    maxSequenceLength_ = 2048;
                } else {
                    maxSequenceLength_ = 512;
                }
            }
            // Respect configured max from EmbeddingConfig if provided
            if (config_.max_sequence_length > 0) {
                maxSequenceLength_ = std::min(maxSequenceLength_, config_.max_sequence_length);
            }

            if (gemmaCoreMLStaticBatchOneEnabled()) {
                providerBatchLimit_.store(1, std::memory_order_relaxed);
                spdlog::info("[ONNX] Enabled static batch_size=1 CoreML path for '{}'; larger "
                             "embedding batches will be chunked per item",
                             modelName_);
            }

            isLoaded_ = true;
            // IoBinding benefits real GPU providers (pre-resolved memory placement);
            // on CPU/CoreML-to-CPU it adds overhead vs direct Session::Run().
            if (actualExecutionProvider_ == "cuda" || actualExecutionProvider_ == "migraphx" ||
                actualExecutionProvider_ == "directml") {
                initIoBinding();
            }
            spdlog::info(
                "[ONNX] Session ready for '{}' (inputs={}, outputs={}, dim={}, max_seq_len={})",
                modelName_.c_str(), numInputs, numOutputs, embeddingDim_, maxSequenceLength_);
            return Result<void>();

        } catch (const Ort::Exception& e) {
            spdlog::error("[ONNX] Failed to load '{}' (Ort::Exception): {}", modelName_.c_str(),
                          e.what());
            return Error{ErrorCode::InternalError,
                         std::string("Failed to load ONNX model: ") + e.what()};
        } catch (const std::system_error& e) {
            spdlog::error("[ONNX] Failed to load '{}' (std::system_error): {}", modelName_.c_str(),
                          e.what());
            return Error{ErrorCode::InternalError,
                         std::string("Failed to load ONNX model (system error): ") + e.what()};
        }
    }

    Result<std::vector<float>> generateEmbedding(const std::string& text) {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        if (test_mode_) {
            // Representative mock: deterministic per-text embedding derived from tokens.
            // - Uses same preprocessing (tokenize/truncate/pad + attention mask)
            // - Pseudo-random per-token vector seeded by token id and position
            // - Pooling respects configured pooling_mode_
            // - Optional L2 normalization
            const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;
            auto toks = preprocessor_.tokenize(text);
            toks = preprocessor_.truncateTokens(toks, seq_len);
            toks = preprocessor_.padTokens(toks, seq_len);
            auto mask = preprocessor_.generateAttentionMask(toks);

            const size_t D = embeddingDim_ > 0 ? embeddingDim_ : 0;
            auto prng_val = [](uint32_t seed) -> float {
                // xorshift32 → map to [-1, 1]
                uint32_t x = seed ? seed : 0x9e3779b9u;
                x ^= x << 13;
                x ^= x >> 17;
                x ^= x << 5;
                // Scale to [0,1)
                float u = (x & 0x00FFFFFFu) / static_cast<float>(0x01000000u);
                return 2.0f * u - 1.0f;
            };

            std::vector<float> emb(D, 0.0f);
            if (pooling_mode_ == Pooling::CLS) {
                const int32_t t = toks.empty() ? 0 : toks[0];
                for (size_t d = 0; d < D; ++d) {
                    emb[d] = prng_val(static_cast<uint32_t>(t * 1315423911u + d * 2654435761u));
                }
            } else if (pooling_mode_ == Pooling::MAX) {
                for (size_t i = 0; i < std::min(seq_len, toks.size()); ++i) {
                    if (i >= mask.size() || mask[i] <= 0)
                        continue;
                    const int32_t t = toks[i];
                    for (size_t d = 0; d < D; ++d) {
                        float v = prng_val(static_cast<uint32_t>(t * 2246822519u + i * 3266489917u +
                                                                 d * 668265263u));
                        emb[d] = (i == 0) ? v : std::max(emb[d], v);
                    }
                }
            } else { // MEAN pooling
                double denom = 0.0;
                for (size_t i = 0; i < std::min(seq_len, toks.size()); ++i) {
                    float w = (i < mask.size() && mask[i] > 0) ? 1.0f : 0.0f;
                    if (w <= 0.0f)
                        continue;
                    denom += w;
                    const int32_t t = toks[i];
                    for (size_t d = 0; d < D; ++d) {
                        float v = prng_val(
                            static_cast<uint32_t>(t * 374761393u + i * 1664525u + d * 1013904223u));
                        emb[d] += v * w;
                    }
                }
                if (denom > 0.0) {
                    for (size_t d = 0; d < D; ++d)
                        emb[d] = static_cast<float>(emb[d] / denom);
                }
            }

            if (normalize_) {
                double n2 = 0.0;
                for (auto v : emb) {
                    double dv = static_cast<double>(v);
                    n2 += dv * dv;
                }
                n2 = std::sqrt(n2);
                if (n2 > 1e-8) {
                    for (auto& v : emb)
                        v = static_cast<float>(v / n2);
                } else {
                    std::fill(emb.begin(), emb.end(), 0.0f);
                }
            }
            return emb;
        }

        if (!isLoaded_) {
            if (auto result = loadModel(); !result) {
                return result.error();
            }
        }

        // Try real ONNX path first (preferred)
        if (auto r = withCoreMLFallback([&]() { return runOnnx(text); }); r) {
            return r.value();
        }

        return Error{ErrorCode::InternalError,
                     "ONNX embedding inference failed and no fallback backend is available"};
    }

    // Public batch API that uses the optimized batched ONNX path
    Result<std::vector<std::vector<float>>>
    generateBatchEmbeddings(const std::vector<std::string>& texts) {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        if (test_mode_) {
            std::vector<std::vector<float>> res;
            res.reserve(texts.size());
            for (const auto& t : texts) {
                auto one = generateEmbedding(t);
                if (!one)
                    return one.error();
                res.emplace_back(std::move(one.value()));
            }
            return res;
        }
        if (!isLoaded_) {
            if (auto result = loadModel(); !result) {
                return result.error();
            }
        }
        return withCoreMLFallback([&]() { return runOnnxBatch(texts); });
    }

    bool isValid() const {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        return test_mode_ ? isLoaded_ : (isLoaded_ && session_ != nullptr);
    }

    std::pair<int, int> getThreading() const {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        return {configuredIntraThreads_, configuredInterThreads_};
    }

    size_t getLearnedBatchLimit() const {
        return providerBatchLimit_.load(std::memory_order_relaxed);
    }

    bool prefersExplicitSentenceEmbeddingOutput() const {
        return modelName_.find("embeddinggemma") != std::string::npos ||
               modelName_.find("gemma") != std::string::npos;
    }

    bool prefersStaticShapesOnCoreML() const {
        return actualExecutionProvider_ == "coreml" &&
               (modelName_.find("embeddinggemma") != std::string::npos ||
                modelName_.find("gemma") != std::string::npos);
    }

    bool gemmaCoreMLStaticBatchOneEnabled() const { return prefersStaticShapesOnCoreML(); }

    bool shouldApplyCoreMLFreeDimensionOverrides() const {
        return gemmaCoreMLStaticBatchOneEnabled();
    }

    int64_t resolveCoreMLStaticSequenceLength() const {
        if (config_.max_sequence_length > 0) {
            return std::max<int64_t>(1, std::min<int64_t>(config_.max_sequence_length, 2048));
        }
        if (maxSequenceLength_ > 512) {
            return static_cast<int64_t>(maxSequenceLength_);
        }
        return int64_t{2048};
    }

    void applyCoreMLFreeDimensionOverrides(Ort::SessionOptions& options) {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        if (!shouldApplyCoreMLFreeDimensionOverrides()) {
            return;
        }

        const int64_t effectiveBatchSize = 1;
        const int64_t sequenceLength = resolveCoreMLStaticSequenceLength();
        const int64_t tripleSequenceLength =
            sequenceLength > (std::numeric_limits<int64_t>::max() / 3) ? sequenceLength
                                                                       : (sequenceLength * 3);

        Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(options, "batch_size",
                                                                       effectiveBatchSize));
        Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(
            options, "total_sequence_length", sequenceLength));
        Ort::ThrowOnError(Ort::GetApi().AddFreeDimensionOverrideByName(
            options, "3*total_sequence_length", tripleSequenceLength));
        spdlog::info("[ONNX] Applied Gemma CoreML static-shape overrides for '{}': batch_size={} "
                     "total_sequence_length={} triple_total_sequence_length={}",
                     modelName_, effectiveBatchSize, sequenceLength, tripleSequenceLength);
    }

    bool supportsCpuFallbackForCurrentModel() const {
        return supportsCoreMLCpuFallbackModel(modelName_);
    }

    bool isCoreMLMaskFeatureFailure(const std::string& message) const {
        return message.find("_model_st_pool_0_mask_unsqueeze_output_0") != std::string::npos ||
               message.find("required but not specified") != std::string::npos;
    }

    Result<void> reinitializeSessionOnCpu() {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        if (!supportsCpuFallbackForCurrentModel()) {
            return Error{ErrorCode::InternalError, "CPU fallback is not enabled for this model"};
        }

        if (actualExecutionProvider_ == "cpu") {
            return Result<void>();
        }

        try {
            session_.reset();
            sessionOptions_ = std::make_unique<Ort::SessionOptions>();
            actualExecutionProvider_ = "cpu";
            providerBatchLimit_.store(0, std::memory_order_relaxed);

            const int hwThreads = static_cast<int>(std::thread::hardware_concurrency());
            const int fallbackIntra = resolveCpuFallbackIntraThreads(configuredIntraThreads_);
            const int fallbackInter = resolveCpuFallbackInterThreads(configuredInterThreads_);
            sessionOptions_->SetIntraOpNumThreads(fallbackIntra);
            sessionOptions_->SetInterOpNumThreads(fallbackInter);
            configuredIntraThreads_ = fallbackIntra;
            configuredInterThreads_ = fallbackInter;

            const bool allowSpinning = resolveCpuFallbackAllowSpinning();
            sessionOptions_->AddConfigEntry("session.intra_op.allow_spinning",
                                            allowSpinning ? "1" : "0");
            sessionOptions_->AddConfigEntry("session.inter_op.allow_spinning",
                                            allowSpinning ? "1" : "0");

#ifdef _WIN32
            sessionOptions_->SetExecutionMode(ExecutionMode::ORT_SEQUENTIAL);
#endif

            GraphOptimizationLevel optLevel = resolveGraphOptimizationLevel(true);
            sessionOptions_->SetGraphOptimizationLevel(optLevel);

            // Graph cache is on by default; disable with YAMS_ONNX_GRAPH_CACHE=0
            if (!std::getenv("YAMS_ONNX_GRAPH_CACHE") ||
                envTruthy(std::getenv("YAMS_ONNX_GRAPH_CACHE"))) {
                auto optPath = resolveOptimizedModelPath(modelPath_, optLevel);
                std::error_code ec;
                fs::create_directories(optPath.parent_path(), ec);
                if (!ec) {
                    sessionOptions_->SetOptimizedModelFilePath(optPath.c_str());
                }
            }

            sessionOptions_->EnableMemPattern();
            sessionOptions_->EnableCpuMemArena();

            auto options = sessionOptions_->Clone();
#ifdef _WIN32
            options.SetIntraOpNumThreads(1);
            options.SetInterOpNumThreads(1);
#endif
            session_ = std::make_unique<Ort::Session>(
                *env_, std::filesystem::path(modelPath_).c_str(), options);
            // CPU fallback — no IoBinding (adds overhead on CPU)
            useIoBinding_ = false;
            ioBinding_.reset();
            OnnxModelSession::denylistCoreML(modelName_);
            const size_t fallbackBatchCap = applyCpuFallbackSafetyCap(
                modelName_, actualExecutionProvider_, providerBatchLimit_);
            spdlog::warn("[ONNX] Reinitialized '{}' on CPU after CoreML inference failure "
                         "(denylisted for future instances, intra_threads={}, inter_threads={}, "
                         "batch_cap={}, spinning={}, hw={})",
                         modelName_, fallbackIntra, fallbackInter, fallbackBatchCap,
                         allowSpinning ? 1 : 0, hwThreads);
            return Result<void>();
        } catch (const Ort::Exception& e) {
            return Error{ErrorCode::InternalError,
                         std::string("CPU fallback session init failed: ") + e.what()};
        } catch (const std::exception& e) {
            return Error{ErrorCode::InternalError,
                         std::string("CPU fallback session init failed: ") + e.what()};
        }
    }

    template <typename Fn> auto withCoreMLFallback(Fn&& fn) -> decltype(fn()) {
        auto result = fn();
        if (result || actualExecutionProvider_ != "coreml" ||
            !supportsCpuFallbackForCurrentModel()) {
            return result;
        }

        const auto& message = result.error().message;
        if (!isCoreMLMaskFeatureFailure(message)) {
            return result;
        }

        auto fallback = reinitializeSessionOnCpu();
        if (!fallback) {
            return decltype(fn())(fallback.error());
        }
        return fn();
    }

    std::vector<const char*> resolveOutputNames() const {
        if (outputNames_.empty()) {
            static const char* default_out = "sentence_embedding";
            return {default_out};
        }

        if (prefersExplicitSentenceEmbeddingOutput()) {
            for (const auto& name : outputNames_) {
                if (name == "sentence_embedding") {
                    return {name.c_str()};
                }
            }
        }

        std::vector<const char*> output_names_c;
        output_names_c.reserve(outputNames_.size());
        for (const auto& n : outputNames_) {
            output_names_c.push_back(n.c_str());
        }
        return output_names_c;
    }

    Result<void> setThreading(int intraThreads, int interThreads) {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        auto valid = [](int value) { return value == -1 || (value >= 1 && value <= 64); };
        if (!valid(intraThreads) || !valid(interThreads)) {
            return Error{ErrorCode::InvalidArgument,
                         "Threading values must be -1 or in range [1, 64]"};
        }

        if (intraThreads > 0) {
            config_.num_threads = intraThreads;
        }
        if (interThreads > 0) {
            config_.inter_op_threads = interThreads;
        }

        configuredIntraThreads_ =
            (config_.num_threads > 0) ? config_.num_threads : configuredIntraThreads_;
        configuredInterThreads_ =
            (config_.inter_op_threads > 0) ? config_.inter_op_threads : configuredInterThreads_;
        if (sessionOptions_) {
            sessionOptions_->SetIntraOpNumThreads(configuredIntraThreads_);
            sessionOptions_->SetInterOpNumThreads(configuredInterThreads_);
        }

        // Force lazy reload on next inference with updated threading.
        session_.reset();
        isLoaded_ = false;

        return Result<void>();
    }

    size_t getEmbeddingDim() const {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        return embeddingDim_;
    }
    size_t getMaxSequenceLength() const {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        return maxSequenceLength_;
    }
    std::string getExecutionProvider() const {
        std::lock_guard<std::recursive_mutex> lock(stateMutex_);
        return actualExecutionProvider_;
    }

private:
    void appendGpuExecutionProvider() {
        std::string cacheDir;
        if (!modelPath_.empty()) {
            auto parentDir = fs::path(modelPath_).parent_path();
            if (fs::exists(parentDir)) {
                cacheDir = parentDir.string();
            }
        }
        actualExecutionProvider_ = onnx_util::appendGpuProvider(*sessionOptions_, cacheDir);
    }

    struct GpuInferenceGuard {
        std::unique_lock<std::recursive_mutex> threadLock;
        std::unique_ptr<ScopedProcessFileLock> processLock;
        bool depthIncremented{false};

        GpuInferenceGuard() = default;
        GpuInferenceGuard(const GpuInferenceGuard&) = delete;
        GpuInferenceGuard& operator=(const GpuInferenceGuard&) = delete;
        GpuInferenceGuard(GpuInferenceGuard&&) noexcept = default;
        GpuInferenceGuard& operator=(GpuInferenceGuard&&) noexcept = default;

        ~GpuInferenceGuard() {
            if (depthIncremented) {
                auto& depth = gpuProcessLockDepth();
                if (depth > 0) {
                    --depth;
                }
            }
        }
    };

    GpuInferenceGuard acquireGpuInferenceLock() const {
        GpuInferenceGuard guard;
        std::string providerLower = actualExecutionProvider_;
        std::transform(providerLower.begin(), providerLower.end(), providerLower.begin(),
                       [](unsigned char c) { return static_cast<char>(std::tolower(c)); });

        bool usesGpuProvider = false;
        if (providerLower.find("migraphx") != std::string::npos ||
            providerLower.find("cuda") != std::string::npos ||
            providerLower.find("dml") != std::string::npos ||
            providerLower.find("directml") != std::string::npos) {
            usesGpuProvider = true;
        }
        if (!usesGpuProvider) {
            return guard;
        }

        static std::recursive_mutex* g_gpu_infer_mutex = new std::recursive_mutex();
        guard.threadLock = std::unique_lock<std::recursive_mutex>(*g_gpu_infer_mutex);

        if (!gpuProcessLockEnabled()) {
            return guard;
        }

        auto& depth = gpuProcessLockDepth();
        ++depth;
        guard.depthIncremented = true;
        if (depth > 1) {
            return guard;
        }

        const auto processLockPath = resolveGpuProcessLockPath();
        guard.processLock = std::make_unique<ScopedProcessFileLock>(processLockPath);
        if (!guard.processLock->locked()) {
            spdlog::warn("[ONNX] Failed to acquire GPU process inference lock: {}",
                         processLockPath.string());
            guard.processLock.reset();
        }
        return guard;
    }

    // Align a sequence length up to a SIMD-friendly multiple (AVX: 8 x int32, NEON: 4).
    static size_t alignSeqLen(size_t len, size_t alignment = 8) {
        return ((len + alignment - 1) / alignment) * alignment;
    }

    void initIoBinding() {
        if (!session_)
            return;
        try {
            ioBinding_ = std::make_unique<Ort::IoBinding>(*session_);
            useIoBinding_ = true;
            spdlog::debug("[ONNX] IoBinding created for '{}'", modelName_);
        } catch (const std::exception& e) {
            spdlog::warn("[ONNX] IoBinding creation failed for '{}': {}; using Run()", modelName_,
                         e.what());
            useIoBinding_ = false;
        }
    }

    // Run inference using IoBinding if available, otherwise fall back to Session::Run().
    std::vector<Ort::Value> runSession(std::vector<Ort::Value>& inputs) {
        if (useIoBinding_ && ioBinding_) {
            ioBinding_->ClearBoundInputs();
            ioBinding_->ClearBoundOutputs();
            for (size_t i = 0; i < inputs.size() && i < inputNames_.size(); ++i) {
                ioBinding_->BindInput(inputNames_[i].c_str(), inputs[i]);
            }
            for (const auto& name : resolveOutputNames()) {
                ioBinding_->BindOutput(name, *memoryInfo_);
            }
            session_->Run(Ort::RunOptions{nullptr}, *ioBinding_);
            return ioBinding_->GetOutputValues();
        }
        // Fallback: direct Run with name arrays
        std::vector<const char*> in_names;
        for (auto& n : inputNames_)
            in_names.push_back(n.c_str());
        auto out_names = resolveOutputNames();
        return session_->Run(Ort::RunOptions{nullptr}, in_names.data(), inputs.data(),
                             inputs.size(), out_names.data(), out_names.size());
    }

    Result<std::vector<float>> runOnnx(const std::string& t) {
        auto gpuInferenceLock = acquireGpuInferenceLock();

        const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;
        auto tokens = tokenizer_.isLoaded() ? tokenizer_.encode(t) : preprocessor_.tokenize(t);
        tokens = preprocessor_.truncateTokens(tokens, seq_len);

        // Dynamic padding: pad to aligned actual length instead of model max
        size_t effective_seq_len = seq_len;
        if (dynamicPaddingEnabled_ && !hasFixedInputShape_ && !prefersStaticShapesOnCoreML()) {
            effective_seq_len = std::min(alignSeqLen(tokens.size()), seq_len);
            if (effective_seq_len == 0)
                effective_seq_len = 1;
            spdlog::debug("[ONNX] Dynamic padding (single): actual={} effective={} (model_max={})",
                          tokens.size(), effective_seq_len, seq_len);
        } else if (prefersStaticShapesOnCoreML()) {
            spdlog::debug("[ONNX] Static padding forced for '{}' on CoreML: seq_len={}", modelName_,
                          effective_seq_len);
        }

        tokens = preprocessor_.padTokens(tokens, effective_seq_len);
        auto attention_mask = preprocessor_.generateAttentionMask(tokens);

        std::vector<int64_t> input_shape = {1, static_cast<int64_t>(effective_seq_len)};
        const size_t input_tensor_size = effective_seq_len;

        // Reuse thread-local buffers to avoid per-call heap allocations.
        auto& bufs = getThreadLocalBuffers();
        bufs.tokens.resize(effective_seq_len);
        bufs.masks.resize(effective_seq_len);

        if (tokenizer_.isLoaded()) {
            for (size_t i = 0; i < effective_seq_len && i < tokens.size(); ++i)
                bufs.tokens[i] = static_cast<int64_t>(tokens[i]);
            for (size_t i = tokens.size(); i < effective_seq_len; ++i)
                bufs.tokens[i] = 0;
        } else {
            const int64_t MAX_TOKEN_ID = 30527;
            const int64_t UNK_TOKEN_ID = 100;
            for (size_t i = 0; i < effective_seq_len && i < tokens.size(); ++i) {
                auto v = tokens[i];
                bufs.tokens[i] =
                    (v < 0 || v > MAX_TOKEN_ID) ? UNK_TOKEN_ID : static_cast<int64_t>(v);
            }
            for (size_t i = tokens.size(); i < effective_seq_len; ++i)
                bufs.tokens[i] = UNK_TOKEN_ID;
        }
        for (size_t i = 0; i < effective_seq_len && i < attention_mask.size(); ++i)
            bufs.masks[i] = static_cast<int64_t>(attention_mask[i]);
        for (size_t i = attention_mask.size(); i < effective_seq_len; ++i)
            bufs.masks[i] = 0;

        std::vector<Ort::Value> input_tensors;
        input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
            *memoryInfo_, bufs.tokens.data(), input_tensor_size, input_shape.data(), 2));
        input_tensors.push_back(Ort::Value::CreateTensor<int64_t>(
            *memoryInfo_, bufs.masks.data(), input_tensor_size, input_shape.data(), 2));
        if (inputNames_.size() >= 3) {
            bufs.token_type_ids.assign(effective_seq_len, 0);
            input_tensors.push_back(
                Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.token_type_ids.data(),
                                                  input_tensor_size, input_shape.data(), 2));
        }

        auto outputs = runSession(input_tensors);
        if (outputs.empty()) {
            return Error{ErrorCode::InternalError, "ONNX session returned no outputs"};
        }

        auto& out = outputs[0];
        auto out_shape = out.GetTensorTypeAndShapeInfo().GetShape();
        const size_t hidden_dim = embeddingDim_;

        if (out_shape.size() == 3) {
            const size_t seq = static_cast<size_t>(out_shape[1]);
            const size_t dim = static_cast<size_t>(out_shape[2]);
            if (dim != hidden_dim) {
                return Error{ErrorCode::InvalidData, "Output dim mismatch"};
            }
            const float* data = out.GetTensorData<float>();
            std::vector<float> embedding(hidden_dim, 0.0f);
            if (pooling_mode_ == Pooling::CLS) {
                // CLS token (first token)
                const size_t off = 0;
                for (size_t d = 0; d < dim; ++d)
                    embedding[d] = data[off + d];
            } else if (pooling_mode_ == Pooling::MAX) {
                for (size_t i = 0; i < std::min(seq, attention_mask.size()); ++i) {
                    if (attention_mask[i] <= 0)
                        continue;
                    const size_t off = i * dim;
                    for (size_t d = 0; d < dim; ++d)
                        embedding[d] = std::max(embedding[d], data[off + d]);
                }
            } else { // MEAN
                float total_attn = 0.0f;
                for (size_t i = 0; i < std::min(seq, attention_mask.size()); ++i) {
                    const float w = static_cast<float>(attention_mask[i]);
                    if (w <= 0.0f)
                        continue;
                    total_attn += w;
                    const size_t off = i * dim;
                    for (size_t d = 0; d < dim; ++d)
                        embedding[d] += data[off + d] * w;
                }
                if (total_attn > 0.0f) {
                    for (auto& v : embedding)
                        v = static_cast<float>(v / total_attn);
                }
            }
            if (normalize_) {
                double norm = 0.0;
                for (auto v : embedding) {
                    double dv = static_cast<double>(v);
                    norm += dv * dv;
                }
                norm = std::sqrt(norm);
                if (norm > 1e-8) {
                    for (auto& v : embedding)
                        v = static_cast<float>(v / norm);
                } else {
                    std::fill(embedding.begin(), embedding.end(), 0.0f);
                }
            }
            return embedding;
        }

        if (out_shape.size() == 2) {
            const size_t dim = static_cast<size_t>(out_shape[1]);
            if (dim != hidden_dim) {
                return Error{ErrorCode::InvalidData, "Output dim mismatch"};
            }
            const float* data = out.GetTensorData<float>();
            std::vector<float> embedding(data, data + dim);
            if (normalize_) {
                double norm = 0.0;
                for (auto v : embedding) {
                    double dv = static_cast<double>(v);
                    norm += dv * dv;
                }
                norm = std::sqrt(norm);
                if (norm > 1e-8) {
                    for (auto& v : embedding)
                        v = static_cast<float>(v / norm);
                } else {
                    std::fill(embedding.begin(), embedding.end(), 0.0f);
                }
            }
            return embedding;
        }

        return Error{ErrorCode::InvalidData, "Unexpected output tensor rank"};
    }

    Result<std::vector<std::vector<float>>> runOnnxBatch(const std::vector<std::string>& texts) {
        auto gpuInferenceLock = acquireGpuInferenceLock();

        if (texts.empty())
            return std::vector<std::vector<float>>{};

        if (providerBatchLimit_.load(std::memory_order_relaxed) == 0) {
            const size_t sharedCap = loadSharedBatchLimit(modelName_, actualExecutionProvider_);
            if (sharedCap > 0) {
                providerBatchLimit_.store(sharedCap, std::memory_order_relaxed);
                const auto currentCap = TuneAdvisor::getEmbedDocCap();
                if (currentCap == 0 || currentCap > sharedCap) {
                    TuneAdvisor::setEmbedDocCap(sharedCap);
                }
            }
        }

        const size_t requestedBatchSize = texts.size();
        size_t learnedBatchCap = providerBatchLimit_.load(std::memory_order_relaxed);
        if (gemmaCoreMLStaticBatchOneEnabled() && requestedBatchSize > 1 && learnedBatchCap == 0) {
            providerBatchLimit_.store(1, std::memory_order_relaxed);
            learnedBatchCap = 1;
        }
        if (learnedBatchCap > 0 && requestedBatchSize > learnedBatchCap) {
            std::vector<std::vector<float>> merged;
            merged.reserve(requestedBatchSize);
            for (size_t offset = 0; offset < requestedBatchSize; offset += learnedBatchCap) {
                const size_t count =
                    std::min(learnedBatchCap, static_cast<size_t>(requestedBatchSize - offset));
                std::vector<std::string> chunk;
                chunk.reserve(count);
                for (size_t i = 0; i < count; ++i) {
                    chunk.push_back(texts[offset + i]);
                }
                auto chunkResult = runOnnxBatch(chunk);
                if (!chunkResult) {
                    return chunkResult.error();
                }
                auto& vectors = chunkResult.value();
                merged.insert(merged.end(), std::make_move_iterator(vectors.begin()),
                              std::make_move_iterator(vectors.end()));
            }
            return merged;
        }

        const size_t seq_len = maxSequenceLength_ > 0 ? maxSequenceLength_ : 512;

        // Pass 1: tokenize + truncate (no padding yet)
        std::vector<std::vector<int32_t>> token_seqs;
        token_seqs.reserve(texts.size());
        for (const auto& t : texts) {
            auto tokens = tokenizer_.isLoaded() ? tokenizer_.encode(t) : preprocessor_.tokenize(t);
            tokens = preprocessor_.truncateTokens(tokens, seq_len);
            token_seqs.push_back(std::move(tokens));
        }

        // Determine effective padding length
        size_t effective_seq_len = seq_len;
        if (dynamicPaddingEnabled_ && !hasFixedInputShape_ && !prefersStaticShapesOnCoreML()) {
            size_t max_actual = 0;
            for (const auto& toks : token_seqs)
                max_actual = std::max(max_actual, toks.size());
            effective_seq_len = std::min(alignSeqLen(max_actual), seq_len);
            if (effective_seq_len == 0)
                effective_seq_len = 1;
            spdlog::debug("[ONNX] Dynamic padding: B={} max_actual={} effective={} (model_max={})",
                          texts.size(), max_actual, effective_seq_len, seq_len);
        } else if (prefersStaticShapesOnCoreML()) {
            spdlog::debug("[ONNX] Static padding forced for '{}' batch on CoreML: B={} seq_len={}",
                          modelName_, texts.size(), effective_seq_len);
        }

        // Pass 2: pad to effective_seq_len + generate masks
        std::vector<std::vector<int32_t>> masks;
        masks.reserve(texts.size());
        for (auto& toks : token_seqs) {
            toks = preprocessor_.padTokens(toks, effective_seq_len);
            masks.push_back(preprocessor_.generateAttentionMask(toks));
        }

        // Prepare batched tensors [B, S]
        const size_t B = texts.size();
        std::vector<int64_t> input_shape = {static_cast<int64_t>(B),
                                            static_cast<int64_t>(effective_seq_len)};
        const size_t tensor_size = B * effective_seq_len;

        // Reuse thread-local buffers to avoid per-call heap allocations.
        auto& bufs = getThreadLocalBuffers();
        bufs.tokens.resize(tensor_size);
        bufs.masks.resize(tensor_size);

        const bool useRealTokenizer = tokenizer_.isLoaded();
        const int64_t MAX_TOKEN_ID = 30527;
        const int64_t UNK_TOKEN_ID = 100;

        size_t offset = 0;
        for (size_t i = 0; i < B; ++i) {
            const auto& toks = token_seqs[i];
            const auto& m = masks[i];
            for (size_t j = 0; j < effective_seq_len && j < toks.size(); ++j) {
                int64_t v = toks[j];
                if (useRealTokenizer) {
                    bufs.tokens[offset] = v;
                } else {
                    bufs.tokens[offset] = (v < 0 || v > MAX_TOKEN_ID) ? UNK_TOKEN_ID : v;
                }
                bufs.masks[offset] = static_cast<int64_t>(m[j]);
                ++offset;
            }
            for (size_t j = toks.size(); j < effective_seq_len; ++j) {
                bufs.tokens[offset] = useRealTokenizer ? 0 : UNK_TOKEN_ID;
                bufs.masks[offset] = 0;
                ++offset;
            }
        }

        std::vector<Ort::Value> inputs;
        inputs.push_back(Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.tokens.data(),
                                                           tensor_size, input_shape.data(), 2));
        inputs.push_back(Ort::Value::CreateTensor<int64_t>(*memoryInfo_, bufs.masks.data(),
                                                           tensor_size, input_shape.data(), 2));
        if (inputNames_.size() >= 3) {
            bufs.token_type_ids.assign(tensor_size, 0);
            inputs.push_back(Ort::Value::CreateTensor<int64_t>(
                *memoryInfo_, bufs.token_type_ids.data(), tensor_size, input_shape.data(), 2));
        }

        // Wrap ONNX Run in try/catch to prevent exceptions from bubbling through ABI boundary
        // On Windows, ORT can throw std::system_error("resource deadlock would occur") when
        // thread pool resources are exhausted (PBI-1c1)
        std::vector<Ort::Value> outputs;
        int retries = 3;
        while (retries-- > 0) {
            try {
                outputs = runSession(inputs);
                break; // synchronous success
            } catch (const Ort::Exception& e) {
                if (retries == 0) {
                    fprintf(stderr, "[ONNX Plugin] Ort::Exception in Run: %s\n", e.what());
                    return Error{ErrorCode::InternalError,
                                 std::string("ONNX runtime error: ") + e.what()};
                }
                std::this_thread::sleep_for(std::chrono::milliseconds(50));
            } catch (const std::system_error& e) {
                if (std::string(e.what()).find("resource deadlock") != std::string::npos ||
                    retries > 0) {
                    spdlog::warn("[ONNX Plugin] Deadlock/system error caught, retrying... ({})",
                                 e.what());
                    std::this_thread::sleep_for(std::chrono::milliseconds(100));
                    if (retries == 0) {
                        fprintf(stderr, "[ONNX Plugin] std::system_error in Run (exhausted): %s\n",
                                e.what());
                        return Error{ErrorCode::InternalError,
                                     std::string("System error during ONNX inference: ") +
                                         e.what()};
                    }
                    continue;
                }
                fprintf(stderr, "[ONNX Plugin] std::system_error in Run: %s\n", e.what());
                return Error{ErrorCode::InternalError,
                             std::string("System error during ONNX inference: ") + e.what()};
            } catch (const std::exception& e) {
                fprintf(stderr, "[ONNX Plugin] std::exception in Run: %s\n", e.what());
                return Error{ErrorCode::InternalError,
                             std::string("Exception during ONNX inference: ") + e.what()};
            }
        }
        if (outputs.empty())
            return Error{ErrorCode::InternalError, "ONNX session returned no outputs"};

        auto& out = outputs[0];
        auto shape = out.GetTensorTypeAndShapeInfo().GetShape();
        // Some execution providers can leave output dims unknown until the first Run.
        // If we don't know the embedding dim yet, infer it from the first output tensor.
        size_t hidden_dim = embeddingDim_;
        std::vector<std::vector<float>> result;
        if (shape.size() == 3) {
            // [B, S, D]
            const int64_t b0 = shape[0];
            const int64_t s1 = shape[1];
            const int64_t d2 = shape[2];
            const size_t SS = (s1 > 0) ? static_cast<size_t>(s1) : maxSequenceLength_;
            const size_t DD = (d2 > 0) ? static_cast<size_t>(d2) : 0;
            if (hidden_dim == 0 && DD > 0) {
                hidden_dim = DD;
                embeddingDim_ = hidden_dim;
            }

            if (b0 > 0 && static_cast<size_t>(b0) < B && (hidden_dim == 0 || DD == hidden_dim)) {
                const size_t observedCap = static_cast<size_t>(b0);
                const size_t mergedCap =
                    mergeSharedBatchLimit(modelName_, actualExecutionProvider_, observedCap);
                if (mergedCap > 0) {
                    providerBatchLimit_.store(mergedCap, std::memory_order_relaxed);
                    const auto currentCap = TuneAdvisor::getEmbedDocCap();
                    if (currentCap == 0 || currentCap > mergedCap) {
                        TuneAdvisor::setEmbedDocCap(mergedCap);
                    }
                }
                spdlog::warn(
                    "[ONNX] Provider returned smaller batch than requested: requested_B={} "
                    "output_B={} (learned_cap={}), processing partial result + chunking remainder",
                    B, observedCap, mergedCap);

                std::vector<std::vector<float>> partial(observedCap,
                                                        std::vector<float>(hidden_dim, 0.0f));
                const float* data = out.GetTensorData<float>();
                for (size_t b = 0; b < observedCap; ++b) {
                    auto& emb = partial[b];
                    if (pooling_mode_ == Pooling::CLS) {
                        const size_t off = b * SS * DD;
                        for (size_t d = 0; d < DD; ++d)
                            emb[d] = data[off + d];
                    } else if (pooling_mode_ == Pooling::MAX) {
                        for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
                            if (masks[b][i] <= 0)
                                continue;
                            const size_t off = b * SS * DD + i * DD;
                            for (size_t d = 0; d < DD; ++d)
                                emb[d] = std::max(emb[d], data[off + d]);
                        }
                    } else {
                        float total = 0.0f;
                        for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
                            float w = static_cast<float>(masks[b][i]);
                            if (w <= 0.0f) {
                                continue;
                            }
                            total += w;
                            const size_t off = b * SS * DD + i * DD;
                            for (size_t d = 0; d < DD; ++d)
                                emb[d] += data[off + d] * w;
                        }
                        if (total > 0.0f) {
                            for (auto& v : emb)
                                v = static_cast<float>(v / total);
                        }
                    }
                    if (normalize_) {
                        double norm = 0.0;
                        for (auto v : emb) {
                            const double dv = static_cast<double>(v);
                            norm += dv * dv;
                        }
                        norm = std::sqrt(norm);
                        if (norm > 1e-8) {
                            for (auto& v : emb)
                                v = static_cast<float>(v / norm);
                        } else {
                            std::fill(emb.begin(), emb.end(), 0.0f);
                        }
                    }
                }

                std::vector<std::string> remainder;
                remainder.reserve(B - observedCap);
                for (size_t idx = observedCap; idx < B; ++idx) {
                    remainder.push_back(texts[idx]);
                }
                if (!remainder.empty()) {
                    auto rem = runOnnxBatch(remainder);
                    if (!rem) {
                        return rem.error();
                    }
                    auto& remVals = rem.value();
                    partial.insert(partial.end(), std::make_move_iterator(remVals.begin()),
                                   std::make_move_iterator(remVals.end()));
                }
                return partial;
            }

            if (hidden_dim > 0 && d2 > 0 && DD != hidden_dim) {
                std::ostringstream oss;
                oss << "Output shape mismatch: expected [B,S,D] with B=" << B;
                if (hidden_dim > 0) {
                    oss << " D=" << hidden_dim;
                }
                oss << ", got [" << b0 << "," << s1 << "," << d2 << "]";
                return Error{ErrorCode::InvalidData, oss.str()};
            }
            if (b0 > 0 && static_cast<size_t>(b0) > B) {
                spdlog::debug("[ONNX] Provider returned padded batch output: requested_B={} "
                              "output_B={} (using first B)",
                              B, b0);
            }
            const float* data = out.GetTensorData<float>();
            result.resize(B, std::vector<float>(hidden_dim, 0.0f));
            for (size_t b = 0; b < B; ++b) {
                auto& emb = result[b];
                if (pooling_mode_ == Pooling::CLS) {
                    const size_t off = b * SS * DD;
                    for (size_t d = 0; d < DD; ++d)
                        emb[d] = data[off + d];
                } else if (pooling_mode_ == Pooling::MAX) {
                    for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
                        if (masks[b][i] <= 0)
                            continue;
                        const size_t off = b * SS * DD + i * DD;
                        for (size_t d = 0; d < DD; ++d)
                            emb[d] = std::max(emb[d], data[off + d]);
                    }
                } else {
                    float total = 0.0f;
                    for (size_t i = 0; i < std::min(SS, masks[b].size()); ++i) {
                        float w = static_cast<float>(masks[b][i]);
                        if (w <= 0.0f) {
                            continue;
                        }
                        total += w;
                        const size_t off = b * SS * DD + i * DD;
                        for (size_t d = 0; d < DD; ++d)
                            emb[d] += data[off + d] * w;
                    }
                    if (total > 0.0f) {
                        for (auto& v : emb)
                            v = static_cast<float>(v / total);
                    }
                }
                if (normalize_) {
                    double norm = 0.0;
                    for (auto v : emb) {
                        double dv = static_cast<double>(v);
                        norm += dv * dv;
                    }
                    norm = std::sqrt(norm);
                    if (norm > 1e-8) {
                        for (auto& v : emb)
                            v = static_cast<float>(v / norm);
                    } else {
                        std::fill(emb.begin(), emb.end(), 0.0f);
                    }
                }
            }
            return result;
        }
        if (shape.size() == 2) {
            // [B, D]
            const int64_t b0 = shape[0];
            const int64_t d1 = shape[1];
            const size_t DD = (d1 > 0) ? static_cast<size_t>(d1) : 0;
            if (hidden_dim == 0 && DD > 0) {
                hidden_dim = DD;
                embeddingDim_ = hidden_dim;
            }

            if (b0 > 0 && static_cast<size_t>(b0) < B && (hidden_dim == 0 || DD == hidden_dim)) {
                const size_t observedCap = static_cast<size_t>(b0);
                const size_t mergedCap =
                    mergeSharedBatchLimit(modelName_, actualExecutionProvider_, observedCap);
                if (mergedCap > 0) {
                    providerBatchLimit_.store(mergedCap, std::memory_order_relaxed);
                    const auto currentCap = TuneAdvisor::getEmbedDocCap();
                    if (currentCap == 0 || currentCap > mergedCap) {
                        TuneAdvisor::setEmbedDocCap(mergedCap);
                    }
                }
                spdlog::warn(
                    "[ONNX] Provider returned smaller batch than requested: requested_B={} "
                    "output_B={} (learned_cap={}), processing partial result + chunking remainder",
                    B, observedCap, mergedCap);

                std::vector<std::vector<float>> partial;
                partial.reserve(B);
                const float* data = out.GetTensorData<float>();
                for (size_t b = 0; b < observedCap; ++b) {
                    partial.emplace_back(data + b * DD, data + (b + 1) * DD);
                    if (normalize_) {
                        double norm = 0.0;
                        for (auto v : partial.back()) {
                            const double dv = static_cast<double>(v);
                            norm += dv * dv;
                        }
                        norm = std::sqrt(norm);
                        if (norm > 1e-8) {
                            for (auto& v : partial.back())
                                v = static_cast<float>(v / norm);
                        } else {
                            std::fill(partial.back().begin(), partial.back().end(), 0.0f);
                        }
                    }
                }

                std::vector<std::string> remainder;
                remainder.reserve(B - observedCap);
                for (size_t idx = observedCap; idx < B; ++idx) {
                    remainder.push_back(texts[idx]);
                }
                if (!remainder.empty()) {
                    auto rem = runOnnxBatch(remainder);
                    if (!rem) {
                        return rem.error();
                    }
                    auto& remVals = rem.value();
                    partial.insert(partial.end(), std::make_move_iterator(remVals.begin()),
                                   std::make_move_iterator(remVals.end()));
                }
                return partial;
            }

            if (hidden_dim > 0 && d1 > 0 && DD != hidden_dim) {
                std::ostringstream oss;
                oss << "Output shape mismatch: expected [B,D] with B=" << B;
                if (hidden_dim > 0) {
                    oss << " D=" << hidden_dim;
                }
                oss << ", got [" << b0 << "," << d1 << "]";
                return Error{ErrorCode::InvalidData, oss.str()};
            }
            if (b0 > 0 && static_cast<size_t>(b0) > B) {
                spdlog::debug("[ONNX] Provider returned padded batch output: requested_B={} "
                              "output_B={} (using first B)",
                              B, b0);
            }
            const float* data = out.GetTensorData<float>();
            result.resize(B);
            for (size_t b = 0; b < B; ++b) {
                result[b].assign(data + b * DD, data + (b + 1) * DD);
                if (normalize_) {
                    double norm = 0.0;
                    for (auto v : result[b]) {
                        double dv = static_cast<double>(v);
                        norm += dv * dv;
                    }
                    norm = std::sqrt(norm);
                    if (norm > 1e-8) {
                        for (auto& v : result[b])
                            v = static_cast<float>(v / norm);
                    } else {
                        std::fill(result[b].begin(), result[b].end(), 0.0f);
                    }
                }
            }
            return result;
        }
        return Error{ErrorCode::InvalidData, "Unexpected output tensor rank (batch)"};
    }

    std::string modelPath_;
    std::string modelName_;
    OnnxTextConfig config_;
    OnnxTextPreprocessor preprocessor_;
    yams::vector::HuggingFaceTokenizer tokenizer_;

    Ort::Env* env_ = nullptr; // shared global env
    std::unique_ptr<Ort::SessionOptions> sessionOptions_;
    std::unique_ptr<Ort::Session> session_;

    std::string actualExecutionProvider_{"cpu"};

    std::vector<std::string> inputNames_;
    std::vector<std::string> outputNames_;

    size_t embeddingDim_ = 0;
    size_t maxSequenceLength_ = 512;
    bool hasFixedInputShape_ = false; // true when ONNX model input dim[1] > 0
    bool isLoaded_ = false;
    bool test_mode_ = false;
    bool dynamicPaddingEnabled_ = true; // controlled by YAMS_ONNX_DYNAMIC_PADDING
    std::atomic<size_t> providerBatchLimit_{0};
    int configuredIntraThreads_{1};
    int configuredInterThreads_{1};

    enum class Pooling { MEAN, CLS, MAX };
    Pooling pooling_mode_ = Pooling::MEAN;
    bool normalize_ = true;

    // Cached MemoryInfo to avoid re-creating per inference call.
    std::unique_ptr<Ort::MemoryInfo> memoryInfo_;

    // I/O Binding for optimized inference dispatch (pre-resolved name lookups,
    // better GPU memory placement). Falls back to Session::Run() if unavailable.
    std::unique_ptr<Ort::IoBinding> ioBinding_;
    bool useIoBinding_ = false;
    mutable std::recursive_mutex stateMutex_;

    void parseModelConfigHints() {
        try {
            namespace fs = std::filesystem;
            fs::path mp(modelPath_);
            std::vector<fs::path> candidates;
            if (mp.has_parent_path()) {
                candidates.push_back(mp.parent_path() / "sentence_bert_config.json");
                candidates.push_back(mp.parent_path() / "config.json");
                if (mp.parent_path().has_parent_path()) {
                    auto up = mp.parent_path().parent_path();
                    candidates.push_back(up / "sentence_bert_config.json");
                    candidates.push_back(up / "config.json");
                }
            }
            for (const auto& c : candidates) {
                if (!fs::exists(c))
                    continue;
                std::ifstream in(c);
                if (!in.good())
                    continue;
                nlohmann::json j;
                in >> j;
                if (j.contains("max_seq_length") && j["max_seq_length"].is_number_integer()) {
                    maxSequenceLength_ = static_cast<size_t>(j["max_seq_length"].get<int>());
                } else if (j.contains("model_max_length") &&
                           j["model_max_length"].is_number_integer()) {
                    maxSequenceLength_ = static_cast<size_t>(j["model_max_length"].get<int>());
                }
                // Common HF key for context length
                if (j.contains("max_position_embeddings") &&
                    j["max_position_embeddings"].is_number_integer()) {
                    maxSequenceLength_ =
                        static_cast<size_t>(j["max_position_embeddings"].get<int>());
                }
                if (j.contains("pooling_mode_cls_token") &&
                    j["pooling_mode_cls_token"].is_boolean() &&
                    j["pooling_mode_cls_token"].get<bool>()) {
                    pooling_mode_ = Pooling::CLS;
                }
                if (j.contains("pooling_mode_max_tokens") &&
                    j["pooling_mode_max_tokens"].is_boolean() &&
                    j["pooling_mode_max_tokens"].get<bool>()) {
                    pooling_mode_ = Pooling::MAX;
                }
                if (j.contains("pooling_mode_mean_tokens") &&
                    j["pooling_mode_mean_tokens"].is_boolean()) {
                    if (j["pooling_mode_mean_tokens"].get<bool>())
                        pooling_mode_ = Pooling::MEAN;
                }
                if (j.contains("normalize_embeddings") && j["normalize_embeddings"].is_boolean()) {
                    normalize_ = j["normalize_embeddings"].get<bool>();
                }
                // Only use config hints for embedding dimension when we could not infer it from
                // the ONNX graph output shape. The ONNX output tensor is authoritative.
                if (embeddingDim_ == 0) {
                    if (j.contains("embedding_dimension") &&
                        j["embedding_dimension"].is_number_integer()) {
                        embeddingDim_ = static_cast<size_t>(j["embedding_dimension"].get<int>());
                    }
                    // Hidden size often present in HF config (standard models use hidden_size,
                    // Nomic models use n_embd)
                    if (j.contains("hidden_size") && j["hidden_size"].is_number_integer()) {
                        embeddingDim_ = static_cast<size_t>(j["hidden_size"].get<int>());
                    } else if (j.contains("n_embd") && j["n_embd"].is_number_integer()) {
                        embeddingDim_ = static_cast<size_t>(j["n_embd"].get<int>());
                    }
                }
                // Nomic-specific hints via architectures
                try {
                    if (j.contains("architectures") && j["architectures"].is_array()) {
                        for (const auto& a : j["architectures"]) {
                            if (a.is_string()) {
                                auto s = a.get<std::string>();
                                if (s.find("Nomic") != std::string::npos) {
                                    pooling_mode_ = Pooling::MEAN;
                                    normalize_ = true;
                                    // If very long context supported and not set, prefer 8192
                                    if (maxSequenceLength_ == 0 || maxSequenceLength_ < 2048) {
                                        maxSequenceLength_ = 8192;
                                    }
                                    break;
                                }
                                // EmbeddingGemma: mean pool + L2 normalize, 2048 context
                                if (s.find("Gemma") != std::string::npos) {
                                    pooling_mode_ = Pooling::MEAN;
                                    normalize_ = true;
                                    if (maxSequenceLength_ == 0) {
                                        maxSequenceLength_ = 2048;
                                    }
                                    spdlog::info("[ONNX] Detected Gemma architecture: "
                                                 "mean pooling, L2 normalize, max_seq={}",
                                                 maxSequenceLength_);
                                    break;
                                }
                            }
                        }
                    }
                } catch (...) {
                }
                break; // stop after first readable config
            }
        } catch (...) {
        }
    }
};

// ============================================================================
// OnnxModelSession Public Interface
// ============================================================================

OnnxModelSession::OnnxModelSession(const std::string& modelPath, const std::string& modelName,
                                   const OnnxTextConfig& config)
    : pImpl(std::make_unique<Impl>(modelPath, modelName, config)) {
    info_.name = modelName;
    info_.path = modelPath;
    info_.loadTime = std::chrono::system_clock::now();

    // Get file size for memory estimation
    if (fs::exists(modelPath)) {
        info_.memoryUsageBytes = fs::file_size(modelPath);
    }

    // Lazy loading: model will be loaded on first inference request
    // This prevents slow session creation during pool initialization
    spdlog::debug("[ONNX] Created session for '{}' - lazy loading (load on first inference)",
                  modelName);

    // Dimension info will be populated on first inference or explicit warm-up
    info_.embeddingDim = 0;
    info_.maxSequenceLength = 0;
}

OnnxModelSession::~OnnxModelSession() = default;

Result<std::vector<float>> OnnxModelSession::generateEmbedding(const std::string& text) {
    info_.requestCount++;

    auto result = pImpl->generateEmbedding(text);
    if (!result) {
        info_.errorCount++;
        return result;
    }

    // Lazy load dimension info on first successful inference
    if (info_.embeddingDim == 0 && pImpl->isValid()) {
        info_.embeddingDim = pImpl->getEmbeddingDim();
        info_.maxSequenceLength = pImpl->getMaxSequenceLength();
        spdlog::debug("[ONNX] Lazy loaded dimensions for '{}': dim={}, max_seq_len={}", info_.name,
                      info_.embeddingDim, info_.maxSequenceLength);
    }

    return result;
}

Result<std::vector<std::vector<float>>>
OnnxModelSession::generateBatchEmbeddings(const std::vector<std::string>& texts) {
    info_.requestCount++;

    auto result = pImpl->generateBatchEmbeddings(texts);

    if (!result) {
        info_.errorCount++;
        return result;
    }

    // Lazy load dimension info on first successful inference
    if (info_.embeddingDim == 0 && pImpl->isValid()) {
        info_.embeddingDim = pImpl->getEmbeddingDim();
        info_.maxSequenceLength = pImpl->getMaxSequenceLength();
        spdlog::debug("[ONNX] Lazy loaded dimensions for '{}' (batch): dim={}, max_seq_len={}",
                      info_.name, info_.embeddingDim, info_.maxSequenceLength);
    }

    return result;
}

bool OnnxModelSession::isValid() const {
    return pImpl->isValid();
}

std::string OnnxModelSession::getExecutionProvider() const {
    return pImpl->getExecutionProvider();
}

size_t OnnxModelSession::getLearnedBatchLimit() const {
    return pImpl->getLearnedBatchLimit();
}

Result<void> OnnxModelSession::setThreading(int intraThreads, int interThreads) {
    return pImpl->setThreading(intraThreads, interThreads);
}

std::pair<int, int> OnnxModelSession::getThreading() const {
    return pImpl->getThreading();
}

// ============================================================================
// OnnxModelPool Implementation
// ============================================================================

OnnxModelPool::OnnxModelPool(const ModelPoolConfig& config) : config_(config) {
#if defined(YAMS_TESTING) || defined(YAMS_TEST_PROVIDER_CONTROLS)
    // In mock mode, avoid initializing any ONNX Runtime state to prevent
    // platform-specific crashes on hosts without compatible runtimes.
    if (std::getenv("YAMS_USE_MOCK_PROVIDER") || std::getenv("YAMS_SKIP_MODEL_LOADING") ||
        std::getenv("YAMS_TEST_MODE")) {
        spdlog::warn("[ONNX] Mock provider mode enabled; skipping global ONNX environment init");
        return;
    }
#endif

    if (config_.asyncLoading) {
        size_t threads = config_.loadWorkerThreads > 0 ? config_.loadWorkerThreads : 1;
        loadPool_ = std::make_unique<boost::asio::thread_pool>(threads);
        spdlog::info("[ONNX] Async model loading enabled (threads={})", threads);
    }

    // NOTE: We do NOT create a separate Ort::Env here.
    // Instead, OnnxModelSession::Impl uses the global get_global_ort_env() function.
    // Creating multiple Ort::Env instances causes thread pool conflicts that lead to
    // "resource deadlock would occur" errors (EDEADLK, code 36).
    // The ortEnv_ and sessionOptions_ members are kept for API compatibility but not used.

    try {
        // Log that we're using the shared global environment
        // NOTE: We do NOT create a separate Ort::Env here - OnnxModelSession::Impl
        // uses the global get_global_ort_env() function to avoid thread pool conflicts.
        spdlog::info(
            "[ONNX] Pool: Using shared global Ort::Env (sessions use single-threaded execution)");

        runtimeGpuEnabled_ = false;
        if (config.enableGPU) {
            const auto& gpu = yams::daemon::resource::detectGpu();
            if (gpu.detected) {
                runtimeGpuEnabled_ = true;
                spdlog::info("[ONNX] GPU enabled: {} ({:.1f} GB VRAM, provider={})", gpu.name,
                             static_cast<double>(gpu.vramBytes) / (1024.0 * 1024.0 * 1024.0),
                             gpu.provider);
            } else {
                spdlog::info("[ONNX] GPU requested but no GPU detected (will use CPU)");
            }
        }
    } catch (const std::exception& e) {
        // Do not crash the daemon if runtime init fails; log and continue.
        spdlog::warn("[ONNX] Runtime environment init failed: {} — continuing (mock may be used)",
                     e.what());
    } catch (...) {
        spdlog::warn("[ONNX] Runtime environment init failed with unknown error — continuing");
    }
}

OnnxModelPool::~OnnxModelPool() {
    try {
        shutdown();
    } catch (const std::exception& e) {
        try {
            spdlog::warn("[OnnxModelPool] destructor exception: {}", e.what());
        } catch (...) {
        }
    } catch (...) {
        try {
            spdlog::warn("[OnnxModelPool] destructor unknown exception");
        } catch (...) {
        }
    }
}

Result<void> OnnxModelPool::initialize() {
    std::lock_guard<std::recursive_mutex> lock(mutex_);

    if (initialized_.load(std::memory_order_acquire)) {
        return Result<void>();
    }

    spdlog::info("Initializing ONNX model pool with max {} models", config_.maxLoadedModels);

    const auto& runtimeInfo = yams::onnx_util::OrtRuntimeLoader::instance().ensureLoaded();
    if (!runtimeInfo.available) {
        return Error{ErrorCode::InternalError,
                     runtimeInfo.errorMessage.empty()
                         ? std::string("ONNX Runtime unavailable")
                         : "ONNX Runtime unavailable: " + runtimeInfo.errorMessage};
    }

    // Eagerly initialize the global Ort::Env NOW, before any concurrent access.
    // This prevents "resource deadlock would occur" errors that can happen when
    // multiple threads try to initialize ONNX Runtime simultaneously.
    try {
        spdlog::info("[ONNX Pool] Eagerly initializing global Ort::Env...");
        (void)get_global_ort_env();
        spdlog::info("[ONNX Pool] Global Ort::Env ready");
    } catch (const std::exception& e) {
        spdlog::error("[ONNX Pool] Failed to initialize Ort::Env: {}", e.what());
        return Error{ErrorCode::InternalError, std::string("ONNX Env init failed: ") + e.what()};
    }

    // Start background preloading if configured (non-blocking)
    if (!config_.lazyLoading && !config_.preloadModels.empty()) {
        // Launch preloading in background thread to avoid blocking
        // Store the thread so we can join it during shutdown
        preloadThread_ = std::thread([this]() {
            try {
                spdlog::info("Starting background model preloading");
                // Check shutdown flag periodically during preload
                for (const auto& modelName : config_.preloadModels) {
                    if (shutdown_.load(std::memory_order_acquire)) {
                        spdlog::info("Background preloading interrupted by shutdown");
                        return;
                    }
                    auto result = loadModel(modelName);
                    if (!result) {
                        spdlog::warn("Failed to preload model {}: {}", modelName,
                                     result.error().message);
                    } else {
                        // Warm up the model to trigger lazy loading
                        // This ensures the first real inference request is fast
                        spdlog::info("Warming up preloaded model: {}", modelName);
                        auto warmupResult = warmupModel(modelName);
                        if (!warmupResult) {
                            spdlog::warn("Failed to warm up model {}: {}", modelName,
                                         warmupResult.error().message);
                        }
                    }
                }
                spdlog::info("Background model preloading completed");
            } catch (const std::exception& e) {
                try {
                    spdlog::error("[ONNX] Background preload crashed: {}", e.what());
                } catch (...) {
                }
            } catch (...) {
                try {
                    spdlog::error("[ONNX] Background preload crashed with unknown exception");
                } catch (...) {
                }
            }
        });
    }

    initialized_.store(true, std::memory_order_release);
    return Result<void>();
}

void OnnxModelPool::shutdown() {
    if (shutdown_.exchange(true, std::memory_order_acq_rel)) {
        return;
    }

    if (preloadThread_.joinable()) {
        try {
            spdlog::info("Waiting for background preload thread to finish...");
        } catch (...) {
        }
        try {
            preloadThread_.join();
            try {
                spdlog::info("Background preload thread joined");
            } catch (...) {
            }
        } catch (const std::exception& e) {
            try {
                spdlog::warn("Failed to join preload thread: {}", e.what());
            } catch (...) {
            }
        } catch (...) {
        }
    }

    if (loadPool_) {
        loadPool_->join();
        loadPool_.reset();
    }

    try {
        std::lock_guard<std::recursive_mutex> lock(mutex_);

        try {
            spdlog::info("Shutting down ONNX model pool");
        } catch (...) {
        }

        loadingFutures_.clear();

        for (auto& [name, entry] : models_) {
            if (entry.pool) {
                try {
                    entry.pool->shutdown();
                } catch (const std::exception& e) {
                    try {
                        spdlog::warn("Failed to shutdown model pool '{}': {}", name, e.what());
                    } catch (...) {
                    }
                } catch (...) {
                }
            }
        }

        models_.clear();
        initialized_.store(false, std::memory_order_release);
    } catch (const std::exception& e) {
        try {
            spdlog::warn("ONNX model pool shutdown exception: {}", e.what());
        } catch (...) {
        }
    } catch (...) {
        try {
            spdlog::warn("ONNX model pool shutdown unknown exception");
        } catch (...) {
        }
    }
}

Result<OnnxModelPool::ModelHandle> OnnxModelPool::acquireModel(const std::string& modelName,
                                                               std::chrono::milliseconds timeout) {
    if (!initialized_.load(std::memory_order_acquire)) {
        if (auto result = initialize(); !result) {
            return result.error();
        }
    }

    totalRequests_++;

    // Check if model is already loaded
    // IMPORTANT: Get pool pointer under lock, then release lock before calling acquire()
    // to avoid deadlock with ResourcePool's internal locking (PBI-1c1 fix)
    std::shared_ptr<ResourcePool<OnnxModelSession>> poolPtr;
    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);

        auto it = models_.find(modelName);
        if (it != models_.end() && it->second.pool) {
            updateAccessStats(modelName);
            cacheHits_++;
            poolPtr = it->second.pool;
        }
    }

    // Acquire outside of lock to avoid deadlock
    if (poolPtr) {
        return poolPtr->acquire(timeout);
    }

    cacheMisses_++;

    // Model not loaded, need to load it
    if (config_.asyncLoading) {
        if (auto result = loadModel(modelName, /*wait=*/false); !result) {
            return result.error();
        }
        {
            std::lock_guard<std::recursive_mutex> lock(mutex_);
            auto it = loadingFutures_.find(modelName);
            if (it != loadingFutures_.end()) {
                if (it->second.wait_for(std::chrono::seconds(0)) != std::future_status::ready) {
                    return Error{ErrorCode::OperationInProgress, "Model loading in background"};
                }
                const auto loadOk = static_cast<bool>(it->second.get());
                std::optional<Error> loadError;
                if (!loadOk) {
                    loadError = it->second.get().error();
                }
                loadingFutures_.erase(it);
                if (loadError) {
                    return std::move(*loadError);
                }
            }
        }
    } else {
        if (auto result = loadModel(modelName); !result) {
            return result.error();
        }
    }

    // Try again after loading - get pool pointer while holding lock
    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);
        auto it = models_.find(modelName);
        if (it != models_.end() && it->second.pool) {
            updateAccessStats(modelName);
            poolPtr = it->second.pool;
        }
    }

    // Acquire outside of lock
    if (poolPtr) {
        return poolPtr->acquire(timeout);
    }

    return Error{ErrorCode::NotFound, "Failed to load model: " + modelName};
}

Result<void> OnnxModelPool::loadModel(const std::string& modelName, bool wait) {
    if (!config_.asyncLoading || wait) {
        return loadModelSync(modelName);
    }
    return scheduleLoad(modelName);
}

Result<void> OnnxModelPool::scheduleLoad(const std::string& modelName) {
    if (!loadPool_) {
        return loadModelSync(modelName);
    }

    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);
        auto it = models_.find(modelName);
        if (it != models_.end() && it->second.pool) {
            return Result<void>();
        }
        if (loadingFutures_.count(modelName) > 0) {
            return Result<void>();
        }

        auto prom = std::make_shared<std::promise<Result<void>>>();
        auto fut = prom->get_future().share();
        loadingFutures_[modelName] = fut;

        boost::asio::post(*loadPool_, [this, modelName, prom]() {
            auto res = loadModelSync(modelName);
            prom->set_value(res);
        });
    }

    return Result<void>();
}

bool OnnxModelPool::isModelLoaded(const std::string& modelName) const {
    std::lock_guard<std::recursive_mutex> lock(mutex_);
    auto it = models_.find(modelName);
    return it != models_.end() && it->second.pool != nullptr;
}

std::vector<std::string> OnnxModelPool::getLoadedModels() const {
    std::lock_guard<std::recursive_mutex> lock(mutex_);

    std::vector<std::string> loaded;
    for (const auto& [name, entry] : models_) {
        if (entry.pool) {
            loaded.push_back(name);
        }
    }

    return loaded;
}

Result<void> OnnxModelPool::loadModelSync(const std::string& modelName) {
    spdlog::info("[ONNX Plugin] loadModel() called for: {}", modelName);

    // In test environments, avoid heavy model loading entirely
    if (yams::test::shouldSkipModelLoading()) {
        return Error{ErrorCode::NotFound, "Model loading skipped in test mode"};
    }

    // Single-flight pattern: only one thread loads a given model at a time.
    // Other threads wait for the loading thread to finish.
    {
        std::unique_lock<std::recursive_mutex> lock(mutex_);

        // Check if already loaded
        auto it = models_.find(modelName);
        if (it != models_.end() && it->second.pool) {
            spdlog::info("[ONNX Plugin] Model already loaded: {}", modelName);
            return Result<void>();
        }

        // Check if another thread is currently loading this model
        if (loadingModels_.count(modelName) > 0) {
            spdlog::info("[ONNX Plugin] Model '{}' is being loaded by another thread, waiting...",
                         modelName);
            // Wait for the other thread to finish loading
            loadingCv_.wait(lock, [this, &modelName]() {
                // Wake up when either: model is loaded, or no longer in loading set
                auto it = models_.find(modelName);
                bool loaded = (it != models_.end() && it->second.pool);
                bool notLoading = (loadingModels_.count(modelName) == 0);
                return loaded || notLoading;
            });

            // Check if model was loaded successfully by the other thread
            it = models_.find(modelName);
            if (it != models_.end() && it->second.pool) {
                spdlog::info("[ONNX Plugin] Model '{}' loaded by another thread", modelName);
                return Result<void>();
            }
            // If not loaded and not loading, we'll try to load it ourselves
            spdlog::warn("[ONNX Plugin] Model '{}' load failed by other thread, retrying",
                         modelName);
        }

        // Mark this model as being loaded by us
        loadingModels_.insert(modelName);
    }

    // RAII guard to remove from loadingModels_ and notify waiters
    struct LoadingGuard {
        OnnxModelPool* pool;
        std::string name;
        ~LoadingGuard() {
            std::lock_guard<std::recursive_mutex> lock(pool->mutex_);
            pool->loadingModels_.erase(name);
            pool->loadingCv_.notify_all();
        }
    } guard{this, modelName};

    // Resolve model path WITHOUT holding the lock
    std::string modelPath = resolveModelPath(modelName);
    spdlog::info("[ONNX Plugin] Resolved path: {}", modelPath);

    // Direct filesystem check
    spdlog::info("[ONNX Plugin] Checking file existence: {}", modelPath);
    bool fileExists = false;
    try {
        fileExists = fs::exists(modelPath) && fs::is_regular_file(modelPath);
    } catch (const std::exception& e) {
        spdlog::warn("Error checking model file existence: {}", e.what());
    }

    if (!fileExists) {
        spdlog::info("Model file not found: {} (searched for: {})", modelName, modelPath);
        const bool offline = []() {
            if (const char* s = std::getenv("YAMS_ONNX_OFFLINE")) {
                std::string v(s);
                std::transform(v.begin(), v.end(), v.begin(), ::tolower);
                return v == "1" || v == "true" || v == "yes" || v == "on";
            }
            return false;
        }();
        if (offline) {
            return Error{ErrorCode::NotFound,
                         "Model file not found in offline mode. Pre-place artifacts in cache or "
                         "~/.yams/models, or unset YAMS_ONNX_OFFLINE"};
        }
        spdlog::info("Download the model with: yams model --download {}", modelName);
        return Error{ErrorCode::NotFound,
                     "Model file not found: " + modelPath +
                         ". Download with: yams model --download " + modelName +
                         " or set YAMS_ALLOW_MODEL_DOWNLOAD=1 for automatic downloads"};
    }

    // Check file size
    size_t modelSize = 0;
    try {
        modelSize = fs::file_size(modelPath);
    } catch (const std::exception& e) {
        spdlog::warn("Error getting model file size: {}", e.what());
        modelSize = 100 * 1024 * 1024; // Assume 100MB if can't determine
    }

    // Configure pool for this model (no lock needed - local variables)
    // Use TuneAdvisor for GPU-aware pool sizing:
    //   GPU mode: larger pool (GPU handles inference, more sessions for throughput)
    //   CPU mode: smaller pool (avoid CPU saturation from parallel inference)
    const bool gpuEnabled = runtimeGpuEnabled_ && !OnnxModelSession::isCoreMLDenylisted(modelName);
    const size_t embedConcurrencyCap =
        std::max<size_t>(1u, std::min<size_t>(TuneAdvisor::postEmbedConcurrent(),
                                              TuneAdvisor::getEmbedMaxConcurrency()));
    const size_t dynamic_max = std::max<size_t>(
        1u, std::min<size_t>(TuneAdvisor::onnxSessionsPerModel(gpuEnabled), embedConcurrencyCap));
    const size_t dynamic_min = std::min(size_t{2}, dynamic_max);

    spdlog::info("[ONNX Plugin] Configuring resource pool (gpu={}, pool_max={}, embed_cap={})",
                 gpuEnabled ? "enabled" : "disabled", dynamic_max, embedConcurrencyCap);
    PoolConfig<OnnxModelSession> poolConfig;
    poolConfig.minSize = dynamic_min;
    poolConfig.maxSize = dynamic_max;
    // Keep maxIdle = maxSize to prevent session churn. With maxIdle < maxSize,
    // excess sessions would be destroyed on return and recreated on next acquire.
    poolConfig.maxIdle = dynamic_max;
    poolConfig.idleTimeout = config_.modelIdleTimeout;
    // CRITICAL: Disable preCreateResources to avoid creating ONNX session during
    // ResourcePool construction. Sessions will be created lazily on first acquire().
    // This prevents EDEADLK errors from ONNX Runtime's internal thread pool.
    poolConfig.preCreateResources = false;

    // Create embedding config
    OnnxTextConfig embConfig;
    embConfig.model_path = modelPath;
    embConfig.model_name = modelName;
    embConfig.enable_gpu = gpuEnabled;
    embConfig.num_threads = config_.numThreads;
    // Leave max_sequence_length unset so session/model metadata can determine the native limit.
    // The OnnxTextConfig default of 512 would otherwise silently clamp models like Gemma.
    embConfig.max_sequence_length = 0;

    // Create ResourcePool WITHOUT holding mutex_ - the factory lambda will be called
    // later during acquire(), also without holding mutex_.
    spdlog::debug("[ONNX Plugin] Creating ResourcePool for '{}'", modelName);
    auto t0 = std::chrono::steady_clock::now();

    std::shared_ptr<ResourcePool<OnnxModelSession>> pool;
    try {
        spdlog::debug("[ONNX Plugin] About to construct ResourcePool...");
        pool = std::make_shared<ResourcePool<OnnxModelSession>>(
            poolConfig,
            [modelPath, modelName, embConfig](const std::string&) -> Result<ModelSessionPtr> {
                try {
                    spdlog::debug("[ONNX Plugin] Factory: creating OnnxModelSession for '{}'",
                                  modelName);
                    return std::make_shared<OnnxModelSession>(modelPath, modelName, embConfig);
                } catch (const std::exception& e) {
                    spdlog::error("[ONNX Plugin] Factory: failed to create session: {}", e.what());
                    return Error{ErrorCode::InternalError,
                                 std::string("Failed to create model session: ") + e.what()};
                } catch (...) {
                    spdlog::error("[ONNX Plugin] Factory: unknown error creating session");
                    return Error{ErrorCode::InternalError, "Unknown error creating ONNX session"};
                }
            },
            [](const OnnxModelSession& session) { return session.isValid(); });
        spdlog::debug("[ONNX Plugin] ResourcePool constructed successfully");
    } catch (const std::system_error& e) {
        spdlog::error("[ONNX Plugin] ResourcePool construction threw system_error: {} (code={})",
                      e.what(), e.code().value());
        return Error{ErrorCode::InternalError,
                     std::string("ResourcePool construction failed: ") + e.what()};
    } catch (const std::exception& e) {
        spdlog::error("[ONNX Plugin] ResourcePool construction threw exception: {}", e.what());
        return Error{ErrorCode::InternalError,
                     std::string("ResourcePool construction failed: ") + e.what()};
    }

    // Now acquire lock briefly to update shared state
    spdlog::debug("[ONNX Plugin] Acquiring mutex for model registration...");
    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);

        // Check if we need to evict models
        size_t loadedCount = 0;
        for (const auto& [name, entry] : models_) {
            if (entry.pool)
                loadedCount++;
        }

        if (loadedCount >= config_.maxLoadedModels) {
            evictLRU(1);
        }

        // Check memory constraints
        if (!canLoadModel(modelSize)) {
            return Error{ErrorCode::ResourceExhausted, "Insufficient memory to load model"};
        }

        // Create/update model entry
        spdlog::info("[ONNX Plugin] Registering model entry for: {}", modelName);
        ModelEntry& entry = models_[modelName];
        entry.name = modelName;
        entry.path = modelPath;
        entry.lastAccess = std::chrono::steady_clock::now();
        entry.pool = std::move(pool);

        // Update cached model count for non-blocking status queries
        loadedModelCount_.fetch_add(1, std::memory_order_relaxed);
    }

    auto t1 = std::chrono::steady_clock::now();
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
    spdlog::info("Registered ONNX model: {} ({}MB) load_ms={} path={}", modelName,
                 modelSize / (1024 * 1024), ms, modelPath);

    spdlog::info("[ONNX Plugin] Model '{}' registered successfully (session created on first use)",
                 modelName);
    return Result<void>();
}

Result<void> OnnxModelPool::unloadModel(const std::string& modelName) {
    std::lock_guard<std::recursive_mutex> lock(mutex_);

    auto it = models_.find(modelName);
    if (it == models_.end()) {
        return Error{ErrorCode::NotFound, "Model not found: " + modelName};
    }

    if (it->second.isHot) {
        return Error{ErrorCode::InvalidOperation, "Cannot unload hot model: " + modelName};
    }

    if (it->second.pool) {
        auto stats = it->second.pool->getStats();
        if (stats.inUseResources > 0) {
            return Error{ErrorCode::OperationInProgress,
                         "Cannot unload model with in-use sessions: " + modelName};
        }
        it->second.pool->shutdown();
        it->second.pool.reset();

        // Update cached model count for non-blocking status queries
        loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
    }

    spdlog::info("Unloaded ONNX model: {}", modelName);

    return Result<void>();
}

Result<void> OnnxModelPool::preloadModels() {
    for (const auto& modelName : config_.preloadModels) {
        spdlog::info("Preloading model: {}", modelName);

        // Load model directly - no async wrapper needed
        auto result = loadModel(modelName);
        if (!result) {
            spdlog::warn("Failed to preload model {}: {}", modelName, result.error().message);
        } else {
            // Warm up the model to trigger lazy loading
            spdlog::info("Warming up preloaded model: {}", modelName);
            auto warmupResult = warmupModel(modelName);
            if (!warmupResult) {
                spdlog::warn("Failed to warm up model {}: {}", modelName,
                             warmupResult.error().message);
            }

            // Mark as hot model
            std::lock_guard<std::recursive_mutex> lock(mutex_);
            auto it = models_.find(modelName);
            if (it != models_.end()) {
                it->second.isHot = true;
            }
        }
    }

    return Result<void>();
}

Result<void> OnnxModelPool::warmupModel(const std::string& modelName) {
    spdlog::info("[ONNX Plugin] Warming up model: {}", modelName);
    auto t0 = std::chrono::steady_clock::now();

    // Acquire a session from the pool
    auto handleResult = acquireModel(modelName, std::chrono::seconds(30));
    if (!handleResult) {
        spdlog::error("[ONNX Plugin] Failed to acquire session for warmup: {}",
                      handleResult.error().message);
        return handleResult.error();
    }

    auto handle = std::move(handleResult.value());

    // Do a batch inference to trigger lazy loading and warm the batch code path
    std::vector<std::string> warmupTexts = {"warmup text one", "warmup text two",
                                            "warmup text three", "warmup text four"};
    auto embedResult = handle->generateBatchEmbeddings(warmupTexts);
    if (!embedResult) {
        spdlog::error("[ONNX Plugin] Warmup inference failed for {}: {}", modelName,
                      embedResult.error().message);
        return embedResult.error();
    }

    auto t1 = std::chrono::steady_clock::now();
    auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(t1 - t0).count();
    spdlog::info("[ONNX Plugin] Model '{}' warmed up in {}ms (batch={}, embedding dim={})",
                 modelName, ms, warmupTexts.size(),
                 embedResult.value().empty() ? 0 : embedResult.value()[0].size());

    // Auto-tune batch size if not explicitly set and auto-tuning not disabled
    auto autotuneDisabled = []() {
        if (const char* s = std::getenv("YAMS_EMBED_AUTOTUNE")) {
            std::string v(s);
            std::transform(v.begin(), v.end(), v.begin(), ::tolower);
            return v == "0" || v == "false" || v == "no" || v == "off";
        }
        return false;
    };
    if (TuneAdvisor::getEmbedDocCap() == 0 && !autotuneDisabled()) {
        spdlog::info("[ONNX Plugin] Auto-tuning batch size for '{}'...", modelName);

        std::vector<std::string> probeTexts;
        probeTexts.reserve(64);
        for (uint32_t i = 0; i < 64; ++i)
            probeTexts.push_back("probe text number " + std::to_string(i) + " for auto-tuning");

        std::vector<uint32_t> candidates = {4, 8, 16, 32, 64};
        double bestTps = 0.0;
        uint32_t bestBatch = 16;

        for (uint32_t bs : candidates) {
            std::vector<std::string> batch(probeTexts.begin(),
                                           probeTexts.begin() + static_cast<ptrdiff_t>(bs));
            auto pt0 = std::chrono::steady_clock::now();
            auto r = handle->generateBatchEmbeddings(batch);
            auto elapsed = std::chrono::steady_clock::now() - pt0;
            if (!r)
                break;
            double probeMs =
                std::chrono::duration_cast<std::chrono::microseconds>(elapsed).count() / 1000.0;
            double tps = static_cast<double>(bs) / (probeMs / 1000.0);
            spdlog::debug("[ONNX Plugin] Autotune probe: batch={} {:.1f}ms ({:.1f} texts/sec)", bs,
                          probeMs, tps);
            if (tps > bestTps) {
                bestTps = tps;
                bestBatch = bs;
            }
        }

        TuneAdvisor::setEmbedDocCap(bestBatch);
        spdlog::info("[ONNX Plugin] Auto-tuned batch size: {} ({:.1f} texts/sec)", bestBatch,
                     bestTps);
    }

    // Pre-warm additional sessions to avoid cold-start penalty on concurrent requests.
    // GPU mode: up to 3 sessions (GPU provider init is expensive).
    // CPU mode: up to 2 sessions (eliminates ~1788ms cold start for second request).
    // The initial handle is still held; additional sessions are acquired and immediately
    // released back to the pool so they are pre-initialized for future use.
    {
        const std::string activeProvider = handle->getExecutionProvider();
        const bool providerGpuEnabled = isGpuExecutionProviderName(activeProvider);
        const bool cpuFallbackPinned = !providerGpuEnabled && runtimeGpuEnabled_ &&
                                       OnnxModelSession::isCoreMLDenylisted(modelName);

        size_t targetWarm =
            std::min(static_cast<size_t>(TuneAdvisor::onnxSessionsPerModel(providerGpuEnabled)),
                     providerGpuEnabled ? size_t{3} : size_t{2});
        if (cpuFallbackPinned) {
            targetWarm = std::min<size_t>(targetWarm, 1u);
        }

        // Keep warm session count aligned with current embedding throughput budget
        // to avoid retaining idle ONNX sessions under low-concurrency settings.
        const size_t embedConcurrencyCap =
            std::max<size_t>(1u, std::min<size_t>(TuneAdvisor::postEmbedConcurrent(),
                                                  TuneAdvisor::getEmbedMaxConcurrency()));
        targetWarm = std::min(targetWarm, embedConcurrencyCap);

        // Optional explicit cap for test/benchmark scenarios.
        if (const char* s = std::getenv("YAMS_ONNX_PREWARM_SESSIONS")) {
            try {
                size_t forced = static_cast<size_t>(std::stoull(s));
                if (forced >= 1 && forced <= 8) {
                    targetWarm = std::min(targetWarm, forced);
                }
            } catch (...) {
            }
        }

        for (size_t i = 1; i < targetWarm; ++i) {
            auto extra = acquireModel(modelName, std::chrono::seconds(60));
            if (!extra) {
                spdlog::debug("[ONNX Plugin] Pre-warm: could not acquire session {} for {}", i + 1,
                              modelName);
                break;
            }
            // Session returned to pool when extra goes out of scope
        }
        if (targetWarm > 1) {
            spdlog::info("[ONNX Plugin] Pre-warmed {} sessions for '{}' (provider={} gpu={})",
                         targetWarm, modelName, activeProvider, providerGpuEnabled ? "yes" : "no");
        }
    }

    return Result<void>();
}

Result<void> OnnxModelPool::setModelHot(const std::string& modelName, bool hot) {
    std::lock_guard<std::recursive_mutex> lock(mutex_);
    auto it = models_.find(modelName);
    if (it == models_.end()) {
        return Error{ErrorCode::NotFound, "Model not found: " + modelName};
    }
    it->second.isHot = hot;
    spdlog::info("[ONNX Plugin] Model '{}' hot={} ", modelName, hot ? "true" : "false");
    return Result<void>();
}

Result<void> OnnxModelPool::setModelThreading(const std::string& modelName, int intraThreads,
                                              int interThreads, bool applyNow) {
    auto valid = [](int value) { return value == -1 || (value >= 1 && value <= 64); };
    if (modelName.empty() || !valid(intraThreads) || !valid(interThreads)) {
        return Error{ErrorCode::InvalidArgument,
                     "setModelThreading requires model name and valid thread values"};
    }

    std::shared_ptr<ResourcePool<OnnxModelSession>> pool;
    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);
        threadingOverrides_[modelName] = {intraThreads, interThreads};
        auto it = models_.find(modelName);
        if (it != models_.end()) {
            pool = it->second.pool;
        }
    }

    if (!applyNow || !pool) {
        return Result<void>();
    }

    if (pool->getStats().inUseResources > 0) {
        return Error{ErrorCode::ResourceBusy,
                     "Model is currently in use; threading override will apply on next reload"};
    }

    auto handle = pool->acquire(std::chrono::milliseconds(500));
    if (!handle) {
        return Error{ErrorCode::Timeout,
                     "Timed out acquiring model session to apply threading override"};
    }

    return handle.value()->setThreading(intraThreads, interThreads);
}

void OnnxModelPool::evictLRU(size_t numToEvict) {
    if (runtimeGpuEnabled_ && !allowGpuAutoUnload()) {
        spdlog::debug(
            "Skipping LRU model eviction on GPU path (set YAMS_ONNX_GPU_AUTO_UNLOAD=1 to enable)");
        return;
    }

    // Find least recently used models that are not hot
    std::vector<std::pair<std::string, std::chrono::steady_clock::time_point>> candidates;

    for (const auto& [name, entry] : models_) {
        if (entry.pool && !entry.isHot) {
            candidates.emplace_back(name, entry.lastAccess);
        }
    }

    // Sort by last access time
    std::sort(candidates.begin(), candidates.end(),
              [](const auto& a, const auto& b) { return a.second < b.second; });

    // Evict the oldest models
    for (size_t i = 0; i < std::min(numToEvict, candidates.size()); ++i) {
        const auto& modelName = candidates[i].first;
        auto it = models_.find(modelName);
        if (it != models_.end() && it->second.pool) {
            auto stats = it->second.pool->getStats();
            if (stats.inUseResources > 0) {
                spdlog::debug("Skipping LRU eviction for model '{}' (in_use={})", modelName,
                              stats.inUseResources);
                continue;
            }
            spdlog::info("Evicting model due to LRU: {}", modelName);
            it->second.pool->shutdown();
            it->second.pool.reset();
            const size_t count = loadedModelCount_.load(std::memory_order_relaxed);
            if (count > 0) {
                loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
            }
        }
    }
}

OnnxModelPool::PoolStats OnnxModelPool::getStats() const {
    std::lock_guard<std::recursive_mutex> lock(mutex_);

    PoolStats stats;
    stats.totalRequests = totalRequests_.load();
    stats.cacheHits = cacheHits_.load();
    stats.cacheMisses = cacheMisses_.load();
    stats.hitRate =
        stats.totalRequests > 0 ? static_cast<double>(stats.cacheHits) / stats.totalRequests : 0.0;

    stats.loadedModels = 0;
    stats.totalMemoryBytes = 0;

    for (const auto& [name, entry] : models_) {
        if (entry.pool) {
            stats.loadedModels++;

            // Get model info (approximate from first session in pool)
            // Note: This is simplified - in production, track actual memory usage
            if (fs::exists(entry.path)) {
                size_t modelSize = fs::file_size(entry.path);
                stats.totalMemoryBytes += modelSize;
            }

            // Populate per-model stats for status reporting (copyable ModelInfo)
            ModelInfo mi;
            mi.name = name;
            mi.path = entry.path;
            mi.embeddingDim = 384;      // default when unknown
            mi.maxSequenceLength = 512; // default when unknown
            try {
                if (fs::exists(entry.path)) {
                    mi.memoryUsageBytes = fs::file_size(entry.path);
                }
            } catch (...) {
                mi.memoryUsageBytes = 0;
            }
            mi.loadTime = std::chrono::system_clock::now();
            mi.requestCount = 0;
            mi.errorCount = 0;
            stats.modelStats.emplace(name, std::move(mi));
        }
    }

    return stats;
}

size_t OnnxModelPool::getMemoryUsage() const {
    return getStats().totalMemoryBytes;
}

void OnnxModelPool::performMaintenance() {
    if (runtimeGpuEnabled_ && !allowGpuAutoUnload()) {
        return;
    }

    std::lock_guard<std::recursive_mutex> lock(mutex_);

    auto now = std::chrono::steady_clock::now();

    for (auto& [name, entry] : models_) {
        if (!entry.pool) {
            continue;
        }

        auto stats = entry.pool->getStats();
        if (stats.inUseResources > 0) {
            continue;
        }

        if (entry.isHot) {
            // Keep the model loaded, but still evict expired idle sessions from its pool.
            entry.pool->evictExpired();
            continue;
        }

        // Check if model has been idle too long
        auto idleTime = now - entry.lastAccess;
        if (idleTime > config_.modelIdleTimeout) {
            spdlog::info("Unloading idle model: {}", name);
            entry.pool->shutdown();
            entry.pool.reset();
            const size_t count = loadedModelCount_.load(std::memory_order_relaxed);
            if (count > 0) {
                loadedModelCount_.fetch_sub(1, std::memory_order_relaxed);
            }
        } else if (entry.pool) {
            // Clean up expired resources in the pool
            entry.pool->evictExpired();
        }
    }
}

#ifdef YAMS_TESTING
std::optional<OnnxModelPool::TestingModelPoolStats>
OnnxModelPool::testingModelPoolStats(const std::string& modelName) const {
    std::lock_guard<std::recursive_mutex> lock(mutex_);

    auto it = models_.find(modelName);
    if (it == models_.end()) {
        return std::nullopt;
    }

    TestingModelPoolStats out;
    out.loaded = (it->second.pool != nullptr);
    out.isHot = it->second.isHot;

    if (it->second.pool) {
        auto stats = it->second.pool->getStats();
        out.totalResources = stats.totalResources;
        out.availableResources = stats.availableResources;
        out.inUseResources = stats.inUseResources;
    }

    return out;
}
#endif

std::string OnnxModelPool::resolveModelPath(const std::string& modelName) const {
    auto is_hf_like = [](const std::string& s) {
        if (s.rfind("hf://", 0) == 0)
            return true;
        if (s.rfind("https://huggingface.co/", 0) == 0)
            return true;
        return (s.find('/') != std::string::npos) && (s.find('/') == s.rfind('/')) &&
               s.find(".onnx") == std::string::npos && s.find('/') > 0;
    };
    auto to_repo_id = [](const std::string& s) {
        if (s.rfind("hf://", 0) == 0)
            return s.substr(5);
        const std::string prefix = "https://huggingface.co/";
        if (s.rfind(prefix, 0) == 0) {
            auto rest = s.substr(prefix.size());
            auto slash = rest.find('/');
            if (slash != std::string::npos) {
                auto slash2 = rest.find('/', slash + 1);
                if (slash2 != std::string::npos)
                    return rest.substr(0, slash2);
            }
            return rest;
        }
        return s;
    };
    // Get home directory (cross-platform)
    std::string homeDir;
    if (const char* home = std::getenv("HOME")) {
        homeDir = home;
    } else if (const char* userprofile = std::getenv("USERPROFILE")) {
        // Windows: USERPROFILE is the user's home directory
        homeDir = userprofile;
    }

    // Windows-specific paths
    std::string localAppData;
    std::string appData;
#ifdef _WIN32
    if (const char* lad = std::getenv("LOCALAPPDATA")) {
        localAppData = lad;
    }
    if (const char* ad = std::getenv("APPDATA")) {
        appData = ad;
    }
#endif

    // Build search paths with priority:
    // 1) Treat modelName as a full path
    // 2) Configured modelsRoot (if set), with common filename layouts
    // 3) User models (~/.yams/models), local models/, system share models
    std::vector<std::string> searchPaths;
    searchPaths.push_back(modelName);

    if (is_hf_like(modelName)) {
        std::string repo = to_repo_id(modelName);
        // Check for revision hint
        std::string preferredRev;
        {
            std::lock_guard<std::recursive_mutex> lk(mutex_);
            auto it = modelHints_.find(modelName);
            if (it != modelHints_.end())
                preferredRev = it->second.hfRevision;
        }
        auto cache_roots = [&]() {
            std::vector<std::string> roots;
            if (const char* env = std::getenv("YAMS_ONNX_HF_CACHE"))
                roots.push_back(env);
            if (const char* hf = std::getenv("HF_HOME")) {
                roots.push_back(std::string(hf));
                roots.push_back(std::string(hf) + "/hub");
            }
            if (!homeDir.empty()) {
                roots.push_back(homeDir + "/.cache/huggingface");
                roots.push_back(homeDir + "/.cache/huggingface/hub");
            }
            return roots;
        }();
        std::string repo_dir = "models--";
        for (auto c : repo)
            repo_dir += (c == '/' ? std::string("--") : std::string(1, c));
        for (const auto& root : cache_roots) {
            try {
                fs::path base = fs::path(root) / repo_dir;
                if (!fs::exists(base))
                    continue;
                fs::path found;
                if (!preferredRev.empty()) {
                    fs::path snapdir;
                    fs::path ref = base / "refs" / preferredRev;
                    if (fs::exists(ref)) {
                        try {
                            std::ifstream in(ref);
                            std::string s;
                            std::getline(in, s);
                            if (!s.empty())
                                snapdir = base / "snapshots" / s;
                        } catch (...) {
                        }
                    }
                    if (snapdir.empty())
                        snapdir = base / "snapshots" / preferredRev;
                    if (fs::exists(snapdir)) {
                        for (auto& p : fs::recursive_directory_iterator(snapdir)) {
                            if (!p.is_regular_file())
                                continue;
                            if (p.path().filename() == "model.onnx") {
                                found = p.path();
                                break;
                            }
                        }
                    }
                }
                if (found.empty()) {
                    for (auto& p : fs::recursive_directory_iterator(base)) {
                        if (!p.is_regular_file())
                            continue;
                        if (p.path().filename() == "model.onnx") {
                            found = p.path();
                            break;
                        }
                    }
                }
                if (!found.empty()) {
                    spdlog::debug("[ONNX] Resolved HF {} to {}", repo, found.string());
                    return found.string();
                }
            } catch (...) {
            }
        }
        spdlog::debug("[ONNX] HF repo '{}' not found in cache; falling back to filesystem search",
                      repo);
    }

    // Expand configured modelsRoot if present (supports '~')
    if (!config_.modelsRoot.empty()) {
        std::string root = config_.modelsRoot;
        if (!root.empty() && root[0] == '~') {
            if (!homeDir.empty()) {
                root = homeDir + root.substr(1);
            }
        }
        // Common layouts under configured root
        searchPaths.push_back(root + "/" + modelName + "/model.onnx");
        searchPaths.push_back(root + "/" + modelName + "/" + modelName + ".onnx");
        searchPaths.push_back(root + "/" + modelName + ".onnx");
    }

    // Default locations (XDG standard first, then fallbacks)
    // Priority: XDG_DATA_HOME > ~/.local/share/yams > ~/.yams > relative > system
    std::string xdg_data_home;
    if (const char* xdg = std::getenv("XDG_DATA_HOME")) {
        xdg_data_home = xdg;
    } else if (!homeDir.empty()) {
        xdg_data_home = homeDir + "/.local/share";
    }

    // XDG paths (highest priority)
    if (!xdg_data_home.empty()) {
        searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + "/model.onnx");
        searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + "/" + modelName +
                              ".onnx");
        searchPaths.push_back(xdg_data_home + "/yams/models/" + modelName + ".onnx");
    }

    // Legacy ~/.yams/models (for backward compatibility)
    if (!homeDir.empty()) {
        searchPaths.push_back(homeDir + "/.yams/models/" + modelName + "/model.onnx");
        searchPaths.push_back(homeDir + "/.yams/models/" + modelName + "/" + modelName + ".onnx");
        searchPaths.push_back(homeDir + "/.yams/models/" + modelName + ".onnx");
    }

#ifdef _WIN32
    // Windows: Check LOCALAPPDATA/yams/models (primary Windows location)
    if (!localAppData.empty()) {
        searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + "\\model.onnx");
        searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + "\\" + modelName +
                              ".onnx");
        searchPaths.push_back(localAppData + "\\yams\\models\\" + modelName + ".onnx");
    }
    // Windows: Check APPDATA/yams/models (roaming profile)
    if (!appData.empty()) {
        searchPaths.push_back(appData + "\\yams\\models\\" + modelName + "\\model.onnx");
        searchPaths.push_back(appData + "\\yams\\models\\" + modelName + "\\" + modelName +
                              ".onnx");
        searchPaths.push_back(appData + "\\yams\\models\\" + modelName + ".onnx");
    }
#endif

    // Relative paths
    searchPaths.push_back("models/" + modelName + "/model.onnx");
    searchPaths.push_back("models/" + modelName + "/" + modelName + ".onnx");
    searchPaths.push_back("models/" + modelName + ".onnx");

#ifdef _WIN32
    // Windows: System paths
    searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + "\\model.onnx");
    searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + "\\" + modelName +
                          ".onnx");
    searchPaths.push_back("C:\\ProgramData\\yams\\models\\" + modelName + ".onnx");
#else
    // Unix: System paths
    searchPaths.push_back("/usr/local/share/yams/models/" + modelName + "/model.onnx");
    searchPaths.push_back("/usr/local/share/yams/models/" + modelName + "/" + modelName + ".onnx");
    searchPaths.push_back("/usr/local/share/yams/models/" + modelName + ".onnx");
#endif

    // Log search paths on first attempt for debugging
    spdlog::info("[ONNX Plugin] Resolving model path for '{}', searching {} locations", modelName,
                 searchPaths.size());
#ifdef _WIN32
    spdlog::debug("[ONNX Plugin] Windows paths: LOCALAPPDATA={}, APPDATA={}, HOME={}", localAppData,
                  appData, homeDir);
#endif

    // Direct filesystem checks - no async, no timeouts needed for local files
    for (const auto& path : searchPaths) {
        if (path.empty())
            continue;

        try {
            if (fs::exists(path)) {
                spdlog::info("[ONNX Plugin] Found model at: {}", path);
                return path;
            }
        } catch (const std::exception& e) {
            spdlog::debug("[ONNX Plugin] Error checking path {}: {}", path, e.what());
        }
    }

    spdlog::warn("[ONNX Plugin] Model '{}' not found in any search path. Searched:", modelName);
    for (const auto& path : searchPaths) {
        if (!path.empty()) {
            spdlog::warn("[ONNX Plugin]   - {}", path);
        }
    }
    // Default to models directory
    return "models/" + modelName + ".onnx";
}

Result<OnnxModelPool::ModelSessionPtr>
OnnxModelPool::createModelSession(const std::string& modelName) {
    std::string modelPath = resolveModelPath(modelName);

    OnnxTextConfig config;
    config.model_path = modelPath;
    config.model_name = modelName;
    config.enable_gpu = config_.enableGPU;
    config.num_threads = config_.numThreads;
    config.max_sequence_length = 0;

    {
        std::lock_guard<std::recursive_mutex> lock(mutex_);
        auto it = threadingOverrides_.find(modelName);
        if (it != threadingOverrides_.end()) {
            config.num_threads = it->second.first;
            config.inter_op_threads = it->second.second;
        }
    }

    try {
        auto session = std::make_shared<OnnxModelSession>(modelPath, modelName, config);
        return session;
    } catch (const std::exception& e) {
        return Error{ErrorCode::InternalError,
                     std::string("Failed to create model session: ") + e.what()};
    }
}

void OnnxModelPool::updateAccessStats(const std::string& modelName) {
    auto it = models_.find(modelName);
    if (it != models_.end()) {
        it->second.lastAccess = std::chrono::steady_clock::now();
        it->second.accessCount++;
    }
}

bool OnnxModelPool::canLoadModel(size_t estimatedSize) const {
    size_t currentUsage = getMemoryUsage();
    size_t maxBytes = config_.maxMemoryGB * 1024 * 1024 * 1024;

    return (currentUsage + estimatedSize) <= maxBytes;
}

void OnnxModelPool::setResolutionHints(const std::string& modelName, const ResolutionHints& hints) {
    std::lock_guard<std::recursive_mutex> lock(mutex_);
    modelHints_[modelName] = hints;
}

} // namespace yams::daemon