package org.lp.medicalai.controller;

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSONObject;
import org.lp.medicalai.component.MemoryUserRecordSpace;
import org.lp.medicalai.component.XfXhStreamClient;
import org.lp.medicalai.config.XfXhConfig;
import org.lp.medicalai.dto.InteractMsg;
import org.lp.medicalai.dto.MsgDTO;
import org.lp.medicalai.dto.RecordsArray;
import org.lp.medicalai.listener.XfXhWebSocketListener;
import okhttp3.WebSocket;
import org.lp.medicalai.pojo.Result;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;

@RestController
@RequestMapping("/medical-ai")
public class AiController {
    @Autowired
    private XfXhConfig xfXhConfig;

    @Autowired
    private XfXhStreamClient xfXhStreamClient;

    @Autowired
    private MemoryUserRecordSpace memoryUserRecordSpace;

    // 使用 id 作为唯一用户的标识（区分不同用户）
    @GetMapping("/sendQuestion")
    public Result question(@RequestAttribute("id") Integer id,String question) throws InterruptedException {

        if (StrUtil.isBlank(question)) {
            return Result.error("NO_QUESTION");//无效问题，请重新输入
        }

        // 尝试锁定用户
        if (!memoryUserRecordSpace.tryLock(id)) {
            return Result.error("WAIT");//正在处理上次问题，请稍后再试
        }

        // 获取连接令牌
        if(!xfXhStreamClient.operateToken(XfXhStreamClient.GET_TOKEN_STATUS)){
            // 释放锁
            memoryUserRecordSpace.unLock(id);
            return Result.error("BUSY");//当前大模型连接数过多，请稍后再试
        }

        MsgDTO msgDTO = MsgDTO.createUserMsg(question);
        String tip = "你是一位专业且耐心的医生，专注于提供医疗建议。非医疗问题请礼貌拒绝。";

        MsgDTO tipDTO = MsgDTO.createSystemMsg(tip);

        XfXhWebSocketListener listener = new XfXhWebSocketListener();

        // 组装上下文内容发送
        List<MsgDTO> msgList = memoryUserRecordSpace.getAllInteractMsg(id);

        msgList.add(tipDTO);
        msgList.add(msgDTO);

        WebSocket webSocket = xfXhStreamClient.sendMsg(UUID.randomUUID().toString().substring(0, 10), msgList, listener);
        if (webSocket == null) {
            // 归还令牌
            xfXhStreamClient.operateToken(XfXhStreamClient.BACK_TOKEN_STATUS);
            // 释放锁
            memoryUserRecordSpace.unLock(id);
            return Result.error("ERROR");//系统内部错误，请联系管理员
        }
        try {
            int count = 0;
            // 为了避免死循环，设置循环次数来定义超时时长
            int maxCount = xfXhConfig.getMaxResponseTime() * 5;
            while (count <= maxCount) {
                Thread.sleep(200);
                if (listener.isWsCloseFlag()) {
                    break;
                }
                count++;
            }
            if (count > maxCount) {
                return Result.error("ERROR");//大模型响应超时，请联系管理员
            }
            // 将记录添加到 memoryUserRecordSpace
            String answer = listener.getAnswer().toString();
            memoryUserRecordSpace.storeInteractMsg(id, new InteractMsg(MsgDTO.createUserMsg(question), MsgDTO.createAssistantMsg(answer)));
            return Result.success(answer);//相应成功
        } finally {
            // 关闭连接
            webSocket.close(1000, "");
            // 释放锁
            memoryUserRecordSpace.unLock(id);
            // 归还令牌
            xfXhStreamClient.operateToken(XfXhStreamClient.BACK_TOKEN_STATUS);
        }
    }

    // 测试使用，查看内存空间中所有的用户记录信息
    @GetMapping("/spaceInfo")
    public List<JSONObject> spaceInfo(@RequestAttribute("id") Integer id){
        if(id == 1){
            ConcurrentHashMap<Long, RecordsArray> userRecordMap = memoryUserRecordSpace.getUserRecordMap();
            ArrayList<JSONObject> infoList = new ArrayList<>(userRecordMap.size());
            for (Map.Entry<Long, RecordsArray> entry : userRecordMap.entrySet()) {
                RecordsArray recordsArray = entry.getValue();
                JSONObject data = new JSONObject();
                data.put("id",entry.getKey());
                data.put("allInteractMsg",recordsArray.getAllInteractMsg());
                data.put("status",recordsArray.getStatus());
                data.put("lock",recordsArray.isLock());
                infoList.add(data);
            }
            return infoList;
        }
        return null;
    }
}