// ============================================================================
// OvLoader.cpp — OpenVINO DLL discovery and pre-loading
//
// Loading order (dependency chain):
//   1. Intel TBB (tbb12.dll / tbbmalloc.dll)  — required by openvino.dll
//   2. openvino.dll                            — core runtime
//   3. openvino_c.dll                          — C API (optional, version query)
//
// Plugins (openvino_intel_cpu_plugin, openvino_intel_gpu_plugin, etc.) and
// frontends (openvino_ir_frontend, openvino_onnx_frontend, etc.) are loaded
// on-demand by OpenVINO's internal plugin system.  They are discovered
// automatically because the shared directory is injected into the DLL
// search path by ANSLibsLoader before OvLoader::Initialize() runs.
// ============================================================================

#include "OvLoader.h"
#include "DynLibUtils.h"

#include <iostream>
#include <mutex>

using namespace ANSCENTER::DynLib;

// ─────────────────────────────────────────────────────────────────────────────
// Static member definitions
// ─────────────────────────────────────────────────────────────────────────────
std::mutex OvLoader::s_mutex;
bool       OvLoader::s_initialized = false;
OvInfo     OvLoader::s_info;
LibHandle  OvLoader::s_hTbb       = nullptr;
LibHandle  OvLoader::s_hTbbMalloc = nullptr;
LibHandle  OvLoader::s_hOvCore    = nullptr;
LibHandle  OvLoader::s_hOvC       = nullptr;

// ─────────────────────────────────────────────────────────────────────────────
// Candidate lists — search ANSCENTER shared dir first, then system PATH
// ─────────────────────────────────────────────────────────────────────────────

std::vector<std::string> OvLoader::TbbCandidates(const std::string& shared_dir)
{
    std::vector<std::string> v;

    // ── Priority 1: ANSCENTER shared directory ────────────────────────────────
#ifdef _WIN32
    v.push_back(JoinPath(shared_dir, "tbb12.dll"));
    v.push_back(JoinPath(shared_dir, "tbb.dll"));
#else
    v.push_back(JoinPath(shared_dir, "libtbb.so.12"));
    v.push_back(JoinPath(shared_dir, "libtbb.so.2"));
    v.push_back(JoinPath(shared_dir, "libtbb.so"));
#endif

    // ── Priority 2: System PATH / LD_LIBRARY_PATH ─────────────────────────────
#ifdef _WIN32
    v.push_back("tbb12.dll");
    v.push_back("tbb.dll");
#else
    v.push_back("libtbb.so.12");
    v.push_back("libtbb.so.2");
    v.push_back("libtbb.so");
#endif

    return v;
}

std::vector<std::string> OvLoader::TbbMallocCandidates(const std::string& shared_dir)
{
    std::vector<std::string> v;

#ifdef _WIN32
    v.push_back(JoinPath(shared_dir, "tbbmalloc.dll"));
    v.push_back("tbbmalloc.dll");
#else
    v.push_back(JoinPath(shared_dir, "libtbbmalloc.so.2"));
    v.push_back(JoinPath(shared_dir, "libtbbmalloc.so"));
    v.push_back("libtbbmalloc.so.2");
    v.push_back("libtbbmalloc.so");
#endif

    return v;
}

std::vector<std::string> OvLoader::OvCoreCandidates(const std::string& shared_dir)
{
    std::vector<std::string> v;

#ifdef _WIN32
    v.push_back(JoinPath(shared_dir, "openvino.dll"));
    v.push_back("openvino.dll");
#else
    v.push_back(JoinPath(shared_dir, "libopenvino.so"));
    v.push_back("libopenvino.so");
#endif

    return v;
}

std::vector<std::string> OvLoader::OvCApiCandidates(const std::string& shared_dir)
{
    std::vector<std::string> v;

#ifdef _WIN32
    v.push_back(JoinPath(shared_dir, "openvino_c.dll"));
    v.push_back("openvino_c.dll");
#else
    v.push_back(JoinPath(shared_dir, "libopenvino_c.so"));
    v.push_back("libopenvino_c.so");
#endif

    return v;
}

// ─────────────────────────────────────────────────────────────────────────────
// Version extraction via OpenVINO C API
//
// The C API exports ov_get_openvino_version() which fills an ov_version_t
// struct with a buildNumber string (e.g. "2024.6.0-17404-4c0f47d2335").
// We replicate the struct layout locally to avoid including OpenVINO headers.
// ─────────────────────────────────────────────────────────────────────────────
std::string OvLoader::TryGetVersion(LibHandle hOvC)
{
    if (!hOvC) return {};

    // Replicate the ov_version_t layout from openvino/c/ov_common.h
    struct ov_version_local {
        const char* buildNumber;
        const char* description;
    };

    using fn_get_version  = int  (*)(ov_version_local*);
    using fn_free_version = void (*)(ov_version_local*);

    fn_get_version  pfnGet  = nullptr;
    fn_free_version pfnFree = nullptr;

    Bind(hOvC, "ov_get_openvino_version", pfnGet);
    Bind(hOvC, "ov_version_free",         pfnFree);

    if (!pfnGet) return {};

    ov_version_local ver = {};
    if (pfnGet(&ver) == 0 && ver.buildNumber) {
        std::string result = ver.buildNumber;
        if (pfnFree) pfnFree(&ver);
        return result;
    }
    return {};
}

// ─────────────────────────────────────────────────────────────────────────────
// Initialize
// ─────────────────────────────────────────────────────────────────────────────
const OvInfo& OvLoader::Initialize(const std::string& shared_dir, bool verbose)
{
    std::lock_guard<std::mutex> lock(s_mutex);
    if (s_initialized) return s_info;

    // Clean up any partially-loaded handles from a previous failed attempt
    if (s_hOvC)       { FreeLib(s_hOvC);       s_hOvC       = nullptr; }
    if (s_hOvCore)    { FreeLib(s_hOvCore);    s_hOvCore    = nullptr; }
    if (s_hTbbMalloc) { FreeLib(s_hTbbMalloc); s_hTbbMalloc = nullptr; }
    if (s_hTbb)       { FreeLib(s_hTbb);       s_hTbb       = nullptr; }

    if (verbose) {
        std::cout << "==================================================" << std::endl;
        std::cout << "[OvLoader] Discovering OpenVINO libraries..."       << std::endl;
        std::cout << "  Shared dir : " << shared_dir                      << std::endl;
        std::cout << "==================================================" << std::endl;
    }

    // ── STEP 1: Inject shared dir into search path ───────────────────────────
    // Ensures plugins/frontends are discoverable by OpenVINO's internal loader.
    InjectDllSearchPath(shared_dir);

    // ── STEP 2: Pre-load Intel TBB (dependency of openvino.dll) ──────────────
    std::string tbbPath;
    s_hTbb = FindAndLoad(TbbCandidates(shared_dir), tbbPath, verbose);
    if (s_hTbb) {
        if (verbose)
            std::cout << "[OvLoader] TBB loaded: " << tbbPath << std::endl;
    } else if (verbose) {
        std::cout << "[OvLoader] WARNING: TBB not found "
                  << "(openvino.dll may fail to load)." << std::endl;
    }

    std::string tbbMallocPath;
    s_hTbbMalloc = FindAndLoad(TbbMallocCandidates(shared_dir),
                               tbbMallocPath, verbose);
    if (!s_hTbbMalloc && verbose) {
        std::cout << "[OvLoader] tbbmalloc not found (optional)." << std::endl;
    }

    // ── STEP 3: Pre-load openvino.dll (core runtime) ─────────────────────────
    std::string corePath;
    s_hOvCore = FindAndLoad(OvCoreCandidates(shared_dir), corePath, verbose);
    if (s_hOvCore) {
        s_info.dllPath = corePath;
        s_info.loaded  = true;

        if (verbose)
            std::cout << "[OvLoader] OpenVINO core loaded: " << corePath
                      << std::endl;
    } else {
        std::cout << "[OvLoader] WARNING: openvino.dll not found in "
                  << shared_dir << " or system PATH." << std::endl;
        std::cout << "[OvLoader] OpenVINO will rely on import-lib resolution "
                  << "at load time." << std::endl;
    }

    // ── STEP 4: Pre-load openvino_c.dll + query version ──────────────────────
    std::string cApiPath;
    s_hOvC = FindAndLoad(OvCApiCandidates(shared_dir), cApiPath, verbose);
    if (s_hOvC) {
        s_info.version = TryGetVersion(s_hOvC);
        if (!s_info.version.empty() && verbose) {
            std::cout << "[OvLoader] OpenVINO version = " << s_info.version
                      << std::endl;
        }
    } else if (verbose) {
        std::cout << "[OvLoader] openvino_c.dll not found "
                  << "(version detection skipped)." << std::endl;
    }

    s_initialized = true;

    if (verbose) {
        std::cout << "[OvLoader] "
                  << (s_info.loaded
                        ? "initialised successfully."
                        : "WARNING — openvino.dll not pre-loaded.")
                  << std::endl;
        std::cout << "==================================================" << std::endl;
    }
    return s_info;
}

// ─────────────────────────────────────────────────────────────────────────────
// Shutdown
// ─────────────────────────────────────────────────────────────────────────────
void OvLoader::Shutdown()
{
    std::lock_guard<std::mutex> lock(s_mutex);

    // Release in reverse load order
    if (s_hOvC)       { FreeLib(s_hOvC);       s_hOvC       = nullptr; }
    if (s_hOvCore)    { FreeLib(s_hOvCore);    s_hOvCore    = nullptr; }
    if (s_hTbbMalloc) { FreeLib(s_hTbbMalloc); s_hTbbMalloc = nullptr; }
    if (s_hTbb)       { FreeLib(s_hTbb);       s_hTbb       = nullptr; }

    s_info = OvInfo{};
    s_initialized = false;
}