#!/usr/bin/env node
import { opendir, readFile, stat } from 'node:fs/promises';
import path from 'node:path';
import { pathToFileURL } from 'node:url';
import { pipeline } from '@xenova/transformers';
import { DEFAULT_EMBEDDING_BATCH_SIZE, DEFAULT_EMBEDDING_MODEL, INDEX_PATH, RAG_SOURCE_DIR } from './constants.js';
import { normalizeVectors, tensorToVectors } from './embeddings.js';
import { saveEmbeddingIndex } from './storage.js';
const MAX_CHUNK_CHARS = 2500;
const MIN_CHUNK_CHARS = 400;
const MIN_PERSISTED_CHARS = 64;
const SENTENCE_SEPARATOR = /(?<=[.!?。？！])\s+/u;
const ALLOWED_EXTENSIONS = new Set(['.md', '.txt']);
const toPosixPath = (input) => input.split(path.sep).join('/');
const collectFiles = async (dir) => {
    const entries = await stat(dir).catch(() => {
        throw new Error(`RAG source directory does not exist: ${dir}`);
    });
    if (!entries.isDirectory()) {
        throw new Error(`RAG source path is not a directory: ${dir}`);
    }
    const files = [];
    const walk = async (current) => {
        const dirEntries = await opendir(current);
        for await (const entry of dirEntries) {
            const fullPath = path.join(current, entry.name);
            if (entry.isDirectory()) {
                await walk(fullPath);
                continue;
            }
            if (!entry.isFile()) {
                continue;
            }
            const extension = path.extname(entry.name).toLowerCase();
            if (!ALLOWED_EXTENSIONS.has(extension)) {
                continue;
            }
            files.push(fullPath);
        }
    };
    await walk(dir);
    return files.sort();
};
const normalizeText = (value) => value.replace(/\r\n/g, '\n').trim();
const splitByLength = (text) => {
    const sanitized = text.trim();
    if (!sanitized) {
        return [];
    }
    if (sanitized.length <= MAX_CHUNK_CHARS) {
        return [sanitized];
    }
    const segments = [];
    for (let index = 0; index < sanitized.length; index += MAX_CHUNK_CHARS) {
        const slice = sanitized.slice(index, index + MAX_CHUNK_CHARS).trim();
        if (slice) {
            segments.push(slice);
        }
    }
    return segments;
};
const splitLongParagraph = (paragraph) => {
    const sanitized = paragraph.trim();
    if (!sanitized) {
        return [];
    }
    if (sanitized.length <= MAX_CHUNK_CHARS) {
        return [sanitized];
    }
    const sentences = sanitized.split(SENTENCE_SEPARATOR).filter(Boolean);
    if (sentences.length <= 1) {
        return splitByLength(sanitized);
    }
    const segments = [];
    let current = '';
    const pushCurrent = () => {
        if (current.trim()) {
            segments.push(current.trim());
            current = '';
        }
    };
    for (const sentence of sentences) {
        const candidate = current ? `${current} ${sentence}` : sentence;
        if (candidate.length > MAX_CHUNK_CHARS) {
            pushCurrent();
            if (sentence.length > MAX_CHUNK_CHARS) {
                segments.push(...splitByLength(sentence));
                continue;
            }
            current = sentence;
            continue;
        }
        current = candidate;
    }
    pushCurrent();
    return segments;
};
const mergeSmallSegments = (segments) => {
    if (segments.length === 0) {
        return [];
    }
    const merged = [];
    for (const segment of segments) {
        const trimmed = segment.trim();
        if (!trimmed) {
            continue;
        }
        const lastIndex = merged.length - 1;
        if (lastIndex >= 0 &&
            trimmed.length < MIN_CHUNK_CHARS &&
            merged[lastIndex].length + 2 + trimmed.length <= MAX_CHUNK_CHARS) {
            merged[lastIndex] = `${merged[lastIndex]}\n\n${trimmed}`;
            continue;
        }
        merged.push(trimmed);
    }
    return merged;
};
const splitContentBlock = (content) => {
    const paragraphs = content
        .split(/\n{2,}/u)
        .map((paragraph) => paragraph.trim())
        .filter(Boolean);
    if (paragraphs.length === 0) {
        return splitByLength(content);
    }
    const segments = [];
    let current = '';
    const pushCurrent = () => {
        if (current.trim()) {
            segments.push(current.trim());
            current = '';
        }
    };
    for (const paragraph of paragraphs) {
        const candidate = current ? `${current}\n\n${paragraph}` : paragraph;
        if (candidate.length > MAX_CHUNK_CHARS) {
            pushCurrent();
            segments.push(...splitLongParagraph(paragraph));
            continue;
        }
        current = candidate;
    }
    pushCurrent();
    return mergeSmallSegments(segments);
};
const chunkDocument = (relativePath, content) => {
    const normalized = normalizeText(content);
    const lines = normalized.split('\n');
    const rawChunks = [];
    let currentLines = [];
    let currentHeading;
    const flush = () => {
        const joined = currentLines.join('\n').trim();
        currentLines = [];
        if (!joined) {
            return;
        }
        const segments = splitContentBlock(joined);
        segments.forEach((segment, segmentIndex) => {
            rawChunks.push({
                heading: currentHeading,
                content: segment,
                includeHeading: segmentIndex === 0,
            });
        });
    };
    for (const line of lines) {
        const headingMatch = line.match(/^(#{1,6})\s+(.*)$/u);
        if (headingMatch) {
            flush();
            currentHeading = headingMatch[2].trim();
            continue;
        }
        currentLines.push(line);
    }
    flush();
    const chunks = [];
    let chunkIndex = 0;
    for (const rawChunk of rawChunks) {
        const { heading, content: chunkContent, includeHeading } = rawChunk;
        const trimmed = chunkContent.trim();
        if (trimmed.length < MIN_PERSISTED_CHARS && chunks.length > 0) {
            const last = chunks[chunks.length - 1];
            if (last.text.length + 2 + trimmed.length <= MAX_CHUNK_CHARS) {
                last.text = `${last.text}\n\n${trimmed}`;
                continue;
            }
        }
        let text = trimmed;
        if (heading && includeHeading && !trimmed.startsWith(heading)) {
            text = `${heading}\n\n${trimmed}`;
        }
        const id = `${toPosixPath(relativePath)}:chunk-${chunkIndex.toString().padStart(4, '0')}`;
        chunks.push({
            id,
            text,
            metadata: {
                source: toPosixPath(relativePath),
                heading,
            },
        });
        chunkIndex += 1;
    }
    return chunks;
};
const embedChunks = async (chunks, model, batchSize) => {
    if (chunks.length === 0) {
        return [];
    }
    const extractor = (await pipeline('feature-extraction', model));
    const embedded = [];
    for (let index = 0; index < chunks.length; index += batchSize) {
        const batch = chunks.slice(index, index + batchSize);
        const inputs = batch.map((chunk) => chunk.text);
        const tensor = await extractor(inputs, { pooling: 'mean' });
        const values = normalizeVectors(tensorToVectors(tensor));
        if (values.length !== batch.length) {
            throw new Error('Embedding batch size mismatch');
        }
        batch.forEach((chunk, batchIndex) => {
            embedded.push({
                ...chunk,
                embedding: values[batchIndex],
            });
        });
        process.stdout.write(`Embedded ${Math.min(index + batch.length, chunks.length)} / ${chunks.length} chunks\r`);
    }
    process.stdout.write('\n');
    return embedded;
};
const buildEmbeddingIndex = async () => {
    const model = process.env.DOCS_MCP_EMBEDDING_MODEL ?? DEFAULT_EMBEDDING_MODEL;
    const parsedBatch = Number.parseInt(process.env.DOCS_MCP_EMBEDDING_BATCH ?? '', 10);
    const batchSize = Number.isNaN(parsedBatch) || parsedBatch < 1 ? DEFAULT_EMBEDDING_BATCH_SIZE : parsedBatch;
    const files = await collectFiles(RAG_SOURCE_DIR);
    if (files.length === 0) {
        throw new Error(`No source documents found in ${RAG_SOURCE_DIR}`);
    }
    const allChunks = [];
    for (const filePath of files) {
        const relativePath = path.relative(RAG_SOURCE_DIR, filePath);
        const contents = await readFile(filePath, 'utf-8');
        const chunks = chunkDocument(relativePath, contents);
        allChunks.push(...chunks);
    }
    if (allChunks.length === 0) {
        throw new Error('No document chunks generated. Check chunking configuration.');
    }
    console.log(`Embedding ${allChunks.length} chunks from ${files.length} documents using ${model}`);
    const embeddedChunks = await embedChunks(allChunks, model, batchSize);
    if (embeddedChunks.length === 0) {
        throw new Error('Failed to generate any embeddings');
    }
    const dimension = embeddedChunks[0].embedding.length;
    const index = {
        model,
        dimension,
        createdAt: new Date().toISOString(),
        chunks: embeddedChunks,
    };
    await saveEmbeddingIndex(index, INDEX_PATH);
    console.log(`Saved embedding index with ${embeddedChunks.length} vectors to ${INDEX_PATH}`);
};
const invokedFromCli = () => {
    const entry = process.argv[1];
    if (!entry) {
        return false;
    }
    return import.meta.url === pathToFileURL(entry).href;
};
if (invokedFromCli()) {
    buildEmbeddingIndex().catch((error) => {
        console.error('Failed to build embedding index:', error);
        process.exitCode = 1;
    });
}
//# sourceMappingURL=build-index.js.map