#!/usr/bin/env node
/**
 * embedding-worker.js — Standalone Node.js child process for embedding generation.
 *
 * Used as a fallback when Bun's ONNX runtime cannot load @huggingface/transformers
 * natively. Spawned by EmbeddingProvider when the native path fails.
 *
 * IPC protocol (JSON-lines over stdin/stdout):
 *   Startup:  sends  { ready: true }  after pipeline is loaded
 *             sends  { error: "..." } if startup fails
 *   Request:  parent → worker: { texts: ["...", "..."] }\n
 *   Response: worker → parent: { vectors: [[...], ...] }\n
 *                           or { error: "..." }\n
 *
 * Vectors are raw (un-normalised) Float32Arrays serialised as plain number arrays.
 * Normalisation is performed by EmbeddingProvider after receiving the response.
 *
 * This script is intentionally standalone — no imports from the CAW codebase.
 * It runs in Node.js, not Bun, to avoid the ONNX runtime incompatibility.
 */

import { createInterface } from 'readline';

const pad = (n) => String(n).padStart(2, '0');
const pad3 = (n) => String(n).padStart(3, '0');
const formatTimestamp = () => {
    const now = new Date();
    const h = now.getHours();
    const ampm = h >= 12 ? 'pm' : 'am';
    const h12 = h % 12 || 12;
    return `${now.getFullYear()}-${pad(now.getMonth() + 1)}-${pad(now.getDate())}` +
           ` ${pad(h12)}:${pad(now.getMinutes())}:${pad(now.getSeconds())}.${pad3(now.getMilliseconds())} ${ampm}`;
};

const MODEL_NAME = process.env.EMBEDDING_MODEL || 'nomic-ai/nomic-embed-text-v1.5';
const CACHE_DIR = process.env.EMBEDDING_CACHE_DIR || null;

const sendLine = (obj) => {
    process.stdout.write(JSON.stringify(obj) + '\n');
};

const run = async () => {
    let pipe;

    try {
        const { pipeline, env } = await import('@huggingface/transformers');

        if (CACHE_DIR) {
            env.cacheDir = CACHE_DIR;
        }

        // Log download progress to stderr (visible in parent process output)
        pipe = await pipeline('feature-extraction', MODEL_NAME, {
            dtype: 'q8',
            device: 'cpu',
            session_options: {
                enableCpuMemArena: false,
                enableMemPattern: false,
                intraOpNumThreads: 1,
                interOpNumThreads: 1,
            },
            progress_callback: (progress) => {
                if (progress.status === 'download') {
                    process.stderr.write(`${formatTimestamp()} [EmbeddingWorker] Downloading: ${progress.file} (${Math.round((progress.loaded / progress.total) * 100)}%)\n`);
                } else if (progress.status === 'ready') {
                    const suffix = progress.file ? `: ${progress.file}` : '';
                    process.stderr.write(`${formatTimestamp()} [EmbeddingWorker] Ready${suffix}\n`);
                }
            }
        });

        // Log loaded model details and baseline memory
        const rss = Math.round(process.memoryUsage().rss / 1024 / 1024);
        process.stderr.write(`${formatTimestamp()} [EmbeddingWorker] Model loaded. RSS: ${rss}MB\n`);

        sendLine({ ready: true });
    } catch (err) {
        sendLine({ error: `Pipeline load failed: ${err.message}` });
        process.exit(1);
    }

    // Process requests from parent line by line
    const rl = createInterface({ input: process.stdin });
    let embedCallCount = 0;

    rl.on('line', async (line) => {
        let req;
        try {
            req = JSON.parse(line);
        } catch {
            sendLine({ error: `Invalid JSON request: ${line}` });
            return;
        }

        const { texts, maxLength } = req;
        if (!Array.isArray(texts)) {
            sendLine({ error: 'Request must have a texts array' });
            return;
        }

        const extractVectors = (output, count) => {
            const dims = output.dims[output.dims.length - 1];
            const vectors = [];
            for (let i = 0; i < count; i++) {
                const start = i * dims;
                vectors.push(Array.from(output.data.slice(start, start + dims)));
            }
            output.dispose();
            return vectors;
        };

        try {
            let output;
            if (maxLength !== undefined) {
                const paddingOptions = { pooling: 'mean', normalize: false, padding: 'max_length', max_length: maxLength };
                try {
                    output = await pipe(texts, paddingOptions);
                } catch (_paddingErr) {
                    process.stderr.write(`${formatTimestamp()} [EmbeddingWorker] padding options unsupported, retrying without padding\n`);
                    output = await pipe(texts, { pooling: 'mean', normalize: false });
                }
            } else {
                output = await pipe(texts, { pooling: 'mean', normalize: false });
            }
            const vectors = extractVectors(output, texts.length);
            embedCallCount++;
            const embedRss = Math.round(process.memoryUsage().rss / 1024 / 1024);
            const shapeStr = maxLength !== undefined ? `[${texts.length}×${maxLength}]` : `[${texts.length}×var]`;
            process.stderr.write(`${formatTimestamp()} [EmbeddingWorker] Embed #${embedCallCount}: shape=${shapeStr}, RSS: ${embedRss}MB\n`);
            sendLine({ vectors });
        } catch (err) {
            sendLine({ error: `Embedding failed: ${err.message}` });
        }
    });

    rl.on('close', () => {
        process.exit(0);
    });
};

run();
