package at.procon.dip.embedding.service;

import at.procon.dip.domain.document.entity.DocumentTextRepresentation;
import at.procon.dip.domain.document.repository.DocumentTextRepresentationRepository;
import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.job.entity.EmbeddingJob;
import at.procon.dip.embedding.job.service.EmbeddingJobService;
import at.procon.dip.embedding.model.EmbeddingJobType;
import at.procon.dip.embedding.model.EmbeddingModelDescriptor;
import at.procon.dip.embedding.model.EmbeddingProviderResult;
import at.procon.dip.embedding.model.EmbeddingUseCase;
import at.procon.dip.embedding.policy.EmbeddingProfile;
import at.procon.dip.embedding.policy.EmbeddingSelectionPolicy;
import at.procon.dip.embedding.registry.EmbeddingModelRegistry;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CompletionException;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.stereotype.Service;

@Service
@Slf4j
public class RepresentationEmbeddingOrchestrator {

    private final EmbeddingJobService jobService;
    private final EmbeddingExecutionService executionService;
    private final EmbeddingJobExecutionPersistenceService executionPersistenceService;
    private final DocumentTextRepresentationRepository representationRepository;
    private final EmbeddingSelectionPolicy selectionPolicy;
    private final EmbeddingModelRegistry modelRegistry;
    private final EmbeddingProperties embeddingProperties;
    private final Executor embeddingJobProcessingExecutor;

    public RepresentationEmbeddingOrchestrator(EmbeddingJobService jobService,
                                               EmbeddingExecutionService executionService,
                                               EmbeddingJobExecutionPersistenceService executionPersistenceService,
                                               DocumentTextRepresentationRepository representationRepository,
                                               EmbeddingSelectionPolicy selectionPolicy,
                                               EmbeddingModelRegistry modelRegistry,
                                               EmbeddingProperties embeddingProperties,
                                               @Qualifier("embeddingJobProcessingExecutor") Executor embeddingJobProcessingExecutor) {
        this.jobService = jobService;
        this.executionService = executionService;
        this.executionPersistenceService = executionPersistenceService;
        this.representationRepository = representationRepository;
        this.selectionPolicy = selectionPolicy;
        this.modelRegistry = modelRegistry;
        this.embeddingProperties = embeddingProperties;
        this.embeddingJobProcessingExecutor = embeddingJobProcessingExecutor;
    }

    public List<EmbeddingJob> enqueueDocument(UUID documentId) {
        return jobService.enqueueForDocument(documentId);
    }

    public List<EmbeddingJob> enqueueDocument(UUID documentId, String modelKey) {
        return jobService.enqueueForDocument(documentId, modelKey);
    }

    public List<EmbeddingJob> enqueueDocument(UUID documentId, String modelKey, EmbeddingProfile profile) {
        var model = modelRegistry.getRequired(modelKey);
        return selectionPolicy.selectRepresentations(documentId, model, profile).stream()
                .map(representation -> enqueueRepresentation(documentId, representation.getId(), modelKey))
                .toList();
    }

    public EmbeddingJob enqueueRepresentation(UUID documentId, UUID representationId, String modelKey) {
        return jobService.enqueueForRepresentation(documentId, representationId, modelKey, EmbeddingJobType.DOCUMENT_EMBED);
    }

    public int processNextReadyBatch() {
        if (!embeddingProperties.isEnabled() || !embeddingProperties.getJobs().isEnabled()) {
            log.debug("New embedding subsystem jobs are disabled");
            return 0;
        }

        int batchSize = Math.max(1, embeddingProperties.getJobs().getBatchSize());
        int parallelBatchCount = Math.max(1, embeddingProperties.getJobs().getParallelBatchCount());

        return parallelBatchCount == 1
                ? processWithSingleBatchSlot(batchSize)
                : processWithSlidingParallelBatchWindow(batchSize, parallelBatchCount);
    }

    public void processClaimedJob(EmbeddingJob job) {
        EmbeddingModelDescriptor model = modelRegistry.getRequired(job.getModelKey());
        PreparedEmbedding prepared = prepareEmbedding(job, model);
        if (prepared == null) {
            return;
        }

        try {
            EmbeddingProviderResult result = executionService.embedTexts(
                    job.getModelKey(),
                    EmbeddingUseCase.DOCUMENT,
                    List.of(prepared.text())
            );
            executionPersistenceService.completeJob(
                    prepared.embeddingId(),
                    result,
                    job.getId(),
                    result.providerRequestId()
            );
        } catch (RuntimeException ex) {
            executionPersistenceService.failJob(
                    prepared.embeddingId(),
                    job.getId(),
                    ex.getMessage(),
                    true
            );
            throw ex;
        }
    }

    private int processWithSingleBatchSlot(int batchSize) {
        int claimedTotal = 0;
        while (true) {
            List<EmbeddingJob> claimedBatch = jobService.claimNextReadyJobs(batchSize);
            if (claimedBatch.isEmpty()) {
                return claimedTotal;
            }
            claimedTotal += claimedBatch.size();
            processClaimedJobBatch(claimedBatch);
        }
    }

    private int processWithSlidingParallelBatchWindow(int batchSize, int parallelBatchCount) {
        ExecutorCompletionService<Integer> completionService = new ExecutorCompletionService<>(embeddingJobProcessingExecutor);
        int claimedTotal = 0;
        int activeBatchCount = 0;
        boolean noMoreReadyJobs = false;

        while (activeBatchCount < parallelBatchCount && !noMoreReadyJobs) {
            List<EmbeddingJob> claimedBatch = jobService.claimNextReadyJobs(batchSize);
            if (claimedBatch.isEmpty()) {
                noMoreReadyJobs = true;
            } else {
                claimedTotal += claimedBatch.size();
                submitClaimedBatch(completionService, claimedBatch);
                activeBatchCount++;
            }
        }

        while (activeBatchCount > 0) {
            waitForBatchCompletion(takeCompletedBatch(completionService));
            activeBatchCount--;

            if (!noMoreReadyJobs) {
                List<EmbeddingJob> claimedBatch = jobService.claimNextReadyJobs(batchSize);
                if (claimedBatch.isEmpty()) {
                    noMoreReadyJobs = true;
                } else {
                    claimedTotal += claimedBatch.size();
                    submitClaimedBatch(completionService, claimedBatch);
                    activeBatchCount++;
                }
            }
        }

        return claimedTotal;
    }

    private void submitClaimedBatch(ExecutorCompletionService<Integer> completionService,
                                    List<EmbeddingJob> claimedBatch) {
        completionService.submit(() -> {
            processClaimedJobBatch(claimedBatch);
            return claimedBatch.size();
        });
    }

    private Future<Integer> takeCompletedBatch(ExecutorCompletionService<Integer> completionService) {
        try {
            return completionService.take();
        } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("Interrupted while waiting for an embedding batch slot to become free", ex);
        }
    }

    private void waitForBatchCompletion(Future<Integer> future) {
        try {
            future.get();
        } catch (InterruptedException ex) {
            Thread.currentThread().interrupt();
            throw new IllegalStateException("Interrupted while waiting for embedding batch completion", ex);
        } catch (java.util.concurrent.ExecutionException ex) {
            Throwable cause = ex.getCause();
            if (cause instanceof RuntimeException runtimeException) {
                throw runtimeException;
            }
            throw new CompletionException(cause);
        }
    }

    private void processClaimedJobBatch(List<EmbeddingJob> jobs) {
        if (embeddingProperties.getJobs().isProcessInBatches()) {
            processClaimedJobsInExecutionBatches(jobs);
        } else {
            jobs.forEach(this::processClaimedJobSafely);
        }
    }

    private void processClaimedJobsInExecutionBatches(List<EmbeddingJob> jobs) {
        LinkedHashMap<String, List<EmbeddingJob>> jobsByModelKey = new LinkedHashMap<>();
        for (EmbeddingJob job : jobs) {
            jobsByModelKey.computeIfAbsent(job.getModelKey(), ignored -> new ArrayList<>()).add(job);
        }

        int executionBatchSize = Math.max(1, embeddingProperties.getJobs().getExecutionBatchSize());
        for (var entry : jobsByModelKey.entrySet()) {
            EmbeddingModelDescriptor model = modelRegistry.getRequired(entry.getKey());
            if (!model.supportsBatch()) {
                entry.getValue().forEach(this::processClaimedJobSafely);
                continue;
            }

            List<EmbeddingJob> sameModelJobs = entry.getValue();
            for (int start = 0; start < sameModelJobs.size(); start += executionBatchSize) {
                List<EmbeddingJob> partition = sameModelJobs.subList(start, Math.min(start + executionBatchSize, sameModelJobs.size()));
                if (partition.size() == 1) {
                    processClaimedJobSafely(partition.getFirst());
                } else {
                    processClaimedBatchSafely(partition, model);
                }
            }
        }
    }


    private void processClaimedBatchSafely(List<EmbeddingJob> jobs, EmbeddingModelDescriptor model) {
        try {
            processClaimedBatch(jobs, model);
        } catch (RuntimeException ex) {
            log.warn("Failed to process embedding batch for model {} ({} jobs): {}",
                    model.modelKey(), jobs.size(), ex.getMessage(), ex);
        }
    }

    private void processClaimedJobSafely(EmbeddingJob job) {
        try {
            processClaimedJob(job);
        } catch (RuntimeException ex) {
            log.warn("Failed to process embedding job {} for representation {}: {}",
                    job.getId(), job.getRepresentationId(), ex.getMessage(), ex);
        }
    }

    private void processClaimedBatch(List<EmbeddingJob> jobs, EmbeddingModelDescriptor model) {
        List<PreparedEmbedding> preparedItems = new ArrayList<>(jobs.size());
        for (EmbeddingJob job : jobs) {
            PreparedEmbedding prepared = prepareEmbedding(job, model);
            if (prepared != null) {
                preparedItems.add(prepared);
            }
        }

        if (preparedItems.isEmpty()) {
            return;
        }

        try {
            EmbeddingProviderResult result = executionService.embedTexts(
                    model.modelKey(),
                    EmbeddingUseCase.DOCUMENT,
                    preparedItems.stream().map(PreparedEmbedding::text).toList()
            );

            if (result.vectors() == null || result.vectors().size() != preparedItems.size()) {
                throw new IllegalStateException(
                        "Embedding provider returned %d vectors for %d batch items"
                                .formatted(result.vectors() == null ? 0 : result.vectors().size(), preparedItems.size())
                );
            }

            for (int i = 0; i < preparedItems.size(); i++) {
                PreparedEmbedding prepared = preparedItems.get(i);
                executionPersistenceService.completeJob(
                        prepared.embeddingId(),
                        result.vectors().get(i),
                        null,
                        prepared.job().getId(),
                        result.providerRequestId(),
                        result.prefixMode(),
                        result.appliedPrefix()
                );
            }
        } catch (RuntimeException ex) {
            for (PreparedEmbedding prepared : preparedItems) {
                executionPersistenceService.failJob(
                        prepared.embeddingId(),
                        prepared.job().getId(),
                        ex.getMessage(),
                        true
                );
            }
            throw ex;
        }
    }

    private PreparedEmbedding prepareEmbedding(EmbeddingJob job, EmbeddingModelDescriptor model) {
        DocumentTextRepresentation representation = representationRepository.findById(job.getRepresentationId())
                .orElseThrow(() -> new IllegalArgumentException("Unknown representation id: " + job.getRepresentationId()));

        String text = representation.getTextBody();
        if (text == null || text.isBlank()) {
            jobService.markFailed(job.getId(), "No text representation available", false);
            return null;
        }

        int maxChars = model.maxInputChars() != null
                ? model.maxInputChars()
                : embeddingProperties.getIndexing().getFallbackMaxInputChars();
        if (text.length() > maxChars) {
            text = text.substring(0, maxChars);
        }

        UUID embeddingId = executionPersistenceService.startProcessing(representation.getId(), job.getModelKey());
        return new PreparedEmbedding(job, embeddingId, text);
    }

    private record PreparedEmbedding(EmbeddingJob job, UUID embeddingId, String text) {
    }
}
