obsidian_ollama/src/vault-indexer.ts

// src/vault-indexer.ts

import { Vault, TFile } from 'obsidian';
import { Logger } from './utils';
import { Cache } from './cache';
import { VaultVectorStore } from './vault-vector-store';
import { VaultIndexEntry } from './types';

interface ParsedFrontmatter {
  title?: string;
  tags?: string;
  [key: string]: unknown;
}

interface TokenizedContent {
  title: string;
  headings: string[];
  frontmatter: ParsedFrontmatter;
  firstParagraph: string;
  content: string;
  basename: string;
}

interface VaultEntry {
  file: TFile;
  title: string;
  frontmatter: ParsedFrontmatter;
  headings: string[];
  content: string;
  basename: string;
  score: number;
}

export class InMemoryCache implements Cache {
  private store = new Map<string, string>();

  get(key: string): Promise<string | null> {
    return Promise.resolve(this.store.get(key) ?? null);
  }

  put(key: string, value: string): Promise<void> {
    this.store.set(key, value);
    return Promise.resolve();
  }
}

const STOP_WORDS = new Set([
  'a',
  'an',
  'the',
  'is',
  'it',
  'in',
  'on',
  'at',
  'to',
  'for',
  'of',
  'and',
  'or',
  'but',
  'with',
  'by',
  'from',
  'up',
  'about',
  'into',
  'this',
  'that',
  'these',
  'those',
  'be',
  'been',
  'being',
  'have',
  'has',
  'had',
  'do',
  'does',
  'did',
  'will',
  'would',
  'could',
  'should',
  'may',
  'might',
  'can',
  'are',
  'was',
  'were',
  'as',
  'so',
  'if',
  'not',
  'no',
  'my',
  'your',
  'our',
  'its',
  'we',
  'you',
  'he',
  'she',
  'they',
]);

const CONTENT_PREVIEW_LENGTH = 500;

export class VaultIndexer {
  private vault: Vault;
  private cache?: Cache;
  private vectorStore?: VaultVectorStore;
  private readonly SCORING_WEIGHTS = {
    TITLE: 5,
    FRONTMATTER_TITLE: 4,
    FRONTMATTER_TAGS: 3,
    HEADINGS: 2,
    CONTENT: 1,
  };

  constructor(vault: Vault, cache?: Cache, vectorStore?: VaultVectorStore) {
    this.vault = vault;
    this.cache = cache;
    this.vectorStore = vectorStore;
  }

  setVectorStore(vectorStore: VaultVectorStore | undefined): void {
    this.vectorStore = vectorStore;
  }

  tokenize(text: string): string[] {
    return text
      .toLowerCase()
      .replace(/[^\w\s]/g, '')
      .split(/\s+/)
      .filter((token) => token.length > 1 && !STOP_WORDS.has(token));
  }

  tokenizeContent(content: string, file: TFile): TokenizedContent {
    const parsed = this.parseMarkdown(content);
    const bodyWithoutFrontmatter = content.replace(/^---[\s\S]*?---\n/, '');
    const paragraphs = bodyWithoutFrontmatter
      .split(/\n\n+/)
      .map((p) => p.trim())
      .filter((p) => p && !p.startsWith('#'));
    const firstParagraph = paragraphs[0] || '';
    return {
      title: parsed.title || file.basename,
      headings: parsed.headings,
      frontmatter: parsed.frontmatter,
      firstParagraph,
      content: parsed.content,
      basename: file.basename,
    };
  }

  calculateWeightedScore(tokenized: TokenizedContent, queryTokens: string[]): { score: number } {
    let score = 0;
    for (const token of queryTokens) {
      if (tokenized.frontmatter?.title && this.exactMatch(tokenized.frontmatter.title, token)) {
        score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE;
      }
      if (tokenized.basename && this.exactMatch(tokenized.basename, token)) {
        score += this.SCORING_WEIGHTS.FRONTMATTER_TITLE;
      }
      if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, token)) {
        score += this.SCORING_WEIGHTS.FRONTMATTER_TAGS;
      }
      if (tokenized.headings.some((h) => h.toLowerCase().includes(token.toLowerCase()))) {
        score += this.SCORING_WEIGHTS.HEADINGS;
      }
      if (tokenized.content.toLowerCase().includes(token.toLowerCase())) {
        score += this.SCORING_WEIGHTS.CONTENT;
      }
      if (tokenized.title && this.exactMatch(tokenized.title, token)) {
        score += this.SCORING_WEIGHTS.TITLE;
      }
    }
    return { score };
  }

  async getVaultEntries(): Promise<VaultEntry[]> {
    const files = this.vault.getMarkdownFiles();
    const entries: VaultEntry[] = [];
    for (const file of files) {
      try {
        const content =
          typeof this.vault.cachedRead === 'function'
            ? await this.vault.cachedRead(file)
            : await this.vault.read(file);
        const parsed = this.parseMarkdown(content);
        entries.push({
          file: file,
          title: parsed.frontmatter.title || file.basename,
          frontmatter: parsed.frontmatter,
          headings: parsed.headings,
          content: parsed.content.slice(0, CONTENT_PREVIEW_LENGTH),
          basename: file.basename,
          score: 0,
        });
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        Logger.warn(`Failed to read file ${file.path}: ${errorMessage}`, 'vault-indexer');
      }
    }
    return entries;
  }

  async searchVault(query: string, limit = 3): Promise<VaultIndexEntry[]> {
    if (!query || !query.trim()) {
      return [];
    }

    // Try semantic search first if vector store is available
    if (this.vectorStore) {
      try {
        const semanticResults = await this.vectorStore.search(query, limit);
        if (semanticResults.length > 0) {
          return semanticResults;
        }
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        Logger.warn(
          `Semantic search failed, falling back to keyword search: ${errorMessage}`,
          'vault-indexer'
        );
      }
    }

    const cacheKey = `query:${query.trim()}:limit:${limit}`;
    if (this.cache) {
      let cachedResults: string | null = null;
      try {
        cachedResults = await this.cache.get(cacheKey);
      } catch {
        cachedResults = null;
      }
      if (cachedResults) {
        try {
          // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
          const parsedResults: VaultIndexEntry[] = JSON.parse(cachedResults);
          return parsedResults.slice(0, limit);
        } catch {
          // ignore parse errors
        }
      }
    }

    const queryTokens = this.tokenize(query);
    if (queryTokens.length === 0) {
      return [];
    }

    const entries = await this.getVaultEntries();
    const scored = entries
      .map((entry) => {
        const { score } = this.calculateWeightedScore(
          {
            title: entry.title,
            headings: entry.headings,
            frontmatter: entry.frontmatter,
            firstParagraph: '',
            content: entry.content,
            basename: entry.basename,
          },
          queryTokens
        );
        return { ...entry, score };
      })
      .filter((e) => e.score > 0);

    scored.sort((a, b) => b.score - a.score);
    const results: VaultIndexEntry[] = scored.slice(0, limit).map((e) => ({
      path: e.file.path,
      title: e.title,
      content: e.content,
      score: e.score,
      tags: e.frontmatter?.tags,
    }));

    if (this.cache) {
      try {
        await this.cache.put(cacheKey, JSON.stringify(results));
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : String(error);
        Logger.warn(
          `Failed to cache results for query "${query}": ${errorMessage}`,
          'vault-indexer'
        );
      }
    }

    return results;
  }

  private stemToken(token: string): string {
    if (token.endsWith('ing') && token.length > 4) return token.slice(0, -3);
    if (token.endsWith('ed') && token.length > 3) return token.slice(0, -2);
    if (token.endsWith('s') && token.length > 2) return token.slice(0, -1);
    return token;
  }

  private exactMatch(text: string | undefined, queryToken: string): boolean {
    if (!text) return false;
    const textLower = text.toLowerCase();
    const queryLower = queryToken.toLowerCase();
    const queryStem = this.stemToken(queryLower);
    return textLower.includes(queryLower) || textLower.includes(queryStem);
  }

  private parseMarkdown(content: string) {
    const frontmatterRegex = /^---\n([\s\S]*?)\n---\n/;
    const frontmatterMatch = content.match(frontmatterRegex);
    const frontmatter: ParsedFrontmatter = {};
    if (frontmatterMatch) {
      try {
        const lines = frontmatterMatch[1].trim().split('\n');
        for (const line of lines) {
          const [key, ...valueParts] = line.split(':');
          if (!key) continue;
          const value = valueParts.join(':').trim();
          if (key.trim() === 'title' && value) frontmatter.title = value;
          else if (key.trim() === 'tags' && value) frontmatter.tags = value;
        }
      } catch {
        Logger.warn('Failed to parse frontmatter', 'vault-indexer');
      }
    }

    const titleMatch = content.match(/^# (.+)$/m);
    const title = titleMatch ? titleMatch[1] : '';

    const headings: string[] = [];
    const headingRegex = /^#{1,6} (.+)$/gm;
    let headingMatch;
    while ((headingMatch = headingRegex.exec(content)) !== null) {
      headings.push(headingMatch[1]);
    }

    const bodyWithoutFrontmatter = frontmatterMatch
      ? content.substring(frontmatterMatch[0].length)
      : content;
    const bodyText = bodyWithoutFrontmatter
      .replace(/#{1,6} .+/g, '')
      .replace(/^\s*[\r\n]/gm, '')
      .trim();

    return { frontmatter, title, headings, content: bodyText };
  }
}