// SF Memory Embeddings — embedding, reranking, and vector ranking
//
// OpenAI-shaped /v1/embeddings and /v1/rerank against a configurable
// gateway endpoint (opt-in via SF_LLM_GATEWAY_KEY). All paths degrade
// gracefully to keyword-only ranking when no gateway is configured.
//
// Read pipeline used by `getRelevantMemoriesRanked` (memory-store.js):
//   1. cosine ranking via rankMemoriesByEmbedding (uses loadEmbeddingMap)
//   2. relation-boost via applyRelationBoost (uses memory_relations edges)
//   3. optional rerank via rerankCandidates
//
// When no gateway is configured (or the worker is offline), all three
// pipeline stages soft-degrade and `getRelevantMemoriesRanked` falls back
// to static (confidence × hit_count) ranking.
import { existsSync, readFileSync } from "node:fs";
import { join } from "node:path";
import {
	_getAdapter,
	deleteMemoryEmbedding,
	isDbAvailable,
	upsertMemoryEmbedding,
} from "./sf-db.js";
import { sfHome } from "./sf-home.js";
import { logWarning } from "./workflow-logger.js";

/** Read the llm-gateway entry from ~/.sf/agent/auth.json if present.
 *  Falls back to SF_LLM_GATEWAY_KEY env var when auth.json is missing
 *  or has no key — enables CI and containerised deployments without
 *  writing secrets to disk. */
function readGatewayFromAuthJson() {
	try {
		const authPath = join(sfHome(), "agent", "auth.json");
		if (existsSync(authPath)) {
			const data = JSON.parse(readFileSync(authPath, "utf8"));
			const entry = data["llm-gateway"];
			if (entry?.key) {
				return {
					key: entry.key,
					url: entry.url || null,
					source: "auth.json:llm-gateway",
				};
			}
		}
	} catch {
		/* fall through to env */
	}
	const envKey = process.env.SF_LLM_GATEWAY_KEY?.trim();
	if (envKey) {
		return {
			key: envKey,
			url: process.env.SF_LLM_GATEWAY_URL?.trim() || null,
			source: "env:SF_LLM_GATEWAY_KEY",
		};
	}
	return null;
}

// ─── Gateway config ──────────────────────────────────────────────────────────
const DEFAULT_TIMEOUT_MS = 30_000;
const ENV_EMBED_MODEL = "SF_LLM_GATEWAY_EMBED_MODEL";
const ENV_RERANK_MODEL = "SF_LLM_GATEWAY_RERANK_MODEL";
const ENV_EMBED_QUERY_INSTRUCTION = "SF_LLM_GATEWAY_EMBED_QUERY_INSTRUCTION";
const DEFAULT_EMBEDDING_MODEL = "Qwen/Qwen3-Embedding-4B";
const DEFAULT_RERANK_MODEL = "Qwen/Qwen3-Reranker-0.6B";
const DEFAULT_QUERY_INSTRUCTION =
	"Instruct: Retrieve relevant software engineering memories, facts, and project decisions for the given query\nQuery: ";
function firstEnvValue(key) {
	return process.env[key]?.trim() || "";
}
/** Read gateway config from auth.json only. Returns null when not configured — opt-in. */
export function loadGatewayConfigFromEnv() {
	const fromAuth = readGatewayFromAuthJson();
	if (!fromAuth) return null;
	const embeddingModel =
		firstEnvValue(ENV_EMBED_MODEL) || DEFAULT_EMBEDDING_MODEL;
	const rerankModel = firstEnvValue(ENV_RERANK_MODEL) || DEFAULT_RERANK_MODEL;
	const queryInstruction =
		firstEnvValue(ENV_EMBED_QUERY_INSTRUCTION) || DEFAULT_QUERY_INSTRUCTION;
	return {
		url: fromAuth.url ?? "https://llm-gateway.centralcloud.com/v1",
		apiKey: fromAuth.key,
		keySource: fromAuth.source ?? "auth.json:llm-gateway",
		urlSource: fromAuth.url
			? (fromAuth.source ?? "auth.json:llm-gateway")
			: "default",
		embeddingModel,
		rerankModel,
		queryInstruction,
	};
}

// ─── Circuit breaker ─────────────────────────────────────────────────────────
const EMBED_CIRCUIT_THRESHOLD = 3;
const EMBED_CIRCUIT_OPEN_MS = 60_000;
const embedCircuit = { failures: 0, openUntil: 0, lastLogAt: 0 };
function embedCircuitIsOpen() {
	return embedCircuit.openUntil > Date.now();
}
function onEmbedSuccess() {
	embedCircuit.failures = 0;
	embedCircuit.openUntil = 0;
}
function onEmbedFailure() {
	embedCircuit.failures += 1;
	if (embedCircuit.failures >= EMBED_CIRCUIT_THRESHOLD) {
		embedCircuit.openUntil = Date.now() + EMBED_CIRCUIT_OPEN_MS;
		const now = Date.now();
		if (now - embedCircuit.lastLogAt >= EMBED_CIRCUIT_OPEN_MS) {
			embedCircuit.lastLogAt = now;
			logWarning(
				"memory-embeddings",
				`gateway /embeddings circuit open after ${EMBED_CIRCUIT_THRESHOLD} failures; ` +
					`skipping for ${EMBED_CIRCUIT_OPEN_MS / 1000}s — falling back to keyword ranking`,
			);
		}
	}
}

// ─── Rerank throttle ─────────────────────────────────────────────────────────
let lastRerankUnavailableLogAt = 0;
const RERANK_UNAVAILABLE_LOG_THROTTLE_MS = 60_000;
function logRerankUnavailable(msg) {
	const now = Date.now();
	if (now - lastRerankUnavailableLogAt < RERANK_UNAVAILABLE_LOG_THROTTLE_MS)
		return;
	lastRerankUnavailableLogAt = now;
	logWarning("memory-embeddings", msg);
}

// ─── Gateway HTTP ─────────────────────────────────────────────────────────────
/** Build an embed function that posts to <url>/embeddings with Bearer auth. */
export function createGatewayEmbedFn(config, opts) {
	return async (texts) => {
		if (texts.length === 0) return [];
		if (embedCircuitIsOpen()) return [];
		const controller = new AbortController();
		const timeout = setTimeout(
			() => controller.abort(),
			config.timeoutMs ?? DEFAULT_TIMEOUT_MS,
		);
		try {
			const body = { model: config.embeddingModel, input: texts };
			if (opts?.instruction) body.instruction = opts.instruction;
			const res = await fetch(`${config.url}/embeddings`, {
				method: "POST",
				headers: {
					Authorization: `Bearer ${config.apiKey}`,
					"Content-Type": "application/json",
				},
				body: JSON.stringify(body),
				signal: controller.signal,
			});
			if (!res.ok) {
				const bodyText = await res.text().catch(() => "");
				throw new Error(
					`gateway /embeddings ${res.status}: ${bodyText.slice(0, 200)}`,
				);
			}
			const json = await res.json();
			if (!Array.isArray(json.data))
				throw new Error("gateway /embeddings: missing data array");
			const sorted = [...json.data].sort((a, b) => a.index - b.index);
			const result = sorted.map((d) => Float32Array.from(d.embedding));
			onEmbedSuccess();
			return result;
		} catch (err) {
			onEmbedFailure();
			throw err;
		} finally {
			clearTimeout(timeout);
		}
	};
}

/** Score candidates against a query via /rerank. Returns null when unavailable. */
export async function rerankCandidates(config, query, candidates) {
	if (!config.rerankModel) return null;
	if (candidates.length === 0) return [];
	const controller = new AbortController();
	const timeout = setTimeout(
		() => controller.abort(),
		config.timeoutMs ?? DEFAULT_TIMEOUT_MS,
	);
	try {
		const res = await fetch(`${config.url}/rerank`, {
			method: "POST",
			headers: {
				Authorization: `Bearer ${config.apiKey}`,
				"Content-Type": "application/json",
			},
			body: JSON.stringify({
				model: config.rerankModel,
				query,
				documents: candidates.map((c) => c.text),
			}),
			signal: controller.signal,
		});
		if (res.status === 503 || res.status === 404) {
			logRerankUnavailable(
				`gateway /rerank unavailable (${res.status}); falling back to cosine order`,
			);
			return null;
		}
		const bodyText = await res.text().catch(() => "");
		if (/no worker.*rerank/i.test(bodyText)) {
			logRerankUnavailable("gateway /rerank: no worker capability available");
			return null;
		}
		if (!res.ok)
			throw new Error(
				`gateway /rerank ${res.status}: ${bodyText.slice(0, 200)}`,
			);
		let json;
		try {
			json = JSON.parse(bodyText);
		} catch {
			throw new Error(
				`gateway /rerank: malformed JSON (${bodyText.slice(0, 200)})`,
			);
		}
		if (!Array.isArray(json.results))
			throw new Error("gateway /rerank: missing results array");
		return json.results.map((r) => ({
			id: candidates[r.index]?.id ?? String(r.index),
			score: r.relevance_score,
		}));
	} finally {
		clearTimeout(timeout);
	}
}

// ─── Vector (de)serialization ───────────────────────────────────────────────
export function packFloat32(vec) {
	return new Uint8Array(vec.buffer, vec.byteOffset, vec.byteLength);
}
export function unpackFloat32(blob) {
	if (!blob) return null;
	try {
		if (blob instanceof Float32Array) return blob;
		let view;
		if (blob instanceof Uint8Array) {
			view = blob;
		} else if (blob instanceof ArrayBuffer) {
			view = new Uint8Array(blob);
		} else if (blob.buffer && blob.byteLength != null) {
			const buf = blob;
			view = new Uint8Array(buf.buffer, buf.byteOffset, buf.byteLength);
		} else if (Array.isArray(blob)) {
			return new Float32Array(blob);
		} else {
			return null;
		}
		if (view.byteLength % 4 !== 0) return null;
		// Copy into an aligned buffer — BLOBs may arrive at odd byte offsets.
		const aligned = new ArrayBuffer(view.byteLength);
		new Uint8Array(aligned).set(view);
		return new Float32Array(aligned);
	} catch {
		return null;
	}
}
// ─── Math ───────────────────────────────────────────────────────────────────
export function cosineSimilarity(a, b) {
	if (a.length === 0 || a.length !== b.length) return 0;
	let dot = 0;
	let normA = 0;
	let normB = 0;
	for (let i = 0; i < a.length; i++) {
		const x = a[i];
		const y = b[i];
		dot += x * y;
		normA += x * x;
		normB += y * y;
	}
	if (normA === 0 || normB === 0) return 0;
	return dot / (Math.sqrt(normA) * Math.sqrt(normB));
}
// ─── Read helpers ───────────────────────────────────────────────────────────
export function getEmbeddingForMemory(memoryId) {
	if (!isDbAvailable()) return null;
	const adapter = _getAdapter();
	if (!adapter) return null;
	try {
		const row = adapter
			.prepare(
				"SELECT memory_id, model, dim, vector FROM memory_embeddings WHERE memory_id = :id",
			)
			.get({ ":id": memoryId });
		if (!row) return null;
		const vector = unpackFloat32(row["vector"]);
		if (!vector) return null;
		return {
			memoryId: row["memory_id"],
			model: row["model"],
			dim: row["dim"],
			vector,
		};
	} catch {
		return null;
	}
}
export function loadAllEmbeddings() {
	if (!isDbAvailable()) return [];
	const adapter = _getAdapter();
	if (!adapter) return [];
	try {
		const rows = adapter
			.prepare(`SELECT e.memory_id, e.model, e.dim, e.vector
         FROM memory_embeddings e
         JOIN memories m ON m.id = e.memory_id
         WHERE m.superseded_by IS NULL`)
			.all();
		const out = [];
		for (const row of rows) {
			const vector = unpackFloat32(row["vector"]);
			if (!vector) continue;
			out.push({
				memoryId: row["memory_id"],
				model: row["model"],
				dim: row["dim"],
				vector,
			});
		}
		return out;
	} catch {
		return [];
	}
}
// ─── Write helpers ──────────────────────────────────────────────────────────
export function saveEmbedding(memoryId, vector, model) {
	if (!isDbAvailable()) return false;
	try {
		upsertMemoryEmbedding({
			memoryId,
			model,
			dim: vector.length,
			vector: packFloat32(vector),
			updatedAt: new Date().toISOString(),
		});
		return true;
	} catch {
		return false;
	}
}
export function removeEmbedding(memoryId) {
	if (!isDbAvailable()) return false;
	try {
		return deleteMemoryEmbedding(memoryId);
	} catch {
		return false;
	}
}
// ─── Orchestration ──────────────────────────────────────────────────────────
/**
 * Embed each memory's content via `embedFn` (if provided) and persist the
 * resulting vectors. Returns the number of successfully embedded memories.
 * Safe to call with embedFn=null — it becomes a no-op.
 */
export async function embedMemories(memories, embedFn, model) {
	if (!embedFn || memories.length === 0) return 0;
	try {
		const vectors = await embedFn(memories.map((m) => m.content));
		let count = 0;
		for (let i = 0; i < memories.length && i < vectors.length; i++) {
			const vector = vectors[i];
			if (!vector || vector.length === 0) continue;
			if (saveEmbedding(memories[i].id, vector, model)) count++;
		}
		return count;
	} catch (err) {
		logWarning("memory-embeddings", `embed failed: ${err.message}`);
		return 0;
	}
}
// ─── Query-aware ranking ──────────────────────────────────────────────────
/** Combine static rank (confidence × hit-count boost) with semantic similarity
 *  to the embedded query. When no embedding rows or no query embedding are
 *  available, returns the input list unchanged so callers can rely on the
 *  static order as a fallback. */
export function rankMemoriesByEmbedding(
	memories,
	queryVector,
	memoryEmbeddings,
	options,
) {
	const alpha = options?.semanticWeight ?? 0.6;
	if (!queryVector || memoryEmbeddings.size === 0) {
		return memories.map((m) => ({
			id: m.id,
			combinedScore: m.staticScore,
			cosine: 0,
		}));
	}
	return memories
		.map((m) => {
			const vec = memoryEmbeddings.get(m.id);
			const cosine = vec ? cosineSimilarity(queryVector, vec) : 0;
			// Static score in [0, ~1.5+]; cosine in [-1, 1] but typically [0, 1].
			// Blend so a perfect static + perfect cosine ≈ 2× a static-only top hit,
			// and a memory with no embedding at all still ranks by static score.
			const combinedScore = m.staticScore * (1 + alpha * cosine);
			return { id: m.id, combinedScore, cosine };
		})
		.sort((a, b) => b.combinedScore - a.combinedScore);
}
/** Embed `query` via the configured gateway and return its Float32Array, or
 *  null when no gateway is configured / the embed call fails. Best-effort.
 *
 *  Passes `config.queryInstruction` so Qwen3-Embedding projects the query into
 *  the correct asymmetric retrieval region of the embedding space. Document
 *  embeddings (backfill) are created without an instruction — the two must be
 *  consistent for cosine similarity to be meaningful. */
export async function embedQueryViaGateway(query) {
	if (!query.trim()) return null;
	try {
		const cfg = loadGatewayConfigFromEnv();
		if (!cfg) return null;
		const embedFn = createGatewayEmbedFn(cfg, {
			instruction: cfg.queryInstruction,
		});
		const vectors = await embedFn([query]);
		return vectors[0] ?? null;
	} catch (err) {
		logWarning("memory-embeddings", `query embed failed: ${err.message}`);
		return null;
	}
}
/** Map of memoryId → vector for all active memories that have embeddings.
 *  Used by the ranker to look up vectors without N+1 queries. */
export function loadEmbeddingMap() {
	const map = new Map();
	for (const row of loadAllEmbeddings()) {
		map.set(row.memoryId, row.vector);
	}
	return map;
}
/** Apply a one-pass graph-boost to an embedding-ranked list. For each
 *  edge (A, B) where both endpoints are in the pool, B's combined score
 *  gains `score(A) × edge.confidence × damping`. Same for the reverse
 *  direction (relations are symmetric for ranking purposes — if A pulls
 *  B, B should also pull A by the same amount). The boost is additive
 *  and capped at one iteration so an N-memory pool never explodes
 *  beyond O(N + |edges|) work.
 *
 *  When the pool has no edges (the empty-graph case while no writers
 *  exist), returns the input unchanged. */
export function applyRelationBoost(ranked, edges, options) {
	const damping = options?.damping ?? 0.4;
	if (edges.length === 0 || ranked.length === 0) return ranked;
	const scoreById = new Map(ranked.map((r) => [r.id, r.combinedScore]));
	const boosts = new Map();
	const bumpScore = (id, contribution) => {
		boosts.set(id, (boosts.get(id) ?? 0) + contribution);
	};
	for (const edge of edges) {
		const fromScore = scoreById.get(edge.from);
		const toScore = scoreById.get(edge.to);
		// Only boost edges where BOTH endpoints are in the pool — keeps the
		// pass intra-pool and prevents a hot edge from pulling in unrelated
		// low-static memories that would otherwise never reach the cosine
		// pool. Cross-pool edges may matter eventually but require a
		// separate "expand pool by relation" step that's out of scope here.
		if (fromScore == null || toScore == null) continue;
		bumpScore(edge.to, fromScore * edge.confidence * damping);
		bumpScore(edge.from, toScore * edge.confidence * damping);
	}
	return ranked
		.map((r) => ({
			...r,
			combinedScore: r.combinedScore + (boosts.get(r.id) ?? 0),
		}))
		.sort((a, b) => b.combinedScore - a.combinedScore);
}
// ─── Auto-engagement / backfill driver ────────────────────────────────────
/** Find active memories (not superseded) that don't yet have an embedding row.
 *  Used by the backfill driver to know what to embed next. */
export function listUnembeddedMemoryIds(limit = 50) {
	if (!isDbAvailable()) return [];
	const adapter = _getAdapter();
	if (!adapter) return [];
	try {
		const rows = adapter
			.prepare(`SELECT m.id, m.content
         FROM memories m
         LEFT JOIN memory_embeddings e ON e.memory_id = m.id
         WHERE m.superseded_by IS NULL AND e.memory_id IS NULL
         ORDER BY m.seq ASC
         LIMIT :lim`)
			.all({ ":lim": limit });
		return rows.map((r) => ({
			id: r["id"],
			content: r["content"],
		}));
	} catch {
		return [];
	}
}
let backfillInFlight = false;
let lastUnavailableLogAt = 0;
/** Best-effort embedding backfill. Probes the gateway by attempting to embed
 *  the first unembedded batch — on success, persists vectors and continues
 *  until either the limit is reached or the queue is empty; on failure (no
 *  worker, network error, missing config), logs once-per-minute and returns
 *  zero so callers can keep firing without spam.
 *
 *  Safe to call from a hook on every turn — guarded against re-entry via an
 *  in-flight flag and bounded by `maxPerInvocation`. */
export async function runEmbeddingBackfill(opts) {
	if (backfillInFlight) return 0;
	const max = opts?.maxPerInvocation ?? 50;
	const batchSize = opts?.batchSize ?? 16;
	const cfg = loadGatewayConfigFromEnv();
	if (!cfg) return 0;
	backfillInFlight = true;
	let embedded = 0;
	try {
		const embedFn = createGatewayEmbedFn(cfg);
		while (embedded < max) {
			const batch = listUnembeddedMemoryIds(
				Math.min(batchSize, max - embedded),
			);
			if (batch.length === 0) break;
			let count = 0;
			try {
				count = await embedMemories(batch, embedFn, cfg.embeddingModel);
			} catch (err) {
				// Throttle "unavailable" log to once per minute so we don't spam
				// a journal when the gateway worker is offline.
				const now = Date.now();
				if (now - lastUnavailableLogAt > 60_000) {
					logWarning(
						"memory-embeddings",
						`backfill: gateway embed failed (${err.message}); will retry next turn`,
					);
					lastUnavailableLogAt = now;
				}
				return embedded;
			}
			if (count === 0) break; // Stop early to avoid loops on transient saves
			embedded += count;
		}
		// Successes are routine and fire on every agent_end — silent. The
		// caller (`runEmbeddingBackfill` consumer) gets the count via the
		// return value; failures still log via the throttled path above.
		return embedded;
	} finally {
		backfillInFlight = false;
	}
}
