alexiel-rag

import "./suppressExperimentalWarnings";
import { DatabaseSync } from "node:sqlite";
import fs from "node:fs";
import path from "node:path";
import os from "node:os";

export interface Chunk {
  file_path: string;
  chunk_index: number;
  content: string;
}

const DB_PATH = process.env.RAG_DB_PATH ?? path.join(os.homedir(), ".alexiel-rag", "rag.db");

let dbInstance: DatabaseSync | null = null;
let ftsAvailable = true;

export function getDb(): DatabaseSync {
  if (dbInstance) return dbInstance;

  const dir = path.dirname(DB_PATH);
  if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });

  const db = new DatabaseSync(DB_PATH);
  db.exec("PRAGMA journal_mode = WAL;");

  db.exec(`
    CREATE TABLE IF NOT EXISTS chunks (
      id INTEGER PRIMARY KEY AUTOINCREMENT,
      file_path TEXT NOT NULL,
      chunk_index INTEGER NOT NULL,
      content TEXT NOT NULL,
      indexed_at TEXT NOT NULL
    );
    CREATE INDEX IF NOT EXISTS idx_chunks_file_path ON chunks(file_path);
  `);

  try {
    db.exec(`
      CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
        file_path, content, content='chunks', content_rowid='id'
      );
      CREATE TRIGGER IF NOT EXISTS chunks_ai AFTER INSERT ON chunks BEGIN
        INSERT INTO chunks_fts(rowid, file_path, content) VALUES (new.id, new.file_path, new.content);
      END;
      CREATE TRIGGER IF NOT EXISTS chunks_ad AFTER DELETE ON chunks BEGIN
        INSERT INTO chunks_fts(chunks_fts, rowid, file_path, content) VALUES ('delete', old.id, old.file_path, old.content);
      END;
    `);
  } catch (error) {
    ftsAvailable = false;
    console.error("FTS5 unavailable, falling back to LIKE-based search:", error);
  }

  dbInstance = db;
  return db;
}

/** Resolve an allowlist string (one path per line) into normalized absolute paths. */
function parseAllowedFolders(raw: string): string[] {
  return raw
    .split(/\r?\n/)
    .map((line) => line.trim())
    .filter(Boolean)
    .map((line) => path.resolve(line));
}

/** True only if `filePath` resolves to somewhere inside `folder`. Blocks ../ escapes. */
function isInside(filePath: string, folder: string): boolean {
  const rel = path.relative(folder, path.resolve(filePath));
  return rel === "" || (!rel.startsWith("..") && !path.isAbsolute(rel));
}

function walkFolder(
  folder: string,
  extensions: Set<string>,
  maxFileSizeBytes: number,
  out: string[],
): void {
  let entries: fs.Dirent[];
  try {
    entries = fs.readdirSync(folder, { withFileTypes: true });
  } catch {
    return; // unreadable folder — skip silently, don't crash the whole index run
  }
  for (const entry of entries) {
    const full = path.join(folder, entry.name);
    if (entry.isSymbolicLink()) continue; // never follow symlinks — could point outside the allowlist
    if (entry.isDirectory()) {
      if (entry.name === "node_modules" || entry.name.startsWith(".")) continue;
      walkFolder(full, extensions, maxFileSizeBytes, out);
    } else if (entry.isFile()) {
      const ext = path.extname(entry.name).toLowerCase();
      if (!extensions.has(ext)) continue;
      try {
        const stat = fs.statSync(full);
        if (stat.size > maxFileSizeBytes) continue;
      } catch {
        continue;
      }
      out.push(full);
    }
  }
}

function chunkText(content: string, chunkSize: number, overlap: number): string[] {
  if (content.length <= chunkSize) return content.trim() ? [content] : [];
  const chunks: string[] = [];
  const step = Math.max(1, chunkSize - overlap);
  for (let start = 0; start < content.length; start += step) {
    const chunk = content.slice(start, start + chunkSize);
    if (chunk.trim()) chunks.push(chunk);
    if (start + chunkSize >= content.length) break;
  }
  return chunks;
}

export interface IndexResult {
  filesIndexed: number;
  chunksIndexed: number;
  filesSkipped: number;
  folders: string[];
}

export function reindex(
  allowedFoldersRaw: string,
  extensionsRaw: string,
  chunkSize: number,
  chunkOverlap: number,
  maxFileSizeBytes: number,
): IndexResult {
  const db = getDb();
  const allowedFolders = parseAllowedFolders(allowedFoldersRaw);
  const extensions = new Set(
    extensionsRaw
      .split(",")
      .map((e) => e.trim().toLowerCase())
      .filter(Boolean)
      .map((e) => (e.startsWith(".") ? e : `.${e}`)),
  );

  if (allowedFolders.length === 0) {
    return { filesIndexed: 0, chunksIndexed: 0, filesSkipped: 0, folders: [] };
  }

  const files: string[] = [];
  for (const folder of allowedFolders) {
    if (!fs.existsSync(folder)) continue;
    walkFolder(folder, extensions, maxFileSizeBytes, files);
  }

  // Wipe and rebuild — reindex is an explicit, user-triggered action, so a full rebuild keeps
  // the logic simple and guarantees stale/deleted files never linger in the index.
  db.exec("DELETE FROM chunks;");

  const insertStmt = db.prepare(
    `INSERT INTO chunks (file_path, chunk_index, content, indexed_at) VALUES (?, ?, ?, ?)`,
  );

  let filesIndexed = 0;
  let filesSkipped = 0;
  let chunksIndexed = 0;
  const now = new Date().toISOString();

  for (const file of files) {
    // Defense in depth: re-verify containment even though walkFolder only descends from
    // allowed roots and never follows symlinks.
    if (!allowedFolders.some((folder) => isInside(file, folder))) {
      filesSkipped++;
      continue;
    }
    let content: string;
    try {
      content = fs.readFileSync(file, "utf-8");
    } catch {
      filesSkipped++;
      continue;
    }
    const chunks = chunkText(content, chunkSize, chunkOverlap);
    chunks.forEach((chunk, idx) => {
      insertStmt.run(file, idx, chunk, now);
      chunksIndexed++;
    });
    if (chunks.length > 0) filesIndexed++;
    else filesSkipped++;
  }

  return { filesIndexed, chunksIndexed, filesSkipped, folders: allowedFolders };
}

function toFtsQuery(raw: string): string {
  const tokens = raw
    .split(/\s+/)
    .map((t) => t.replace(/["]/g, "").trim())
    .filter((t) => t.length > 1);
  if (tokens.length === 0) return '""';
  return tokens.map((t) => `"${t}"*`).join(" OR ");
}

export function search(query: string, limit: number): Chunk[] {
  const db = getDb();
  if (ftsAvailable) {
    try {
      const ftsQuery = toFtsQuery(query);
      return db
        .prepare(
          `SELECT c.file_path, c.chunk_index, c.content
           FROM chunks_fts f JOIN chunks c ON c.id = f.rowid
           WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?`,
        )
        .all(ftsQuery, limit) as unknown as Chunk[];
    } catch {
      // fall through
    }
  }
  const likeTerm = `%${query.trim()}%`;
  return db
    .prepare(
      `SELECT file_path, chunk_index, content FROM chunks WHERE content LIKE ? LIMIT ?`,
    )
    .all(likeTerm, limit) as unknown as Chunk[];
}

export function listSources(): Array<{ file_path: string; chunks: number }> {
  const db = getDb();
  return db
    .prepare(
      `SELECT file_path, COUNT(*) as chunks FROM chunks GROUP BY file_path ORDER BY file_path`,
    )
    .all() as unknown as Array<{ file_path: string; chunks: number }>;
}

export function getDbPath(): string {
  return DB_PATH;
}