package com.ch.jedge.glm;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.ch.jedge.glm.zhipuai.*;
import com.changhong.jedge.JMgbusModual;
import com.changhong.jedge.JMgbusService;
import com.changhong.jedge.JMgbusUtil;
import com.changhong.jedge.MgbusApi;
import com.changhong.qlib.QData;
import com.changhong.qlib.intf.QIData;
import com.changhong.qlib.util.NumberUtils;
import com.changhong.qlib.util.StringUtils;
import com.changhong.qlib.util.sync.SystemUtils;
import com.google.gson.Gson;
import okhttp3.*;
import okhttp3.internal.sse.RealEventSource;
import org.apache.commons.math3.linear.ArrayRealVector;
import org.apache.commons.math3.linear.RealVector;

import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;

import static com.ch.jedge.jbot.JEJBotLLMConst.ival_mufis_text_message_gap_time_;
import static com.ch.jedge.jbot.JEJBotLLMConst.sval_glm_config_filename;

public class JedgeGlmBridgeService extends JMgbusService {

//    static final QIData QIData_DefaultConfig = QData.fromFile(sval_glm_config_filename);

    private ConfigV3 zhipuConfig;
    private ICache tokenCache;
    private TokenManagerV3 tokenManagerV3;
    private OkHttpClient okHttpClient;

    private static final String sseModelUrl = "https://open.bigmodel.cn/api/paas/v3/model-api/{model}/sse-invoke";
    private static final String reqModelUrl = "https://open.bigmodel.cn/api/paas/v3/model-api/{model}/invoke";

    private static final String embeddingURL = "https://open.bigmodel.cn/api/paas/v3/model-api/text_embedding/invoke";
    private static final String roleAssistant = "assistant";
    private static final String roleUser = "user";
    private static final String sseFormat = "data";

    private String sval_glm_model_pro = "chatglm_pro";
    private String sval_glm_model_std = "chatglm_std";
    private String sval_glm_model_lite = "chatglm_std";
    private String sval_glm_model_lite32 = "chatglm_lite_32k";
    private String intentModel = sval_glm_model_pro;


    private final Queue<QIData> ret_msgs_ = new LinkedList<>();
    private boolean quit_ = false;
    private boolean active_ = false;
    private static final float fval_temperature = 0.9f;
    private static final float fval_top_p = 0.7f;

    public JedgeGlmBridgeService(JMgbusModual holder) {
        super(holder);
        auto_remove_call_info = false;
        watchLocalEvent("OnModuleConnected", (w, keyVal, msg) -> {
            String host = msg.getString("host");
            int port = msg.getInteger("port");
            return true;
        });
        watchLocalEvent("OnModuleDisconnected", (w, keyVal, msg) -> {
            String host = msg.getString("host");
            int port = msg.getInteger("port");
            return true;
        });
    }

    @Override
    public void onServiceStart() {
        super.onServiceStart();
        quit_ = false;

        zhipuConfig = new ConfigV3();
        zhipuConfig.setApiKey("297ba886b74433514543eac6ff485ddf"); // QIData_DefaultConfig.getString("ApiKey");
        zhipuConfig.setApiSecret("yKwpPFiUmreQd5Vq");              // QIData_DefaultConfig.getString("ApiSecret");
        zhipuConfig.setCache(LocalCache.getInstance());
        zhipuConfig.setDisableTokenCache(false);

        tokenCache = LocalCache.getInstance();
        tokenManagerV3 = new TokenManagerV3(tokenCache);
        GlobalTokenManager.setTokenManager(tokenManagerV3);

        okHttpClient = new OkHttpClient.Builder()
                .connectTimeout(3, TimeUnit.SECONDS)
                .readTimeout(3, TimeUnit.MINUTES)
                .build();

        module.postThread(this::PostMgbusMessage);
    }

    private void PostMgbusMessage() {
        active_ = true;
        module.warnLog("Posting thread start ok.");
        while(!quit_) {
            try {
                QIData msg = null;
                {
                    if(!ret_msgs_.isEmpty()) {
                        msg = ret_msgs_.poll();
                    }
                }
                if(msg!= null) {
                    String jeMod = msg.getString("~m");
                    msg.removeKey("~m");
                    postServiceMessage(jeMod, "/gpt/LocalStream", msg);
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
            SystemUtils.tryWaitforSingal(ret_msgs_, ival_mufis_text_message_gap_time_);
        }
        module.warnLog("Posting thread end.");
        active_ = false;
    }

    @Override
    public void onServiceStop() {
        quit_ = true;
        SystemUtils.notifyAllObject(ret_msgs_);
        //等待退出完成
        while(active_) {
            SystemUtils.tryWaitforSingal(ret_msgs_, 20);
        }
        super.onServiceStop();
    }

    @MgbusApi
    public QIData stream(QIData eMsg) {
        String user_msg = eMsg.getString("query");
        System.out.printf("[SSSSSSSS]Posting>>%s\n", user_msg);
        String content;
        if(StringUtils.isValidStr(user_msg))
            content = user_msg;
        else
            return JMgbusUtil.MgbusResult(401,"Empty query string");
        String sid = String.valueOf(eMsg.getInteger("sid", NumberUtils.randomInt(100000)));
        String srcMod = JMgbusUtil.GetMgModSource(eMsg);

        // url
        Map<String, String> pathParams = new HashMap<>();
        pathParams.put("model", intentModel);
        String url = ZhipuStringUtils.formatString(sseModelUrl, "{", "}", pathParams);

        // headers
        String token = GlobalTokenManager.getTokenManagerV3().getToken(zhipuConfig);
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", token);
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        headers.put("Accept", "text/event-stream");
        headers.put("sid", sid);
        headers.put("srcMod", srcMod);

        // body
        Map<String, Object> body = new HashMap<>();

        Prompt prompt = new Prompt(roleUser, content);
        List<Prompt> prompts = new ArrayList<>();
        prompts.add(prompt);
        body.put("prompt", prompts);
        body.put("request_id", sid);
        // incremental暂不支持配置，强制为true（增量返回）
        body.put("incremental", true);
        body.put("temperature", fval_temperature);
        body.put("top_p", fval_top_p);
        body.put("sseFormat",sseFormat);
        String reqBodyStr = new Gson().toJson(body);
        RequestBody formBody = RequestBody.create(MediaType.get("application/json; charset=utf-8"), reqBodyStr);

        Request request = new Request.Builder()
                .url(url)
                .post(formBody)
                .headers(Headers.of(headers))
                .build();

        RealEventSource eventSource = new RealEventSource(request, new StandardEventSourceListener(this));
        eventSource.connect(okHttpClient);
        return JMgbusUtil.MgbusResult(200,"ok");
    }

    @MgbusApi
    public QIData chat(QIData eMsg) {
        String user_msg = eMsg.getString("query");
        System.out.printf("[CCCCCCC]Posting>>%s\n", user_msg);
        String content;
        if(StringUtils.isValidStr(user_msg))
            content = user_msg;
        else
            return JMgbusUtil.MgbusResult(401,"Empty query string");

        String sid = String.valueOf(eMsg.getInteger("sid", NumberUtils.randomInt(100000)));


        // url
        Map<String, String> pathParams = new HashMap<>();
        pathParams.put("model", intentModel);
        String url = ZhipuStringUtils.formatString(reqModelUrl, "{", "}", pathParams);

        // headers
        String token = GlobalTokenManager.getTokenManagerV3().getToken(zhipuConfig);
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", token);
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        headers.put("Accept", "application/json; charset=utf-8");

        // body
        Map<String, Object> body = new HashMap<>();

        Prompt prompt = new Prompt(roleUser, content);
        List<Prompt> prompts = new ArrayList<>();
        prompts.add(prompt);
        body.put("prompt", prompts);
        body.put("request_id", sid);
        body.put("temperature", fval_temperature);
        body.put("top_p", fval_top_p);

        String reqBodyStr = new Gson().toJson(body);
        RequestBody formBody = RequestBody.create(MediaType.get("application/json; charset=utf-8"), reqBodyStr);

        Request request = new Request.Builder()
                .url(url)
                .post(formBody)
                .headers(Headers.of(headers))
                .build();

        Response okHttpResp = null;
        String resultDataStr = null;
        try {
            long time = new Date().getTime();
            okHttpResp = okHttpClient.newCall(request).execute();
            time = new Date().getTime() - time;
            module.infoLog(String.format("%d ms cost.", time));
            ResponseBody respBody = okHttpResp.body();
            resultDataStr = respBody.string();
        } catch (IOException e) {
            e.printStackTrace();
        }

        if(StringUtils.isNotValidStr(resultDataStr))
            return JMgbusUtil.MgbusResult(504, "Http request err");

        QIData qData = new QData();
        QIData resultMap = QData.fromString(resultDataStr);
        int code = resultMap.getInteger("code");
        String msg = resultMap.getString("msg");
        if (code == 200) {
            QIData data = resultMap.getData("data");
            String status = (String)data.get("task_status");
            if(StringUtils.isValidStr(status) && "success".equals(status.toLowerCase(Locale.ROOT))) {
                JSONArray contents = data.getJsonArray("choices");
                if(contents!=null && !contents.isEmpty()) {
                    Object item = contents.get(0);
                    if(item instanceof JSONObject) {
                        String cc = ((JSONObject) item).getString("content");
//                        module.printClientLog(String.format("GLM response:>>%s", cc));
                        if(StringUtils.isValidStr(cc))
                            cc = cc.substring(1,cc.length()-1);
                        module.printClientLog(String.format("GLM response:>>%s", cc));
                        qData.put("r", cc);
                        qData.put("code", 200);
                        return qData;
                    }
                }
            }

        }

        qData.put("code", 201);
        qData.put("msg", msg);
        return qData;
    }

    @MgbusApi
    public QIData embedding(QIData eMsg) throws IOException {
        String text = eMsg.getString("text");
        if(text.isEmpty())
            return JMgbusUtil.MgbusResult(201,"text is empty");
        if(text.length() > 512)
            return JMgbusUtil.MgbusResult(201,"text length should be less than 512");

        String url = embeddingURL;
        // headers
        String token = GlobalTokenManager.getTokenManagerV3().getToken(zhipuConfig);
        Map<String, String> headers = new HashMap<>();
        headers.put("Authorization", token);
        headers.put("Content-Type", "application/json; charset=utf-8");
        headers.put("User-Agent", "Mozilla/4.0 (compatible; MSIE 5.0; Windows NT; DigExt)");
        headers.put("Accept", "application/json; charset=utf-8");

        // body
        Map<String, Object> body = new HashMap<>();
        body.put("prompt", text);
        body.put("request_id", String.valueOf(System.currentTimeMillis()));
        String reqBodyStr = new Gson().toJson(body);
        RequestBody formBody = RequestBody.create(MediaType.get("application/json; charset=utf-8"), reqBodyStr);

        Request request = new Request.Builder()
                .url(url)
                .post(formBody)
                .headers(Headers.of(headers))
                .build();

        OkHttpClient okHttpClient = new OkHttpClient.Builder()
                .connectTimeout(3, TimeUnit.SECONDS)
                .readTimeout(3, TimeUnit.MINUTES)
                .build();

        Response okHttpResp = okHttpClient.newCall(request).execute();
        ResponseBody respBody = okHttpResp.body();
        String resultDataStr = respBody.string();

        QIData qData = new QData();
        Gson gson = new Gson();
        Map<?,?> resultMap = gson.fromJson(resultDataStr, Map.class);
        Double code = (Double) resultMap.get("code");
        String msg = (String) resultMap.get("msg");
        if (code.intValue() == 200) {
            Map<?,?> data = (Map<?,?>)resultMap.get("data");
            String status = (String)data.get("task_status");
            if(StringUtils.isValidStr(status) && "success".equals(status.toLowerCase(Locale.ROOT))) {
                ArrayList<Double> embd = (ArrayList<Double>)data.get("embedding");
                double[] doubleArray = embd.stream()
                        .mapToDouble(Double::doubleValue)
                        .toArray();
                RealVector vector = new ArrayRealVector(doubleArray);
                double norm = vector.getNorm();
                if (norm != 0.0) {
                    vector.mapDivideToSelf(norm);
                }
                qData.put("embedding", vector.toArray());
                qData.put("code", 200);
                return qData;
            }

        }

        qData.put("code", 201);
        qData.put("msg", msg);
        return qData;
    }

    public void putMgbusMessage(QIData msg) {
        synchronized (ret_msgs_) {
            ret_msgs_.add(msg);
        }
    }

    public static class Prompt {
        private String role;
        private String content;

        public Prompt(String role, String content) {
            this.role = role;
            this.content = content;
        }

        public String getRole() {
            return role;
        }

        public void setRole(String role) {
            this.role = role;
        }

        public String getContent() {
            return content;
        }

        public void setContent(String content) {
            this.content = content;
        }
    }

}
