package tech.easyflow.chatlog.service.impl;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.Assert;
import org.junit.Test;
import tech.easyflow.chatlog.domain.dto.ChatRoundRecord;
import tech.easyflow.chatlog.domain.command.ChatSessionUpsertCommand;
import tech.easyflow.chatlog.domain.dto.ChatMessageRecord;
import tech.easyflow.chatlog.domain.dto.ChatSessionExtPayload;
import tech.easyflow.chatlog.domain.dto.ChatSessionSummary;
import tech.easyflow.chatlog.service.ChatPersistDispatcher;
import tech.easyflow.chatlog.service.ChatRoundOperateService;
import tech.easyflow.chatlog.service.ChatRoundQueryService;
import tech.easyflow.chatlog.service.ChatSessionQueryService;
import tech.easyflow.chatlog.support.ChatJsonSupport;
import tech.easyflow.core.runtime.ChatRuntimeContext;
import tech.easyflow.core.runtime.ChatRuntimeExtKeys;

import java.math.BigInteger;
import java.util.Date;
import java.util.List;

/**
 * {@link ChatlogRuntimeListener} 单元测试。
 */
public class ChatlogRuntimeListenerTest {

    /**
     * 会话准备阶段应把额外知识库选择写入 ext_json。
     */
    @Test
    public void onSessionPreparedShouldWriteExtraKnowledgeIdsToExtJson() {
        CapturingChatPersistDispatcher dispatcher = new CapturingChatPersistDispatcher();
        ChatlogRuntimeListener listener = new ChatlogRuntimeListener(
                dispatcher,
                new NoopChatRoundOperateService(),
                new NoopChatRoundQueryService(),
                new NoopChatSessionQueryService(),
                new ChatJsonSupport(new ObjectMapper())
        );
        ChatRuntimeContext context = new ChatRuntimeContext();
        context.setSessionId(BigInteger.valueOf(1001));
        context.setTenantId(BigInteger.ONE);
        context.setDeptId(BigInteger.TEN);
        context.setUserId(BigInteger.valueOf(7));
        context.setUserAccount("admin");
        context.setAssistantId(BigInteger.valueOf(88));
        context.setAssistantCode("bot-88");
        context.setAssistantName("测试助手");
        context.setSessionTitle("你好");
        context.getExt().put(
                ChatRuntimeExtKeys.EXTRA_KNOWLEDGE_IDS,
                List.of(BigInteger.valueOf(11), BigInteger.valueOf(22))
        );

        listener.onSessionPrepared(context);

        Assert.assertNotNull(dispatcher.captured);
        ChatSessionExtPayload payload = new ChatJsonSupport(new ObjectMapper())
                .fromJson(dispatcher.captured.getExtJson(), ChatSessionExtPayload.class);
        Assert.assertEquals(
                List.of(BigInteger.valueOf(11), BigInteger.valueOf(22)),
                payload.getExtraKnowledgeIds()
        );
    }

    /**
     * 重新生成时历史上下文应排除当前轮旧问题和旧答案。
     */
    @Test
    public void loadMessagesShouldExcludeRegenerateRoundHistory() {
        ChatlogRuntimeListener listener = new ChatlogRuntimeListener(
                null,
                new NoopChatRoundOperateService(),
                new NoopChatRoundQueryService(),
                new TailChatSessionQueryService(List.of(
                        record(4, 2, "assistant", "旧答案"),
                        record(3, 2, "user", "当前问题"),
                        record(2, 1, "assistant", "上一轮答案"),
                        record(1, 1, "user", "上一轮问题")
                )),
                new ChatJsonSupport(new ObjectMapper())
        );
        ChatRuntimeContext context = new ChatRuntimeContext();
        context.setSessionId(BigInteger.valueOf(1001));
        context.getExt().put(ChatRuntimeExtKeys.REGENERATE_ROUND_ID, BigInteger.valueOf(2));

        List<tech.easyflow.core.runtime.ChatRuntimeMessage> messages = listener.loadMessages(context, 10);

        Assert.assertEquals(2, messages.size());
        Assert.assertEquals("上一轮问题", messages.get(0).getContentText());
        Assert.assertEquals("上一轮答案", messages.get(1).getContentText());
    }

    private static ChatMessageRecord record(long id, int roundId, String role, String text) {
        ChatMessageRecord record = new ChatMessageRecord();
        record.setId(BigInteger.valueOf(id));
        record.setSessionId(BigInteger.valueOf(1001));
        record.setRoundId(BigInteger.valueOf(roundId));
        record.setSenderRole(role);
        record.setContentType("TEXT");
        record.setContentText(text);
        record.setCreated(new Date(id));
        return record;
    }

    private static class CapturingChatPersistDispatcher extends ChatPersistDispatcher {

        private ChatSessionUpsertCommand captured;

        private CapturingChatPersistDispatcher() {
            super(null, null, null, null);
        }

        @Override
        public ChatSessionSummary createOrTouchSession(ChatSessionUpsertCommand command) {
            this.captured = command;
            return new ChatSessionSummary();
        }
    }

    private static class NoopChatRoundOperateService implements ChatRoundOperateService {

        @Override
        public ChatRoundRecord requireRegeneratableRound(BigInteger sessionId, BigInteger roundId) {
            return null;
        }

        @Override
        public List<tech.easyflow.chatlog.domain.dto.ChatMessageRecord> listVariants(BigInteger sessionId, BigInteger roundId) {
            return List.of();
        }

        @Override
        public tech.easyflow.chatlog.domain.dto.ChatMessageRecord selectVariant(BigInteger sessionId, BigInteger roundId, Integer variantIndex, BigInteger operatorId) {
            return null;
        }
    }

    private static class NoopChatRoundQueryService implements ChatRoundQueryService {

        @Override
        public ChatRoundRecord getLatestRound(BigInteger sessionId) {
            return null;
        }

        @Override
        public ChatRoundRecord getRound(BigInteger sessionId, BigInteger roundId) {
            return null;
        }

        @Override
        public List<tech.easyflow.chatlog.domain.dto.ChatMessageRecord> listRoundVariants(BigInteger sessionId, BigInteger roundId) {
            return List.of();
        }

        @Override
        public tech.easyflow.chatlog.domain.dto.ChatMessageRecord getRoundVariant(BigInteger sessionId, BigInteger roundId, Integer variantIndex) {
            return null;
        }

        @Override
        public boolean hasRounds(BigInteger sessionId) {
            return false;
        }
    }

    private static class NoopChatSessionQueryService implements ChatSessionQueryService {

        @Override
        public List<ChatSessionSummary> listSessions(BigInteger userId, BigInteger assistantId, tech.easyflow.chatlog.domain.query.ChatPageQuery query) {
            return List.of();
        }

        @Override
        public long countSessions(BigInteger userId, BigInteger assistantId) {
            return 0;
        }

        @Override
        public tech.easyflow.chatlog.domain.dto.ChatSessionPage pageSessions(BigInteger userId, BigInteger assistantId, tech.easyflow.chatlog.domain.query.ChatPageQuery query) {
            return new tech.easyflow.chatlog.domain.dto.ChatSessionPage();
        }

        @Override
        public ChatSessionSummary getSessionSummary(BigInteger sessionId) {
            return null;
        }

        @Override
        public tech.easyflow.chatlog.domain.dto.ChatHistoryPage pageMainlineMessages(BigInteger sessionId, tech.easyflow.chatlog.domain.query.ChatPageQuery query) {
            return new tech.easyflow.chatlog.domain.dto.ChatHistoryPage();
        }

        @Override
        public List<tech.easyflow.chatlog.domain.dto.ChatMessageRecord> listMainlineMessages(BigInteger sessionId) {
            return List.of();
        }

        @Override
        public List<tech.easyflow.chatlog.domain.dto.ChatMessageRecord> getRecentTail(BigInteger sessionId, int limit) {
            return List.of();
        }
    }

    private static class TailChatSessionQueryService extends NoopChatSessionQueryService {

        private final List<ChatMessageRecord> records;

        private TailChatSessionQueryService(List<ChatMessageRecord> records) {
            this.records = records;
        }

        @Override
        public List<ChatMessageRecord> getRecentTail(BigInteger sessionId, int limit) {
            return records.subList(0, Math.min(records.size(), limit));
        }
    }
}
