package at.procon.dip.clustering.service;

import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.createdAtFrom;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.createdAtTo;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.hasAlgorithm;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.hasDocumentType;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.hasExecutionBackend;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.hasStatus;
import static at.procon.dip.clustering.repository.EmbeddingClusterRunSpecifications.nameContains;

import at.procon.dip.clustering.ClusterRunStatus;
import at.procon.dip.clustering.ClusteringAlgorithm;
import at.procon.dip.clustering.ClusteringExecutionBackend;
import at.procon.dip.clustering.ReductionMethod;
import at.procon.dip.clustering.PythonRequestMode;
import at.procon.dip.clustering.client.PythonClusteringClient;
import at.procon.dip.clustering.config.ClusteringPhaseBProperties;
import at.procon.dip.clustering.dto.ClusterAssignmentViewResponse;
import at.procon.dip.clustering.dto.ClusterMembersResponse;
import at.procon.dip.clustering.dto.ClusterResponse;
import at.procon.dip.clustering.dto.ClusterRunResponse;
import at.procon.dip.clustering.dto.ClusteringEngineAssignment;
import at.procon.dip.clustering.dto.ClusteringEngineCluster;
import at.procon.dip.clustering.dto.ClusteringEngineRequest;
import at.procon.dip.clustering.dto.ClusteringEngineResult;
import at.procon.dip.clustering.dto.CreateClusterRunRequest;
import at.procon.dip.clustering.dto.EmbeddingSelectionSpec;
import at.procon.dip.clustering.dto.PythonClusteringRequest;
import at.procon.dip.clustering.dto.PythonClusteringResponse;
import at.procon.dip.clustering.dto.PythonRunExecutionRequest;
import at.procon.dip.clustering.dto.ReductionConfig;
import at.procon.dip.clustering.dto.SelectedEmbeddingRow;
import at.procon.dip.clustering.entity.EmbeddingClusterRun;
import at.procon.dip.clustering.entity.EmbeddingClusterSet;
import at.procon.dip.clustering.repository.EmbeddingClusterRepository;
import at.procon.dip.clustering.repository.EmbeddingClusterResultQueryRepository;
import at.procon.dip.clustering.repository.EmbeddingClusterRunRepository;
import at.procon.dip.clustering.repository.EmbeddingClusterSetRepository;
import at.procon.dip.clustering.spi.EmbeddingClusteringEngine;
import at.procon.dip.domain.document.DocumentType;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import java.time.OffsetDateTime;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.UUID;
import lombok.RequiredArgsConstructor;
import org.springframework.data.domain.Sort;
import org.springframework.data.jpa.domain.Specification;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.server.ResponseStatusException;

@Service
@RequiredArgsConstructor
public class EmbeddingClusterRunService {

    private static final int DEFAULT_PREVIEW_LENGTH = 1000;
    private static final TypeReference<Map<String, Object>> PARAMETERS_TYPE = new TypeReference<>() {
    };
    private static final Set<ClusterRunStatus> STARTABLE_STATUSES = Set.of(ClusterRunStatus.CREATED);
    private static final Set<ClusterRunStatus> FINAL_STATUSES = Set.of(
            ClusterRunStatus.COMPLETED, ClusterRunStatus.FAILED, ClusterRunStatus.CANCELLED);

    private final EmbeddingClusterRunRepository runRepository;
    private final EmbeddingClusterSetRepository clusterSetRepository;
    private final EmbeddingClusterRepository clusterRepository;
    private final EmbeddingClusterResultQueryRepository resultQueryRepository;
    private final EmbeddingSelectionService selectionService;
    private final EmbeddingClusterPersistenceService persistenceService;
    private final List<EmbeddingClusteringEngine> clusteringEngines;
    private final ClusteringPhaseBProperties pythonProperties;
    private final Optional<PythonClusteringClient> pythonClusteringClient;
    private final ObjectMapper objectMapper;

    @Transactional
    public ClusterRunResponse createRun(CreateClusterRunRequest request) {
        EmbeddingClusterRun run = EmbeddingClusterRun.builder()
                .clusterSet(resolveClusterSet(request.clusterSetCode()))
                .name(request.name())
                .status(ClusterRunStatus.CREATED)
                .algorithm(request.algorithm())
                .algorithmVersion("phase-e-dual-python")
                .executionBackend(resolveExecutionBackend(request))
                .reductionMethod(resolveReductionMethod(request.reduction()))
                .reductionDimensions(resolveReductionDimensions(request.reduction()))
                .selectionJson(writeJson(request.selection()))
                .parametersJson(writeJson(request.resolvedParameters()))
                .documentType(firstOrNull(request.selection() == null ? null : request.selection().documentTypes()))
                .documentFamily(firstOrNull(request.selection() == null ? null : request.selection().documentFamilies()))
                .representationType(firstOrNull(request.selection() == null ? null : request.selection().representationTypes()))
                .builderKey(firstOrNullString(request.selection() == null ? null : request.selection().builderKeys()))
                .build();
        return toResponse(runRepository.save(run));
    }

    public ClusterRunResponse executeRun(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        return run.getExecutionBackend() == ClusteringExecutionBackend.PYTHON_REMOTE
                ? executeRunRemote(run)
                : executeRunLocal(run);
    }

    @Transactional
    public ClusterRunResponse queueRun(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        if (!STARTABLE_STATUSES.contains(run.getStatus())) {
            throw new ResponseStatusException(HttpStatus.CONFLICT,
                    "Cluster run cannot be started from status " + run.getStatus());
        }
        run.setStatus(ClusterRunStatus.QUEUED);
        run.setStartedAt(null);
        run.setFinishedAt(null);
        run.setErrorMessage(null);
        run.setItemCount(null);
        run.setClusterCount(null);
        run.setNoiseCount(null);
        return toResponse(runRepository.save(run));
    }

    public ClusterRunResponse executeQueuedRun(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        if (run.getStatus() == ClusterRunStatus.CANCELLED) {
            return toResponse(run);
        }
        if (run.getStatus() != ClusterRunStatus.QUEUED && run.getStatus() != ClusterRunStatus.CANCEL_REQUESTED) {
            throw new ResponseStatusException(HttpStatus.CONFLICT,
                    "Cluster run is not queued for execution: " + runId);
        }

        transitionToRunning(runId);

        if (loadRun(runId).getExecutionBackend() == ClusteringExecutionBackend.PYTHON_REMOTE) {
            return executeQueuedRemoteRun(runId);
        }
        return executeQueuedLocalRun(runId);
    }

    public ClusterRunResponse getRun(UUID runId) {
        return toResponse(loadRun(runId));
    }

    public List<ClusterRunResponse> listRuns(
            ClusterRunStatus status,
            ClusteringAlgorithm algorithm,
            ClusteringExecutionBackend executionBackend,
            DocumentType documentType,
            String nameLike,
            OffsetDateTime createdFrom,
            OffsetDateTime createdTo) {

        Specification<EmbeddingClusterRun> spec = Specification
                .where(hasStatus(status))
                .and(hasAlgorithm(algorithm))
                .and(hasExecutionBackend(executionBackend))
                .and(hasDocumentType(documentType))
                .and(nameContains(nameLike))
                .and(createdAtFrom(createdFrom))
                .and(createdAtTo(createdTo));

        return runRepository.findAll(spec, Sort.by(Sort.Direction.DESC, "createdAt")).stream()
                .map(this::toResponse)
                .toList();
    }

    @Transactional
    public ClusterRunResponse requestCancellation(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        if (FINAL_STATUSES.contains(run.getStatus())) {
            return toResponse(run);
        }
        if (run.getStatus() == ClusterRunStatus.CREATED || run.getStatus() == ClusterRunStatus.QUEUED) {
            run.setStatus(ClusterRunStatus.CANCELLED);
            run.setFinishedAt(OffsetDateTime.now());
            run.setErrorMessage("Cluster run was cancelled");
            return toResponse(runRepository.save(run));
        }
        if (run.getStatus() == ClusterRunStatus.RUNNING) {
            run.setStatus(ClusterRunStatus.CANCEL_REQUESTED);
            run.setErrorMessage("Cluster run cancellation requested");
            return toResponse(runRepository.save(run));
        }
        return toResponse(run);
    }

    public List<ClusterResponse> listClusters(UUID runId) {
        ensureRunExists(runId);
        return clusterRepository.findByClusterRun_IdOrderByClusterLabelAsc(runId).stream()
                .map(cluster -> new ClusterResponse(
                        cluster.getId(),
                        cluster.getClusterLabel(),
                        cluster.getDisplayName(),
                        cluster.getItemCount(),
                        cluster.isNoiseCluster(),
                        cluster.getSummaryText()))
                .toList();
    }

    public List<ClusterAssignmentViewResponse> listAssignments(UUID runId, boolean includeText) {
        ensureRunExists(runId);
        return resultQueryRepository.findAssignments(runId, includeText, DEFAULT_PREVIEW_LENGTH);
    }

    public List<ClusterMembersResponse> listClusterMembers(UUID runId, UUID clusterId, boolean includeText) {
        ensureRunExists(runId);
        return resultQueryRepository.findClusterMembers(runId, clusterId, includeText, DEFAULT_PREVIEW_LENGTH);
    }

    private ClusterRunResponse executeRunRemote(EmbeddingClusterRun run) {
        EmbeddingSelectionSpec selection = readSelection(run.getSelectionJson());
        List<SelectedEmbeddingRow> selected = null;
        long itemCount;

        if (useCompactPythonRunMode()) {
            itemCount = selectionService.count(selection);
        } else {
            selected = selectionService.load(selection);
            itemCount = selected.size();
        }

        if (itemCount == 0L) {
            failRun(run.getId(), "Selection contains no completed embeddings");
            return getRun(run.getId());
        }

        run.setStatus(ClusterRunStatus.RUNNING);
        run.setStartedAt(OffsetDateTime.now());
        run.setItemCount(itemCount);
        run.setErrorMessage(null);
        run.setFinishedAt(null);
        run.setClusterCount(null);
        run.setNoiseCount(null);
        runRepository.save(run);

        try {
            ClusteringEngineResult result = executePythonClustering(run, selected);
            persistenceService.persist(run, result);
            run.setStatus(ClusterRunStatus.COMPLETED);
            run.setFinishedAt(OffsetDateTime.now());
            run.setClusterCount(result.clusters().stream().filter(cluster -> !cluster.noiseCluster()).count());
            run.setNoiseCount(result.noiseCount());
            runRepository.save(run);
            return toResponse(run);
        } catch (Exception ex) {
            failRun(run.getId(), ex.getMessage());
            return getRun(run.getId());
        }
    }

    private ClusterRunResponse executeRunLocal(EmbeddingClusterRun run) {
        EmbeddingSelectionSpec selection = readSelection(run.getSelectionJson());
        List<SelectedEmbeddingRow> selected = selectionService.load(selection);
        if (selected.isEmpty()) {
            failRun(run.getId(), "Selection contains no completed embeddings");
            return getRun(run.getId());
        }

        run.setStatus(ClusterRunStatus.RUNNING);
        run.setStartedAt(OffsetDateTime.now());
        run.setItemCount((long) selected.size());
        run.setErrorMessage(null);
        run.setFinishedAt(null);
        run.setClusterCount(null);
        run.setNoiseCount(null);
        runRepository.save(run);

        try {
            ClusteringEngineResult result = executeLocalClustering(run, selected);
            persistenceService.persist(run, result);
            run.setStatus(ClusterRunStatus.COMPLETED);
            run.setFinishedAt(OffsetDateTime.now());
            run.setClusterCount(result.clusters().stream().filter(cluster -> !cluster.noiseCluster()).count());
            run.setNoiseCount(result.noiseCount());
            runRepository.save(run);
            return toResponse(run);
        } catch (Exception ex) {
            failRun(run.getId(), ex.getMessage());
            return getRun(run.getId());
        }
    }

    private ClusterRunResponse executeQueuedRemoteRun(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        EmbeddingSelectionSpec selection = readSelection(run.getSelectionJson());
        List<SelectedEmbeddingRow> selected = null;
        long itemCount;

        if (useCompactPythonRunMode()) {
            itemCount = selectionService.count(selection);
        } else {
            selected = selectionService.load(selection);
            itemCount = selected.size();
        }

        if (itemCount == 0L) {
            failRun(runId, "Selection contains no completed embeddings");
            return getRun(runId);
        }
        updateItemCount(runId, itemCount);

        if (isCancellationRequested(runId)) {
            cancelRunNow(runId, "Cluster run was cancelled before clustering started");
            return getRun(runId);
        }

        try {
            ClusteringEngineResult result = executePythonClustering(loadRun(runId), selected);
            if (isCancellationRequested(runId)) {
                cancelRunNow(runId, "Cluster run was cancelled before results were persisted");
                return getRun(runId);
            }
            persistenceService.persist(loadRun(runId), result);
            completeRun(runId, result);
            return getRun(runId);
        } catch (Exception ex) {
            failRun(runId, ex.getMessage());
            return getRun(runId);
        }
    }

    private ClusterRunResponse executeQueuedLocalRun(UUID runId) {
        EmbeddingSelectionSpec selection = readSelection(loadRun(runId).getSelectionJson());
        List<SelectedEmbeddingRow> selected = selectionService.load(selection);
        if (selected.isEmpty()) {
            failRun(runId, "Selection contains no completed embeddings");
            return getRun(runId);
        }
        updateItemCount(runId, (long) selected.size());

        if (isCancellationRequested(runId)) {
            cancelRunNow(runId, "Cluster run was cancelled before clustering started");
            return getRun(runId);
        }

        try {
            ClusteringEngineResult result = executeLocalClustering(loadRun(runId), selected);
            if (isCancellationRequested(runId)) {
                cancelRunNow(runId, "Cluster run was cancelled before results were persisted");
                return getRun(runId);
            }
            persistenceService.persist(loadRun(runId), result);
            completeRun(runId, result);
            return getRun(runId);
        } catch (Exception ex) {
            failRun(runId, ex.getMessage());
            return getRun(runId);
        }
    }

    private void transitionToRunning(UUID runId) {
        EmbeddingClusterRun run = loadRun(runId);
        if (run.getStatus() == ClusterRunStatus.CANCELLED) {
            return;
        }
        run.setStatus(ClusterRunStatus.RUNNING);
        run.setStartedAt(OffsetDateTime.now());
        run.setFinishedAt(null);
        run.setErrorMessage(null);
        runRepository.save(run);
    }

    private void updateItemCount(UUID runId, Long itemCount) {
        EmbeddingClusterRun run = loadRun(runId);
        run.setItemCount(itemCount);
        runRepository.save(run);
    }

    private void completeRun(UUID runId, ClusteringEngineResult result) {
        EmbeddingClusterRun run = loadRun(runId);
        run.setStatus(ClusterRunStatus.COMPLETED);
        run.setFinishedAt(OffsetDateTime.now());
        run.setClusterCount(result.clusters().stream().filter(cluster -> !cluster.noiseCluster()).count());
        run.setNoiseCount(result.noiseCount());
        run.setErrorMessage(null);
        runRepository.save(run);
    }

    private void cancelRunNow(UUID runId, String message) {
        EmbeddingClusterRun run = loadRun(runId);
        run.setStatus(ClusterRunStatus.CANCELLED);
        run.setFinishedAt(OffsetDateTime.now());
        run.setErrorMessage(message);
        runRepository.save(run);
    }

    private boolean isCancellationRequested(UUID runId) {
        ClusterRunStatus status = loadRun(runId).getStatus();
        return status == ClusterRunStatus.CANCEL_REQUESTED || status == ClusterRunStatus.CANCELLED;
    }

    private ClusteringEngineResult executePythonClustering(EmbeddingClusterRun run, List<SelectedEmbeddingRow> selected) {
        PythonClusteringClient client = pythonClusteringClient.orElseThrow(() -> new ResponseStatusException(
                HttpStatus.BAD_REQUEST,
                "Cluster run requires PYTHON_REMOTE backend but no Python client is configured"));

        PythonClusteringResponse response = useCompactPythonRunMode()
                ? client.clusterRun(new PythonRunExecutionRequest(run.getId()))
                : client.cluster(new PythonClusteringRequest(
                        run.getAlgorithm(),
                        readParameters(run.getParametersJson()),
                        run.getReductionMethod(),
                        run.getReductionDimensions(),
                        selected == null ? List.of() : selected.stream()
                                .map(item -> new PythonClusteringRequest.PythonClusteringItem(
                                        item.embeddingId(),
                                        item.documentId(),
                                        item.representationId(),
                                        item.embeddingVector()))
                                .toList()));

        return mapPythonResponse(response);
    }

    private boolean useCompactPythonRunMode() {
        return pythonProperties.effectiveRequestMode() == PythonRequestMode.RUN_ID;
    }

    private ClusteringEngineResult executeLocalClustering(EmbeddingClusterRun run, List<SelectedEmbeddingRow> selected) {
        EmbeddingClusteringEngine engine = resolveEngine(run.getAlgorithm());
        return engine.cluster(selected, new ClusteringEngineRequest(readParameters(run.getParametersJson())));
    }

    private ClusteringEngineResult mapPythonResponse(PythonClusteringResponse response) {
        long noiseCount = response.noiseCount() == null
                ? response.assignments().stream().filter(assignment -> Boolean.TRUE.equals(assignment.noise())).count()
                : response.noiseCount();

        return new ClusteringEngineResult(
                response.clusters().stream()
                        .map(cluster -> new ClusteringEngineCluster(
                                cluster.clusterLabel(),
                                cluster.itemCount(),
                                Boolean.TRUE.equals(cluster.noiseCluster())))
                        .toList(),
                response.assignments().stream()
                        .map(assignment -> new ClusteringEngineAssignment(
                                assignment.embeddingId(),
                                assignment.documentId(),
                                assignment.representationId(),
                                assignment.clusterLabel(),
                                assignment.distanceToCentroid(),
                                assignment.membershipScore(),
                                Boolean.TRUE.equals(assignment.noise())))
                        .toList(),
                noiseCount);
    }

    private EmbeddingClusterSet resolveClusterSet(String clusterSetCode) {
        if (clusterSetCode == null || clusterSetCode.isBlank()) {
            return null;
        }
        return clusterSetRepository.findByCode(clusterSetCode)
                .orElseThrow(() -> new ResponseStatusException(HttpStatus.BAD_REQUEST,
                        "Cluster set not found: " + clusterSetCode));
    }

    private void ensureRunExists(UUID runId) {
        if (!runRepository.existsById(runId)) {
            throw new ResponseStatusException(HttpStatus.NOT_FOUND, "Cluster run not found: " + runId);
        }
    }

    private void failRun(UUID runId, String message) {
        EmbeddingClusterRun run = loadRun(runId);
        run.setStatus(ClusterRunStatus.FAILED);
        run.setFinishedAt(OffsetDateTime.now());
        run.setErrorMessage(message);
        runRepository.save(run);
    }

    private EmbeddingClusteringEngine resolveEngine(ClusteringAlgorithm algorithm) {
        return clusteringEngines.stream()
                .filter(engine -> engine.algorithm() == algorithm)
                .findFirst()
                .orElseThrow(() -> new ResponseStatusException(HttpStatus.BAD_REQUEST,
                        "No clustering engine registered for algorithm " + algorithm));
    }

    private EmbeddingSelectionSpec readSelection(String json) {
        try {
            return objectMapper.readValue(json, EmbeddingSelectionSpec.class);
        } catch (JsonProcessingException e) {
            throw new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR,
                    "Cannot parse stored clustering selection", e);
        }
    }

    private Map<String, Object> readParameters(String json) {
        if (json == null || json.isBlank()) {
            return Map.of();
        }
        try {
            Map<String, Object> values = objectMapper.readValue(json, PARAMETERS_TYPE);
            return values == null ? Map.of() : values;
        } catch (JsonProcessingException e) {
            throw new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR,
                    "Cannot parse clustering parameters", e);
        }
    }

    private String writeJson(Object value) {
        try {
            return objectMapper.writeValueAsString(value);
        } catch (JsonProcessingException e) {
            throw new ResponseStatusException(HttpStatus.INTERNAL_SERVER_ERROR,
                    "Cannot serialize clustering payload", e);
        }
    }

    private ClusteringExecutionBackend resolveExecutionBackend(CreateClusterRunRequest request) {
        if (request.executionBackend() != null) {
            if (request.executionBackend() == ClusteringExecutionBackend.JAVA_LOCAL
                    && request.reduction() != null
                    && request.reduction().method() != null
                    && request.reduction().method() != ReductionMethod.NONE) {
                throw new ResponseStatusException(HttpStatus.BAD_REQUEST,
                        "JAVA_LOCAL backend does not support PCA/UMAP reduction; use PYTHON_REMOTE");
            }
            return request.executionBackend();
        }
        if (request.reduction() != null && request.reduction().method() != null && request.reduction().method() != ReductionMethod.NONE) {
            return ClusteringExecutionBackend.PYTHON_REMOTE;
        }
        return switch (request.algorithm()) {
            case KMEANS -> ClusteringExecutionBackend.JAVA_LOCAL;
            default -> ClusteringExecutionBackend.PYTHON_REMOTE;
        };
    }

    private ReductionMethod resolveReductionMethod(ReductionConfig reduction) {
        if (reduction == null || reduction.method() == null) {
            return ReductionMethod.NONE;
        }
        return reduction.method();
    }

    private Integer resolveReductionDimensions(ReductionConfig reduction) {
        return reduction == null ? null : reduction.targetDimensions();
    }

    private <T> T firstOrNull(Iterable<T> values) {
        if (values == null) {
            return null;
        }
        for (T value : values) {
            return value;
        }
        return null;
    }

    private String firstOrNullString(Iterable<String> values) {
        return firstOrNull(values);
    }

    private ClusterRunResponse toResponse(EmbeddingClusterRun run) {
        return new ClusterRunResponse(
                run.getId(),
                run.getName(),
                run.getStatus(),
                run.getAlgorithm(),
                run.getExecutionBackend(),
                run.getReductionMethod(),
                run.getReductionDimensions(),
                run.getItemCount(),
                run.getClusterCount(),
                run.getNoiseCount(),
                run.getStartedAt(),
                run.getFinishedAt(),
                run.getErrorMessage());
    }

    private EmbeddingClusterRun loadRun(UUID runId) {
        return runRepository.findById(runId)
                .orElseThrow(() -> new ResponseStatusException(HttpStatus.NOT_FOUND, "Cluster run not found: " + runId));
    }
}
