/**
 * Vector similarity search over stored snippet embeddings.
 *
 * SQLite does not natively support vector operations, so cosine similarity is
 * computed in JavaScript after loading candidate embeddings from the
 * snippet_embeddings table.
 *
 * Performance note: For repositories with > 50k snippets, pre-filtering by
 * FTS5 candidates before computing cosine similarity is recommended. For v1,
 * in-memory computation is acceptable.
 */

import type Database from 'better-sqlite3';

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

export interface VectorSearchResult {
	snippetId: string;
	score: number;
}

/** Raw DB row from snippet_embeddings joined with snippets. */
interface RawEmbeddingRow {
	snippet_id: string;
	embedding: Buffer;
}

// ---------------------------------------------------------------------------
// Math helpers
// ---------------------------------------------------------------------------

/**
 * Compute cosine similarity between two Float32Array vectors.
 *
 * Returns a value in [-1, 1] where 1 is identical direction. Returns 0 when
 * either vector has zero magnitude to avoid division by zero.
 */
export function cosineSimilarity(a: Float32Array, b: Float32Array): number {
	if (a.length !== b.length) {
		throw new Error(
			`Embedding dimension mismatch: ${a.length} vs ${b.length}`
		);
	}

	let dot = 0;
	let normA = 0;
	let normB = 0;

	for (let i = 0; i < a.length; i++) {
		dot += a[i] * b[i];
		normA += a[i] * a[i];
		normB += b[i] * b[i];
	}

	const denom = Math.sqrt(normA) * Math.sqrt(normB);
	if (denom === 0) return 0;
	return dot / denom;
}

// ---------------------------------------------------------------------------
// VectorSearch class
// ---------------------------------------------------------------------------

export class VectorSearch {
	private readonly stmt: Database.Statement<[string], RawEmbeddingRow>;

	constructor(private readonly db: Database.Database) {
		// Prepare once — reused for every call.
		this.stmt = this.db.prepare<[string], RawEmbeddingRow>(`
			SELECT se.snippet_id, se.embedding
			FROM snippet_embeddings se
			JOIN snippets s ON s.id = se.snippet_id
			WHERE s.repository_id = ?
		`);
	}

	/**
	 * Search stored embeddings by cosine similarity to the query embedding.
	 *
	 * @param queryEmbedding - The embedded representation of the search query.
	 * @param repositoryId   - Scope the search to a single repository.
	 * @param limit          - Maximum number of results to return. Default: 50.
	 * @returns Results sorted by descending cosine similarity score.
	 */
	vectorSearch(
		queryEmbedding: Float32Array,
		repositoryId: string,
		limit = 50
	): VectorSearchResult[] {
		const rows = this.stmt.all(repositoryId);

		const scored: VectorSearchResult[] = rows.map((row) => {
			const embedding = new Float32Array(
				row.embedding.buffer,
				row.embedding.byteOffset,
				row.embedding.byteLength / 4
			);
			return {
				snippetId: row.snippet_id,
				score: cosineSimilarity(queryEmbedding, embedding)
			};
		});

		return scored.sort((a, b) => b.score - a.score).slice(0, limit);
	}
}
