package com.ch.jedge.jbot.session;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.ch.jedge.jbot.JedgeBotGptService;
import com.ch.jedge.jbot.agent.JedgeGdeSession;
import com.ch.jedge.jbot.knbase.JedgeGptApp;
import com.changhong.qlib.util.StringUtils;

import java.util.*;

import static com.ch.jedge.jbot.JEJBotLLMConst.*;

public class JedgeGdeKeywordPool {

    private static class knItem {
        public String content;
        int ttl= 600;
        public knItem(String content) {
            this.content = content;
        }

        public void resetTTL() {
            ttl = 600;
        }
    }

    private static class keyItem {
        public final String key;
        public int ttl = 180;
        public int qaRange = 0; //对话距离，使得对话超过一定距离的词不再纳入知识计算。

        public keyItem(String mk) {
            this.key = mk;
        }
    }

    private final JedgeGdeSession gdeSession;
    private final Map<String, knItem> mKnItemMap = new HashMap<>();

    //记录最近和次进的主key
    private final Map<String, keyItem> mMainKeymap = new HashMap<>();
    private final Map<String, keyItem> mRelKeymap = new HashMap<>();
    private final Map<String, keyItem> mJoinKeymap = new HashMap<>();

    private keyItem lastMainKey,secondMainKey,lastJoinKey;

    public void refreshQARange() {
        synchronized (mKnItemMap) {
            for(keyItem item : mMainKeymap.values()) {
                item.qaRange ++;
            }
            for(keyItem item : mRelKeymap.values()) {
                item.qaRange ++;
            }
            for(keyItem item : mJoinKeymap.values()) {
                item.qaRange ++;
            }
        }
    }


    public JedgeGdeKeywordPool(JedgeGdeSession jedgeGdeSession) {
        gdeSession = jedgeGdeSession;
    }

    public void pushKeyList(final Map<String, keyItem> keyMap, List<String> knKeyList) {
        synchronized (keyMap) {
            for (String mk : knKeyList) {
                keyItem theItem = keyMap.get(mk);
                if (theItem==null) {
                    theItem = new keyItem(mk);
                    keyMap.put(mk, theItem);
                } else {
                    theItem.qaRange = 0;
                }
                //滚动MainKey
                if(keyMap==mMainKeymap) {
                    secondMainKey = lastMainKey;
                    lastMainKey = theItem;
                } else if(keyMap == mJoinKeymap) {
                    lastJoinKey = theItem;
                }
            }
        }
    }

    //获取一个组合Key
    //是否记录ttl和Key的距离信息，只针对特定距离的key进行取值
    //   拼装的知识引用的过程。
    //

    /**
     *
     //  1、得到一个主Key的qa距离；
     //  2、检查主key以内有无JoinKey： 如果有，抽取第二个最近的主Key
     //  3、抽取主Key的qa距离以内的所有的relKey
     //  4、如何组合两个主key及其关系？
     *
     *   如果当前对话与主key和当前的缓存数据，没有任何关系，也无法检索出任何内容，则应当提示帮助
     */

    public List<String> getAvailableKnList() {
        List<String> rr = new ArrayList<>();

        int maxRange;
        if(lastMainKey!=null) {
            rr.add(lastMainKey.key);
            maxRange = lastMainKey.qaRange;
        } else {
            maxRange = 99;      //很远距离没有相关的距离信息
        }

        if(lastJoinKey!=null && lastJoinKey.qaRange < maxRange) {
            rr.add(secondMainKey.key);
            maxRange = secondMainKey.qaRange;
        }

        //应该是多重subKey的组合
        List<String> relKeySrc = new ArrayList<>();

        //只提取qaRange小于maxRange的数据
        //如果当前对话，既没有mainKey, 又没有relKey，则表现为无关联的对话
        boolean hasRelKey = false;
        for (keyItem relItem : mRelKeymap.values()) {
            if(relItem.qaRange > maxRange)
                continue;
            if(relItem.qaRange<1)
                hasRelKey = true;
            relKeySrc.add(relItem.key);
        }

        String relKey = hasRelKey?
                StringUtils.ListToString(mRelKeymap.keySet(),'$'):null;

        //组合
        if(StringUtils.isValidStr(relKey)) {
            for (Map.Entry<String, keyItem> baseKey : mMainKeymap.entrySet()) {
                if(baseKey.getValue().qaRange > maxRange)
                    continue;
                rr.add(baseKey.getKey() + '$' + relKey);
            }
        }

        if(lastJoinKey!=null && lastJoinKey.qaRange < maxRange) {
            rr.add(lastJoinKey.key + '$' + lastMainKey.key + '$' + relKey);
            rr.add(lastJoinKey.key + '$' + secondMainKey.key + '$' + relKey);
            rr.add(lastMainKey.key + '$' + relKey);
            rr.add(secondMainKey.key + '$' + relKey);
        } else if(StringUtils.isValidStr(relKey)) {
            if(lastMainKey!=null)
                rr.add(lastMainKey.key + '$' + relKey);
            // else 提供默认的MainKey
        } else if(maxRange!=0) {
            rr.clear();
        }
        //找到mainKey的帮助信息

        //只提取qa距离为0的 joinKey
        return rr;
    }

    public String getComposedPrompt(JedgeServiceSession ss, String prompt, boolean helping) {
        JedgeBotGptService s = ss.getGptBotService();
        //分词，得到词性的组合
        //将这些分词的属性记录到会话缓存中。
        JSONArray results = s.cutWords(prompt);

        //分词后，进行联想，扩充这个词库
        List<String> knKeyList = getTypeKeyItems(results, val_verb_type_user_defined_noun);
        List<String> relKeyList = getTypeKeyItems(results, sval_verb_type_user_defined_relation_noun);
        List<String> joinKeyList = getTypeKeyItems(results, sval_verb_type_user_defined_join_noun);

        //分解的词汇汇集到缓存池
        List<String> realKnList = null;
        synchronized (mKnItemMap) {
            pushKeyList(mMainKeymap, knKeyList);
            pushKeyList(mRelKeymap, relKeyList);
            pushKeyList(mJoinKeymap, joinKeyList);
            //获取所有可能的key组合
            realKnList = getAvailableKnList();
        }

        //如果得到多个关键字， 如何组合去查询数据？
        if(realKnList!=null && !realKnList.isEmpty()) {
            StringBuilder result = new StringBuilder();
            for (String knKey : realKnList) {
                //先检索本地缓存中是否已有这个key的知识，如果有，则从缓存中读取，否则从知识库中查询
                String content = getKnowledgeFromMilvus(ss, knKey, helping);
                if(JedgeGptApp.isDebugMode())
                    gdeSession.getGdeApp().getMgModule().exMarkLog(String.format("查询知识:%s => %s", knKey, content));
                if(StringUtils.isValidStr(content))
                    result.append(content);
            }
            if(result.length()>0)
                return String.format(s.getKnAskPrompt(), "咨询答案", result, "TestGPT智能助手", prompt);
        }
        return null;
    }

    public void onKnItemHeartbeat() {
        refreshTTL(mMainKeymap);
        refreshTTL(mRelKeymap);
        refreshTTL(mJoinKeymap);
        refreshKnTTL(mKnItemMap);
    }

    private void refreshKnTTL(final Map<String, knItem> keyMap) {
        List<String> toDel = new ArrayList<>();
        synchronized (keyMap) {
            Set<String> keys = keyMap.keySet();
            for(String k : keys) {
                knItem kt = keyMap.get(k);
                if(kt.ttl<0) continue;
                kt.ttl--;
                if(kt.ttl < 1) {
                    if(JedgeGptApp.isDebugMode()) {
                        gdeSession.getGdeApp().getMgModule().highLog(String.format("Key word removed : %s", k));
                    }
                    toDel.add(k);
                }
            }
            if(!toDel.isEmpty()) {
                for (String del : toDel)
                    keyMap.remove(del);
                if(JedgeGptApp.isDebugMode()) {
                    gdeSession.getGdeApp().getMgModule().highLog(String.format("KnItem Left : %s", StringUtils.ListToString(keyMap.keySet(), ',')));
                }
            }
        }
    }

    private void refreshTTL(final Map<String, keyItem> keyMap) {
        List<String> toDel = new ArrayList<>();
        synchronized (keyMap) {
            Set<String> keys = keyMap.keySet();
            for(String k : keys) {
                keyItem kt = keyMap.get(k);
                if(kt.ttl<0) continue;
                kt.ttl--;
                if(kt.ttl < 1) {
                    if(JedgeGptApp.isDebugMode()) {
                        gdeSession.getGdeApp().getMgModule().highLog(String.format("Key word removed : %s", k));
                    }
                    toDel.add(k);
                }
            }
            if(!toDel.isEmpty()) {
                for (String del : toDel)
                    keyMap.remove(del);
                if(JedgeGptApp.isDebugMode()) {
                    gdeSession.getGdeApp().getMgModule().highLog(String.format("Key Left : %s", StringUtils.ListToString(keyMap.keySet(), ',')));
                }
            }
        }
    }


    public String getKnowledgeFromMilvus(JedgeServiceSession ss, String knKey, boolean helping) {
        knItem knContent;
        synchronized (mKnItemMap) {
            knContent = mKnItemMap.get(knKey);
        }
        if(knContent!=null) {
            knContent.resetTTL();
            return knContent.content;
        }
        JedgeBotGptService s = ss.getGptBotService();
        String content = s.getKnowledgeFromMilvus(ss, knKey, helping);
        if(StringUtils.isValidStr(content)) {
            synchronized (mKnItemMap) {
                mKnItemMap.put(knKey, new knItem(content));
            }
        }
        return content;
    }

    //生成已对接的key
    //  主key：
    //  分类key：
    // 主key或，分类keyAnd
    public List<String> getTypeKeyItems(JSONArray results, String keyType) {
        List<String> re = new ArrayList<>();
        if(results.isEmpty())
            return re;
        for(Object obj : results) {
            if(obj instanceof JSONObject && !((JSONObject) obj).isEmpty()) {
                String t = ((JSONObject) obj).getString("t");
                if(keyType.equals(t) && !re.contains(keyType))
                    re.add(((JSONObject) obj).getString("w"));
            }
        }

        return re;
    }

}
