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

import at.procon.dip.embedding.model.EmbeddingInputPrefixing;
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.EmbeddingUseCase;
import at.procon.dip.embedding.model.ResolvedEmbeddingProviderConfig;
import at.procon.dip.embedding.provider.EmbeddingProvider;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.io.IOException;
import java.net.http.HttpResponse;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.springframework.stereotype.Component;

/**
 * HTTP provider for vector APIs.
 *
 * Supported endpoints:
 *   POST {baseUrl}/vector-sync      - single text
 *   POST {baseUrl}/vectorize-batch  - multiple texts
 */
@Component
public class VectorSyncHttpEmbeddingProvider extends AbstractHttpEmbeddingProviderSupport implements EmbeddingProvider {
    private static final String PROVIDER_TYPE = "http-vector-sync";

    private static final boolean DEFAULT_TRUNCATE_TEXT = false;
    private static final int DEFAULT_TRUNCATE_LENGTH = 512;
    private static final int DEFAULT_CHUNK_SIZE = 20;

    private static final List<String> TRUNCATE_TEXT_KEYS = List.of(
            "vectorize-batch.truncate-text",
            "vectorize-batch.truncate_text",
            "truncate_text",
            "truncate-text",
            "truncateText"
    );

    private static final List<String> TRUNCATE_LENGTH_KEYS = List.of(
            "vectorize-batch.truncate-length",
            "vectorize-batch.truncate_length",
            "truncate_length",
            "truncate-length",
            "truncateLength"
    );

    private static final List<String> CHUNK_SIZE_KEYS = List.of(
            "vectorize-batch.chunk-size",
            "vectorize-batch.chunk_size",
            "chunk_size",
            "chunk-size",
            "chunkSize"
    );

    public VectorSyncHttpEmbeddingProvider(ObjectMapper objectMapper) {
        super(objectMapper);
    }

    @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, EmbeddingUseCase.DOCUMENT);
    }

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

    private EmbeddingProviderResult execute(ResolvedEmbeddingProviderConfig providerConfig,
                                            EmbeddingModelDescriptor model,
                                            EmbeddingRequest request,
                                            EmbeddingUseCase useCase) {
        if (request.texts() == null || request.texts().isEmpty()) {
            throw new IllegalArgumentException("Embedding request texts must not be empty");
        }

        EmbeddingInputPrefixing.PrefixedTexts prefixedTexts = EmbeddingInputPrefixing.apply(model, useCase, request.texts());

        try {
            return prefixedTexts.texts().size() == 1
                    ? executeSingle(providerConfig, model, prefixedTexts)
                    : executeBatch(providerConfig, model, request, prefixedTexts);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("Embedding provider call interrupted", e);
        } catch (IOException e) {
            throw new IllegalStateException("Failed to call embedding provider", e);
        }
    }

    private EmbeddingProviderResult executeSingle(ResolvedEmbeddingProviderConfig providerConfig,
                                                  EmbeddingModelDescriptor model,
                                                  EmbeddingInputPrefixing.PrefixedTexts prefixedTexts) throws IOException, InterruptedException {
        HttpResponse<String> response = postJson(
                providerConfig,
                "/vector-sync",
                new VectorSyncRequest(model.providerModelKey(), prefixedTexts.texts().getFirst())
        );

        VectorSyncResponse parsed = objectMapper.readValue(response.body(), VectorSyncResponse.class);
        float[] vector = extractVector(parsed.vector, parsed.combinedVector, model);

        return new EmbeddingProviderResult(
                model,
                List.of(vector),
                List.of(),
                null,
                parsed.tokenCount,
                prefixedTexts.prefixMode(),
                prefixedTexts.appliedPrefix()
        );
    }

    private EmbeddingProviderResult executeBatch(ResolvedEmbeddingProviderConfig providerConfig,
                                                 EmbeddingModelDescriptor model,
                                                 EmbeddingRequest request,
                                                 EmbeddingInputPrefixing.PrefixedTexts prefixedTexts) throws IOException, InterruptedException {
        BatchRequestSettings settings = resolveBatchRequestSettings(providerConfig, request.providerOptions());

        if (settings.truncateLength() <= 0) {
            throw new IllegalArgumentException("Batch truncate length must be > 0");
        }
        if (settings.chunkSize() <= 0) {
            throw new IllegalArgumentException("Batch chunk size must be > 0");
        }

        List<String> requestOrder = new ArrayList<>(prefixedTexts.texts().size());
        List<VectorizeBatchItemRequest> items = new ArrayList<>(prefixedTexts.texts().size());

        for (String text : prefixedTexts.texts()) {
            String id = UUID.randomUUID().toString();
            requestOrder.add(id);
            items.add(new VectorizeBatchItemRequest(id, text));
        }

        HttpResponse<String> response = postJson(
                providerConfig,
                "/vectorize-batch",
                new VectorizeBatchRequest(
                        model.providerModelKey(),
                        settings.truncateText(),
                        settings.truncateLength(),
                        settings.chunkSize(),
                        true,
                        items
                )
        );

        VectorizeBatchResponse parsed = objectMapper.readValue(response.body(), VectorizeBatchResponse.class);
        if (parsed.results == null || parsed.results.isEmpty()) {
            throw new IllegalStateException("Vectorize-batch provider returned no results");
        }

        Map<String, VectorizeBatchItemResponse> resultById = new HashMap<>();
        for (VectorizeBatchItemResponse result : parsed.results) {
            resultById.put(result.id, result);
        }

        List<float[]> vectors = new ArrayList<>(prefixedTexts.texts().size());
        int totalTokenCount = 0;
        boolean hasAnyTokenCount = false;

        for (String id : requestOrder) {
            VectorizeBatchItemResponse item = resultById.get(id);
            if (item == null) {
                throw new IllegalStateException("Vectorize-batch provider response is missing item for id " + id);
            }

            vectors.add(extractVector(item.vector, item.combinedVector, model));

            if (item.tokenCount != null) {
                totalTokenCount += item.tokenCount;
                hasAnyTokenCount = true;
            }
        }

        return new EmbeddingProviderResult(
                model,
                vectors,
                List.of(),
                null,
                hasAnyTokenCount ? totalTokenCount : null,
                prefixedTexts.prefixMode(),
                prefixedTexts.appliedPrefix()
        );
    }

    private BatchRequestSettings resolveBatchRequestSettings(ResolvedEmbeddingProviderConfig providerConfig,
                                                             Map<String, Object> providerOptions) {
        boolean truncateText = resolveBooleanOption(
                providerOptions,
                TRUNCATE_TEXT_KEYS,
                providerConfig.batchTruncateText() != null ? providerConfig.batchTruncateText() : DEFAULT_TRUNCATE_TEXT
        );
        int truncateLength = resolveIntOption(
                providerOptions,
                TRUNCATE_LENGTH_KEYS,
                providerConfig.batchTruncateLength() != null ? providerConfig.batchTruncateLength() : DEFAULT_TRUNCATE_LENGTH
        );
        int chunkSize = resolveIntOption(
                providerOptions,
                CHUNK_SIZE_KEYS,
                providerConfig.batchChunkSize() != null ? providerConfig.batchChunkSize() : DEFAULT_CHUNK_SIZE
        );
        return new BatchRequestSettings(truncateText, truncateLength, chunkSize);
    }

    private boolean resolveBooleanOption(Map<String, Object> providerOptions,
                                         List<String> keys,
                                         boolean defaultValue) {
        Object raw = resolveOption(providerOptions, keys);
        if (raw == null) {
            return defaultValue;
        }
        if (raw instanceof Boolean booleanValue) {
            return booleanValue;
        }
        String normalized = String.valueOf(raw).trim();
        if (normalized.isEmpty()) {
            return defaultValue;
        }
        return Boolean.parseBoolean(normalized);
    }

    private int resolveIntOption(Map<String, Object> providerOptions,
                                 List<String> keys,
                                 int defaultValue) {
        Object raw = resolveOption(providerOptions, keys);
        if (raw == null) {
            return defaultValue;
        }
        if (raw instanceof Number number) {
            return number.intValue();
        }
        String normalized = String.valueOf(raw).trim();
        if (normalized.isEmpty()) {
            return defaultValue;
        }
        return Integer.parseInt(normalized);
    }

    private Object resolveOption(Map<String, Object> providerOptions, List<String> keys) {
        if (providerOptions == null || providerOptions.isEmpty()) {
            return null;
        }
        for (String key : keys) {
            if (providerOptions.containsKey(key)) {
                return providerOptions.get(key);
            }
        }
        return null;
    }

    private float[] extractVector(List<Float> vector,
                                  List<Float> combinedVector,
                                  EmbeddingModelDescriptor model) {
        float[] resolved;

        if (combinedVector != null && !combinedVector.isEmpty()) {
            resolved = toArray(combinedVector);
        } else if (vector != null && !vector.isEmpty()) {
            resolved = toArray(vector);
        } else {
            throw new IllegalStateException("Embedding provider returned no vector");
        }

        if (model.dimensions() > 0 && resolved.length != model.dimensions()) {
            throw new IllegalStateException(
                    "Embedding provider returned dimension %d for model %s, expected %d"
                            .formatted(resolved.length, model.modelKey(), model.dimensions())
            );
        }

        return resolved;
    }

    private record BatchRequestSettings(boolean truncateText, int truncateLength, int chunkSize) {
    }

    private record VectorSyncRequest(
            @JsonProperty("model") String model,
            @JsonProperty("text") String text
    ) {
    }

    private record VectorizeBatchRequest(
            @JsonProperty("model") String model,
            @JsonProperty("truncate_text") boolean truncateText,
            @JsonProperty("truncate_length") int truncateLength,
            @JsonProperty("chunk_size") int chunkSize,
            @JsonProperty("balance_chunks") boolean balanceChunks,
            @JsonProperty("items") List<VectorizeBatchItemRequest> items
    ) {
    }

    private record VectorizeBatchItemRequest(
            @JsonProperty("id") String id,
            @JsonProperty("text") String text
    ) {
    }

    static class VectorSyncResponse {
        @JsonProperty("runtime_ms")
        public Double runtimeMs;

        @JsonProperty("vector")
        public List<Float> vector;

        @JsonProperty("incomplete")
        public Boolean incomplete;

        @JsonProperty("combined_vector")
        public List<Float> combinedVector;

        @JsonProperty("token_count")
        public Integer tokenCount;

        @JsonProperty("model")
        public String model;

        @JsonProperty("max_seq_length")
        public Integer maxSeqLength;
    }

    static class VectorizeBatchResponse {
        @JsonProperty("model")
        public String model;

        @JsonProperty("count")
        public Integer count;

        @JsonProperty("results")
        public List<VectorizeBatchItemResponse> results;
    }

    static class VectorizeBatchItemResponse {
        @JsonProperty("id")
        public String id;

        @JsonProperty("vector")
        public List<Float> vector;

        @JsonProperty("token_count")
        public Integer tokenCount;

        @JsonProperty("runtime_ms")
        public Double runtimeMs;

        @JsonProperty("incomplete")
        public Boolean incomplete;

        @JsonProperty("combined_vector")
        public List<Float> combinedVector;

        @JsonProperty("truncated")
        public Boolean truncated;

        @JsonProperty("truncate_length")
        public Integer truncateLength;

        @JsonProperty("model")
        public String model;

        @JsonProperty("max_seq_length")
        public Integer maxSeqLength;
    }
}
