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

import static org.assertj.core.api.Assertions.assertThat;

import at.procon.dip.domain.document.DistanceMetric;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.model.EmbeddingPrefixMode;
import at.procon.dip.embedding.model.EmbeddingRequest;
import at.procon.dip.embedding.model.EmbeddingUseCase;
import at.procon.dip.embedding.model.ResolvedEmbeddingProviderConfig;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpServer;
import java.io.IOException;
import java.io.InputStream;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;

class VectorSyncHttpEmbeddingProviderTest {

    private final ObjectMapper objectMapper = new ObjectMapper();
    private HttpServer server;
    private final AtomicReference<String> lastBatchBody = new AtomicReference<>();
    private final AtomicReference<String> lastSingleBody = new AtomicReference<>();

    @AfterEach
    void tearDown() {
        if (server != null) {
            server.stop(0);
        }
    }

    @Test
    void shouldCallVectorSyncEndpointAndParseVector() throws Exception {
        server = HttpServer.create(new InetSocketAddress(0), 0);
        server.createContext("/vector-sync", this::handleVectorSync);
        server.start();

        var provider = new VectorSyncHttpEmbeddingProvider(objectMapper);
        var config = ResolvedEmbeddingProviderConfig.builder()
                .key("vector-sync-local")
                .providerType("http-vector-sync")
                .baseUrl("http://localhost:" + server.getAddress().getPort())
                .readTimeout(Duration.ofSeconds(5))
                .headers(Map.of("X-Client", "dip-test"))
                .batchTruncateText(false)
                .batchTruncateLength(512)
                .batchChunkSize(20)
                .build();
        var model = new EmbeddingModelDescriptor(
                "e5-default",
                "vector-sync-local",
                "intfloat/multilingual-e5-large",
                3,
                DistanceMetric.COSINE,
                true,
                false,
                8192,
                true
        );
        var request = EmbeddingRequest.builder()
                .modelKey("e5-default")
                .useCase(EmbeddingUseCase.DOCUMENT)
                .texts(List.of("This is a sample text to vectorize"))
                .providerOptions(Map.of())
                .build();

        var result = provider.embedDocuments(config, model, request);

        assertThat(result.vectors()).hasSize(1);
        assertThat(result.vectors().getFirst()).containsExactly(0.1f, 0.2f, 0.3f);
        assertThat(result.tokenCount()).isEqualTo(9);
        assertThat(result.prefixMode()).isEqualTo(EmbeddingPrefixMode.OFF);
        assertThat(result.appliedPrefix()).isNull();
    }

    @Test
    void shouldPrefixTextsInClientModeForDocuments() throws Exception {
        server = HttpServer.create(new InetSocketAddress(0), 0);
        server.createContext("/vector-sync", this::handleVectorSync);
        server.start();

        var provider = new VectorSyncHttpEmbeddingProvider(objectMapper);
        var config = ResolvedEmbeddingProviderConfig.builder()
                .key("vector-sync-local")
                .providerType("http-vector-sync")
                .baseUrl("http://localhost:" + server.getAddress().getPort())
                .readTimeout(Duration.ofSeconds(5))
                .headers(Map.of("X-Client", "dip-test"))
                .batchTruncateText(false)
                .batchTruncateLength(512)
                .batchChunkSize(20)
                .build();
        var model = new EmbeddingModelDescriptor(
                "e5-default",
                "vector-sync-local",
                "intfloat/multilingual-e5-large",
                3,
                DistanceMetric.COSINE,
                true,
                false,
                8192,
                true,
                EmbeddingPrefixMode.CLIENT,
                "query: ",
                "passage: "
        );
        var request = EmbeddingRequest.builder()
                .modelKey("e5-default")
                .useCase(EmbeddingUseCase.DOCUMENT)
                .texts(List.of("This is a sample text to vectorize"))
                .providerOptions(Map.of())
                .build();

        var result = provider.embedDocuments(config, model, request);

        assertThat(result.prefixMode()).isEqualTo(EmbeddingPrefixMode.CLIENT);
        assertThat(result.appliedPrefix()).isEqualTo("passage: ");
        assertThat(lastSingleBody.get()).contains("passage: This is a sample text to vectorize");
    }

    @Test
    void shouldCallVectorizeBatchEndpointUsingConfiguredBatchRequestSettings() throws Exception {
        server = HttpServer.create(new InetSocketAddress(0), 0);
        server.createContext("/vectorize-batch", this::handleVectorizeBatch);
        server.start();

        var provider = new VectorSyncHttpEmbeddingProvider(objectMapper);
        var config = ResolvedEmbeddingProviderConfig.builder()
                .key("vector-sync-local")
                .providerType("http-vector-sync")
                .baseUrl("http://localhost:" + server.getAddress().getPort())
                .readTimeout(Duration.ofSeconds(5))
                .headers(Map.of("X-Client", "dip-test"))
                .batchTruncateText(true)
                .batchTruncateLength(768)
                .batchChunkSize(33)
                .build();
        var model = new EmbeddingModelDescriptor(
                "e5-default",
                "vector-sync-local",
                "intfloat/multilingual-e5-large",
                3,
                DistanceMetric.COSINE,
                true,
                true,
                8192,
                true
        );
        var request = EmbeddingRequest.builder()
                .modelKey("e5-default")
                .useCase(EmbeddingUseCase.DOCUMENT)
                .texts(List.of("First text", "Second text"))
                .providerOptions(Map.of())
                .build();

        var result = provider.embedDocuments(config, model, request);

        assertThat(result.vectors()).hasSize(2);
        assertThat(result.vectors().get(0)).containsExactly(0.1f, 0.2f, 0.3f);
        assertThat(result.vectors().get(1)).containsExactly(0.4f, 0.5f, 0.6f);
        assertThat(result.tokenCount()).isEqualTo(12);
        assertThat(result.prefixMode()).isEqualTo(EmbeddingPrefixMode.OFF);

        JsonNode requestBody = objectMapper.readTree(lastBatchBody.get());
        assertThat(requestBody.get("truncate_text").asBoolean()).isTrue();
        assertThat(requestBody.get("truncate_length").asInt()).isEqualTo(768);
        assertThat(requestBody.get("chunk_size").asInt()).isEqualTo(33);
    }

    @Test
    void shouldAllowPerRequestBatchOverrides() throws Exception {
        server = HttpServer.create(new InetSocketAddress(0), 0);
        server.createContext("/vectorize-batch", this::handleVectorizeBatch);
        server.start();

        var provider = new VectorSyncHttpEmbeddingProvider(objectMapper);
        var config = ResolvedEmbeddingProviderConfig.builder()
                .key("vector-sync-local")
                .providerType("http-vector-sync")
                .baseUrl("http://localhost:" + server.getAddress().getPort())
                .readTimeout(Duration.ofSeconds(5))
                .batchTruncateText(false)
                .batchTruncateLength(512)
                .batchChunkSize(20)
                .build();
        var model = new EmbeddingModelDescriptor(
                "e5-default",
                "vector-sync-local",
                "intfloat/multilingual-e5-large",
                3,
                DistanceMetric.COSINE,
                true,
                true,
                8192,
                true
        );
        var request = EmbeddingRequest.builder()
                .modelKey("e5-default")
                .useCase(EmbeddingUseCase.DOCUMENT)
                .texts(List.of("First text", "Second text"))
                .providerOptions(Map.of(
                        "truncate_text", true,
                        "truncate_length", 1024,
                        "chunk_size", 44
                ))
                .build();

        provider.embedDocuments(config, model, request);

        JsonNode requestBody = objectMapper.readTree(lastBatchBody.get());
        assertThat(requestBody.get("truncate_text").asBoolean()).isTrue();
        assertThat(requestBody.get("truncate_length").asInt()).isEqualTo(1024);
        assertThat(requestBody.get("chunk_size").asInt()).isEqualTo(44);
    }

    private void handleVectorSync(HttpExchange exchange) throws IOException {
        String body;
        try (InputStream in = exchange.getRequestBody()) {
            body = new String(in.readAllBytes(), StandardCharsets.UTF_8);
        }

        lastSingleBody.set(body);
        assertThat(exchange.getRequestMethod()).isEqualTo("POST");
        assertThat(body).contains("\"model\":\"intfloat/multilingual-e5-large\"");
        assertThat(body).contains("This is a sample text to vectorize");
        assertThat(exchange.getRequestHeaders().getFirst("X-Client")).isEqualTo("dip-test");

        respondJson(exchange, """
                {
                  "runtime_ms": 472.49,
                  "vector": [0.1, 0.2, 0.3],
                  "incomplete": false,
                  "combined_vector": null,
                  "token_count": 9,
                  "model": "intfloat/multilingual-e5-large",
                  "max_seq_length": 512
                }
                """);
    }

    private void handleVectorizeBatch(HttpExchange exchange) throws IOException {
        String body;
        try (InputStream in = exchange.getRequestBody()) {
            body = new String(in.readAllBytes(), StandardCharsets.UTF_8);
        }
        lastBatchBody.set(body);

        assertThat(exchange.getRequestMethod()).isEqualTo("POST");
        JsonNode requestBody = objectMapper.readTree(body);
        assertThat(requestBody.get("model").asText()).isEqualTo("intfloat/multilingual-e5-large");
        assertThat(requestBody.get("items")).hasSize(2);

        String firstId = requestBody.get("items").get(0).get("id").asText();
        String secondId = requestBody.get("items").get(1).get("id").asText();

        respondJson(exchange, """
                {
                  "model": "intfloat/multilingual-e5-large",
                  "count": 2,
                  "results": [
                    {
                      "id": "%s",
                      "vector": [0.1, 0.2, 0.3],
                      "token_count": 5,
                      "runtime_ms": 0.0,
                      "incomplete": false,
                      "combined_vector": null,
                      "truncated": false,
                      "truncate_length": 512,
                      "model": "intfloat/multilingual-e5-large",
                      "max_seq_length": 512
                    },
                    {
                      "id": "%s",
                      "vector": [0.4, 0.5, 0.6],
                      "token_count": 7,
                      "runtime_ms": 0.0,
                      "incomplete": false,
                      "combined_vector": null,
                      "truncated": false,
                      "truncate_length": 512,
                      "model": "intfloat/multilingual-e5-large",
                      "max_seq_length": 512
                    }
                  ]
                }
                """.formatted(firstId, secondId));
    }

    private void respondJson(HttpExchange exchange, String json) throws IOException {
        byte[] response = json.getBytes(StandardCharsets.UTF_8);
        exchange.getResponseHeaders().add("Content-Type", "application/json");
        exchange.sendResponseHeaders(200, response.length);
        exchange.getResponseBody().write(response);
        exchange.close();
    }
}
