defmodule Vecto do
  @moduledoc """
  Hybrid Search with Postgres and Ecto
  """
  import Ecto.Query
  import Pgvector.Ecto.Query

  @doc """
  Keyword Search

  Uses Postgres full-text search (tsvector) to find documents that match the query string, ranked by relevance.
  """
  def keyword_search(schema, query_field, query_string, limit_by \\ 100, select_columns \\ [])
      when is_atom(query_field) and is_binary(query_string) do
    select_columns = build_select_columns(schema, select_columns)

    schema
    |> join(:inner, [s], q in fragment("websearch_to_tsquery('english', ?)", ^query_string),
      on: fragment("? @@ ?", field(s, ^query_field), q)
    )
    |> order_by([s, q], fragment("ts_rank_cd(?, ?) DESC", field(s, ^query_field), q))
    |> limit(^limit_by)
    |> select([s, q], %{
      id: s.id,
      rank: fragment("RANK() OVER (ORDER BY ts_rank_cd(?, ?) DESC)", field(s, ^query_field), q)
    })
    |> select_merge([t], map(t, ^select_columns))
  end

  @doc """
  Semantic Search

  Uses Postgres vector search (pg_vector) to find documents that are semantically similar to the query embedding, ranked by similarity.
  """
  def semantic_search(schema, query_field, query_embedding, limit_by \\ 100, select_columns \\ [])
      when is_atom(query_field) and is_list(query_embedding) do
    select_columns = build_select_columns(schema, select_columns)

    schema
    |> order_by([s], cosine_distance(field(s, ^query_field), ^query_embedding))
    |> limit(^limit_by)
    |> select([s], %{
      id: s.id,
      rank:
        fragment(
          "RANK() OVER (ORDER BY ?)",
          cosine_distance(field(s, ^query_field), ^query_embedding)
        )
    })
    |> select_merge([t], map(t, ^select_columns))
  end

  @doc """
    Postgres Hybrid Search

    Uses Reciprocal Rank Fusion (RRF) approach to combine full-text search and vector search.

    Assumes your primary key is "id", vector column is "embedding" and tsvector column is "content",
    but is flexible enough to work with any table or column name
  """
  def hybrid_search(schema, query_embedding, query_string, opts \\ [])
      when is_list(query_embedding) and is_binary(query_string) do
    opts = Keyword.merge(hybrid_search_default_opts(), opts)
    max_limit = opts[:limit] * 2
    select_columns = build_select_columns(schema, opts[:select_columns])

    semantic_search_cte =
      semantic_search(schema, opts[:vector_column], query_embedding, max_limit, select_columns)

    keyword_search_cte =
      keyword_search(schema, opts[:tsvector_column], query_string, max_limit, select_columns)

    base_query =
      from(s in schema)
      |> with_cte("semantic_search", as: ^semantic_search_cte)
      |> join(:full, [s], vs in "semantic_search", on: vs.id == s.id)
      |> with_cte("keyword_search", as: ^keyword_search_cte)
      |> join(:full, [s, vs], ss in "keyword_search", on: ss.id == s.id)
      |> select([s, vs, ss], %{
        score:
          fragment(
            "(COALESCE(1.0 / (? + ?), 0.0) * ?) + (COALESCE(1.0 / (? + ?), 0.0) * ?)",
            ^opts[:rrf_k],
            vs.rank,
            ^opts[:vector_weight],
            ss.rank,
            ^opts[:rrf_k],
            ^opts[:tsvector_weight]
          ),
        id: coalesce(vs.id, ss.id)
      })
      |> select_merge([t, _t], map(t, ^select_columns))

    base_query
    |> subquery()
    |> order_by([bq], desc: bq.score)
    |> limit(^opts[:limit])
    |> select([bq], %{score: bq.score, id: bq.id})
    |> select_merge([t, _t], map(t, ^select_columns))
  end

  defp hybrid_search_default_opts do
    [
      rrf_k: 50,
      vector_column: :embedding,
      vector_weight: 1.0,
      tsvector_column: :content,
      tsvector_weight: 1.0,
      select_columns: [],
      limit: 100
    ]
  end

  defp build_select_columns(schema, []), do: schema.__schema__(:fields)
  defp build_select_columns(_schema, columns), do: columns
end