%%%-------------------------------------------------------------------
%%% @doc ColBERT late interaction embedding provider
%%%
%%% Uses erlang_python with ColBERT models for multi-vector embeddings. Each
%%% document produces multiple vectors (one per token) for fine-grained matching.
%%%
%%% == Requirements ==
%%% ```
%%% pip install transformers torch
%%% '''
%%%
%%% == Configuration ==
%%% ```
%%% Config = #{
%%%     venv => "/path/to/.venv",              %% Virtualenv path (recommended)
%%%     model => "colbert-ir/colbertv2.0",     %% Model name (default, 128 dims)
%%%     timeout => 120000                      %% Timeout in ms (default)
%%% }.
%%% '''
%%%
%%% When `venv' is specified, the provider uses the venv's Python executable
%%% and properly activates the venv environment.
%%%
%%% == Multi-Vector Format ==
%%% Unlike single-vector embeddings, ColBERT produces a list of vectors:
%%% ```
%%% [[0.1, 0.2, ...], [0.3, 0.4, ...], ...]  %% One vector per token
%%% '''
%%%
%%% == Late Interaction ==
%%% ColBERT scoring uses MaxSim:
%%% ```
%%% Score(Q, D) = sum(max(qi · dj for all dj in D) for all qi in Q)
%%% '''
%%% This enables fine-grained token-level matching.
%%%
%%% == Supported Models ==
%%% - `"colbert-ir/colbertv2.0"' - Default, 128 dimensions
%%% - `"answerdotai/answerai-colbert-small-v1"' - 96 dimensions, smaller
%%% - `"jinaai/jina-colbert-v2"' - 128 dimensions, long context (8192 tokens)
%%%
%%% == Use Cases ==
%%% - Fine-grained semantic matching
%%% - Passage retrieval with token-level scoring
%%% - Question answering
%%%
%%% @end
%%%-------------------------------------------------------------------
-module(barrel_embed_colbert).
-behaviour(barrel_embed_provider).

%% Behaviour callbacks
-export([
    embed/2,
    embed_batch/2,
    dimension/1,
    name/0,
    init/1,
    available/1
]).

%% Multi-vector API
-export([
    embed_multi/2,
    embed_batch_multi/2,
    maxsim_score/2
]).

-define(DEFAULT_MODEL, "colbert-ir/colbertv2.0").
-define(DEFAULT_TIMEOUT, 120000).
-define(DEFAULT_DIMENSION, 128).
-define(PROVIDER, <<"colbert">>).

%% Multi-vector type: list of token vectors
-type multi_vector() :: [[float()]].

-export_type([multi_vector/0]).

%%====================================================================
%% Behaviour Callbacks
%%====================================================================

%% @doc Provider name.
-spec name() -> atom().
name() -> colbert.

%% @doc Get dimension for this provider.
-spec dimension(map()) -> pos_integer().
dimension(Config) ->
    maps:get(dimension, Config, ?DEFAULT_DIMENSION).

%% @doc Initialize the provider.
-spec init(map()) -> {ok, map()} | {error, term()}.
init(Config) ->
    Model = maps:get(model, Config, ?DEFAULT_MODEL),
    Timeout = maps:get(timeout, Config, ?DEFAULT_TIMEOUT),
    Venv = maps:get(venv, Config, undefined),

    %% Validate model (warning only)
    validate_model(Model),

    %% Initialize Python environment
    PyConfig = case Venv of
        undefined -> #{};
        _ -> #{venv => Venv}
    end,

    case barrel_embed_py:init(PyConfig) of
        ok ->
            ModelBin = ensure_binary(Model),
            case barrel_embed_py:load_model(?PROVIDER, ModelBin) of
                {ok, #{dimensions := Dims}} ->
                    {ok, Config#{
                        dimension => Dims,
                        model => ModelBin,
                        provider => ?PROVIDER,
                        timeout => Timeout,
                        initialized => true
                    }};
                {ok, _} ->
                    %% No dimensions in response, use default
                    {ok, Config#{
                        dimension => ?DEFAULT_DIMENSION,
                        model => ModelBin,
                        provider => ?PROVIDER,
                        timeout => Timeout,
                        initialized => true
                    }};
                {error, Reason} ->
                    {error, Reason}
            end;
        {error, Reason} ->
            {error, {init_failed, Reason}}
    end.

%% @doc Check if provider is available.
-spec available(map()) -> boolean().
available(#{initialized := true}) ->
    true;
available(_Config) ->
    false.

%% @doc Generate single-vector embedding (mean pooling of token vectors).
%% Note: For ColBERT, use embed_multi/2 to get full multi-vector output.
-spec embed(binary(), map()) -> {ok, [float()]} | {error, term()}.
embed(Text, Config) ->
    case embed_multi(Text, Config) of
        {ok, MultiVec} ->
            {ok, mean_pool(MultiVec)};
        {error, _} = Error ->
            Error
    end.

%% @doc Generate single-vector embeddings for batch (mean pooling).
-spec embed_batch([binary()], map()) -> {ok, [[float()]]} | {error, term()}.
embed_batch(Texts, Config) ->
    case embed_batch_multi(Texts, Config) of
        {ok, MultiVecs} ->
            {ok, [mean_pool(MV) || MV <- MultiVecs]};
        {error, _} = Error ->
            Error
    end.

%%====================================================================
%% Multi-Vector API
%%====================================================================

%% @doc Generate multi-vector embedding for a single text.
%% Returns a list of token vectors.
-spec embed_multi(binary(), map()) -> {ok, multi_vector()} | {error, term()}.
embed_multi(Text, Config) ->
    case embed_batch_multi([Text], Config) of
        {ok, [MultiVec]} -> {ok, MultiVec};
        {error, _} = Error -> Error
    end.

%% @doc Generate multi-vector embeddings for multiple texts.
-spec embed_batch_multi([binary()], map()) -> {ok, [multi_vector()]} | {error, term()}.
embed_batch_multi(Texts, #{model := Model, provider := Provider, initialized := true}) ->
    TextsBin = [ensure_binary(T) || T <- Texts],
    barrel_embed_py:embed_multi(Provider, Model, TextsBin);
embed_batch_multi(_Texts, _Config) ->
    {error, not_initialized}.

%% @doc Calculate MaxSim score between query and document multi-vectors.
%% This is the standard ColBERT scoring function.
%% Score = sum(max(qi · dj for all dj in D) for all qi in Q)
-spec maxsim_score(multi_vector(), multi_vector()) -> float().
maxsim_score(QueryVecs, DocVecs) ->
    lists:sum([max_dot_product(QVec, DocVecs) || QVec <- QueryVecs]).

%%====================================================================
%% Internal Functions
%%====================================================================

ensure_binary(B) when is_binary(B) -> B;
ensure_binary(L) when is_list(L) -> unicode:characters_to_binary(L).

%% @private
validate_model(Model) ->
    ModelBin = ensure_binary(Model),
    case is_known_model(ModelBin) of
        true -> ok;
        false ->
            error_logger:warning_msg(
                "Model ~s is not in the known list. "
                "It may still work if it's a valid ColBERT model.~n",
                [ModelBin]
            )
    end.

%% @private
is_known_model(<<"colbert-ir/colbertv2.0">>) -> true;
is_known_model(<<"answerdotai/answerai-colbert-small-v1">>) -> true;
is_known_model(<<"jinaai/jina-colbert-v2">>) -> true;
is_known_model(_) -> false.

%% @private
%% Mean pooling of token vectors to get single vector
mean_pool([]) -> [];
mean_pool(Vectors) ->
    N = length(Vectors),
    Dim = length(hd(Vectors)),
    %% Sum all vectors element-wise
    Sums = lists:foldl(
        fun(Vec, Acc) ->
            lists:zipwith(fun(A, B) -> A + B end, Vec, Acc)
        end,
        lists:duplicate(Dim, 0.0),
        Vectors
    ),
    %% Divide by N
    [S / N || S <- Sums].

%% @private
%% Find maximum dot product between query vector and all doc vectors
max_dot_product(QueryVec, DocVecs) ->
    DotProducts = [dot_product(QueryVec, DocVec) || DocVec <- DocVecs],
    lists:max(DotProducts).

%% @private
dot_product(V1, V2) ->
    lists:sum(lists:zipwith(fun(A, B) -> A * B end, V1, V2)).