package tech.easyflow.agent.runtime.agui;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.ArgumentCaptor;

import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

/**
 * AG-UI 批量事件日志测试。
 */
public class AgentAguiRunJournalTest {

    /**
     * 验证相邻同消息 delta 会在一次 Redis 批次中合并且顺序不变。
     */
    @Test
    public void shouldMergeAdjacentDeltasBeforeRedisFlush() {
        AgentAguiRunStore store = mock(AgentAguiRunStore.class);
        AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
        journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_CONTENT\",\"messageId\":\"m1\",\"delta\":\"你\"}");
        journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_CONTENT\",\"messageId\":\"m1\",\"delta\":\"好\"}");
        journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_END\",\"messageId\":\"m1\"}");

        journal.flush("run-1");

        @SuppressWarnings("unchecked")
        ArgumentCaptor<List<String>> events = ArgumentCaptor.forClass(List.class);
        verify(store).append(eq("run-1"), anyString(), events.capture());
        Assert.assertEquals(2, events.getValue().size());
        Assert.assertTrue(events.getValue().get(0).contains("\"delta\":\"你好\""));
        Assert.assertTrue(events.getValue().get(1).contains("TEXT_MESSAGE_END"));
    }

    /**
     * 验证同步终态写入失败后，后台刷新仍会保留并重试原批次。
     */
    @Test
    public void shouldRetryTerminalBatchWithoutDroppingEvents() {
        AgentAguiRunStore store = mock(AgentAguiRunStore.class);
        AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
        String terminal = "{\"type\":\"RUN_FINISHED\",\"runId\":\"run-1\",\"threadId\":\"1\"}";
        journal.append("run-1", terminal);
        doThrow(new IllegalStateException("redis down"))
                .doReturn(1L)
                .when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());

        try {
            journal.complete("run-1");
            Assert.fail("首次同步写入应抛出异常");
        } catch (IllegalStateException expected) {
            Assert.assertEquals("redis down", expected.getMessage());
        }
        journal.flushPending();

        ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
        verify(store, times(2)).append(
                eq("run-1"), batchIds.capture(), org.mockito.ArgumentMatchers.anyList());
        Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
    }

    /**
     * 验证结果不确定的批次内容保持冻结，期间到达的新事件使用新批次写入。
     */
    @Test
    public void shouldKeepAmbiguousBatchImmutableAndFlushTailSeparately() {
        AgentAguiRunStore store = mock(AgentAguiRunStore.class);
        AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
        doThrow(new IllegalStateException("ambiguous timeout"))
                .doReturn(1L)
                .doReturn(2L)
                .when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());
        String first = "{\"type\":\"RUN_STARTED\",\"runId\":\"run-1\"}";
        String tail = "{\"type\":\"TEXT_MESSAGE_START\",\"messageId\":\"m1\"}";
        journal.append("run-1", first);

        Assert.assertThrows(IllegalStateException.class, () -> journal.flush("run-1"));
        journal.append("run-1", tail);
        journal.flushPending();

        ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
        @SuppressWarnings("unchecked")
        ArgumentCaptor<List<String>> batches = ArgumentCaptor.forClass(List.class);
        verify(store, times(3)).append(eq("run-1"), batchIds.capture(), batches.capture());
        Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
        Assert.assertNotEquals(batchIds.getAllValues().get(1), batchIds.getAllValues().get(2));
        Assert.assertEquals(List.of(first), batches.getAllValues().get(0));
        Assert.assertEquals(List.of(first), batches.getAllValues().get(1));
        Assert.assertEquals(List.of(tail), batches.getAllValues().get(2));
    }

    /**
     * 验证终态完成会等待在途后台刷盘，并在其失败后接管重试全部剩余批次。
     *
     * @throws Exception 并发任务等待失败时抛出
     */
    @Test
    public void completeShouldWaitForBackgroundFlushAndPersistAllTailEvents() throws Exception {
        AgentAguiRunStore store = mock(AgentAguiRunStore.class);
        AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
        CountDownLatch firstAppendStarted = new CountDownLatch(1);
        CountDownLatch releaseFirstAppend = new CountDownLatch(1);
        CountDownLatch completeStarted = new CountDownLatch(1);
        AtomicInteger appendCalls = new AtomicInteger();
        doAnswer(invocation -> {
            int call = appendCalls.incrementAndGet();
            if (call == 1) {
                firstAppendStarted.countDown();
                Assert.assertTrue(releaseFirstAppend.await(5, TimeUnit.SECONDS));
                throw new IllegalStateException("first append failed");
            }
            return (long) call;
        }).when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());
        String started = "{\"type\":\"RUN_STARTED\",\"runId\":\"run-1\"}";
        String finished = "{\"type\":\"RUN_FINISHED\",\"runId\":\"run-1\",\"threadId\":\"1\"}";
        journal.append("run-1", started);
        ExecutorService executor = Executors.newFixedThreadPool(2);

        try {
            Future<?> background = executor.submit(journal::flushPending);
            Assert.assertTrue(firstAppendStarted.await(5, TimeUnit.SECONDS));
            journal.append("run-1", finished);
            Future<?> completion = executor.submit(() -> {
                completeStarted.countDown();
                journal.complete("run-1");
            });
            Assert.assertTrue(completeStarted.await(5, TimeUnit.SECONDS));
            Assert.assertFalse(completion.isDone());

            releaseFirstAppend.countDown();
            background.get(5, TimeUnit.SECONDS);
            completion.get(5, TimeUnit.SECONDS);
        } finally {
            releaseFirstAppend.countDown();
            executor.shutdownNow();
        }

        @SuppressWarnings("unchecked")
        ArgumentCaptor<List<String>> batches = ArgumentCaptor.forClass(List.class);
        ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
        verify(store, times(3)).append(eq("run-1"), batchIds.capture(), batches.capture());
        Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
        Assert.assertEquals(List.of(started), batches.getAllValues().get(1));
        Assert.assertEquals(List.of(finished), batches.getAllValues().get(2));
    }

    /**
     * 验证容量故障会丢弃未确认普通事件并以独立小批次写入失败终态。
     */
    @Test
    public void shouldReplacePendingEventsWithJournalFailureTerminal() {
        AgentAguiRunStore store = mock(AgentAguiRunStore.class);
        AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
        journal.append("run-1", "{\"type\":\"RUN_STARTED\"}");

        journal.fail("thread-1", "run-1", "AGENT_RUN_JOURNAL_LIMIT", "输出过长");

        @SuppressWarnings("unchecked")
        ArgumentCaptor<List<String>> events = ArgumentCaptor.forClass(List.class);
        verify(store).append(eq("run-1"), anyString(), events.capture());
        Assert.assertEquals(1, events.getValue().size());
        Assert.assertTrue(events.getValue().get(0).contains("AGENT_RUN_JOURNAL_LIMIT"));
        Assert.assertFalse(events.getValue().get(0).contains("RUN_STARTED"));
    }

    /**
     * 验证异常大的单事件在进入共享内存缓冲前即被拒绝。
     */
    @Test
    public void shouldRejectOversizedSingleEvent() {
        AgentAguiRunJournal journal = new AgentAguiRunJournal(
                mock(AgentAguiRunStore.class), new ObjectMapper());
        String oversized = "x".repeat(512 * 1024 + 1);

        Assert.assertThrows(IllegalStateException.class,
                () -> journal.append("run-1", oversized));
    }
}
