obsidian_ollama/src/vault-indexer.ts

// src/vault-indexer.ts

import { VaultIndexEntry } from './types';
import { Logger } from './utils';

interface Cache {
  get(key: string): Promise<string | null>;
  put(key: string, value: string): Promise<void>;
  clear(): Promise<void>;
}

class InMemoryCache implements Cache {
  private store: Map<string, string>;

  constructor() {
    this.store = new Map();
  }

  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();
  }

  clear(): Promise<void> {
    this.store.clear();
    return Promise.resolve();
  }
}

class CancellationToken {
  private cancelled = false;

  cancel(): void {
    this.cancelled = true;
  }

  get isCancelled(): boolean {
    return this.cancelled;
  }
}

interface Frontmatter {
  title?: string;
  tags?: string;
}

interface VaultFile {
  basename: string;
  path: string;
}

interface VaultLike {
  getMarkdownFiles(): VaultFile[];
  read(file: VaultFile): Promise<string>;
}

interface TokenizedContent {
  tokens: string[];
  headings: string[];
  frontmatter: Frontmatter;
  firstParagraph?: string;
}

interface ScoreResult {
  score: number;
  matchedFields: string[];
}

class VaultIndexer {
  private vault: VaultLike | null = null;
  private cache?: Cache;

  constructor(vault: VaultLike, cache?: Cache) {
    this.vault = vault;
    this.cache = cache;
  }

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

    if (!this.vault) {
      throw new Error('Vault-like object not provided to VaultIndexer');
    }

    const cacheKey = `query:${query.trim()}:limit:${limit}`;
    if (this.cache) {
      let cachedResults;
      try {
        cachedResults = await this.cache.get(cacheKey);
      } catch {
        // Ignore cache retrieval errors and continue with normal processing
        cachedResults = null;
      }
      if (cachedResults) {
        try {
          const parsedResults = JSON.parse(cachedResults) as VaultIndexEntry[];
          return parsedResults.slice(0, limit);
        } catch {
          // Ignore cache parse errors and continue with normal processing
        }
      }
    }

    const queryTokens = this.tokenize(query.trim());
    const vault = this.vault;
    const allFiles = vault.getMarkdownFiles();
    const results = await this.processFilesInBatches(vault, allFiles, queryTokens);

    const filteredResults = results
      .filter((result): result is NonNullable<typeof result> => result !== null)
      .sort((a, b) => b.score - a.score)
      .slice(0, limit);

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

    return filteredResults;
  }

  private async processFilesInBatches(
    vault: VaultLike,
    files: VaultFile[],
    queryTokens: string[]
  ): Promise<Array<VaultIndexEntry | null>> {
    const batchSize = 10;
    const results: VaultIndexEntry[] = [];
    const seenPaths = new Set<string>();
    const cancellationToken = new CancellationToken();
    // Removed unused processedCount variable

    // Set up a check for cancellation every 100 files
    const checkInterval = setInterval(() => {
      if (cancellationToken.isCancelled) {
        clearInterval(checkInterval);
      }
    }, 100);

    try {
      for (let i = 0; i < files.length; i += batchSize) {
        if (cancellationToken.isCancelled) {
          break;
        }

        const batch = files.slice(i, i + batchSize);
        const batchResults = await Promise.all(
          batch.map(async (file) => {
            try {
              const content = await vault.read(file);
              const tokenized = this.tokenizeContent(content);
              const scoreResult = this.calculateWeightedScore(tokenized, queryTokens, file);
              if (scoreResult.score > 0) {
                const entry: VaultIndexEntry = {
                  path: file.path,
                  title: file.basename.replace(/\.md$/, ''),
                  content: content.substring(0, 500),
                  score: scoreResult.score,
                };
                if (!seenPaths.has(entry.path)) {
                  seenPaths.add(entry.path);
                  return entry;
                }
                return null;
              }
              return null;
            } catch (error) {
              Logger.warn(
                `Failed to process ${file.path}: ${error instanceof Error ? error.message : String(error)}`,
                'vault-indexer'
              );
              return null;
            }
          })
        );

        const validResults = batchResults.filter(
          (result): result is NonNullable<typeof result> => result !== null
        );
        results.push(...validResults);

        // Continue processing all files to ensure we don't miss higher-scoring results
        // even if we've already found some matches

        // Removed processedCount increment
      }
    } finally {
      clearInterval(checkInterval);
    }

    return results;
  }

  private tokenize(text: string): string[] {
    const stopWords = new Set([
      'the',
      'a',
      'an',
      'and',
      'or',
      'but',
      'is',
      'are',
      'was',
      'were',
      'in',
      'on',
      'at',
      'to',
      'of',
      'for',
      'with',
      'as',
      'by',
      'it',
      'its',
      'that',
      'this',
      'these',
      'those',
    ]);
    return text
      .toLowerCase()
      .split(/\W+/)
      .filter((token) => token.length > 1 && !stopWords.has(token));
  }

  private tokenizeContent(content: string): TokenizedContent {
    const tokens: string[] = [];
    const headings: string[] = [];
    const frontmatter: Frontmatter = {};
    let firstParagraph: string | undefined;

    const frontmatterMatch = content.match(/^---(.*?)---/s);
    if (frontmatterMatch) {
      try {
        const frontmatterContent = frontmatterMatch[1];
        const lines = frontmatterContent.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') {
            if (value) {
              frontmatter.title = value;
            }
          } else if (key.trim() === 'tags') {
            if (value) {
              frontmatter.tags = value;
            }
          }
        }
      } catch {
        Logger.warn('Failed to parse frontmatter', 'vault-indexer');
      }
    }

    const headingMatches = content.match(/^#{1,6} (.*?)$/gm);
    if (headingMatches) {
      headings.push(...headingMatches.map((h: string) => h.replace(/^#{1,6} /, '')));
    }

    const paragraphMatch = content.match(/^([^#]*?)(?=#|\n\n|$)/s);
    if (paragraphMatch) {
      firstParagraph = paragraphMatch[1].trim();
    }

    const allText = content
      .replace(/^---.*?---/s, '')
      .replace(/^#.*?$/gm, '')
      .replace(/```.*?```/gs, '')
      .replace(/`.*?`/g, '')
      .replace(/\[.*?\]\(.*?\)/g, '');
    tokens.push(...this.tokenize(allText));

    return { tokens, headings, frontmatter, firstParagraph };
  }

  private calculateWeightedScore(
    tokenized: TokenizedContent,
    queryTokens: string[],
    file?: VaultFile
  ): ScoreResult {
    let totalScore = 0;
    const matchedTokens: Set<string> = new Set<string>();

    for (const queryToken of queryTokens) {
      let tokenScore = 0;
      const stemmed = this.stemToken(queryToken);
      let matched = false;

      if (
        tokenized.frontmatter?.title &&
        this.exactMatch(tokenized.frontmatter.title, queryToken)
      ) {
        tokenScore += 3;
        matched = true;
      } else if (
        file &&
        file.basename &&
        this.exactMatch(file.basename.replace(/\.md$/, ''), queryToken)
      ) {
        tokenScore += 3;
        matched = true;
      }

      if (tokenized.frontmatter?.tags && this.exactMatch(tokenized.frontmatter.tags, queryToken)) {
        tokenScore += 2.5;
        matched = true;
      }

      if (tokenized.headings.some((heading: string) => heading.toLowerCase().includes(stemmed))) {
        tokenScore += 5;
        matched = true;
      }

      if (tokenized.firstParagraph && tokenized.firstParagraph.toLowerCase().includes(stemmed)) {
        tokenScore += 1.5;
        matched = true;
      }

      if (tokenized.tokens.includes(stemmed)) {
        tokenScore += 1;
        matched = true;
      }

      if (matched) {
        totalScore += tokenScore;
        matchedTokens.add(queryToken);
      }
    }

    return {
      score: totalScore,
      matchedFields: Array.from(matchedTokens),
    };
  }

  private stemToken(token: string): string {
    // Improved stemmer that handles edge cases
    if (token.length <= 3) return token; // Don't stem very short tokens
    if (token.endsWith('s')) return token.slice(0, -1);
    if (token.endsWith('ed') && token.length > 4) return token.slice(0, -2); // Don't stem 3-letter words ending in ed
    if (token.endsWith('ing') && token.length > 5) return token.slice(0, -3); // Don't stem 4-letter words ending in ing
    return token;
  }

  private exactMatch(content: string, token: string): boolean {
    const stemmedToken = this.stemToken(token);
    return content.toLowerCase().includes(stemmedToken);
  }
}

export { VaultIndexer, Cache, InMemoryCache, CancellationToken };

// Convenience method to create a VaultIndexer with an in-memory cache
export function createVaultIndexerWithCache(vault: VaultLike): VaultIndexer {
  return new VaultIndexer(vault, new InMemoryCache());
}