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.List;
import java.util.Map;
import org.springframework.stereotype.Component;

/**
 * Compatibility provider for the currently running external Python vectorization service.
 *
 * It calls /embed with the legacy payload shape { text, isQuery }.
 */
@Component
public class ExternalHttpEmbeddingProvider extends AbstractHttpEmbeddingProviderSupport implements EmbeddingProvider {

    private static final String PROVIDER_TYPE = "http-json";

    public ExternalHttpEmbeddingProvider(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 {
            HttpResponse<String> response = postJson(
                    providerConfig,
                    "/embed",
                    Map.of(
                            "text", prefixedTexts.texts().getFirst(),
                            "isQuery", useCase == EmbeddingUseCase.QUERY
                    )
            );

            EmbedResponse parsed = objectMapper.readValue(response.body(), EmbedResponse.class);
            if (parsed.embedding == null) {
                throw new IllegalStateException("Embedding provider returned no embedding");
            }

            return new EmbeddingProviderResult(
                    model,
                    List.of(parsed.embedding),
                    List.of(),
                    null,
                    parsed.tokenCount,
                    prefixedTexts.prefixMode(),
                    prefixedTexts.appliedPrefix()
            );
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("Embedding provider call interrupted", e);
        } catch (IOException e) {
            throw new IllegalStateException("Failed to call external embedding provider", e);
        }
    }

    public static class EmbedResponse {
        @JsonProperty("embedding")
        public float[] embedding;

        @JsonProperty("dimensions")
        public Integer dimensions;

        @JsonProperty("token_count")
        public Integer tokenCount;
    }
}
