// Picks an embedding model loaded in LM Studio and computes vectors for the
// RAG layer: auto-pick the first loaded embedding model, or honour the
// identifier configured in the plugin settings.

import type { LMStudioClient, EmbeddingDynamicHandle } from "@lmstudio/sdk";

export interface PickedEmbedder {
  model: EmbeddingDynamicHandle;
  identifier: string;
}

export async function pickEmbeddingModel(
  client: LMStudioClient,
  override: string,
): Promise<PickedEmbedder | { error: string }> {
  if (override && override.trim()) {
    try {
      const model = await client.embedding.model(override.trim());
      return { model, identifier: override.trim() };
    } catch (e: unknown) {
      return {
        error:
          `Failed to access the configured embedding model "${override}". ` +
          `Check the identifier in the plugin settings or load the model in LM Studio. ` +
          `(${e instanceof Error ? e.message : String(e)})`,
      };
    }
  }
  const loaded = await client.embedding.listLoaded();
  if (loaded.length === 0) {
    return {
      error:
        "No embedding model is currently loaded in LM Studio. Load one " +
        "(e.g. nomic-embed-text, bge-small/large, all-MiniLM) and try again.",
    };
  }
  const model = loaded[0];
  let identifier = "(loaded embedding model)";
  try {
    const info = await model.getModelInfo();
    identifier = info?.identifier ?? info?.modelKey ?? identifier;
  } catch {
    /* keep fallback */
  }
  return { model, identifier };
}

/** Embed a single string, normalising the SDK's return shape to number[]. */
export async function embedOne(model: EmbeddingDynamicHandle, text: string): Promise<number[]> {
  const result = await model.embed(text);
  const embedding = Array.isArray(result) ? result[0]?.embedding : result?.embedding;
  if (!embedding || !Array.isArray(embedding)) {
    throw new Error("Embedding model returned no vector.");
  }
  return embedding;
}

export interface EmbedManyOptions {
  onProgress?: (done: number, total: number) => void;
  abortSignal?: AbortSignal;
}

/**
 * Embed many strings sequentially (LM Studio serves one request at a time).
 * Returns vectors aligned with the input order.
 */
export async function embedMany(
  model: EmbeddingDynamicHandle,
  texts: string[],
  opts: EmbedManyOptions = {},
): Promise<number[][]> {
  const out: number[][] = [];
  for (let i = 0; i < texts.length; i++) {
    if (opts.abortSignal?.aborted) throw new Error("Embedding aborted.");
    out.push(await embedOne(model, texts[i]));
    opts.onProgress?.(i + 1, texts.length);
  }
  return out;
}

// Picks an embedding model loaded in LM Studio and computes vectors for the
// RAG layer: auto-pick the first loaded embedding model, or honour the
// identifier configured in the plugin settings.

import type { LMStudioClient, EmbeddingDynamicHandle } from "@lmstudio/sdk";

export interface PickedEmbedder {
  model: EmbeddingDynamicHandle;
  identifier: string;
}

export async function pickEmbeddingModel(
  client: LMStudioClient,
  override: string,
): Promise<PickedEmbedder | { error: string }> {
  if (override && override.trim()) {
    try {
      const model = await client.embedding.model(override.trim());
      return { model, identifier: override.trim() };
    } catch (e: unknown) {
      return {
        error:
          `Failed to access the configured embedding model "${override}". ` +
          `Check the identifier in the plugin settings or load the model in LM Studio. ` +
          `(${e instanceof Error ? e.message : String(e)})`,
      };
    }
  }
  const loaded = await client.embedding.listLoaded();
  if (loaded.length === 0) {
    return {
      error:
        "No embedding model is currently loaded in LM Studio. Load one " +
        "(e.g. nomic-embed-text, bge-small/large, all-MiniLM) and try again.",
    };
  }
  const model = loaded[0];
  let identifier = "(loaded embedding model)";
  try {
    const info = await model.getModelInfo();
    identifier = info?.identifier ?? info?.modelKey ?? identifier;
  } catch {
    /* keep fallback */
  }
  return { model, identifier };
}

/** Embed a single string, normalising the SDK's return shape to number[]. */
export async function embedOne(model: EmbeddingDynamicHandle, text: string): Promise<number[]> {
  const result = await model.embed(text);
  const embedding = Array.isArray(result) ? result[0]?.embedding : result?.embedding;
  if (!embedding || !Array.isArray(embedding)) {
    throw new Error("Embedding model returned no vector.");
  }
  return embedding;
}

export interface EmbedManyOptions {
  onProgress?: (done: number, total: number) => void;
  abortSignal?: AbortSignal;
}

/**
 * Embed many strings sequentially (LM Studio serves one request at a time).
 * Returns vectors aligned with the input order.
 */
export async function embedMany(
  model: EmbeddingDynamicHandle,
  texts: string[],
  opts: EmbedManyOptions = {},
): Promise<number[][]> {
  const out: number[][] = [];
  for (let i = 0; i < texts.length; i++) {
    if (opts.abortSignal?.aborted) throw new Error("Embedding aborted.");
    out.push(await embedOne(model, texts[i]));
    opts.onProgress?.(i + 1, texts.length);
  }
  return out;
}