obsidian-promptfire/src/presets.ts

import { App, TFile, CachedMetadata } from 'obsidian';
import { estimateTokens } from './history';

// === Types ===

export type PresetMode = 'manual' | 'auto';

export interface ContextPreset {
	mode?: PresetMode;
	template?: string;
	includeLinked?: boolean;
	linkDepth?: number;
	includeTags?: string[];
	excludePaths?: string[];
	excludeTags?: string[];
	maxTokens?: number;
	includeActiveNote?: boolean;
}

export interface PresetValidationResult {
	valid: boolean;
	preset: ContextPreset | null;
	errors: string[];
	warnings: string[];
}

export interface LinkedNoteResult {
	file: TFile;
	depth: number;
	linkPath: string[];
}

// === Frontmatter Parsing ===

/**
 * Extract and validate ai-context preset from file's frontmatter
 */
export function parsePresetFromFrontmatter(
	cache: CachedMetadata | null
): PresetValidationResult {
	const errors: string[] = [];
	const warnings: string[] = [];

	if (!cache?.frontmatter) {
		return { valid: false, preset: null, errors: ['No frontmatter found'], warnings };
	}

	const aiContext = cache.frontmatter['ai-context'];

	if (!aiContext) {
		return { valid: false, preset: null, errors: ['No ai-context block in frontmatter'], warnings };
	}

	if (typeof aiContext !== 'object' || Array.isArray(aiContext)) {
		return { valid: false, preset: null, errors: ['ai-context must be an object'], warnings };
	}

	const preset: ContextPreset = {};

	// Parse mode
	if (aiContext.mode !== undefined) {
		if (aiContext.mode === 'auto' || aiContext.mode === 'manual') {
			preset.mode = aiContext.mode;
		} else {
			errors.push('mode must be "auto" or "manual"');
		}
	}

	// Parse template
	if (aiContext.template !== undefined) {
		if (typeof aiContext.template === 'string') {
			preset.template = aiContext.template;
		} else {
			errors.push('template must be a string');
		}
	}

	// Parse includeLinked
	if (aiContext['include-linked'] !== undefined) {
		if (typeof aiContext['include-linked'] === 'boolean') {
			preset.includeLinked = aiContext['include-linked'];
		} else {
			errors.push('include-linked must be a boolean');
		}
	}

	// Parse linkDepth
	if (aiContext['link-depth'] !== undefined) {
		const depth = Number(aiContext['link-depth']);
		if (!isNaN(depth) && depth >= 0 && depth <= 10) {
			preset.linkDepth = depth;
		} else {
			errors.push('link-depth must be a number between 0 and 10');
		}
	}

	// Parse includeTags
	if (aiContext['include-tags'] !== undefined) {
		if (Array.isArray(aiContext['include-tags'])) {
			const tags = aiContext['include-tags'].filter((t: unknown) => typeof t === 'string');
			if (tags.length > 0) {
				preset.includeTags = tags.map((t: string) => t.startsWith('#') ? t : `#${t}`);
			}
		} else {
			errors.push('include-tags must be an array of strings');
		}
	}

	// Parse excludePaths
	if (aiContext['exclude-paths'] !== undefined) {
		if (Array.isArray(aiContext['exclude-paths'])) {
			const paths = aiContext['exclude-paths'].filter((p: unknown) => typeof p === 'string');
			if (paths.length > 0) {
				preset.excludePaths = paths;
			}
		} else {
			errors.push('exclude-paths must be an array of strings');
		}
	}

	// Parse excludeTags
	if (aiContext['exclude-tags'] !== undefined) {
		if (Array.isArray(aiContext['exclude-tags'])) {
			const tags = aiContext['exclude-tags'].filter((t: unknown) => typeof t === 'string');
			if (tags.length > 0) {
				preset.excludeTags = tags.map((t: string) => t.startsWith('#') ? t : `#${t}`);
			}
		} else {
			errors.push('exclude-tags must be an array of strings');
		}
	}

	// Parse maxTokens
	if (aiContext['max-tokens'] !== undefined) {
		const tokens = Number(aiContext['max-tokens']);
		if (!isNaN(tokens) && tokens > 0) {
			preset.maxTokens = tokens;
		} else {
			errors.push('max-tokens must be a positive number');
		}
	}

	// Parse includeActiveNote
	if (aiContext['include-active-note'] !== undefined) {
		if (typeof aiContext['include-active-note'] === 'boolean') {
			preset.includeActiveNote = aiContext['include-active-note'];
		} else {
			errors.push('include-active-note must be a boolean');
		}
	}

	// Warn about unknown fields
	const knownFields = [
		'mode', 'template', 'include-linked', 'link-depth', 'include-tags',
		'exclude-paths', 'exclude-tags', 'max-tokens', 'include-active-note'
	];
	for (const key of Object.keys(aiContext)) {
		if (!knownFields.includes(key)) {
			warnings.push(`Unknown field: ${key}`);
		}
	}

	return {
		valid: errors.length === 0,
		preset: errors.length === 0 ? preset : null,
		errors,
		warnings,
	};
}

// === Link Traversal ===

export class LinkTraverser {
	private app: App;
	private visited: Set<string> = new Set();
	private results: LinkedNoteResult[] = [];

	constructor(app: App) {
		this.app = app;
	}

	/**
	 * Traverse links from a starting file up to a given depth
	 */
	async traverseLinks(
		startFile: TFile,
		maxDepth: number,
		excludePaths: string[] = [],
		excludeTags: string[] = []
	): Promise<LinkedNoteResult[]> {
		this.visited.clear();
		this.results = [];

		await this.traverse(startFile, 0, maxDepth, [], excludePaths, excludeTags);

		return this.results;
	}

	private async traverse(
		file: TFile,
		currentDepth: number,
		maxDepth: number,
		linkPath: string[],
		excludePaths: string[],
		excludeTags: string[]
	): Promise<void> {
		// Check if already visited
		if (this.visited.has(file.path)) {
			return;
		}

		// Check excluded paths
		if (this.isPathExcluded(file.path, excludePaths)) {
			return;
		}

		// Check excluded tags
		if (this.hasExcludedTag(file, excludeTags)) {
			return;
		}

		this.visited.add(file.path);

		// Add to results (except starting file at depth 0)
		if (currentDepth > 0) {
			this.results.push({
				file,
				depth: currentDepth,
				linkPath: [...linkPath, file.basename],
			});
		}

		// Stop if max depth reached
		if (currentDepth >= maxDepth) {
			return;
		}

		// Get outgoing links
		const cache = this.app.metadataCache.getFileCache(file);
		if (!cache?.links) {
			return;
		}

		for (const link of cache.links) {
			const linkedFile = this.app.metadataCache.getFirstLinkpathDest(
				link.link,
				file.path
			);

			if (linkedFile instanceof TFile) {
				await this.traverse(
					linkedFile,
					currentDepth + 1,
					maxDepth,
					[...linkPath, file.basename],
					excludePaths,
					excludeTags
				);
			}
		}
	}

	private isPathExcluded(path: string, excludePaths: string[]): boolean {
		for (const excludePath of excludePaths) {
			if (path.startsWith(excludePath) || path.includes(`/${excludePath}`)) {
				return true;
			}
		}
		return false;
	}

	private hasExcludedTag(file: TFile, excludeTags: string[]): boolean {
		if (excludeTags.length === 0) {
			return false;
		}

		const cache = this.app.metadataCache.getFileCache(file);
		if (!cache?.frontmatter?.tags) {
			return false;
		}

		const fileTags: string[] = Array.isArray(cache.frontmatter.tags)
			? cache.frontmatter.tags
			: [cache.frontmatter.tags];

		for (const fileTag of fileTags) {
			const normalizedTag = fileTag.startsWith('#') ? fileTag : `#${fileTag}`;
			if (excludeTags.includes(normalizedTag)) {
				return true;
			}
		}

		return false;
	}
}

// === Tag-based File Collection ===

export class TagCollector {
	private app: App;

	constructor(app: App) {
		this.app = app;
	}

	/**
	 * Collect all files that have any of the specified tags
	 */
	collectFilesWithTags(
		tags: string[],
		excludePaths: string[] = []
	): TFile[] {
		const files: TFile[] = [];
		const normalizedTags = tags.map(t => t.startsWith('#') ? t.substring(1) : t);

		for (const file of this.app.vault.getMarkdownFiles()) {
			// Check excluded paths
			let excluded = false;
			for (const excludePath of excludePaths) {
				if (file.path.startsWith(excludePath) || file.path.includes(`/${excludePath}`)) {
					excluded = true;
					break;
				}
			}
			if (excluded) continue;

			// Check tags
			const cache = this.app.metadataCache.getFileCache(file);
			if (!cache) continue;

			// Check frontmatter tags
			const frontmatterTags = this.getFrontmatterTags(cache);
			// Check inline tags
			const inlineTags = cache.tags?.map(t => t.tag.substring(1)) || [];

			const allTags = [...frontmatterTags, ...inlineTags];

			for (const tag of allTags) {
				// Support hierarchical tag matching
				for (const searchTag of normalizedTags) {
					if (tag === searchTag || tag.startsWith(`${searchTag}/`)) {
						files.push(file);
						break;
					}
				}
			}
		}

		return files;
	}

	private getFrontmatterTags(cache: CachedMetadata): string[] {
		if (!cache.frontmatter?.tags) {
			return [];
		}

		const tags = cache.frontmatter.tags;
		if (Array.isArray(tags)) {
			return tags.map(t => typeof t === 'string' ? t : String(t));
		}
		if (typeof tags === 'string') {
			return [tags];
		}
		return [];
	}
}

// === Preset Executor ===

export interface PresetExecutionResult {
	files: TFile[];
	linkedFiles: LinkedNoteResult[];
	taggedFiles: TFile[];
	truncated: boolean;
	totalTokens: number;
}

export class PresetExecutor {
	private app: App;
	private linkTraverser: LinkTraverser;
	private tagCollector: TagCollector;

	constructor(app: App) {
		this.app = app;
		this.linkTraverser = new LinkTraverser(app);
		this.tagCollector = new TagCollector(app);
	}

	/**
	 * Execute a preset and collect all files to include
	 */
	async execute(
		activeFile: TFile,
		preset: ContextPreset,
		contextFiles: TFile[]
	): Promise<PresetExecutionResult> {
		const allFiles = new Set<TFile>(contextFiles);
		let linkedFiles: LinkedNoteResult[] = [];
		let taggedFiles: TFile[] = [];

		// Include linked notes
		if (preset.includeLinked) {
			const depth = preset.linkDepth ?? 1;
			linkedFiles = await this.linkTraverser.traverseLinks(
				activeFile,
				depth,
				preset.excludePaths || [],
				preset.excludeTags || []
			);

			for (const result of linkedFiles) {
				allFiles.add(result.file);
			}
		}

		// Include notes with specified tags
		if (preset.includeTags && preset.includeTags.length > 0) {
			taggedFiles = this.tagCollector.collectFilesWithTags(
				preset.includeTags,
				preset.excludePaths || []
			);

			for (const file of taggedFiles) {
				allFiles.add(file);
			}
		}

		// Convert to array and apply exclusions
		let files = Array.from(allFiles);

		// Apply path exclusions
		if (preset.excludePaths && preset.excludePaths.length > 0) {
			files = files.filter(f => {
				for (const excludePath of preset.excludePaths!) {
					if (f.path.startsWith(excludePath) || f.path.includes(`/${excludePath}`)) {
						return false;
					}
				}
				return true;
			});
		}

		// Apply tag exclusions
		if (preset.excludeTags && preset.excludeTags.length > 0) {
			files = files.filter(f => {
				const cache = this.app.metadataCache.getFileCache(f);
				if (!cache?.frontmatter?.tags) return true;

				const fileTags: string[] = Array.isArray(cache.frontmatter.tags)
					? cache.frontmatter.tags
					: [cache.frontmatter.tags];

				for (const fileTag of fileTags) {
					const normalizedTag = fileTag.startsWith('#') ? fileTag : `#${fileTag}`;
					if (preset.excludeTags!.includes(normalizedTag)) {
						return false;
					}
				}
				return true;
			});
		}

		// Check token budget
		let totalTokens = 0;
		let truncated = false;

		if (preset.maxTokens) {
			const filesWithTokens: { file: TFile; tokens: number }[] = [];

			for (const file of files) {
				const content = await this.app.vault.read(file);
				const tokens = estimateTokens(content);
				filesWithTokens.push({ file, tokens });
			}

			// Sort by tokens (smallest first) to maximize included files
			filesWithTokens.sort((a, b) => a.tokens - b.tokens);

			const includedFiles: TFile[] = [];
			for (const { file, tokens } of filesWithTokens) {
				if (totalTokens + tokens <= preset.maxTokens) {
					includedFiles.push(file);
					totalTokens += tokens;
				} else {
					truncated = true;
				}
			}

			files = includedFiles;
		} else {
			// Calculate total tokens without truncation
			for (const file of files) {
				const content = await this.app.vault.read(file);
				totalTokens += estimateTokens(content);
			}
		}

		return {
			files,
			linkedFiles,
			taggedFiles,
			truncated,
			totalTokens,
		};
	}
}

// === Helper ===

export function formatPresetErrors(result: PresetValidationResult): string {
	let message = '';

	if (result.errors.length > 0) {
		message += 'Errors:\n' + result.errors.map(e => `  • ${e}`).join('\n');
	}

	if (result.warnings.length > 0) {
		if (message) message += '\n\n';
		message += 'Warnings:\n' + result.warnings.map(w => `  • ${w}`).join('\n');
	}

	return message;
}