import { _embedTexts, _cosineSimilarity } from "agency-lang/stdlib-lib/embedding.js"

/** @module
  @summary Embed text and compare vectors for similarity search.
  Turn text into embedding vectors. An embedding is a list of numbers
that stands for the meaning of a piece of text. Two texts with similar
meaning get vectors that are close together. You can use embeddings to
find related documents, group similar notes, or rank matches for a
query.

  ```ts
  import { embed, embedMany, cosineSimilarity } from "std::embedding"

  node main() {
    const query = embed("a red bicycle in the rain")
    const docs = embedMany(["apples", "oranges", "a bicycle"])
    if (isFailure(query) || isFailure(docs)) { return }
    const score = cosineSimilarity(query.value.vector, docs.value.vectors[2])
    print(score)
  }
  ```

`embed` turns one string into one vector. `embedMany` turns a list of
strings into one vector each, in the same order, with one provider call.
`cosineSimilarity` compares two vectors and returns a number from -1 to
1. Higher scores indicate more similar vector directions.

Provider usage counts toward the branch's cost and token totals, as with
`llm()`. Cost guards apply.

The default model is OpenAI's `text-embedding-3-small`. Pass `model` to
pick another one. The provider is derived from the model name, or you
can set `provider` yourself. Ollama works as a local provider. Only
compare vectors that came from the same model.
*/

/** One vector and the model that produced it. */
export type Embedding = {
  vector: number[];
  model: string
}

/** One vector per input, in input order, and the model that produced them. */
export type Embeddings = {
  vectors: number[][];
  model: string
}

@capabilities(llm)
export def embed(
  text: string,
  model: string = "",
  provider: string = "",
  dimensions: number = 0,
  apiKey: string = "",
  baseUrl: string = "",
): Result<Embedding> {
  """
  Turn one piece of text into an embedding vector.

  @param model - Embedding model (default: text-embedding-3-small). For a local provider (mlx, llama-cpp) this may be a catalog name, alias, or path
  @param provider - Override the provider (normally derived from the model name)
  @param dimensions - Shorten the vector to this length, for models that support it (0 means the model's full length)
  @param apiKey - Override the API key
  @param baseUrl - Base URL for the ollama, deepinfra, litellm, openai-compat, or mlx provider
  """
  const r = _embedTexts([text], model, provider, dimensions, apiKey, baseUrl)
  if (isFailure(r)) {
    return r
  }
  return success({ vector: r.value.vectors[0], model: r.value.model })
}

@capabilities(llm)
export def embedMany(
  texts: string[],
  model: string = "",
  provider: string = "",
  dimensions: number = 0,
  apiKey: string = "",
  baseUrl: string = "",
): Result<Embeddings> {
  """
  Turn a list of texts into embedding vectors in one call, one vector per text in the same order.

  @param model - Embedding model (default: text-embedding-3-small). For a local provider (mlx, llama-cpp) this may be a catalog name, alias, or path
  @param provider - Override the provider (normally derived from the model name)
  @param dimensions - Shorten each vector to this length, for models that support it (0 means the model's full length)
  @param apiKey - Override the API key
  @param baseUrl - Base URL for the ollama, deepinfra, litellm, openai-compat, or mlx provider
  """
  return _embedTexts(texts, model, provider, dimensions, apiKey, baseUrl)
}

export idempotent def cosineSimilarity(a: number[], b: number[]): Result<number> {
  """
  Compare equal-length, nonzero vectors from the same embedding model.
  Returns cosine similarity from -1 to 1; higher means more similar.
  """
  return _cosineSimilarity(a, b)
}