package com.infoepoch.pms.agent.platform.chat.checkpoint;

import com.alibaba.cloud.ai.graph.RunnableConfig;
import com.alibaba.cloud.ai.graph.StateGraph;
import com.alibaba.cloud.ai.graph.checkpoint.BaseCheckpointSaver;
import com.alibaba.cloud.ai.graph.checkpoint.Checkpoint;
import com.alibaba.cloud.ai.graph.serializer.StateSerializer;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.infoepoch.pms.agent.platform.chat.domain.GraphCheckPoint;
import com.infoepoch.pms.agent.platform.chat.domain.GraphThread;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.GraphCheckPointCriteria;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.GraphThreadCriteria;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.IGraphCheckPointRepository;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.IGraphThreadRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

import java.io.IOException;
import java.time.LocalDateTime;
import java.util.Base64;
import java.util.Collection;
import java.util.Comparator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;

/**
 * 基于仓储层的 checkpoint 持久化实现，语义尽量对齐官方 OracleSaver。
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class DatabaseCheckpointStore implements BaseCheckpointSaver {

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

    private final IGraphCheckPointRepository graphCheckPointRepository;
    private final IGraphThreadRepository graphThreadRepository;
    private final StateSerializer stateSerializer = StateGraph.DEFAULT_JACKSON_SERIALIZER;

    @Override
    public Collection<Checkpoint> list(RunnableConfig config) {
        return findActiveThreadByName(resolveThreadName(config))
                .map(GraphThread::getThreadId)
                .map(this::findCheckpointsByThreadId)
                .orElse(List.of())
                .stream()
                .map(this::buildCheckpoint)
                .toList();
    }

    @Override
    public Optional<Checkpoint> get(RunnableConfig config) {
        Optional<GraphThread> activeThread = findActiveThreadByName(resolveThreadName(config));
        if (activeThread.isEmpty()) {
            return Optional.empty();
        }

        String persistedThreadId = activeThread.get().getThreadId();
        if (config.checkPointId().isPresent()) {
            return findCheckpointByThreadIdAndCheckpointId(persistedThreadId, config.checkPointId().get())
                    .map(this::buildCheckpoint);
        }

        return findCheckpointsByThreadId(persistedThreadId)
                .stream()
                .findFirst()
                .map(this::buildCheckpoint);
    }

    @Override
    public RunnableConfig put(RunnableConfig config, Checkpoint checkpoint) throws Exception {
        String threadName = resolveThreadName(config);
        String persistedThreadId = config.checkPointId()
                .map(previousCheckpointId -> resolveThreadIdForReplacement(threadName, previousCheckpointId))
                .orElseGet(() -> ensureActiveThread(threadName).getThreadId());

        GraphCheckPoint entity = new GraphCheckPoint(
                checkpoint.getId(),
                null,
                persistedThreadId,
                checkpoint.getNodeId(),
                checkpoint.getNextNodeId(),
                encodeState(checkpoint.getState()),
                stateSerializer.contentType(),
                LocalDateTime.now()
        );

        if (config.checkPointId().isPresent()) {
            graphCheckPointRepository.replaceCheckpoint(config.checkPointId().get(), entity);
        } else {
            graphCheckPointRepository.insert(entity);
        }

        return RunnableConfig.builder(config)
                .threadId(threadName)
                .checkPointId(checkpoint.getId())
                .build();
    }

    @Override
    public Tag release(RunnableConfig config) throws Exception {
        String threadName = resolveThreadName(config);
        GraphThread activeThread = findActiveThreadByName(threadName)
                .orElseThrow(() -> new IllegalStateException(
                        String.format("Thread '%s' not found or already released", threadName)));
        Collection<Checkpoint> checkpoints = findCheckpointsByThreadId(activeThread.getThreadId())
                .stream()
                .map(this::buildCheckpoint)
                .toList();
        graphThreadRepository.update(new GraphThread(activeThread.getThreadId(), activeThread.getThreadName(), true));
        return new Tag(threadName, checkpoints);
    }

    public List<String> findCandidateThreadIds(String sessionId) {
        if (!StringUtils.hasText(sessionId)) {
            return List.of();
        }
        LinkedHashSet<String> candidateThreadNames = new LinkedHashSet<>();
        candidateThreadNames.add(sessionId);
        Map<String, LocalDateTime> latestSavedAtByThreadName = findAllThreads().stream()
                .filter(thread -> belongsToThreadGroup(sessionId, thread.getThreadName()))
                .filter(thread -> !Boolean.TRUE.equals(thread.getReleased()))
                .collect(java.util.stream.Collectors.toMap(
                        GraphThread::getThreadName,
                        thread -> findLatestSavedAt(thread.getThreadId()),
                        DatabaseCheckpointStore::maxSavedAt
                ));
        candidateThreadNames.addAll(latestSavedAtByThreadName.keySet());
        return candidateThreadNames.stream()
                .sorted(Comparator.comparing(latestSavedAtByThreadName::get,
                        Comparator.nullsLast(Comparator.reverseOrder())))
                .toList();
    }

    public void deleteSession(String sessionId) {
        deleteThreadGroups(List.of(sessionId));
    }

    public void deleteThreadGroups(List<String> rootThreadIds) {
        if (rootThreadIds == null || rootThreadIds.isEmpty()) {
            return;
        }
        LinkedHashSet<String> distinctRootThreadIds = new LinkedHashSet<>(rootThreadIds);
        for (GraphThread thread : findAllThreads()) {
            if (distinctRootThreadIds.stream().noneMatch(rootThreadId ->
                    belongsToThreadGroup(rootThreadId, thread.getThreadName()))) {
                continue;
            }
            findCheckpointsByThreadId(thread.getThreadId()).forEach(checkpoint ->
                    graphCheckPointRepository.delete(checkpoint.getCheckPointId()));
            graphThreadRepository.delete(thread.getThreadId());
        }
    }

    private String resolveThreadName(RunnableConfig config) {
        return config.threadId().filter(StringUtils::hasText).orElse(THREAD_ID_DEFAULT);
    }

    private GraphThread ensureActiveThread(String threadName) {
        return findActiveThreadByName(threadName)
                .orElseGet(() -> {
                    GraphThread newThread = new GraphThread(UUID.randomUUID().toString(), threadName, false);
                    graphThreadRepository.insert(newThread);
                    return newThread;
                });
    }

    private Optional<GraphThread> findActiveThreadByName(String threadName) {
        return findThreadsByName(threadName, false).stream().findFirst();
    }

    private List<GraphThread> findThreadsByName(String threadName, Boolean released) {
        GraphThreadCriteria criteria = new GraphThreadCriteria();
        criteria.setThreadName(threadName);
        if (released != null) {
            criteria.setReleased(released);
        }
        List<GraphThread> threads = graphThreadRepository.selectByCriteria(criteria);
        return threads == null ? List.of() : threads;
    }

    private List<GraphThread> findAllThreads() {
        List<GraphThread> threads = graphThreadRepository.selectByCriteria(null);
        return threads == null ? List.of() : threads;
    }

    private Optional<GraphCheckPoint> findCheckpointByThreadIdAndCheckpointId(String threadId, String checkpointId) {
        GraphCheckPointCriteria criteria = new GraphCheckPointCriteria();
        criteria.setThreadId(threadId);
        criteria.setCheckPointId(checkpointId);
        List<GraphCheckPoint> records = graphCheckPointRepository.selectByCriteria(criteria);
        if (records == null || records.isEmpty()) {
            return Optional.empty();
        }
        return records.stream().findFirst();
    }

    private List<GraphCheckPoint> findCheckpointsByThreadId(String threadId) {
        GraphCheckPointCriteria criteria = new GraphCheckPointCriteria();
        criteria.setThreadId(threadId);
        List<GraphCheckPoint> records = graphCheckPointRepository.selectByCriteria(criteria);
        if (records == null) {
            return List.of();
        }
        return records.stream()
                .sorted(Comparator.comparing(GraphCheckPoint::getSavedAt, Comparator.nullsLast(Comparator.reverseOrder())))
                .toList();
    }

    private LocalDateTime findLatestSavedAt(String threadId) {
        return findCheckpointsByThreadId(threadId).stream()
                .map(GraphCheckPoint::getSavedAt)
                .filter(java.util.Objects::nonNull)
                .max(LocalDateTime::compareTo)
                .orElse(null);
    }

    private boolean belongsToThreadGroup(String rootThreadId, String candidateThreadName) {
        if (!StringUtils.hasText(rootThreadId) || !StringUtils.hasText(candidateThreadName)) {
            return false;
        }
        return rootThreadId.equals(candidateThreadName)
                || candidateThreadName.startsWith(rootThreadId + "_subgraph_");
    }

    private static LocalDateTime maxSavedAt(LocalDateTime left, LocalDateTime right) {
        if (left == null) {
            return right;
        }
        if (right == null) {
            return left;
        }
        return left.isAfter(right) ? left : right;
    }

    private String resolveThreadIdForReplacement(String threadName, String previousCheckpointId) {
        Optional<String> activeThreadId = findActiveThreadByName(threadName).map(GraphThread::getThreadId);
        if (activeThreadId.isPresent()) {
            return activeThreadId.get();
        }

        GraphCheckPoint existingCheckpoint = graphCheckPointRepository.selectById(previousCheckpointId);
        if (existingCheckpoint != null && StringUtils.hasText(existingCheckpoint.getThreadId())) {
            return existingCheckpoint.getThreadId();
        }

        throw new IllegalStateException(String.format("Checkpoint '%s' not found for thread '%s'", previousCheckpointId, threadName));
    }

    private Checkpoint buildCheckpoint(GraphCheckPoint record) {
        return Checkpoint.builder()
                .id(record.getCheckPointId())
                .nodeId(record.getNodeId())
                .nextNodeId(record.getNextNodeId())
                .state(decodeState(record.getStateData(), record.getStateContentType()))
                .build();
    }

    private String encodeState(Map<String, Object> state) throws IOException {
        byte[] binaryData = stateSerializer.dataToBytes(state == null ? Map.of() : state);
        return OBJECT_MAPPER.writeValueAsString(Map.of(
                "binaryPayload", Base64.getEncoder().encodeToString(binaryData)
        ));
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> decodeState(String stateData, String contentType) {
        if (!StringUtils.hasText(stateData)) {
            return Map.of();
        }
        if (!stateSerializer.contentType().equals(contentType)) {
            throw new IllegalStateException(String.format(
                    "Content Type used for store state '%s' is different from one '%s' used for deserialize it",
                    contentType,
                    stateSerializer.contentType()));
        }
        try {
            JsonNode jsonNode = OBJECT_MAPPER.readTree(stateData);
            JsonNode payloadNode = jsonNode.get("binaryPayload");
            if (payloadNode == null || !payloadNode.isTextual()) {
                throw new IllegalStateException("Checkpoint state does not contain binaryPayload");
            }
            byte[] bytes = Base64.getDecoder().decode(payloadNode.asText());
            return (Map<String, Object>) stateSerializer.dataFromBytes(bytes);
        } catch (JsonProcessingException ex) {
            throw new IllegalStateException("Failed to parse checkpoint state JSON", ex);
        } catch (ClassNotFoundException ex) {
            throw new IllegalStateException("Failed to deserialize checkpoint state", ex);
        } catch (IOException ex) {
            throw new IllegalStateException("Failed to deserialize checkpoint state", ex);
        }
    }
}
