package tech.easyflow.chatlog.service.impl;

import com.mybatisflex.core.keygen.impl.SnowFlakeIDKeyGenerator;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;
import tech.easyflow.chatlog.domain.command.ChatAppendMessageCommand;
import tech.easyflow.chatlog.domain.command.ChatRoundUpsertCommand;
import tech.easyflow.chatlog.domain.command.ChatSessionUpsertCommand;
import tech.easyflow.chatlog.domain.dto.ChatRoundRecord;
import tech.easyflow.chatlog.domain.dto.ChatSessionExtPayload;
import tech.easyflow.chatlog.domain.dto.ChatMessageRecord;
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.ChatConstants;
import tech.easyflow.chatlog.support.ChatJsonSupport;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.core.runtime.ChatRuntimeExtKeys;
import tech.easyflow.core.runtime.ChatRuntimeHistoryPayloadHelper;
import tech.easyflow.core.runtime.ChatRuntimeContext;
import tech.easyflow.core.runtime.ChatRuntimeListener;
import tech.easyflow.core.runtime.ChatRuntimeMessage;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;

@Component
@Order(100)
public class ChatlogRuntimeListener implements ChatRuntimeListener {

    private final SnowFlakeIDKeyGenerator idGenerator = new SnowFlakeIDKeyGenerator();

    private final ChatPersistDispatcher chatPersistDispatcher;
    private final ChatRoundOperateService chatRoundOperateService;
    private final ChatRoundQueryService chatRoundQueryService;
    private final ChatSessionQueryService chatSessionQueryService;
    private final ChatJsonSupport chatJsonSupport;

    public ChatlogRuntimeListener(ChatPersistDispatcher chatPersistDispatcher,
                                  ChatRoundOperateService chatRoundOperateService,
                                  ChatRoundQueryService chatRoundQueryService,
                                  ChatSessionQueryService chatSessionQueryService,
                                  ChatJsonSupport chatJsonSupport) {
        this.chatPersistDispatcher = chatPersistDispatcher;
        this.chatRoundOperateService = chatRoundOperateService;
        this.chatRoundQueryService = chatRoundQueryService;
        this.chatSessionQueryService = chatSessionQueryService;
        this.chatJsonSupport = chatJsonSupport;
    }

    @Override
    public void onSessionPrepared(ChatRuntimeContext context) {
        try {
            ChatSessionUpsertCommand command = new ChatSessionUpsertCommand();
            command.setSessionId(context.getSessionId());
            command.setTenantId(defaultNumber(context.getTenantId()));
            command.setDeptId(defaultNumber(context.getDeptId()));
            command.setUserId(defaultNumber(context.getUserId()));
            command.setUserAccount(context.getUserAccount());
            command.setAssistantId(defaultNumber(context.getAssistantId()));
            command.setAssistantCode(context.getAssistantCode());
            command.setAssistantName(context.getAssistantName());
            command.setTitle(context.getSessionTitle());
            command.setExtJson(resolveExtJson(context));
            command.setOperatorId(defaultNumber(context.getUserId()));
            chatPersistDispatcher.createOrTouchSession(command);
        } catch (RuntimeException ex) {
            throw persistFailed(ex);
        }
    }

    @Override
    public void onUserMessage(ChatRuntimeContext context, ChatRuntimeMessage message) {
        try {
            if (prepareRoundContext(context, message)) {
                return;
            }
            chatPersistDispatcher.appendUserMessage(toAppendCommand(context, message));
        } catch (RuntimeException ex) {
            throw persistFailed(ex);
        }
    }

    @Override
    public void onAssistantCompleted(ChatRuntimeContext context, ChatRuntimeMessage message) {
        try {
            applyAssistantRoundMetadata(context, message);
            chatPersistDispatcher.appendAssistantMessage(toAppendCommand(context, message));
            chatPersistDispatcher.createOrTouchRound(buildAssistantCompletedRoundCommand(context, message));
        } catch (RuntimeException ex) {
            throw persistFailed(ex);
        }
    }

    @Override
    public void onChatFailed(ChatRuntimeContext context, Throwable throwable) {
        try {
            BigInteger roundId = resolveNumber(context, ChatRuntimeExtKeys.CURRENT_ROUND_ID);
            if (context == null || context.getSessionId() == null || roundId == null) {
                return;
            }
            ChatRoundRecord currentRound = chatRoundQueryService.getRound(context.getSessionId(), roundId);
            if (currentRound == null) {
                return;
            }
            ChatRoundUpsertCommand command = new ChatRoundUpsertCommand();
            command.setRoundId(currentRound.getId());
            command.setSessionId(currentRound.getSessionId());
            command.setRoundNo(currentRound.getRoundNo());
            command.setUserMessageId(currentRound.getUserMessageId());
            command.setSelectedAssistantMessageId(currentRound.getSelectedAssistantMessageId());
            command.setSelectedVariantIndex(currentRound.getSelectedVariantIndex());
            command.setVariantCount(currentRound.getVariantCount());
            command.setStatus(ChatConstants.ROUND_STATUS_READY);
            command.setOperatorId(defaultNumber(context.getUserId()));
            chatPersistDispatcher.createOrTouchRound(command);
        } catch (RuntimeException ex) {
            throw persistFailed(ex);
        }
    }

    @Override
    public List<ChatRuntimeMessage> loadMessages(ChatRuntimeContext context, int limit) {
        if (context == null || context.getSessionId() == null || limit <= 0) {
            return Collections.emptyList();
        }
        BigInteger regenerateRoundId = resolveNumber(context, ChatRuntimeExtKeys.REGENERATE_ROUND_ID);
        int queryLimit = regenerateRoundId == null ? limit : limit + 4;
        List<ChatMessageRecord> records = new ArrayList<>(chatSessionQueryService.getRecentTail(context.getSessionId(), queryLimit));
        if (regenerateRoundId != null) {
            records.removeIf(record -> regenerateRoundId.equals(record.getRoundId()));
            if (records.size() > limit) {
                records = new ArrayList<>(records.subList(0, limit));
            }
        }
        Collections.reverse(records);
        List<ChatRuntimeMessage> messages = new ArrayList<>(records.size());
        for (ChatMessageRecord record : records) {
            if ((record.getContentText() == null || record.getContentText().isBlank())
                    && !ChatRuntimeHistoryPayloadHelper.hasStructuredHistory(record.getContentPayload())) {
                continue;
            }
            ChatRuntimeMessage message = new ChatRuntimeMessage();
            message.setMessageId(record.getId());
            message.setRole(record.getSenderRole());
            message.setContentType(record.getContentType());
            message.setContentText(record.getContentText());
            message.setContentPayload(record.getContentPayload());
            message.setCreatedAt(record.getCreated());
            message.setSenderId(record.getSenderId());
            message.setSenderName(record.getSenderName());
            messages.add(message);
        }
        return messages;
    }

    private ChatAppendMessageCommand toAppendCommand(ChatRuntimeContext context, ChatRuntimeMessage message) {
        ChatAppendMessageCommand command = new ChatAppendMessageCommand();
        command.setMessageId(message.getMessageId() == null ? BigInteger.valueOf(idGenerator.nextId()) : message.getMessageId());
        command.setTenantId(defaultNumber(context.getTenantId()));
        command.setDeptId(defaultNumber(context.getDeptId()));
        command.setSessionId(context.getSessionId());
        command.setUserId(defaultNumber(context.getUserId()));
        command.setUserAccount(context.getUserAccount());
        command.setAssistantId(defaultNumber(context.getAssistantId()));
        command.setAssistantCode(context.getAssistantCode());
        command.setAssistantName(context.getAssistantName());
        command.setSessionTitle(context.getSessionTitle());
        command.setSenderId(defaultNumber(message.getSenderId()));
        command.setSenderName(message.getSenderName());
        command.setSenderRole(message.getRole());
        command.setContentType(message.getContentType());
        command.setContentText(message.getContentText());
        command.setContentPayload(message.getContentPayload());
        command.setRoundId(message.getRoundId());
        command.setRoundNo(message.getRoundNo());
        command.setMessageKind(message.getMessageKind());
        command.setVariantIndex(message.getVariantIndex());
        command.setCreatedBy(defaultNumber(context.getUserId()));
        command.setCreated(message.getCreatedAt());
        return command;
    }

    private boolean prepareRoundContext(ChatRuntimeContext context, ChatRuntimeMessage message) {
        if (context == null || message == null || context.getSessionId() == null) {
            return false;
        }
        BigInteger regenerateRoundId = resolveNumber(context, ChatRuntimeExtKeys.REGENERATE_ROUND_ID);
        if (regenerateRoundId != null) {
            ChatRoundRecord round = chatRoundOperateService.requireRegeneratableRound(context.getSessionId(), regenerateRoundId);
            context.getExt().put(ChatRuntimeExtKeys.CURRENT_ROUND_ID, round.getId());
            context.getExt().put(ChatRuntimeExtKeys.CURRENT_ROUND_NO, round.getRoundNo());
            context.getExt().put(ChatRuntimeExtKeys.CURRENT_VARIANT_INDEX, Math.max(round.getVariantCount() + 1, 1));
            ChatRoundUpsertCommand command = new ChatRoundUpsertCommand();
            command.setRoundId(round.getId());
            command.setSessionId(round.getSessionId());
            command.setRoundNo(round.getRoundNo());
            command.setUserMessageId(round.getUserMessageId());
            command.setSelectedAssistantMessageId(round.getSelectedAssistantMessageId());
            command.setSelectedVariantIndex(round.getSelectedVariantIndex());
            command.setVariantCount(round.getVariantCount());
            command.setStatus(ChatConstants.ROUND_STATUS_ANSWERING);
            command.setOperatorId(defaultNumber(context.getUserId()));
            chatPersistDispatcher.createOrTouchRound(command);
            return true;
        }
        ChatRoundRecord latestRound = chatRoundQueryService.getLatestRound(context.getSessionId());
        if (latestRound != null && latestRound.getId() != null
                && !ChatConstants.ROUND_STATUS_LOCKED.equalsIgnoreCase(latestRound.getStatus())) {
            ChatRoundUpsertCommand lockCommand = new ChatRoundUpsertCommand();
            lockCommand.setRoundId(latestRound.getId());
            lockCommand.setSessionId(latestRound.getSessionId());
            lockCommand.setRoundNo(latestRound.getRoundNo());
            lockCommand.setUserMessageId(latestRound.getUserMessageId());
            lockCommand.setSelectedAssistantMessageId(latestRound.getSelectedAssistantMessageId());
            lockCommand.setSelectedVariantIndex(latestRound.getSelectedVariantIndex());
            lockCommand.setVariantCount(latestRound.getVariantCount());
            lockCommand.setStatus(ChatConstants.ROUND_STATUS_LOCKED);
            lockCommand.setOperatorId(defaultNumber(context.getUserId()));
            chatPersistDispatcher.createOrTouchRound(lockCommand);
        }
        BigInteger roundId = BigInteger.valueOf(idGenerator.nextId());
        int roundNo = latestRound == null || latestRound.getRoundNo() == null ? 1 : latestRound.getRoundNo() + 1;
        if (message.getMessageId() == null) {
            message.setMessageId(BigInteger.valueOf(idGenerator.nextId()));
        }
        context.getExt().put(ChatRuntimeExtKeys.CURRENT_ROUND_ID, roundId);
        context.getExt().put(ChatRuntimeExtKeys.CURRENT_ROUND_NO, roundNo);
        context.getExt().put(ChatRuntimeExtKeys.CURRENT_VARIANT_INDEX, 1);
        message.setRoundId(roundId);
        message.setRoundNo(roundNo);
        message.setMessageKind(ChatConstants.MESSAGE_KIND_USER_PROMPT);
        message.setVariantIndex(null);
        ChatRoundUpsertCommand command = new ChatRoundUpsertCommand();
        command.setRoundId(roundId);
        command.setSessionId(context.getSessionId());
        command.setRoundNo(roundNo);
        command.setUserMessageId(message.getMessageId());
        command.setSelectedVariantIndex(0);
        command.setVariantCount(0);
        command.setStatus(ChatConstants.ROUND_STATUS_ANSWERING);
        command.setOperatorId(defaultNumber(context.getUserId()));
        chatPersistDispatcher.createOrTouchRound(command);
        return false;
    }

    private void applyAssistantRoundMetadata(ChatRuntimeContext context, ChatRuntimeMessage message) {
        if (message.getMessageId() == null) {
            message.setMessageId(BigInteger.valueOf(idGenerator.nextId()));
        }
        message.setRoundId(resolveNumber(context, ChatRuntimeExtKeys.CURRENT_ROUND_ID));
        message.setRoundNo(resolveInteger(context, ChatRuntimeExtKeys.CURRENT_ROUND_NO));
        message.setVariantIndex(resolveInteger(context, ChatRuntimeExtKeys.CURRENT_VARIANT_INDEX));
        message.setMessageKind(ChatConstants.MESSAGE_KIND_ASSISTANT_VARIANT);
    }

    private ChatRoundUpsertCommand buildAssistantCompletedRoundCommand(ChatRuntimeContext context, ChatRuntimeMessage message) {
        ChatRoundUpsertCommand command = new ChatRoundUpsertCommand();
        command.setRoundId(message.getRoundId());
        command.setSessionId(context.getSessionId());
        command.setRoundNo(message.getRoundNo());
        ChatRoundRecord existing = chatRoundQueryService.getRound(context.getSessionId(), message.getRoundId());
        if (existing != null) {
            command.setUserMessageId(existing.getUserMessageId());
        }
        command.setSelectedAssistantMessageId(message.getMessageId());
        command.setSelectedVariantIndex(message.getVariantIndex());
        command.setVariantCount(message.getVariantIndex());
        command.setStatus(ChatConstants.ROUND_STATUS_READY);
        command.setOperatorId(defaultNumber(context.getUserId()));
        return command;
    }

    private BigInteger resolveNumber(ChatRuntimeContext context, String key) {
        if (context == null || context.getExt() == null || key == null) {
            return null;
        }
        Object value = context.getExt().get(key);
        if (value == null) {
            return null;
        }
        return new BigInteger(String.valueOf(value));
    }

    private Integer resolveInteger(ChatRuntimeContext context, String key) {
        if (context == null || context.getExt() == null || key == null) {
            return null;
        }
        Object value = context.getExt().get(key);
        if (value == null) {
            return null;
        }
        return Integer.parseInt(String.valueOf(value));
    }

    private BigInteger defaultNumber(BigInteger value) {
        return value == null ? BigInteger.ZERO : value;
    }

    private BusinessException persistFailed(RuntimeException ex) {
        if (ex instanceof BusinessException businessException) {
            return businessException;
        }
        return new BusinessException("聊天记录持久化失败，请稍后重试");
    }

    private String resolveExtJson(ChatRuntimeContext context) {
        if (context == null || context.getExt() == null || context.getExt().isEmpty()) {
            return null;
        }
        if (!context.getExt().containsKey(ChatRuntimeExtKeys.EXTRA_KNOWLEDGE_IDS)) {
            return null;
        }
        Object rawExtraKnowledgeIds = context.getExt().get(ChatRuntimeExtKeys.EXTRA_KNOWLEDGE_IDS);
        if (!(rawExtraKnowledgeIds instanceof List<?> rawList)) {
            return null;
        }
        List<BigInteger> extraKnowledgeIds = new ArrayList<>();
        for (Object item : rawList) {
            if (item == null) {
                continue;
            }
            extraKnowledgeIds.add(new BigInteger(String.valueOf(item)));
        }
        ChatSessionExtPayload payload = new ChatSessionExtPayload();
        payload.setExtraKnowledgeIds(extraKnowledgeIds);
        return chatJsonSupport.toJson(payload);
    }
}
