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

import com.seer.teach.common.ResultBean;
import com.seer.teach.common.cache.SseDataCacheService;
import com.seer.teach.common.config.LlmConfig;
import com.seer.teach.common.enums.ResultCodeEnum;
import com.seer.teach.openapi.cache.OpenApiSseLifeCycleCacheService;
import com.seer.teach.openapi.llm.resp.CorrectDTO;
import com.seer.teach.openapi.llm.req.LlmAiCorrectReq;
import com.seer.teach.openapi.service.LlmSseService;
import com.seer.teach.openapi.service.SnCallLogsService;
import com.seer.teach.openapi.uitl.OpenApiUtils;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
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.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<String, Disposable>> TYPE_SSE_MAP = new ConcurrentHashMap<>();

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

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

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

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

    /**
     * 呼叫日志服务实例
     * 用于处理呼叫日志相关的业务逻辑
     */
    private final SnCallLogsService snCallLogsService;

    /**
     * SSE数据缓存服务实例
     * 用于管理SSE数据的缓存操作
     */
    private final SseDataCacheService sseDataCacheService;

    @PostConstruct
    public void initTypeSseMap() {
        TYPE_SSE_MAP.put(OpenApiUtils.SSE_TYPE_CORRECT, CORRECT_SSE_MAP);
        TYPE_SSE_MAP.put(OpenApiUtils.SSE_TYPE_CHAT, CHAT_SSE_MAP);
        TYPE_SSE_MAP.put(OpenApiUtils.SSE_TYPE_STUDY, STUDY_SSE_MAP);
    }

    @Override
    public Flux<ResultBean<CorrectDTO>> aiCorrect(LlmAiCorrectReq params) {
        snCallLogsService.saveSseLog(params.getSn(), "correct", params.getSessionId(), params, params.getImageUrl());
        return sseFactoryService.streamSse(
                        LlmConfig.SSE_AI_CORRECT,
                        params,
                        OpenApiSseLifeCycleCacheService.AI_CORRECT_REDIS_KEY_PREFIX,
                        params.getSn(),
                        params.getSessionId(),
                        20,
                        CorrectDTO.class)
                // 将每条 CorrectDTO 封装到 ResultBean 中
                .map(ResultBean::success)
                .onErrorResume(throwable -> handlerError(params, throwable))
                .doOnNext(data -> sseDataCacheService.appendSseData(params.getSessionId(), data))
                .doFinally(signal -> handleStreamFinally(
                        params.getSessionId(),
                        params.getSn(),
                        OpenApiUtils.SSE_TYPE_CORRECT,
                        signal
                ));
    }

    @NotNull
    private Flux<ResultBean<CorrectDTO>> handlerError(LlmAiCorrectReq params, Throwable throwable) {
        snCallLogsService.updateErrorLog(params.getSessionId(), throwable.getMessage());
        log.info("SSE stream error: {}", throwable.getMessage());
        removeStream(params.getSn(), OpenApiUtils.SSE_TYPE_CORRECT);
        return Flux.just(ResultBean.error(ResultCodeEnum.CORRECT_ERROR));
    }

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

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

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

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


}
