/*
 * Copyright Debezium Authors.
 *
 * Licensed under the Apache Software License version 2.0, available at http://www.apache.org/licenses/LICENSE-2.0
 */
package io.debezium.ai.embeddings;

import static io.debezium.ai.embeddings.FieldToEmbedding.LEGACY_EMBEDDINGS_PREFIX;
import static org.assertj.core.api.Assertions.assertThat;

import java.io.IOException;
import java.time.Duration;
import java.util.List;
import java.util.Map;

import org.apache.kafka.connect.data.Struct;
import org.apache.kafka.connect.source.SourceRecord;
import org.assertj.core.data.Offset;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.testcontainers.ollama.OllamaContainer;
import org.testcontainers.utility.DockerImageName;

/**
 * Integrations tests for {@link FieldToEmbedding} SMT which uses {@link OllamaModelFactory}.
 *
 * @author vjuranek
 */
public class EmbeddingsOllamaIT {
    private static final String OLLAMA_IMAGE_NAME = "mirror.gcr.io/ollama/ollama:0.6.2";
    private static final String OLLAMA_TEST_MODEL = "all-minilm";

    private static final OllamaContainer ollama = new OllamaContainer(
            DockerImageName.parse(OLLAMA_IMAGE_NAME).asCompatibleSubstituteFor("ollama/ollama"))
            .withStartupTimeout(Duration.ofSeconds(180));

    private final FieldToEmbedding<SourceRecord> embeddingSmt = new FieldToEmbedding();

    @BeforeAll
    public static void startDatabase() {
        ollama.start();
    }

    @AfterAll
    public static void stopDatabase() {
        ollama.stop();
    }

    @Test
    public void testOllamaEmbeddings() throws InterruptedException, IOException {
        assertEmbeddingsForConfig(Map.of(
                "field.source", "after.product",
                "field.embedding", "after.prod_embedding",
                "ollama.url", ollama.getEndpoint(),
                "ollama.model.name", OLLAMA_TEST_MODEL,
                "operation.timeout.ms", 20_000));
    }

    @Test
    public void testOllamaEmbeddingsWithLegacyConfig() throws InterruptedException, IOException {
        assertEmbeddingsForConfig(Map.of(
                LEGACY_EMBEDDINGS_PREFIX + "field.source", "after.product",
                LEGACY_EMBEDDINGS_PREFIX + "field.embedding", "after.prod_embedding",
                LEGACY_EMBEDDINGS_PREFIX + "ollama.url", ollama.getEndpoint(),
                LEGACY_EMBEDDINGS_PREFIX + "ollama.model.name", OLLAMA_TEST_MODEL,
                LEGACY_EMBEDDINGS_PREFIX + "operation.timeout.ms", 20_000));
    }

    private void assertEmbeddingsForConfig(Map<String, ?> config) throws InterruptedException, IOException {
        ollama.execInContainer("ollama", "pull", OLLAMA_TEST_MODEL);

        embeddingSmt.configure(config);
        SourceRecord transformedRecord = embeddingSmt.apply(FieldToEmbeddingTest.SOURCE_RECORD);

        Struct payloadStruct = (Struct) transformedRecord.value();
        assertThat(payloadStruct.getStruct("after").getString("product")).contains("a product");
        List<Float> embeddings = payloadStruct.getStruct("after").getArray("prod_embedding");
        assertThat(embeddings.size()).isEqualTo(384);
        final Offset<Float> offset = Offset.offset(0.001f);
        assertThat(embeddings.get(0)).isCloseTo(-0.07157089f, offset);
        assertThat(embeddings.get(1)).isCloseTo(0.022460647f, offset);
        assertThat(embeddings.get(2)).isCloseTo(-0.02369636f, offset);
    }
}
