/**
 * LocalEmbeddingProvider — fully offline embeddings via Transformers.js
 * (ONNX Runtime in-process, no Ollama / Python / external daemon).
 *
 * Why Transformers.js (not Ollama):
 *   - Zero external runtime dependency. Users who don't have Ollama,
 *     LM Studio, or any model runner still get high-quality local
 *     embeddings out of the box.
 *   - Pure `npm install @xenova/transformers` — ONNX Runtime ships
 *     prebuilt binaries for Win/Mac/Linux.
 *   - Runs in-process: no port, no socket, no IPC. Privacy boundary
 *     is the OS process.
 *   - Same model family Ollama would serve (nomic-embed-text), just
 *     loaded directly via ONNX instead of through a server.
 *
 * Model selection:
 *   - DEFAULT: Xenova/nomic-embed-text-v1.5 — 768-dim, 137M params,
 *     MTEB ~62.4, 8K context, ~250 MB one-time download. Quality on
 *     par with text-embedding-3-small. Recommended for most users.
 *   - SPEED:  Xenova/bge-small-en-v1.5 — 384-dim, 33M params, MTEB
 *     ~62.2, 512-token context, ~30 MB. Half the index size, ~3x
 *     faster query. Recommended when corpora are large or hardware
 *     is modest.
 *
 * Both models use task-specific prefixes for queries vs documents
 * ("search_query: " / "search_document: " for nomic;
 *  "Represent this sentence for searching relevant passages: " for BGE
 *  queries — though many users skip this for BGE).
 *
 * @xenova/transformers is declared as an OPTIONAL dependency so cloud-only
 * users don't pay the ~50 MB install cost. If the user selects the local
 * provider without the package installed, they get a clear remediation
 * message (NOT_INITIALIZED with install instructions).
 *
 * Lazy loading: the model is fetched + initialized on first call,
 * not at construction. This means:
 *   - The CLI starts instantly even with local provider configured.
 *   - First-call latency is ~1-3s for model load + ~50-150ms for
 *     the actual inference. The orchestrator can pre-warm via a
 *     dummy call after enable.
 */

import path from 'path';
import {
  EmbeddingProvider,
  EmbeddingError,
  EMBEDDING_ERROR_CODES,
} from '../providerInterface.js';

/**
 * Known-good local model catalog. New models can be added without
 * touching consumer code — they'll get auto-discovered dimensions and
 * the right query/document prefix handling.
 *
 * `prefix.query` / `prefix.document` are strings prepended to the
 * input. `null` means "no prefix" (the model treats both the same).
 */
export const LOCAL_MODELS = Object.freeze({
  'Xenova/nomic-embed-text-v1.5': {
    dimensions: 768,
    maxTokensPerText: 8192,
    prefix: {
      query:    'search_query: ',
      document: 'search_document: ',
    },
    description: 'Best quality, recommended (250 MB)',
  },
  'Xenova/bge-small-en-v1.5': {
    dimensions: 384,
    maxTokensPerText: 512,
    prefix: {
      // BGE-small's "represent for searching" instruction is only used
      // for query embeddings — documents stay raw.
      query:    'Represent this sentence for searching relevant passages: ',
      document: null,
    },
    description: 'Fast, smaller index (30 MB)',
  },
  'Xenova/all-MiniLM-L6-v2': {
    dimensions: 384,
    maxTokensPerText: 256,
    prefix: { query: null, document: null },
    description: 'Classic baseline (25 MB)',
  },
});

export const DEFAULT_LOCAL_MODEL = 'Xenova/nomic-embed-text-v1.5';

export class LocalEmbeddingProvider extends EmbeddingProvider {
  /**
   * @param {object} opts
   * @param {string} [opts.modelId]    - Defaults to nomic-embed-text-v1.5.
   * @param {string} [opts.cacheDir]   - Where Transformers.js stores ONNX
   *        weights. Defaults to the Transformers.js default; pass an
   *        absolute path under `.loxia-state/embeddings/cache/` to keep
   *        downloads scoped to Loxia's state dir.
   * @param {number} [opts.maxBatchSize=32] - Per-call inference batch
   *        size. CPU-bound; bigger batches help throughput but increase
   *        memory pressure. 32 is a safe default on 8 GB RAM laptops.
   * @param {(modelId: string) => Promise<Function>} [opts.pipelineFactory]
   *        - Injectable for tests. When supplied, replaces the real
   *          `@xenova/transformers` dynamic import. The factory must
   *          return an awaitable pipeline function with signature
   *          `(input: string|string[], options) => Promise<{tolist():
   *          number[][] | number[]}>`.
   * @param {object} [opts.logger]
   */
  constructor({
    modelId = DEFAULT_LOCAL_MODEL,
    cacheDir = null,
    maxBatchSize = 32,
    pipelineFactory = null,
    logger = null,
  } = {}) {
    super();
    this._modelId = modelId;
    this._cacheDir = cacheDir;
    this._maxBatchSize = maxBatchSize;
    this._pipelineFactory = pipelineFactory;
    this._logger = logger;

    const known = LOCAL_MODELS[modelId];
    this._dimensions = known?.dimensions ?? 768;       // safe default
    this._maxTokensPerText = known?.maxTokensPerText ?? 512;
    this._prefix = known?.prefix ?? { query: null, document: null };

    // Lazy-load state.
    this._pipeline = null;
    this._loadPromise = null;     // dedup concurrent first-call loads
    this._loadError = null;       // sticky failure so we don't retry every call
  }

  getInfo() {
    return {
      name: 'local',
      modelId: this._modelId,
      dimensions: this._dimensions,
      maxBatchSize: this._maxBatchSize,
      maxTokensPerText: this._maxTokensPerText,
      distinguishesQueryAndDocument: !!(this._prefix.query || this._prefix.document),
      isLocal: true,
    };
  }

  /**
   * Model is loaded? Lightweight — doesn't trigger a load. For first-call
   * latency hiding, callers should explicitly `await provider.warmup()`.
   */
  async health() {
    if (this._loadError) return { ok: false, reason: this._loadError };
    if (!this._pipeline) {
      return { ok: false, reason: `Model "${this._modelId}" not loaded yet. Call warmup() or any embed* method to trigger load.` };
    }
    return { ok: true };
  }

  /**
   * Force the model load now (downloads on first call, then cached).
   * Safe to call multiple times — subsequent calls are no-ops.
   * Throws EmbeddingError on load failure.
   * @returns {Promise<void>}
   */
  async warmup() {
    await this._ensureLoaded();
  }

  async _ensureLoaded() {
    if (this._pipeline) return this._pipeline;
    if (this._loadError) {
      throw new EmbeddingError(
        EMBEDDING_ERROR_CODES.NOT_INITIALIZED,
        this._loadError,
        { provider: 'local', transient: false },
      );
    }
    if (this._loadPromise) return this._loadPromise;

    const start = Date.now();
    this._loadPromise = (async () => {
      try {
        const pipe = this._pipelineFactory
          ? await this._pipelineFactory(this._modelId)
          : await this._loadRealPipeline();
        if (typeof pipe !== 'function') {
          throw new Error('pipelineFactory must return a callable pipeline');
        }
        this._pipeline = pipe;
        this._logger?.info?.('[LocalEmbeddingProvider] Model loaded', {
          modelId: this._modelId,
          ms: Date.now() - start,
        });
        return pipe;
      } catch (err) {
        this._loadError = err?.message || String(err);
        this._loadPromise = null;
        if (err instanceof EmbeddingError) throw err;
        throw new EmbeddingError(
          EMBEDDING_ERROR_CODES.NOT_INITIALIZED,
          `Failed to load local model "${this._modelId}": ${this._loadError}`,
          { provider: 'local', transient: false, cause: err },
        );
      }
    })();
    return this._loadPromise;
  }

  async _loadRealPipeline() {
    let mod;
    try {
      // Dynamic import keeps @xenova/transformers an optional dep —
      // cloud-only users don't pay the install cost.
      mod = await import('@xenova/transformers');
    } catch (err) {
      throw new EmbeddingError(
        EMBEDDING_ERROR_CODES.NOT_INITIALIZED,
        '@xenova/transformers is not installed. Run `npm install @xenova/transformers` to use local embeddings, or switch the embedding provider to cloud.',
        { provider: 'local', transient: false, cause: err },
      );
    }
    const { pipeline, env } = mod;
    if (this._cacheDir) {
      // Keep ONNX weights inside Loxia's state dir, not scattered around
      // the home dir. This makes backups + uninstall clean.
      env.cacheDir = path.resolve(this._cacheDir);
    }
    return await pipeline('feature-extraction', this._modelId);
  }

  async _embed(texts, kind) {
    const pipe = await this._ensureLoaded();
    const prefixed = this._applyPrefix(texts, kind);

    let raw;
    try {
      // Transformers.js feature-extraction pipeline signature:
      //   pipe(input, { pooling, normalize }) -> Tensor
      // We mean-pool to get one vector per input and skip normalize=true
      // because the base class normalizes for us (avoids double work +
      // ensures consistent behavior across providers).
      raw = await pipe(prefixed, { pooling: 'mean', normalize: false });
    } catch (err) {
      throw new EmbeddingError(
        EMBEDDING_ERROR_CODES.PROVIDER_UNAVAILABLE,
        `Local pipeline inference failed: ${err?.message || err}`,
        { provider: 'local', transient: true, cause: err },
      );
    }

    // Tensor → JS arrays. Different Tensor shapes:
    //   - Single input: shape [dim] → tolist() = [n0, n1, ...]
    //   - Batch input:  shape [batch, dim] → tolist() = [[...], [...], ...]
    // We always pass an array (even when texts.length === 1) so the
    // pipeline ALWAYS gives the batch shape — keeps downstream simple.
    let list;
    try {
      list = typeof raw?.tolist === 'function' ? raw.tolist() : raw;
    } catch (err) {
      throw new EmbeddingError(
        EMBEDDING_ERROR_CODES.PROVIDER_UNAVAILABLE,
        `Failed to materialize pipeline output: ${err?.message || err}`,
        { provider: 'local', transient: true, cause: err },
      );
    }

    // Defensive: if the runtime returned a 1-D array even though we
    // batched, wrap it. Happens with some Transformers.js versions.
    if (Array.isArray(list) && list.length > 0 && typeof list[0] === 'number') {
      list = [list];
    }
    return list;
  }

  /**
   * Apply the model-specific task prefix. Falls back to identity when the
   * model doesn't distinguish queries from documents.
   */
  _applyPrefix(texts, kind) {
    const prefix = kind === 'query' ? this._prefix.query : this._prefix.document;
    if (!prefix) return texts;
    return texts.map(t => prefix + t);
  }

  /** Fingerprint includes modelId; vectors get re-embedded on model swap. */
  getModelFingerprint() {
    return `local:${this._modelId}`;
  }
}
