defmodule AI.Embeddings do
  @behaviour AI.Endpoint
  @base_url "https://api.openai.com"

  # Upper bound on attempts to work around token limit exceeded errors.
  @max_attempts 3

  # Short wait between retries to avoid server-side rate limiting.
  @retry_interval 250

  @model "text-embedding-3-large"
  @batch_size 300_000
  @batch_reduction_factor 0.75
  @chunk_size 8192

  @impl AI.Endpoint
  def endpoint_path(), do: "#{@base_url}/v1/embeddings"

  @doc "Returns the embeddings model name."
  @spec model_name() :: String.t()
  def model_name(), do: @model

  @type embedding :: list(float())
  @type embeddings :: list(embedding())

  @type error ::
          {:error, :max_attempts_reached}
          | {:error, :http_error}
          | {:error, :transport_error}
          | {:error, String.t()}

  @type attempt :: non_neg_integer()
  @type inputs :: list(String.t())

  @spec get(String.t()) :: {:ok, embeddings} | error
  def get(input) do
    input
    |> split_into_batches()
    |> get(1, [])
  end

  @spec get(inputs, attempt, embeddings) :: {:ok, embeddings} | error
  defp get(batches, attempt, acc)
  defp get([], _attempt, acc), do: {:ok, acc}

  defp get(_batches, attempt, _acc) when attempt > @max_attempts do
    {:error, :max_attempts_reached}
  end

  defp get([batch | rest], attempt, acc) do
    if attempt > 1 do
      Process.sleep(@retry_interval)
    end

    batch
    |> split_into_chunks(attempt)
    |> endpoint()
    |> case do
      {:ok, embeddings} ->
        get(rest, 1, merge_embeddings(embeddings, acc))

      {:error, :token_limit_exceeded} ->
        get([batch | rest], attempt + 1, acc)

      other ->
        {:error, other}
    end
  end

  @spec split_into_batches(String.t()) :: inputs
  defp split_into_batches(input) do
    tokens = AI.PretendTokenizer.guesstimate_tokens(input)

    if tokens >= @batch_size do
      AI.PretendTokenizer.chunk(input, @batch_size, @batch_reduction_factor)
    else
      [input]
    end
  end

  @spec split_into_chunks(String.t(), attempt) :: inputs
  defp split_into_chunks(batch, attempt) do
    reduction_factor = token_reduction_factor(attempt)
    AI.PretendTokenizer.chunk(batch, @chunk_size, reduction_factor)
  end

  # -----------------------------------------------------------------------------
  # For each dimension, find the maximum value across all embeddings. This
  # isn't necessarily the _most_ accurate, but it selects the highest rating
  # for each dimension found in the file, which should be reasonable for
  # semantic searching.
  # -----------------------------------------------------------------------------
  @spec merge_embeddings(embeddings, embedding) :: embeddings | error
  defp merge_embeddings([], acc), do: acc
  defp merge_embeddings([first | rest], []), do: merge_embeddings(rest, first)

  defp merge_embeddings([first | rest], acc) do
    merged = Enum.zip_with(acc, first, &max/2)
    merge_embeddings(rest, merged)
  end

  @spec get_api_key!() :: String.t()
  defp get_api_key!() do
    ["FNORD_OPENAI_API_KEY", "OPENAI_API_KEY"]
    |> Enum.find_value(fn k -> Util.Env.get_env(k, nil) end)
    |> case do
      nil ->
        raise "Either FNORD_OPENAI_API_KEY or OPENAI_API_KEY environment variable must be set"

      api_key ->
        api_key
    end
  end

  # -----------------------------------------------------------------------------
  # Calculate the token reduction factor based on the number of attempts. This
  # is used to dial back the number of (estimated) tokens sent to the endpoint
  # when retrying requests.
  # -----------------------------------------------------------------------------
  @spec token_reduction_factor(attempt) :: float()
  defp token_reduction_factor(attempt) do
    case attempt do
      1 -> 0.75
      2 -> 0.50
      _ -> 0.25
    end
  end

  @spec token_limit_error?(String.t()) :: boolean
  defp token_limit_error?(body) do
    String.contains?(body, "maximum context length")
  end

  @spec endpoint(inputs) ::
          {:ok, embeddings}
          | {:error, :token_limit_exceeded}
          | {:error, :http_error}
          | {:error, :transport_error}
  defp endpoint(input) do
    endpoint(input, 1)
  end

  @spec endpoint(inputs, attempt) ::
          {:ok, embeddings}
          | {:error, :token_limit_exceeded}
          | {:error, :http_error}
          | {:error, :transport_error}
  defp endpoint(input, _attempt) do
    api_key = get_api_key!()

    headers = [
      {"Authorization", "Bearer #{api_key}"},
      {"Content-Type", "application/json"}
    ]

    payload = %{
      encoding_format: "float",
      model: @model,
      input: input
    }

    AI.Endpoint.post_json(__MODULE__, headers, payload)
    |> case do
      {:ok, %{body: %{"data" => embeddings}}} ->
        {:ok, Enum.map(embeddings, &Map.get(&1, "embedding"))}

      {:http_error, {status_code, body}} ->
        if token_limit_error?(body) do
          {:error, :token_limit_exceeded}
        else
          UI.warn("[AI.Embeddings] Error getting embeddings: #{status_code} - #{body}")
          {:error, :http_error}
        end

      {:transport_error, error} ->
        UI.warn("[AI.Embeddings] Error getting embeddings: #{inspect(error)}")
        {:error, :transport_error}
    end
  end
end