"use strict";
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.RedisEmbeddingCache = exports.MemoryEmbeddingCache = void 0;
exports.loadModel = loadModel;
exports.disposeModel = disposeModel;
exports.disposeAllModels = disposeAllModels;
exports.embed = embed;
exports.embedQuery = embedQuery;
exports.embedBatch = embedBatch;
exports.resetEmbedder = resetEmbedder;
exports.cosineSimilarity = cosineSimilarity;
const transformers_1 = require("@huggingface/transformers");
const crypto_1 = __importDefault(require("crypto"));
const fs_1 = __importDefault(require("fs"));
const path_1 = __importDefault(require("path"));
const logger_1 = require("../lib/logger");
const defaults_1 = require("../lib/defaults");
const embedding_codec_1 = require("../lib/embedding-codec");
const log = (0, logger_1.createLogger)('embedder');
// ---------------------------------------------------------------------------
// LRU cache (in-memory) — default implementation
// ---------------------------------------------------------------------------
const DEFAULT_CACHE_SIZE = defaults_1.DEFAULT_EMBEDDING_CACHE_SIZE;
class LruCache {
    maxSize;
    map = new Map();
    constructor(maxSize) {
        this.maxSize = maxSize;
    }
    get(key) {
        const v = this.map.get(key);
        if (v !== undefined) {
            // Move to end (most recently used)
            this.map.delete(key);
            this.map.set(key, v);
        }
        return v;
    }
    set(key, value) {
        if (this.map.has(key))
            this.map.delete(key);
        this.map.set(key, value);
        if (this.map.size > this.maxSize) {
            // Evict oldest (first key)
            const first = this.map.keys().next().value;
            this.map.delete(first);
        }
    }
    clear() { this.map.clear(); }
}
class MemoryEmbeddingCache {
    lru;
    constructor(maxSize) { this.lru = new LruCache(maxSize); }
    async get(text) { return this.lru.get(text); }
    async set(text, vector) { this.lru.set(text, vector); }
    async clear() { this.lru.clear(); }
}
exports.MemoryEmbeddingCache = MemoryEmbeddingCache;
class RedisEmbeddingCache {
    client;
    prefix;
    ttlSeconds;
    constructor(client, prefix, ttlSeconds) {
        this.client = client;
        this.prefix = prefix;
        this.ttlSeconds = ttlSeconds;
    }
    hashKey(text) {
        const hash = crypto_1.default.createHash('sha256').update(text).digest('hex');
        return `${this.prefix}emb:${hash}`;
    }
    async get(text) {
        const raw = await this.client.get(this.hashKey(text));
        if (!raw)
            return undefined;
        return (0, embedding_codec_1.base64ToFloat32)(raw);
    }
    async set(text, vector) {
        const key = this.hashKey(text);
        const value = (0, embedding_codec_1.float32ToBase64)(vector);
        if (this.ttlSeconds > 0) {
            await this.client.set(key, value, { EX: this.ttlSeconds });
        }
        else {
            await this.client.set(key, value);
        }
    }
    async clear() {
        // Clear by pattern — expensive, but only used in dispose/reset
        const pattern = `${this.prefix}emb:*`;
        for await (const key of this.client.scanIterator({ MATCH: pattern, COUNT: 100 })) {
            await this.client.del(key);
        }
    }
}
exports.RedisEmbeddingCache = RedisEmbeddingCache;
const _models = new Map(); // name → { pipe, model, embedding }
const _pipeCache = new Map(); // "model|dtype" → pipe (dedup)
function validateRemoteUrl(url) {
    let parsed;
    try {
        parsed = new URL(url);
    }
    catch {
        throw new Error(`Invalid remote embedding URL: ${url}`);
    }
    if (!['http:', 'https:'].includes(parsed.protocol)) {
        throw new Error(`Remote embedding URL must use http or https: ${url}`);
    }
}
async function loadModel(model, embedding, modelsDir, name = 'default', cacheOrFactory) {
    const maxChars = embedding.maxChars;
    const cacheSize = embedding.cacheSize ?? DEFAULT_CACHE_SIZE;
    // Cache key by model identity — graphs sharing the same model share one cache
    const modelFingerprint = `${model.name}|${model.dtype ?? 'default'}`;
    const embeddingCache = typeof cacheOrFactory === 'function'
        ? cacheOrFactory(modelFingerprint)
        : cacheOrFactory ?? new MemoryEmbeddingCache(cacheSize);
    // Remote embedding: register proxy, skip ONNX loading
    if (embedding.remote) {
        validateRemoteUrl(embedding.remote);
        _models.set(name, { pipe: null, pipePromise: null, model, embedding, modelsDir, maxChars, cache: embeddingCache, remote: { url: embedding.remote, apiKey: embedding.remoteApiKey, model: embedding.remoteModel } });
        log.info({ name, endpoint: embedding.remote }, 'Model using remote endpoint');
        return;
    }
    // Register for lazy loading — pipeline created on first use
    _models.set(name, { pipe: null, pipePromise: null, model, embedding, modelsDir, maxChars, cache: embeddingCache });
    log.info({ model: model.name, name }, 'Registered model (lazy)');
}
// ---------------------------------------------------------------------------
// Lazy pipeline initialization
// ---------------------------------------------------------------------------
/** ONNX session options to reduce memory footprint. */
const SESSION_OPTIONS = {
    enableCpuMemArena: false,
    enableMemPattern: false,
    executionMode: 'sequential',
};
/**
 * Ensure the ONNX pipeline is loaded for a model entry.
 * Deduplicates by model+dtype via _pipeCache. Concurrent calls share the same promise.
 */
async function ensurePipeline(entry) {
    if (entry.pipe)
        return entry.pipe;
    if (entry.remote)
        throw new Error('ensurePipeline called on remote model');
    const cacheKey = `${entry.model.name}|${entry.model.dtype ?? ''}`;
    // Reuse pipeline if same model+dtype already loaded by another entry
    const cached = _pipeCache.get(cacheKey);
    if (cached) {
        entry.pipe = cached;
        entry.pipePromise = null;
        log.debug({ model: entry.model.name }, 'Reusing model for lazy init');
        return cached;
    }
    // Deduplicate concurrent lazy inits for same entry
    if (entry.pipePromise)
        return entry.pipePromise;
    entry.pipePromise = (async () => {
        transformers_1.env.cacheDir = entry.modelsDir;
        const modelDir = path_1.default.join(entry.modelsDir, entry.model.name.replace('/', path_1.default.sep));
        if (fs_1.default.existsSync(modelDir)) {
            transformers_1.env.allowRemoteModels = false;
            log.info({ model: entry.model.name }, 'Loading local model');
        }
        else {
            transformers_1.env.allowRemoteModels = true;
            log.info({ model: entry.model.name, modelsDir: entry.modelsDir }, 'Downloading model');
        }
        const pipeOpts = { session_options: SESSION_OPTIONS };
        if (entry.model.dtype)
            pipeOpts.dtype = entry.model.dtype;
        const pipe = await (0, transformers_1.pipeline)('feature-extraction', entry.model.name, pipeOpts);
        _pipeCache.set(cacheKey, pipe);
        entry.pipe = pipe;
        entry.pipePromise = null;
        log.info({ model: entry.model.name }, 'Model ready');
        return pipe;
    })();
    return entry.pipePromise;
}
// ---------------------------------------------------------------------------
// Dispose
// ---------------------------------------------------------------------------
/** Dispose a single model entry. Pipeline is only disposed if no other entries reference it. */
async function disposeModel(name) {
    const entry = _models.get(name);
    if (!entry)
        return;
    if (entry.pipe) {
        // Check if any other entry shares this pipeline
        const cacheKey = `${entry.model.name}|${entry.model.dtype ?? ''}`;
        let shared = false;
        for (const [n, e] of _models) {
            if (n !== name && e.pipe === entry.pipe) {
                shared = true;
                break;
            }
        }
        if (!shared) {
            await entry.pipe.dispose();
            _pipeCache.delete(cacheKey);
            log.debug({ model: entry.model.name }, 'Disposed pipeline');
        }
    }
    await entry.cache.clear();
    _models.delete(name);
}
/** Dispose all pipelines and clear registries. */
async function disposeAllModels() {
    const disposed = new Set();
    for (const entry of _models.values()) {
        if (entry.pipe && !disposed.has(entry.pipe)) {
            disposed.add(entry.pipe);
            await entry.pipe.dispose();
        }
        await entry.cache.clear();
    }
    _models.clear();
    _pipeCache.clear();
}
// ---------------------------------------------------------------------------
// Remote embedding HTTP client
// ---------------------------------------------------------------------------
async function remoteEmbed(url, texts, apiKey, remoteModel) {
    const headers = { 'Content-Type': 'application/json' };
    if (apiKey)
        headers['Authorization'] = `Bearer ${apiKey}`;
    const payload = { texts };
    if (remoteModel)
        payload.model = remoteModel;
    const body = JSON.stringify(payload);
    for (let attempt = 0; attempt < defaults_1.REMOTE_MAX_RETRIES; attempt++) {
        let resp;
        try {
            resp = await fetch(url, { method: 'POST', headers, body });
        }
        catch (err) {
            // Network error — retry
            if (attempt < defaults_1.REMOTE_MAX_RETRIES - 1) {
                const delay = defaults_1.REMOTE_BASE_DELAY_MS * 2 ** attempt;
                await new Promise(r => setTimeout(r, delay));
                continue;
            }
            throw new Error(`Remote embed network error after ${defaults_1.REMOTE_MAX_RETRIES} attempts: ${err}`);
        }
        if (resp.ok) {
            const data = await resp.json();
            return data.embeddings;
        }
        // Client errors (4xx) — don't retry
        if (resp.status < 500) {
            const respBody = (await resp.text()).slice(0, defaults_1.ERROR_BODY_LIMIT);
            throw new Error(`Remote embed failed (${resp.status}): ${respBody}`);
        }
        // Server errors (5xx) — retry
        if (attempt < defaults_1.REMOTE_MAX_RETRIES - 1) {
            const delay = defaults_1.REMOTE_BASE_DELAY_MS * 2 ** attempt;
            await new Promise(r => setTimeout(r, delay));
            continue;
        }
        const respBody = (await resp.text()).slice(0, defaults_1.ERROR_BODY_LIMIT);
        throw new Error(`Remote embed failed after ${defaults_1.REMOTE_MAX_RETRIES} attempts (${resp.status}): ${respBody}`);
    }
    throw new Error('Remote embed: unreachable');
}
function getEntry(modelName) {
    const entry = _models.get(modelName);
    if (!entry)
        throw new Error(`Model "${modelName}" not loaded. Call loadModel() first.`);
    return entry;
}
/** Embed a document (indexing). Applies documentPrefix and configured pooling. */
async function embed(title, content, modelName = 'default') {
    const entry = getEntry(modelName);
    const raw = `${title}\n${content}`;
    const text = `${entry.model.documentPrefix}${raw}`.slice(0, entry.maxChars);
    const cached = await entry.cache.get(text);
    if (cached)
        return cached;
    let vec;
    if (entry.remote) {
        [vec] = await remoteEmbed(entry.remote.url, [text], entry.remote.apiKey, entry.remote.model);
    }
    else {
        const pipe = await ensurePipeline(entry);
        const tensor = await pipe._call(text, { pooling: entry.model.pooling, normalize: entry.model.normalize });
        vec = Array.from(tensor.data);
    }
    await entry.cache.set(text, vec);
    return vec;
}
/** Embed a search query. Applies queryPrefix and configured pooling. */
async function embedQuery(query, modelName = 'default') {
    const entry = getEntry(modelName);
    const text = `${entry.model.queryPrefix}${query}`.slice(0, entry.maxChars);
    const cached = await entry.cache.get(text);
    if (cached)
        return cached;
    let vec;
    if (entry.remote) {
        [vec] = await remoteEmbed(entry.remote.url, [text], entry.remote.apiKey, entry.remote.model);
    }
    else {
        const pipe = await ensurePipeline(entry);
        const tensor = await pipe._call(text, { pooling: entry.model.pooling, normalize: entry.model.normalize });
        vec = Array.from(tensor.data);
    }
    await entry.cache.set(text, vec);
    return vec;
}
/** Batch-embed documents (indexing). Applies documentPrefix and configured pooling. */
async function embedBatch(inputs, modelName = 'default') {
    const entry = getEntry(modelName);
    if (inputs.length === 0)
        return [];
    if (inputs.length === 1)
        return [await embed(inputs[0].title, inputs[0].content, modelName)];
    const texts = inputs.map(({ title, content }) => `${entry.model.documentPrefix}${title}\n${content}`.slice(0, entry.maxChars));
    // Check cache: split into hits and misses
    const cachedResults = await Promise.all(texts.map(t => entry.cache.get(t)));
    const result = cachedResults.map(v => v ?? null);
    const missIndices = result.map((v, i) => v === null ? i : -1).filter(i => i >= 0);
    if (missIndices.length === 0)
        return result;
    const missTexts = missIndices.map(i => texts[i]);
    let missVecs;
    if (entry.remote) {
        missVecs = await remoteEmbed(entry.remote.url, missTexts, entry.remote.apiKey, entry.remote.model);
    }
    else {
        const pipe = await ensurePipeline(entry);
        missVecs = [];
        const batchSize = entry.embedding.batchSize;
        for (let start = 0; start < missTexts.length; start += batchSize) {
            const chunk = missTexts.slice(start, start + batchSize);
            const tensor = await pipe._call(chunk, { pooling: entry.model.pooling, normalize: entry.model.normalize });
            const dim = tensor.dims[1];
            const data = tensor.data;
            for (let i = 0; i < chunk.length; i++) {
                missVecs.push(Array.from(data.slice(i * dim, (i + 1) * dim)));
            }
        }
    }
    // Fill results and populate cache
    for (let j = 0; j < missIndices.length; j++) {
        const idx = missIndices[j];
        result[idx] = missVecs[j];
        await entry.cache.set(texts[idx], missVecs[j]);
    }
    return result;
}
async function resetEmbedder() {
    for (const entry of _models.values())
        await entry.cache.clear();
    _models.clear();
    _pipeCache.clear();
}
// Vectors are L2-normalized → dot product = cosine similarity
function cosineSimilarity(a, b) {
    if (a.length !== b.length)
        return 0;
    let dot = 0;
    for (let i = 0; i < a.length; i++)
        dot += a[i] * b[i];
    return dot;
}
