package at.procon.dip.embedding.config;

import at.procon.dip.domain.document.DistanceMetric;
import at.procon.dip.domain.document.DocumentType;
import at.procon.dip.domain.document.RepresentationType;
import at.procon.dip.embedding.model.EmbeddingPrefixMode;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.Map;
import lombok.Data;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Configuration;

@ConfigurationProperties(prefix = "dip.embedding")
@Data
@Configuration
public class EmbeddingProperties {

    private boolean enabled = false;
    private String defaultDocumentModel;
    private String defaultQueryModel;
    private Map<String, ProviderProperties> providers = new LinkedHashMap<>();
    private Map<String, ModelProperties> models = new LinkedHashMap<>();
    private IndexingProperties indexing = new IndexingProperties();
    private JobsProperties jobs = new JobsProperties();
    private StartupProperties startup = new StartupProperties();

    @Data
    public static class ProviderProperties {
        private String type;
        private String baseUrl;
        private String apiKey;
        private Duration connectTimeout = Duration.ofSeconds(5);
        private Duration readTimeout = Duration.ofSeconds(60);
        private Map<String, String> headers = new LinkedHashMap<>();
        private Integer dimensions;
        private BatchRequestProperties batchRequest = new BatchRequestProperties();
    }

    @Data
    public static class BatchRequestProperties {
        private boolean truncateText = false;
        private int truncateLength = 512;
        private int chunkSize = 20;
    }

    @Data
    public static class ModelProperties {
        private String providerConfigKey;
        private String providerModelKey;
        private Integer dimensions;
        private DistanceMetric distanceMetric = DistanceMetric.COSINE;
        private boolean supportsQueryEmbeddingMode = true;
        private boolean supportsBatch = false;
        private Integer maxInputChars;
        private boolean active = true;
        private EmbeddingPrefixMode prefixMode = EmbeddingPrefixMode.OFF;
        private String queryPrefix = "query: ";
        private String documentPrefix = "passage: ";
    }

    @Data
    public static class IndexingProperties {
        private boolean embedSemanticText = true;
        private boolean embedTitleAbstract = true;
        private boolean embedChunks = true;
        private boolean embedFulltext = false;
        private boolean embedSummary = false;
        private int chunkMinLength = 300;
        private int fallbackMaxInputChars = 8192;
    }

    @Data
    public static class JobsProperties {
        private boolean enabled = false;
        private int batchSize = 16;
        private int parallelBatchCount = 1;
        private boolean processInBatches = false;
        private int executionBatchSize = 8;
        private int maxRetries = 5;
        private Duration initialRetryDelay = Duration.ofSeconds(30);
        private Duration maxRetryDelay = Duration.ofHours(6);
        private long schedulerDelayMs = 5000;
    }

    @Data
    public static class StartupProperties {
        private boolean enqueueMissingEnabled = false;
        private boolean processReadyEnabled = false;
        private DocumentType documentType;
        private RepresentationType representationType;
        private String builderKey;
        private boolean primaryOnly = false;
        private String modelKey;
        private boolean force = false;
        private int batchSize = 1000;
        private long maxRepresentationsPerRun = 0;
    }
}
