package at.procon.dip.migration.service;

import at.procon.dip.embedding.config.EmbeddingProperties;
import at.procon.dip.embedding.service.EmbeddingPersistenceService;
import at.procon.dip.embedding.service.RepresentationEmbeddingOrchestrator;
import at.procon.dip.migration.config.LegacyTedEmbeddingBackfillProperties;
import at.procon.dip.migration.repository.LegacyTedEmbeddingTargetQueryRepository;
import at.procon.dip.runtime.condition.ConditionalOnRuntimeMode;
import at.procon.dip.runtime.config.RuntimeMode;
import at.procon.ted.repository.LegacyTedEmbeddingSnapshot;
import at.procon.ted.repository.LegacyTedMigrationCursor;
import at.procon.ted.repository.ProcurementDocumentRepository;
import java.time.Instant;
import java.time.OffsetDateTime;
import java.time.ZoneId;
import java.util.ArrayList;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import org.springframework.util.StringUtils;

@Service
@ConditionalOnRuntimeMode(RuntimeMode.NEW)
@RequiredArgsConstructor
@Slf4j
public class LegacyTedEmbeddingBackfillService {

    private final LegacyTedEmbeddingBackfillProperties properties;
    private final ProcurementDocumentRepository procurementDocumentRepository;
    private final LegacyTedEmbeddingTargetQueryRepository targetQueryRepository;
    private final EmbeddingPersistenceService embeddingPersistenceService;
    private final RepresentationEmbeddingOrchestrator embeddingOrchestrator;
    private final EmbeddingProperties embeddingProperties;

    public void runBackfill() {
        if (!properties.isEnabled()) {
            log.info("Legacy TED embedding-only backfill is disabled");
            return;
        }
        if (!embeddingProperties.isEnabled() || !StringUtils.hasText(embeddingProperties.getDefaultDocumentModel())) {
            log.warn("Skipping legacy TED embedding-only backfill because no default document model is configured");
            return;
        }

        log.info("Starting legacy TED embedding-only backfill (batchSize={}, maxDocumentsPerRun={}, queueMissingEmbeddings={})",
                properties.getBatchSize(), properties.getMaxDocumentsPerRun(), properties.isQueueMissingEmbeddings());

        Instant lastCreatedAt = null;
        UUID lastId = null;
        long inspected = 0;
        long migrated = 0;
        long queued = 0;
        long skippedMissingTarget = 0;
        long skippedMissingVector = 0;

        while (true) {
            int limit = effectiveBatchLimit(inspected);
            if (limit <= 0) {
                break;
            }

            List<LegacyTedMigrationCursor> cursors = loadNextBatch(
                    lastCreatedAt != null ?
                    lastCreatedAt.atZone(ZoneId.systemDefault()).toOffsetDateTime() : null, lastId, limit);
            if (cursors.isEmpty()) {
                break;
            }

            List<UUID> legacyIds = new ArrayList<>(cursors.size());
            for (LegacyTedMigrationCursor cursor : cursors) {
                legacyIds.add(cursor.getId());
            }

            Map<UUID, LegacyTedEmbeddingTarget> targetsByLegacyId = indexTargets(
                    targetQueryRepository.findPrimarySemanticTargetsByLegacyIds(legacyIds)
            );
            Map<UUID, LegacyTedEmbeddingSnapshot> snapshotsByLegacyId = indexSnapshots(
                    procurementDocumentRepository.findEmbeddingSnapshotsByIdsForMigration(legacyIds)
            );

            for (LegacyTedMigrationCursor cursor : cursors) {
                inspected++;
                LegacyTedEmbeddingTarget target = targetsByLegacyId.get(cursor.getId());
                LegacyTedEmbeddingSnapshot snapshot = snapshotsByLegacyId.get(cursor.getId());

                if (target == null) {
                    if (properties.isSkipWhenPrimaryRepresentationMissing()) {
                        skippedMissingTarget++;
                    }
                    lastCreatedAt = cursor.getCreatedAt();
                    lastId = cursor.getId();
                    continue;
                }

                if (snapshot != null && StringUtils.hasText(snapshot.getVectorText())) {
                    float[] vector = parseVectorText(snapshot.getVectorText());
                    if (vector.length > 0) {
                        var embedding = embeddingPersistenceService.ensurePending(target.representationId(), embeddingProperties.getDefaultDocumentModel());
                        embeddingPersistenceService.saveCompleted(embedding.getId(), vector, snapshot.getTokenCount());
                        migrated++;
                    } else {
                        skippedMissingVector++;
                    }
                } else if (properties.isQueueMissingEmbeddings()) {
                    embeddingOrchestrator.enqueueRepresentation(
                            target.documentId(),
                            target.representationId(),
                            embeddingProperties.getDefaultDocumentModel()
                    );
                    queued++;
                } else {
                    skippedMissingVector++;
                }

                lastCreatedAt = cursor.getCreatedAt();
                lastId = cursor.getId();
            }

            log.info("Legacy TED embedding-only backfill progress: inspected={}, migrated={}, queued={}, skippedMissingTarget={}, skippedMissingVector={}",
                    inspected, migrated, queued, skippedMissingTarget, skippedMissingVector);
        }

        log.info("Legacy TED embedding-only backfill finished: inspected={}, migrated={}, queued={}, skippedMissingTarget={}, skippedMissingVector={}",
                inspected, migrated, queued, skippedMissingTarget, skippedMissingVector);
    }

    protected List<LegacyTedMigrationCursor> loadNextBatch(OffsetDateTime lastCreatedAt, UUID lastId, int limit) {
        if (lastCreatedAt == null || lastId == null) {
            return procurementDocumentRepository.findFirstMigrationBatch(limit);
        }
        return procurementDocumentRepository.findNextMigrationBatch(lastCreatedAt, lastId, limit);
    }

    private int effectiveBatchLimit(long processedInThisInvocation) {
        long maxPerRun = properties.getMaxDocumentsPerRun();
        if (maxPerRun <= 0) {
            return Math.max(1, properties.getBatchSize());
        }
        long remaining = maxPerRun - processedInThisInvocation;
        if (remaining <= 0) {
            return 0;
        }
        return (int) Math.max(1L, Math.min(properties.getBatchSize(), remaining));
    }

    private Map<UUID, LegacyTedEmbeddingTarget> indexTargets(Collection<LegacyTedEmbeddingTarget> targets) {
        Map<UUID, LegacyTedEmbeddingTarget> indexed = new LinkedHashMap<>();
        for (LegacyTedEmbeddingTarget target : targets) {
            indexed.putIfAbsent(target.legacyProcurementDocumentId(), target);
        }
        return indexed;
    }

    private Map<UUID, LegacyTedEmbeddingSnapshot> indexSnapshots(Collection<LegacyTedEmbeddingSnapshot> snapshots) {
        Map<UUID, LegacyTedEmbeddingSnapshot> indexed = new LinkedHashMap<>();
        for (LegacyTedEmbeddingSnapshot snapshot : snapshots) {
            indexed.put(snapshot.getId(), snapshot);
        }
        return indexed;
    }

    private float[] parseVectorText(String vectorText) {
        if (!StringUtils.hasText(vectorText)) {
            return new float[0];
        }
        String trimmed = vectorText.trim();
        if (trimmed.startsWith("[")) {
            trimmed = trimmed.substring(1);
        }
        if (trimmed.endsWith("]")) {
            trimmed = trimmed.substring(0, trimmed.length() - 1);
        }
        if (!StringUtils.hasText(trimmed)) {
            return new float[0];
        }
        String[] parts = trimmed.split(",");
        float[] result = new float[parts.length];
        for (int i = 0; i < parts.length; i++) {
            result[i] = Float.parseFloat(parts[i].trim());
        }
        return result;
    }
}
