package com.infoepoch.pms.agent.platform.chat.adapter.out.persistence;

import com.infoepoch.pms.agent.platform.chat.domain.GraphCheckPoint;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.RowMapper;
import org.springframework.test.util.ReflectionTestUtils;

import java.sql.ResultSet;
import java.sql.Timestamp;
import java.time.LocalDateTime;
import java.util.Arrays;
import java.util.Base64;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertInstanceOf;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.when;

@ExtendWith(MockitoExtension.class)
class GraphCheckPointRepositoryTest {

    private static final String CONTENT_TYPE = "application/x-java-serialized-object;base64";

    @Mock
    private JdbcTemplate jdbcTemplate;

    private GraphCheckPointRepository repository;

    @BeforeEach
    void setUp() {
        repository = new GraphCheckPointRepository();
        ReflectionTestUtils.setField(repository, "jdbcTemplate", jdbcTemplate);
    }

    @Test
    void shouldWriteRuntimeCheckpointIntoAiTableWithEncodedPayload() {
        LocalDateTime now = LocalDateTime.now();
        byte[] payload = new byte[]{1, 2, 3};
        GraphCheckPoint checkpoint = new GraphCheckPoint(
                "thread-1",
                "session-1",
                "checkpoint-1",
                "node-1",
                "node-2",
                payload,
                false,
                now
        );
        AtomicReference<String> sqlRef = new AtomicReference<>();
        AtomicReference<Object[]> argsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            sqlRef.set(invocation.getArgument(0));
            argsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any());

        repository.insert(new GraphCheckPoint(
                checkpoint.getCheckPointId(),
                null,
                checkpoint.getThreadId(),
                checkpoint.getNodeId(),
                checkpoint.getNextNodeId(),
                Base64.getEncoder().encodeToString(payload),
                CONTENT_TYPE,
                now
        ));

        assertTrue(sqlRef.get().contains("AI_GRAPHCHECKPOINT"));
        assertFalse(sqlRef.get().contains("platform_agent_checkpoint"));
        assertEquals(Base64.getEncoder().encodeToString(payload), argsRef.get()[5]);
        assertEquals(CONTENT_TYPE, argsRef.get()[6]);
        assertInstanceOf(Timestamp.class, argsRef.get()[7]);
        assertEquals(Timestamp.valueOf(now), argsRef.get()[7]);
    }

    @Test
    void shouldReadRuntimeCheckpointFromAiTableViaGenericCriteria() throws Exception {
        LocalDateTime now = LocalDateTime.now();
        byte[] payload = new byte[]{9, 8, 7};
        AtomicReference<String> sqlRef = new AtomicReference<>();
        when(jdbcTemplate.query(anyString(), any(Object[].class), any(RowMapper.class)))
                .thenAnswer(invocation -> {
                    sqlRef.set(invocation.getArgument(0));
                    Object[] args = invocation.getArgument(1);
                    assertEquals("checkpoint-1", args[0]);
                    assertEquals("thread-1", args[1]);
                    @SuppressWarnings("unchecked")
                    RowMapper<GraphCheckPoint> mapper = invocation.getArgument(2);
                    ResultSet rs = mock(ResultSet.class);
                    when(rs.getString("CHECKPOINT_ID")).thenReturn("checkpoint-1");
                    when(rs.getString("PARENT_CHECKPOINT_ID")).thenReturn(null);
                    when(rs.getString("THREAD_ID")).thenReturn("thread-1");
                    when(rs.getString("NODE_ID")).thenReturn("node-1");
                    when(rs.getString("NEXT_NODE_ID")).thenReturn("node-2");
                    when(rs.getString("STATE_DATA")).thenReturn(Base64.getEncoder().encodeToString(payload));
                    when(rs.getString("STATE_CONTENT_TYPE")).thenReturn(CONTENT_TYPE);
                    when(rs.getTimestamp("SAVED_AT")).thenReturn(Timestamp.valueOf(now));
                    return List.of(mapper.mapRow(rs, 0));
                });

        List<GraphCheckPoint> result = repository.selectByCriteria(new com.infoepoch.pms.agent.platform.chat.domain.irepository.GraphCheckPointCriteria() {{
            setThreadId("thread-1");
            setCheckPointId("checkpoint-1");
        }});

        assertTrue(sqlRef.get().contains("AI_GRAPHCHECKPOINT"));
        assertFalse(sqlRef.get().contains("platform_agent_checkpoint"));
        assertEquals(1, result.size());
        assertEquals(Base64.getEncoder().encodeToString(payload), result.getFirst().getStateData());
    }

    @Test
    void shouldReplaceCheckpointIdWhenUpdatingRuntimeCheckpoint() {
        LocalDateTime now = LocalDateTime.now();
        AtomicReference<String> sqlRef = new AtomicReference<>();
        AtomicReference<Object[]> argsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            sqlRef.set(invocation.getArgument(0));
            argsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any(), any());

        boolean updated = repository.replaceCheckpoint("checkpoint-old", new GraphCheckPoint(
                "checkpoint-new",
                null,
                "thread-1",
                "node-2",
                "node-3",
                "{\"binaryPayload\":\"AQI=\"}",
                "application/json",
                now
        ));

        assertTrue(updated);
        assertTrue(sqlRef.get().contains("SET CHECKPOINT_ID = ?"));
        assertEquals("checkpoint-new", argsRef.get()[0]);
        assertInstanceOf(Timestamp.class, argsRef.get()[7]);
        assertEquals(Timestamp.valueOf(now), argsRef.get()[7]);
        assertEquals("checkpoint-old", argsRef.get()[8]);
    }

    @Test
    void shouldWriteTimestampWhenUpdatingCheckpoint() {
        LocalDateTime now = LocalDateTime.now();
        AtomicReference<Object[]> argsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            argsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any());

        boolean updated = repository.update(new GraphCheckPoint(
                "checkpoint-1",
                "checkpoint-0",
                "thread-1",
                "node-1",
                "node-2",
                "{\"binaryPayload\":\"AQI=\"}",
                "application/json",
                now
        ));

        assertTrue(updated);
        assertInstanceOf(Timestamp.class, argsRef.get()[6]);
        assertEquals(Timestamp.valueOf(now), argsRef.get()[6]);
    }

    @Test
    void shouldAllowNullSavedAtWhenWritingCheckpoint() {
        AtomicReference<Object[]> insertArgsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            insertArgsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any());

        repository.insert(new GraphCheckPoint(
                "checkpoint-1",
                null,
                "thread-1",
                "node-1",
                "node-2",
                "{\"binaryPayload\":\"AQI=\"}",
                "application/json",
                null
        ));

        assertNull(insertArgsRef.get()[7]);

        reset(jdbcTemplate);

        AtomicReference<Object[]> updateArgsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            updateArgsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any());

        repository.update(new GraphCheckPoint(
                "checkpoint-1",
                null,
                "thread-1",
                "node-1",
                "node-2",
                "{\"binaryPayload\":\"AQI=\"}",
                "application/json",
                null
        ));

        assertNull(updateArgsRef.get()[6]);

        reset(jdbcTemplate);

        AtomicReference<Object[]> replaceArgsRef = new AtomicReference<>();
        doAnswer(invocation -> {
            replaceArgsRef.set(Arrays.copyOfRange(invocation.getArguments(), 1, invocation.getArguments().length));
            return 1;
        }).when(jdbcTemplate).update(anyString(), any(), any(), any(), any(), any(), any(), any(), any(), any());

        repository.replaceCheckpoint("checkpoint-old", new GraphCheckPoint(
                "checkpoint-1",
                null,
                "thread-1",
                "node-1",
                "node-2",
                "{\"binaryPayload\":\"AQI=\"}",
                "application/json",
                null
        ));

        assertNull(replaceArgsRef.get()[7]);
    }
}
