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.Checkpoint;
import com.alibaba.cloud.ai.graph.serializer.StateSerializer;
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.IGraphCheckPointRepository;
import com.infoepoch.pms.agent.platform.chat.domain.irepository.IGraphThreadRepository;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;

import java.time.LocalDateTime;
import java.util.Base64;
import java.util.List;
import java.util.Map;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class DatabaseCheckpointStoreTest {

    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
    private static final StateSerializer STATE_SERIALIZER = StateGraph.DEFAULT_JACKSON_SERIALIZER;

    @Mock
    private IGraphCheckPointRepository graphCheckPointRepository;

    @Mock
    private IGraphThreadRepository graphThreadRepository;

    private DatabaseCheckpointStore store;

    @BeforeEach
    void setUp() {
        store = new DatabaseCheckpointStore(graphCheckPointRepository, graphThreadRepository);
    }

    @Test
    void shouldInsertThreadAndCheckpointWhenNoActiveThreadExists() throws Exception {
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of());

        Checkpoint checkpoint = Checkpoint.builder()
                .id("checkpoint-1")
                .nodeId("node-1")
                .nextNodeId("node-2")
                .state(Map.of("user", "u1"))
                .build();

        RunnableConfig result = store.put(RunnableConfig.builder().threadId("thread-1").build(), checkpoint);

        ArgumentCaptor<GraphThread> threadCaptor = ArgumentCaptor.forClass(GraphThread.class);
        verify(graphThreadRepository).insert(threadCaptor.capture());
        GraphThread insertedThread = threadCaptor.getValue();
        assertNotNull(insertedThread.getThreadId());
        assertEquals("thread-1", insertedThread.getThreadName());
        assertEquals(Boolean.FALSE, insertedThread.getReleased());

        ArgumentCaptor<GraphCheckPoint> checkpointCaptor = ArgumentCaptor.forClass(GraphCheckPoint.class);
        verify(graphCheckPointRepository).insert(checkpointCaptor.capture());
        GraphCheckPoint insertedCheckpoint = checkpointCaptor.getValue();
        assertEquals(insertedThread.getThreadId(), insertedCheckpoint.getThreadId());
        assertEquals("checkpoint-1", insertedCheckpoint.getCheckPointId());
        assertEquals("node-1", insertedCheckpoint.getNodeId());
        assertEquals("node-2", insertedCheckpoint.getNextNodeId());
        assertEquals(STATE_SERIALIZER.contentType(), insertedCheckpoint.getStateContentType());
        assertEquals(Map.of("user", "u1"), decodeState(insertedCheckpoint.getStateData(), insertedCheckpoint.getStateContentType()));
        assertEquals("thread-1", result.threadId().orElseThrow());
        assertEquals("checkpoint-1", result.checkPointId().orElseThrow());
    }

    @Test
    void shouldReplaceCheckpointWhenPreviousCheckpointIdExists() throws Exception {
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));

        Checkpoint checkpoint = Checkpoint.builder()
                .id("checkpoint-new")
                .nodeId("node-2")
                .nextNodeId("node-3")
                .state(Map.of("approved", true))
                .build();

        store.put(RunnableConfig.builder().threadId("thread-1").checkPointId("checkpoint-old").build(), checkpoint);

        ArgumentCaptor<GraphCheckPoint> checkpointCaptor = ArgumentCaptor.forClass(GraphCheckPoint.class);
        verify(graphCheckPointRepository).replaceCheckpoint(eq("checkpoint-old"), checkpointCaptor.capture());
        verify(graphThreadRepository, never()).insert(org.mockito.ArgumentMatchers.any());
        GraphCheckPoint replacedCheckpoint = checkpointCaptor.getValue();
        assertEquals("checkpoint-new", replacedCheckpoint.getCheckPointId());
        assertEquals("db-thread-1", replacedCheckpoint.getThreadId());
        assertEquals(Map.of("approved", true), decodeState(replacedCheckpoint.getStateData(), replacedCheckpoint.getStateContentType()));
    }

    @Test
    void shouldGetCheckpointByIdFromActiveThread() {
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "db-thread-1".equals(criteria.getThreadId())
                        && "checkpoint-1".equals(criteria.getCheckPointId()))))
                .thenReturn(List.of(checkpointRecord(
                        "checkpoint-1",
                        "db-thread-1",
                        "node-1",
                        "node-2",
                        Map.of("user", "u1"),
                        LocalDateTime.now()
                )));

        Checkpoint checkpoint = store.get(RunnableConfig.builder()
                        .threadId("thread-1")
                        .checkPointId("checkpoint-1")
                        .build())
                .orElseThrow();

        assertEquals("checkpoint-1", checkpoint.getId());
        assertEquals(Map.of("user", "u1"), checkpoint.getState());
    }

    @Test
    void shouldReturnLatestCheckpointWhenCheckpointIdMissing() {
        LocalDateTime now = LocalDateTime.now();
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "db-thread-1".equals(criteria.getThreadId())
                        && criteria.getCheckPointId() == null)))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-old", "db-thread-1", "node-1", "node-2", Map.of("seq", 1), now.minusMinutes(5)),
                        checkpointRecord("checkpoint-new", "db-thread-1", "node-2", "node-3", Map.of("seq", 2), now)
                ));

        Checkpoint checkpoint = store.get(RunnableConfig.builder().threadId("thread-1").build()).orElseThrow();

        assertEquals("checkpoint-new", checkpoint.getId());
        assertEquals(Map.of("seq", 2), checkpoint.getState());
    }

    @Test
    void shouldListCheckpointsInSavedAtDescendingOrder() {
        LocalDateTime now = LocalDateTime.now();
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "db-thread-1".equals(criteria.getThreadId())
                        && criteria.getCheckPointId() == null)))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-1", "db-thread-1", "node-1", "node-2", Map.of("seq", 1), now.minusMinutes(3)),
                        checkpointRecord("checkpoint-3", "db-thread-1", "node-3", "node-4", Map.of("seq", 3), now),
                        checkpointRecord("checkpoint-2", "db-thread-1", "node-2", "node-3", Map.of("seq", 2), now.minusMinutes(1))
                ));

        List<String> checkpointIds = store.list(RunnableConfig.builder().threadId("thread-1").build())
                .stream()
                .map(Checkpoint::getId)
                .toList();

        assertEquals(List.of("checkpoint-3", "checkpoint-2", "checkpoint-1"), checkpointIds);
    }

    @Test
    void shouldMarkThreadReleased() throws Exception {
        LocalDateTime now = LocalDateTime.now();
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "db-thread-1".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-1", "db-thread-1", "node-1", "node-2", Map.of("seq", 1), now)
                ));

        store.release(RunnableConfig.builder().threadId("thread-1").build());

        ArgumentCaptor<GraphThread> threadCaptor = ArgumentCaptor.forClass(GraphThread.class);
        verify(graphThreadRepository).update(threadCaptor.capture());
        GraphThread releasedThread = threadCaptor.getValue();
        assertEquals("db-thread-1", releasedThread.getThreadId());
        assertEquals("thread-1", releasedThread.getThreadName());
        assertEquals(Boolean.TRUE, releasedThread.getReleased());
    }

    @Test
    void shouldThrowWhenReleaseThreadDoesNotExist() {
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-missing".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of());

        IllegalStateException ex = assertThrows(IllegalStateException.class,
                () -> store.release(RunnableConfig.builder().threadId("thread-missing").build()));

        assertTrue(ex.getMessage().contains("thread-missing"));
    }

    @Test
    void shouldFindCandidateThreadIdsIncludingSubgraphThreadsOrderedByLatestCheckpoint() {
        LocalDateTime now = LocalDateTime.now();
        when(graphThreadRepository.selectByCriteria(null))
                .thenReturn(List.of(
                        new GraphThread("db-thread-root", "session-1", false),
                        new GraphThread("db-thread-sub-1", "session-1_subgraph_a", false),
                        new GraphThread("db-thread-sub-2", "session-1_subgraph_b", false)
                ));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-root".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-root", "db-thread-root", "node-1", "node-2", Map.of("seq", 1), now.minusMinutes(5))
                ));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-sub-1".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-sub-1", "db-thread-sub-1", "node-2", "node-3", Map.of("seq", 2), now)
                ));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-sub-2".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-sub-2", "db-thread-sub-2", "node-3", "node-4", Map.of("seq", 3), now.minusMinutes(1))
                ));

        List<String> threadIds = store.findCandidateThreadIds("session-1");

        assertEquals(List.of("session-1_subgraph_a", "session-1_subgraph_b", "session-1"), threadIds);
    }

    @Test
    void shouldDeleteRootThreadGroupsIncludingSubgraphs() {
        when(graphThreadRepository.selectByCriteria(null))
                .thenReturn(List.of(
                        new GraphThread("db-thread-root", "session-1", false),
                        new GraphThread("db-thread-sub", "session-1_subgraph_a", false)
                ));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-root".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-root", "db-thread-root", "node-1", "node-2", Map.of("seq", 1), LocalDateTime.now())
                ));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-sub".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-sub", "db-thread-sub", "node-2", "node-3", Map.of("seq", 2), LocalDateTime.now())
                ));

        store.deleteThreadGroups(List.of("session-1"));

        verify(graphCheckPointRepository).delete("checkpoint-root");
        verify(graphCheckPointRepository).delete("checkpoint-sub");
        verify(graphThreadRepository).delete("db-thread-root");
        verify(graphThreadRepository).delete("db-thread-sub");
    }

    @Test
    void shouldDeleteSessionByDelegatingToThreadGroupDeletion() {
        when(graphThreadRepository.selectByCriteria(null))
                .thenReturn(List.of(new GraphThread("db-thread-root", "session-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null && "db-thread-root".equals(criteria.getThreadId()))))
                .thenReturn(List.of(
                        checkpointRecord("checkpoint-root", "db-thread-root", "node-1", "node-2", Map.of("seq", 1), LocalDateTime.now())
                ));

        store.deleteSession("session-1");

        verify(graphCheckPointRepository).delete("checkpoint-root");
        verify(graphThreadRepository).delete("db-thread-root");
    }

    @Test
    void shouldThrowWhenContentTypeDoesNotMatchSerializer() {
        when(graphThreadRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "thread-1".equals(criteria.getThreadName())
                        && Boolean.FALSE.equals(criteria.getReleased()))))
                .thenReturn(List.of(new GraphThread("db-thread-1", "thread-1", false)));
        when(graphCheckPointRepository.selectByCriteria(argThat(criteria ->
                criteria != null
                        && "db-thread-1".equals(criteria.getThreadId())
                        && "checkpoint-1".equals(criteria.getCheckPointId()))))
                .thenReturn(List.of(new GraphCheckPoint(
                        "checkpoint-1",
                        null,
                        "db-thread-1",
                        "node-1",
                        "node-2",
                        "{\"binaryPayload\":\"AQI=\"}",
                        "application/custom",
                        LocalDateTime.now()
                )));

        IllegalStateException ex = assertThrows(IllegalStateException.class, () ->
                store.get(RunnableConfig.builder().threadId("thread-1").checkPointId("checkpoint-1").build()));

        assertTrue(ex.getMessage().contains("Content Type"));
    }

    private GraphCheckPoint checkpointRecord(String checkpointId,
                                            String threadId,
                                            String nodeId,
                                            String nextNodeId,
                                            Map<String, Object> state,
                                            LocalDateTime savedAt) {
        return new GraphCheckPoint(
                checkpointId,
                null,
                threadId,
                nodeId,
                nextNodeId,
                encodeState(state),
                STATE_SERIALIZER.contentType(),
                savedAt
        );
    }

    @SuppressWarnings("unchecked")
    private Map<String, Object> decodeState(String stateData, String contentType) throws Exception {
        assertEquals(STATE_SERIALIZER.contentType(), contentType);
        JsonNode jsonNode = OBJECT_MAPPER.readTree(stateData);
        byte[] bytes = Base64.getDecoder().decode(jsonNode.get("binaryPayload").asText());
        return (Map<String, Object>) STATE_SERIALIZER.dataFromBytes(bytes);
    }

    private String encodeState(Map<String, Object> state) {
        try {
            byte[] bytes = STATE_SERIALIZER.dataToBytes(state);
            return OBJECT_MAPPER.writeValueAsString(Map.of(
                    "binaryPayload", Base64.getEncoder().encodeToString(bytes)
            ));
        } catch (Exception ex) {
            throw new IllegalStateException(ex);
        }
    }
}
