package at.procon.dip.embedding.provider.mock;

import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.model.EmbeddingProviderResult;
import at.procon.dip.embedding.model.EmbeddingRequest;
import at.procon.dip.embedding.model.ResolvedEmbeddingProviderConfig;
import at.procon.dip.embedding.provider.EmbeddingProvider;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import org.springframework.stereotype.Component;

@Component
public class MockEmbeddingProvider implements EmbeddingProvider {

    private static final String PROVIDER_TYPE = "mock";

    @Override
    public String providerType() {
        return PROVIDER_TYPE;
    }

    @Override
    public boolean supports(EmbeddingModelDescriptor model, ResolvedEmbeddingProviderConfig providerConfig) {
        return PROVIDER_TYPE.equalsIgnoreCase(providerConfig.providerType());
    }

    @Override
    public EmbeddingProviderResult embedDocuments(ResolvedEmbeddingProviderConfig providerConfig,
                                                  EmbeddingModelDescriptor model,
                                                  EmbeddingRequest request) {
        return execute(providerConfig, model, request);
    }

    @Override
    public EmbeddingProviderResult embedQuery(ResolvedEmbeddingProviderConfig providerConfig,
                                              EmbeddingModelDescriptor model,
                                              EmbeddingRequest request) {
        return execute(providerConfig, model, request);
    }

    private EmbeddingProviderResult execute(ResolvedEmbeddingProviderConfig providerConfig,
                                            EmbeddingModelDescriptor model,
                                            EmbeddingRequest request) {
        int dimensions = model.dimensions() > 0
                ? model.dimensions()
                : providerConfig.dimensions() == null ? 16 : providerConfig.dimensions();

        List<float[]> vectors = new ArrayList<>();
        for (String text : request.texts()) {
            vectors.add(embedDeterministically(text, dimensions));
        }

        return new EmbeddingProviderResult(
                model,
                vectors,
                List.of(),
                "mock-" + UUID.randomUUID(),
                request.texts().stream().mapToInt(text -> text == null ? 0 : text.length()).sum()
        );
    }

    private float[] embedDeterministically(String text, int dimensions) {
        float[] vector = new float[dimensions];
        String value = text == null ? "" : text;
        for (int i = 0; i < value.length(); i++) {
            int bucket = i % dimensions;
            vector[bucket] += ((value.charAt(i) % 31) + 1) / 31.0f;
        }
        return vector;
    }
}
