/**
 * Embedding provider for clankie-memory
 * Supports local CPU embeddings (HuggingFace Transformers) and OpenAI-compatible APIs
 */

import { createHash } from "node:crypto";
import type { EmbeddingConfig } from "./types.ts";

export interface EmbeddingProvider {
	embed(texts: string[]): Promise<number[][]>;
}

// Cached local embedder pipeline
let localEmbedder: ((text: string) => Promise<number[]>) | null = null;

interface CacheEntry {
	embedding: number[];
	timestamp: number;
}

// Simple LRU cache for embeddings
class EmbeddingCache {
	private cache = new Map<string, CacheEntry>();
	private maxSize: number;

	constructor(maxSize = 1000) {
		this.maxSize = maxSize;
	}

	private getHash(text: string): string {
		return createHash("sha256").update(text).digest("hex");
	}

	get(text: string): number[] | undefined {
		const hash = this.getHash(text);
		const entry = this.cache.get(hash);
		if (entry) {
			// Move to front (LRU)
			this.cache.delete(hash);
			this.cache.set(hash, entry);
			return entry.embedding;
		}
		return undefined;
	}

	set(text: string, embedding: number[]): void {
		const hash = this.getHash(text);

		// Evict oldest if at capacity
		if (this.cache.size >= this.maxSize && !this.cache.has(hash)) {
			const firstKey = this.cache.keys().next().value;
			if (firstKey) {
				this.cache.delete(firstKey);
			}
		}

		this.cache.set(hash, { embedding, timestamp: Date.now() });
	}

	clear(): void {
		this.cache.clear();
	}
}

// Global cache instance
const globalCache = new EmbeddingCache(2000);

/** Null provider for text-only mode (no embeddings) */
class NullEmbeddingProvider implements EmbeddingProvider {
	async embed(): Promise<number[][]> {
		throw new Error(
			"No embedding provider configured. " +
				"Set embedding.provider to 'local', 'openai', or 'ollama' in config."
		);
	}
}

/**
 * Create a local embedder using HuggingFace Transformers (ONNX, runs on CPU)
 */
async function createLocalEmbedder(config: EmbeddingConfig): Promise<(text: string) => Promise<number[]>> {
	if (localEmbedder) return localEmbedder;

	const { pipeline } = await import("@huggingface/transformers");
	const cacheDir = config.cacheDir ?? process.env.MEMORY_MODEL_CACHE_DIR;

	// Try user-provided model first, then known-compatible fallbacks
	const candidates = [
		config.model || "Xenova/all-MiniLM-L6-v2",
		"Xenova/all-MiniLM-L6-v2",
		"onnx-community/all-MiniLM-L6-v2",
	].filter((m, i, arr) => arr.indexOf(m) === i);

	let lastError: Error | undefined;
	for (const model of candidates) {
		try {
			console.log(`[memory] Loading local embedding model: ${model}${cacheDir ? ` (cache: ${cacheDir})` : ""}`);
			const extractor = await pipeline("feature-extraction", model, {
				cache_dir: cacheDir,
				local_files_only: false,
			});
			console.log(`[memory] Local embedding model loaded: ${model}`);

			localEmbedder = async (text: string): Promise<number[]> => {
				const output = await extractor(text, { pooling: "mean", normalize: true });
				return Array.from(output.data as Float32Array);
			};

			return localEmbedder;
		} catch (error) {
			lastError = error instanceof Error ? error : new Error(String(error));
			console.warn(`[memory] Failed to load local model '${model}': ${lastError.message}`);
		}
	}

	throw new Error(
		`Failed to load local embedding model. Tried: ${candidates.join(", ")}. ` +
			`Last error: ${lastError?.message ?? "unknown"}`,
	);
}

/**
 * Create an embedding provider from config
 * Returns NullEmbeddingProvider if no provider configured (text-only mode)
 */
export function createEmbeddingProvider(
	config: EmbeddingConfig,
	getApiKey?: (provider: string) => string | undefined,
): EmbeddingProvider {
	// Allow env override for quick debugging
	const provider = (process.env.MEMORY_EMBEDDING_PROVIDER as EmbeddingConfig["provider"]) || config.provider;

	// No provider configured = text-only mode
	if (!provider) {
		return new NullEmbeddingProvider();
	}

	// Local provider uses HuggingFace Transformers
	if (provider === "local") {
		let embedderPromise: Promise<(text: string) => Promise<number[]>> | null = null;

		return {
			async embed(texts: string[]): Promise<number[][]> {
				// Lazy-load the embedder
				if (!embedderPromise) {
					embedderPromise = createLocalEmbedder(config);
				}
				const embedder = await embedderPromise;

				// Check cache first
				const results: (number[] | undefined)[] = texts.map((t) => globalCache.get(t));
				const uncachedIndices: number[] = [];

				for (let i = 0; i < texts.length; i++) {
					if (results[i] === undefined) {
						uncachedIndices.push(i);
					}
				}

				// Generate uncached embeddings
				for (const idx of uncachedIndices) {
					const embedding = await embedder(texts[idx]);
					globalCache.set(texts[idx], embedding);
					results[idx] = embedding;
				}

				return results as number[][];
			},
		};
	}

	// OpenAI/Ollama providers use API
	return {
		async embed(texts: string[]): Promise<number[][]> {
			// Check cache first
			const results: (number[] | undefined)[] = texts.map((t) => globalCache.get(t));
			const uncachedIndices: number[] = [];
			const uncachedTexts: string[] = [];

			for (let i = 0; i < texts.length; i++) {
				if (results[i] === undefined) {
					uncachedIndices.push(i);
					uncachedTexts.push(texts[i]);
				}
			}

			// Fetch uncached embeddings
			if (uncachedTexts.length > 0) {
				const embeddings = await fetchEmbeddings(uncachedTexts, config, getApiKey);

				// Store in cache and fill results
				for (let i = 0; i < uncachedIndices.length; i++) {
					const idx = uncachedIndices[i];
					const embedding = embeddings[i];
					globalCache.set(texts[idx], embedding);
					results[idx] = embedding;
				}
			}

			return results as number[][];
		},
	};
}

/**
 * Fetch embeddings from API with retry
 */
async function fetchEmbeddings(
	texts: string[],
	config: EmbeddingConfig,
	getApiKey?: (provider: string) => string | undefined,
): Promise<number[][]> {
	const baseUrl = config.baseUrl ?? getDefaultBaseUrl(config.provider);
	const apiKey = resolveApiKey(config, getApiKey);

	const url = `${baseUrl}/v1/embeddings`;
	const body = {
		model: config.model,
		input: texts,
	};

	let lastError: Error | undefined;

	// Try up to 2 times
	for (let attempt = 0; attempt < 2; attempt++) {
		try {
			const headers: Record<string, string> = {
				"Content-Type": "application/json",
			};
			if (apiKey) {
				headers.Authorization = `Bearer ${apiKey}`;
			}

			const response = await fetch(url, {
				method: "POST",
				headers,
				body: JSON.stringify(body),
			});

			if (!response.ok) {
				const errorText = await response.text();
				throw new Error(`Embedding API error: ${response.status} ${errorText}`);
			}

			const data = await response.json();

			// Handle different response formats
			if (data.data && Array.isArray(data.data)) {
				// OpenAI format
				return data.data
					.sort((a: { index: number }, b: { index: number }) => a.index - b.index)
					.map((item: { embedding: number[] }) => item.embedding);
			}
			if (Array.isArray(data.embeddings)) {
				// Alternative format
				return data.embeddings;
			}
			if (Array.isArray(data)) {
				// Direct array
				return data;
			}

			throw new Error("Unexpected embedding API response format");
		} catch (error) {
			lastError = error instanceof Error ? error : new Error(String(error));

			// Wait before retry (exponential backoff)
			if (attempt === 0) {
				await new Promise((resolve) => setTimeout(resolve, 1000));
			}
		}
	}

	throw lastError ?? new Error("Failed to fetch embeddings");
}

function getDefaultBaseUrl(provider: string | null | undefined): string {
	switch (provider) {
		case "openai":
			return "https://api.openai.com";
		case "ollama":
			return "http://localhost:11434";
		default:
			return "https://api.openai.com";
	}
}

function resolveApiKey(
	config: EmbeddingConfig,
	getApiKey?: (provider: string) => string | undefined,
): string | undefined {
	// 1. Config literal
	if (config.apiKey && !config.apiKey.startsWith("$")) {
		return config.apiKey;
	}

	// 2. Environment variable reference ($ENV_VAR)
	if (config.apiKey?.startsWith("$")) {
		const envVar = config.apiKey.substring(1);
		return process.env[envVar];
	}

	// 3. From pi's model registry
	if (getApiKey && config.provider) {
		return getApiKey(config.provider);
	}

	// 4. Default env vars
	if (config.provider === "openai") {
		return process.env.OPENAI_API_KEY;
	}

	return undefined;
}

/**
 * Clear the global embedding cache and unload local model
 */
export function clearEmbeddingCache(): void {
	globalCache.clear();
	localEmbedder = null;
}