package com.seer.teach.teacher.service.impl;

import com.seer.teach.common.cache.SseDataCacheService;
import com.seer.teach.common.cache.SseLifeCycleCacheService;
import com.seer.teach.common.cache.dto.CorrectResultDTO;
import com.seer.teach.common.config.LlmConfig;
import com.seer.teach.common.dto.mq.LlmChatDTO;
import com.seer.teach.teacher.common.llm.sse.req.LlmAiCorrectReq;
import com.seer.teach.teacher.common.llm.sse.req.LlmSseChatReq;
import com.seer.teach.teacher.common.llm.sse.req.LlmSseStudyReq;
import com.seer.teach.teacher.common.llm.sse.resp.LlmSseStudyResp;
import com.seer.teach.teacher.common.utils.TeacherUtils;
import com.seer.teach.teacher.service.LlmRequestLogsService;
import com.seer.teach.teacher.service.LlmSseService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.SignalType;

import jakarta.annotation.PostConstruct;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

/**
 * LLM SSE服务实现类，用于管理与AI模型之间的SSE（Server-Sent Events）通信。
 * 提供聊天、批改和学习场景下的流式数据传输功能，并维护用户会话状态。
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class LlmSseServiceImpl implements LlmSseService {

    /**
     * 根据类型映射对应的Disposable容器
     */
    private final ConcurrentHashMap<Integer, ConcurrentHashMap<Integer, Disposable>> TYPE_SSE_MAP = new ConcurrentHashMap<>();

    /**
     * 存储学习类型的SSE连接
     */
    private final ConcurrentHashMap<Integer, Disposable> STUDY_SSE_MAP = new ConcurrentHashMap<>();

    /**
     * 存储批改类型的SSE连接
     */
    private final ConcurrentHashMap<Integer, Disposable> CORRECT_SSE_MAP = new ConcurrentHashMap<>();

    /**
     * 存储聊天类型的SSE连接
     */
    private final ConcurrentHashMap<Integer, Disposable> CHAT_SSE_MAP = new ConcurrentHashMap<>();

    /**
     * SSE工厂服务，提供通用的SSE流构建逻辑
     */
    private final SseFactoryService sseFactoryService;

    private final LlmRequestLogsService llmRequestLogsService;

    private final SseDataCacheService sseDataCacheService;

    /**
     * 初始化TYPE_SSE_MAP结构，在Bean初始化完成后执行
     */
    @PostConstruct
    public void initTypeSseMap() {
        TYPE_SSE_MAP.put(TeacherUtils.SSE_TYPE_CORRECT, CORRECT_SSE_MAP);
        TYPE_SSE_MAP.put(TeacherUtils.SSE_TYPE_CHAT, CHAT_SSE_MAP);
        TYPE_SSE_MAP.put(TeacherUtils.SSE_TYPE_STUDY, STUDY_SSE_MAP);
    }

    @Override
    public Flux<LlmChatDTO> aiChat(LlmSseChatReq llmReq) {
        llmRequestLogsService.saveSseLog(llmReq.getUserId(), "chat", llmReq.getSessionId(), llmReq, llmReq.getMessages());
        return sseFactoryService.streamSse(
                        LlmConfig.SSE_AI_CHAT,
                        llmReq,
                        SseLifeCycleCacheService.AI_CHAT_REDIS_KEY_PREFIX,
                        llmReq.getUserId(),
                        llmReq.getSessionId(),
                        10,
                        LlmChatDTO.class)
                .doOnError(e -> handleStreamError(llmReq.getSessionId(), llmReq.getUserId(), TeacherUtils.SSE_TYPE_CHAT, e.getMessage()))
                .doOnNext(data -> sseDataCacheService.appendSseData(llmReq.getSessionId(), data))
                .doFinally(signal -> handleStreamFinally(llmReq.getSessionId(), llmReq.getUserId(), TeacherUtils.SSE_TYPE_CHAT, signal));
    }

    @Override
    public Flux<CorrectResultDTO> aiCorrect(LlmAiCorrectReq params) {
        llmRequestLogsService.saveSseLog(params.getUserId(), "correct", params.getSessionId(), params, params.getImageUrl());
        return sseFactoryService.streamSse(
                        LlmConfig.SSE_AI_CORRECT,
                        params,
                        SseLifeCycleCacheService.AI_CORRECT_REDIS_KEY_PREFIX,
                        params.getUserId(),
                        params.getSessionId(),
                        20,
                        CorrectResultDTO.class)
                .doOnError(e -> handleStreamError(params.getSessionId(), params.getUserId(), TeacherUtils.SSE_TYPE_CORRECT, e.getMessage()))
                .doOnNext(data -> sseDataCacheService.appendSseData(params.getSessionId(), data))
                .doFinally(signal -> handleStreamFinally(params.getSessionId(), params.getUserId(), TeacherUtils.SSE_TYPE_CORRECT, signal));
    }

    @Override
    public Flux<LlmSseStudyResp> aiStudy(LlmSseStudyReq params) {
        if (params.getMessages() == null) {
            params.setMessages(Collections.emptyList());
        }
        llmRequestLogsService.saveSseLog(params.getUserId(), "study", params.getSessionId(), params, params.getVideoIndex());
        return sseFactoryService.streamSse(
                        LlmConfig.SSE_STUDY,
                        params,
                        SseLifeCycleCacheService.AI_STUDY_REDIS_KEY_PREFIX,
                        params.getUserId(),
                        params.getSessionId(),
                        10,
                        LlmSseStudyResp.class)
                .doOnError(e -> handleStreamError(params.getSessionId(), params.getUserId(), TeacherUtils.SSE_TYPE_STUDY, e.getMessage()))
                .doOnNext(data -> sseDataCacheService.appendSseData(params.getSessionId(), data))
                .doFinally(signal -> handleStreamFinally(params.getSessionId(), params.getUserId(), TeacherUtils.SSE_TYPE_STUDY, signal));
    }

    @Override
    public void registerStream(Integer userId, Integer type, Disposable disposable) {
        closeStream(userId, type);
        ConcurrentHashMap<Integer, Disposable> sseMap = TYPE_SSE_MAP.get(type);
        if (sseMap != null) {
            sseMap.put(userId, disposable);
        }
    }

    @Override
    public void removeStudyStream(Integer userId) {
        ConcurrentHashMap<Integer, Disposable> map = TYPE_SSE_MAP.get(TeacherUtils.SSE_TYPE_STUDY);
        Disposable disposable = map.remove(userId);
        if (disposable != null) {
            disposable.dispose();
            log.info("SSE 学习流已关闭: userId={}", userId);
        }
    }

    @Override
    public void removeChatStream(Integer userId) {
        ConcurrentHashMap<Integer, Disposable> map = TYPE_SSE_MAP.get(TeacherUtils.SSE_TYPE_CHAT);
        Disposable disposable = map.remove(userId);
        if (disposable != null) {
            disposable.dispose();
            log.info("SSE 对话流已关闭: userId={}", userId);
        }
    }

    @Override
    public void removeStream(Integer userId, Integer type) {
        ConcurrentHashMap<Integer, Disposable> map = TYPE_SSE_MAP.get(type);
        Disposable disposable = map.remove(userId);
        if (disposable != null && !disposable.isDisposed()) {
            disposable.dispose();
        }
    }

    public void closeStream(Integer userId, Integer type) {
        ConcurrentHashMap<Integer, Disposable> sseMap = TYPE_SSE_MAP.get(type);
        Disposable disposable = sseMap.remove(userId);
        if (disposable != null && !disposable.isDisposed()) {
            disposable.dispose();
        }
    }

    private void handleStreamError(Integer sessionId, Integer userId, Integer type, String errorMessage) {
        llmRequestLogsService.updateErrorLog(sessionId, errorMessage);
        removeStream(userId, type);
    }

    private void handleStreamFinally(Integer sessionId, Integer userId, Integer type, SignalType signal) {
        if (sseDataCacheService.exists(sessionId)) {
            List<Object> data = sseDataCacheService.getSseData(sessionId);
            llmRequestLogsService.updateResponseLog(sessionId, data);
            sseDataCacheService.clearSseData(sessionId);
        } else {
            llmRequestLogsService.updateEmptyResponseLog(sessionId);
        }
        if (signal == SignalType.CANCEL) {
            removeStream(userId, type);
        }
    }

    @Override
    public Boolean checkUserHasStudyStream(Integer userId) {
        ConcurrentHashMap<Integer, Disposable> hashMap = TYPE_SSE_MAP.get(TeacherUtils.SSE_TYPE_STUDY);
        return hashMap.containsKey(userId);
    }

}
