// Embedding provider for AgenTeam memory system
// Uses @huggingface/transformers for local embeddings (ONNX WASM)

import { pipeline, env } from "@huggingface/transformers"
import type { FeatureExtractionPipeline } from "@huggingface/transformers"
import { debugLog } from "../state-io.ts"

// Configure for local/WASM usage
env.allowLocalModels = true

// Singleton cache — one pipeline instance per model per process
let _cachedPipeline: FeatureExtractionPipeline | null = null
let _cachedModelName: string | null = null
let _pipelinePromise: Promise<FeatureExtractionPipeline> | null = null

/**
 * Lazily initialize and cache the embedding pipeline.
 * Only one instance ever exists per process. The pipeline is NOT loaded
 * at import time — it loads on first call.
 */
export async function getEmbedder(modelName: string): Promise<FeatureExtractionPipeline> {
  // Return cached instance if model matches
  if (_cachedPipeline && _cachedModelName === modelName) {
    return _cachedPipeline
  }

  // If a load is already in progress for this model, await it
  if (_pipelinePromise && _cachedModelName === modelName) {
    return _pipelinePromise
  }

  // If model changed, discard old pipeline
  _cachedPipeline = null
  _cachedModelName = modelName

  debugLog("system", "memory.embedding_loading", {
    model: modelName,
    note: "First run may download model (~22MB)",
  })

  _pipelinePromise = (async () => {
    try {
      const extractor = await pipeline("feature-extraction", modelName, {
        dtype: "fp32",
      }) as FeatureExtractionPipeline

      _cachedPipeline = extractor
      debugLog("system", "memory.embedding_loaded", { model: modelName })
      return extractor
    } catch (err) {
      // Reset so subsequent calls retry
      _cachedPipeline = null
      _cachedModelName = null
      _pipelinePromise = null

      throw new Error(
        `Failed to load embedding model "${modelName}": ${err instanceof Error ? err.message : String(err)}`
      )
    }
  })()

  return _pipelinePromise
}

/**
 * Compute embedding vector for a single text string.
 * Uses mean pooling + L2 normalization.
 */
export async function embed(text: string, modelName: string): Promise<number[]> {
  const extractor = await getEmbedder(modelName)

  try {
    const output = await extractor(text, {
      pooling: "mean",
      normalize: true,
    })

    // Convert typed array to plain number[]
    const data = output.tolist?.() ?? Array.from(output.data as Float32Array)

    // output.tolist() may return nested arrays; flatten if needed
    if (Array.isArray(data) && data.length === 1 && Array.isArray(data[0])) {
      return data[0] as number[]
    }

    return data as number[]
  } catch (err) {
    throw new Error(
      `Embedding computation failed for model "${modelName}": ${err instanceof Error ? err.message : String(err)}`
    )
  }
}