package com.mortals.xhx.webflux.factory;

import cn.hutool.core.util.StrUtil;
import com.mortals.xhx.webflux.common.JsonUtils;
import com.mortals.xhx.webflux.service.FlowChatCallBack;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Subscriber;
import org.reactivestreams.Subscription;
import org.springframework.util.ObjectUtils;
import reactor.core.Disposable;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.util.concurrent.ConcurrentHashMap;

/**
 * 流式问答监听器
 *
 * @author: zxfei
 * @date: 2025/3/13 10:02
 */
@Slf4j
public class FlowChatSubscriber implements Subscriber<String>, Disposable {

    private final FluxSink<String> emitter;
    private Subscription subscription;
    private final FlowChatContext context;
    private final FlowChatCallBack callBack;
    private final ConcurrentHashMap<Long, FlowChatSubscriber> subscriberMap;
    private boolean stopFlag = false;

    public void stop() {
        stopFlag = true;
    }

    public FlowChatSubscriber(FluxSink<String> emitter, FlowChatCallBack callBack,
                              FlowChatContext context, ConcurrentHashMap<Long, FlowChatSubscriber> subscriberMap) {
        this.emitter = emitter;
        this.callBack = callBack;
        this.context = context;
        this.context.setAnswer(new StringBuilder());
        this.context.setDocs(new StringBuilder());
        this.subscriberMap = subscriberMap;
        this.context.setEmitter(emitter);
        log.info("subscriberMap in ChatSubscriber before remove: {}", JsonUtils.toJson(this.subscriberMap));
    }

    @Override
    public void onSubscribe(Subscription subscription) {
        log.info("流式问答订阅开始, sessionId:{}, msgId:{}", context.getRequestParam().getId(), context.getMsgId());
        // 订阅开始，初始化subscription
        this.subscription = subscription;
        // 请求接收一个数据项
        subscription.request(1);
    }

    @Override
    public void onNext(String data) {

        String rsp = "";
        try {
            rsp = callBack.onNext(data, context);
            if (ObjectUtils.isEmpty(rsp)) {
                this.stopFlag = true;
            }
           // rsp="data:"+rsp;
            log.info(" onNext rsp:{}", rsp);
        } catch (Exception e) {
            log.error("流式问答异常:{}", e.getMessage());
        } finally {
            // todo 临时打印日志
            // log.info("=============== data: {}", data);
            if (stopFlag) {
                log.info("stopFlag is true");
                subscription.cancel();
                onComplete();
            } else {
                // emitter.(msg);
                // 将数据发送给前端 如果是起始包 需要要处理一下
                emitter.next(rsp);
                // 继续请求接收下一个数据项
                subscription.request(1);
            }
        }
    }

    @Override
    public void onError(Throwable t) {
        // 处理数据流完成后的回调逻辑
        try {
            callBack.error(context);
            log.error("流式问答异常, sessionId:{}, msgId:{}", context.getRequestParam().getId(), context.getMsgId());
        } catch (Exception e) {
            log.error("流式问答异常", e);
        } finally {
            dispose();
        }
    }

    @Override
    public void onComplete() {
        // 处理数据流完成后的回调逻辑
        try {
            //   callBack.completed(context);
            log.info("流式问答订阅结束, sessionId:{}, msgId:{}", context.getRequestParam().getId(), context.getMsgId());
        } catch (Exception e) {
            log.error("流式问答异常", e);
        } finally {
            dispose();
        }
    }

    @Override
    public void dispose() {
        emitter.complete();
        //subscriberMap.remove(context.getRequestParam().getSessionId());
        log.info("subscriberMap in ChatSubscriber after remove: {}", JsonUtils.toJson(this.subscriberMap));
    }
}
