package com.mortals.xhx.webflux.template;

import com.mortals.framework.util.ThreadPool;
import com.mortals.xhx.thread.ChatMsgRspThread;
import com.mortals.xhx.webflux.common.CommonError;
import com.mortals.xhx.webflux.common.JsonUtils;
import com.mortals.xhx.webflux.factory.*;
import com.mortals.xhx.webflux.service.FlowChatCallBack;
import com.mortals.xhx.webflux.service.IFlowChat;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.core.publisher.Flux;

import javax.annotation.PostConstruct;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 流式问答抽象父类模板
 *
 * @author: zxfei
 * @date: 2025/3/13 10:02
 */
@Slf4j
@Component
public abstract class AbstractFlowChatTemplate implements IFlowChat, FlowChatCallBack {

    private WebClient webClient;

    @PostConstruct
    private void init() {
        webClient = WebClient
                .builder()
                .defaultHeader(HttpHeaders.CONTENT_TYPE, "application/json").build();
    }

    private final ConcurrentHashMap<Long, FlowChatSubscriber> subscriberMap = new ConcurrentHashMap<>();

    /**
     * 问答入口
     */
    @Override
    public Flux<String> request(FlowChatContext context) {
        // 请求大模型问答之前的逻辑处理
        doPreRequest(context);
        // 请求大模型、处理回调逻辑
        return Flux.create(emitter -> {
            Flux<String> response = this.doRequest(context, buildRequest(context), buildRequestHeaders(context));
            log.info("subscriberMap in AbstractChatService before put: {}", JsonUtils.toJson(subscriberMap));
            FlowChatSubscriber subscriber = new FlowChatSubscriber(emitter, this, context, subscriberMap);
            //subscriberMap.put(context.getRequestParam().getSessionId(), subscriber);
            log.info("subscriberMap in AbstractChatService after put: " + JsonUtils.toJson(subscriberMap));
            response.subscribe(subscriber);
            emitter.onDispose(subscriber);

            ChatMsgRspThread chatMsgRspThread = new ChatMsgRspThread(context);

            ThreadPool.getInstance().execute(chatMsgRspThread);

        });
    }

    @Override
    public void stop(Long sessionId) {
        FlowChatSubscriber subscriber = subscriberMap.get(sessionId);
        if (subscriber == null) {
            return;
        }
        subscriber.stop();
    }

    /**
     * 保存对话内容
     */
    protected Long saveChatMsg(FlowChatContext context, String query) {

        // 保存对话query  todo

        return null;
    }

    /**
     * 更新对话内容
     */
    protected void updateChatMsg(FlowChatContext context) {

        // 更新answer  todo

    }

    /**
     * 构建响应参数
     */
    protected Map<String, Object> buildAnswer(String answer) {
        Map<String, Object> res = new HashMap<>(4, 1);
        res.put("answer", answer);
        res.put("docs", Collections.emptyList());
        res.put("process_docs", Collections.emptyList());
        return res;
    }

    /**
     * 请求大模型
     */
    private Flux<String> doRequest(FlowChatContext context, FLowChatRequest request, HttpHeaders headers) {
        log.info("请求大模型开始，URL:{}, 参数:{}", request.getUrl(), request.getJsonBody());

        Flux<String> flux = webClient
                .post()
                .uri(request.getUrl())
                .accept(MediaType.TEXT_EVENT_STREAM)
                .headers(httpHeaders -> httpHeaders.addAll(headers))
                .bodyValue(request.getJsonBody())
                //.bodyValue(bodyStr)
                .retrieve()
                .bodyToFlux(String.class)
                .onErrorResume(WebClientResponseException.class, ex -> {
                    //subscriberMap.remove(context.getRequestParam().getSessionId());
                    log.error("请求大模型接口异常", ex);
                    return Flux.just(JsonUtils.toJson(buildAnswer(CommonError.LLM_ERROR.getMsg())));
                })
                .onErrorResume(Throwable.class, ex -> {
                    //subscriberMap.remove(context.getRequestParam().getSessionId());
                    log.error("系统异常", ex);
                    return Flux.just(JsonUtils.toJson(buildAnswer(CommonError.GLOBAL_ERROR.getMsg())));
                });


        //如果一定时间未收到响应，则模拟响应数据回写

        return flux;
    }

    /**
     * 前置逻辑处理
     */
    protected abstract void doPreRequest(FlowChatContext context);

    /**
     * 构建大模型f请求参数(robot to fastapi)
     */
    protected abstract FLowChatRequest buildRequest(FlowChatContext context);

    protected abstract HttpHeaders buildRequestHeaders(FlowChatContext context);
}
