package com.ch.jedge.glm.zhipuai;

import com.ch.jedge.glm.JedgeGlmBridgeService;
import com.changhong.qlib.QData;
import com.changhong.qlib.intf.QIData;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.CountDownLatch;

/**
 * 监听模型结果，等到所有事件完成后统一返回。效果类似同步调用，如您有其他Listener需求可以联系我们，或自己对Listener进行扩展
 */
public class StandardEventSourceListener extends ModelEventSourceListener {
    public StandardEventSourceListener(JedgeGlmBridgeService service){
        this.service = service;
    }
    private final JedgeGlmBridgeService service;
    private static final Logger logger = LoggerFactory.getLogger(StandardEventSourceListener.class);

    private String outputText = "";

    private boolean incremental;

    /*    private EventStatus status = new EventStatus();*/

    protected CountDownLatch countDownLatch = new CountDownLatch(1);

    private Gson gson = new Gson();

    private SseMeta meta;

    @Override
    public void onOpen(EventSource eventSource, Response response) {
        //logger.info("server start sending events");
    }

    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data, String meta) {
        boolean is_finish = false;
        String formatData = "";
        if ("finish".equals(type) && ZhipuStringUtils.isNotEmpty(meta)) {
            this.meta = gson.fromJson(meta, SseMeta.class);
            System.out.println();
            is_finish = true;
        }
        if("add".equals(type)) {
            formatData = gson.fromJson(data, JsonObject.class).get("data").getAsString();
            System.out.print(formatData);
            is_finish = false;
        }
        QIData qData = new QData();
        String sid = eventSource.request().headers().get("sid");
        String srcMod = eventSource.request().headers().get("srcMod");
        qData.putString("sid", sid);
        qData.putString("~m", srcMod);
        qData.putString("delta", formatData);
        qData.putBoolean("finished", is_finish);
//        service.putMgbusMessage(qData);
        service.postServiceMessage(srcMod, "/gpt/LocalStream", qData);
        if (this.isIncremental()) {
            outputText += data;
        } else {
            outputText = data;
        }
    }

    @Override
    public void onClosed(EventSource eventSource) {
        logger.info("server stream closed");
        try {
/*            synchronized (status) {
                status.setStatus(EventStatus.StatusClosed);
                status.notify();
            }*/
            eventSource.cancel();
        } finally {
            countDownLatch.countDown();
        }
    }

    @Override
    public void onFailure(EventSource eventSource, Throwable t, Response response) {
        try {
            logger.error("sse connection fail");
/*            synchronized (status) {
                status.setStatus(EventStatus.StatusFailed);
                status.notify();
            }*/
            eventSource.cancel();
        } finally {
            countDownLatch.countDown();
        }
    }

    @Override
    public String getOutputText() {
        return outputText;
    }

    public SseMeta getMeta() {
        return this.meta;
    }

    public CountDownLatch getCountDownLatch() {
        return countDownLatch;
    }

    public boolean isIncremental() {
        return incremental;
    }

    public void setIncremental(boolean incremental) {
        this.incremental = incremental;
    }
/*
    private static class EventStatus {

        public static final String StatusClosed = "closed";
        public static final String StatusFailed = "failed";

        private String status;

        public String getStatus() {
            return status;
        }

        public void setStatus(String status) {
            this.status = status;
        }
    }*/
}
