package org.jeecg.modules.adweb.gpt.controller;

import lombok.extern.slf4j.Slf4j;

import org.apache.commons.lang3.tuple.Pair;
import org.apache.shiro.SecurityUtils;
import org.jeecg.common.api.vo.Result;
import org.jeecg.common.system.vo.LoginUser;
import org.jeecg.modules.adweb.gpt.service.ChatService;
import org.jeecg.modules.adweb.gpt.vo.ChatHistoryVO;
import org.jeecg.modules.adweb.quota.service.IResourceQuotaService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

/**
 * @Description: ChatGPT-聊天接口
 * @Author: chenrui
 * @Date: 2024/1/9 16:30
 */
@Controller
@RequestMapping("/ai/chat")
@Slf4j
public class ChatController {

    @Autowired private ChatService chatService;

    @Autowired private IResourceQuotaService resourceQuotaService;

    /** 创建sse连接 */
    @GetMapping(value = "/send")
    public SseEmitter createConnect(
            @RequestParam(name = "topicId", required = false) String topicId,
            @RequestParam(name = "message", required = true) String message) {
        SseEmitter sse = chatService.createChat();
        chatService.sendMessage(topicId, message);
        return sse;
    }

    /**
     * 保存聊天记录
     *
     * @param chatHistoryVO
     * @return
     * @author chenrui
     * @date 2024/2/22 13:54
     */
    @PostMapping(value = "/history/save")
    @ResponseBody
    public Result<?> saveHistory(@RequestBody ChatHistoryVO chatHistoryVO) {
        chatService.saveHistory(chatHistoryVO);
        return Result.OK("保存成功");
    }

    /**
     * 查询聊天记录
     *
     * @return
     * @author chenrui
     * @date 2024/2/22 14:03
     */
    @GetMapping(value = "/history/get")
    @ResponseBody
    public Result<ChatHistoryVO> getHistoryByTopic() {
        return Result.OK(chatService.getHistoryByTopic());
    }

    /** 关闭连接 */
    @GetMapping(value = "/close")
    public void closeConnect() {
        chatService.closeChat();
    }

    /** 检查AI算力额度 */
    @GetMapping(value = "checkQuota")
    @ResponseBody
    public Result<String> checkResourceQuota() {
        String uid = ((LoginUser) SecurityUtils.getSubject().getPrincipal()).getId();

        Pair<Integer, Integer> aiPowerQuota = resourceQuotaService.getAIPowerQuota(uid);
        if (aiPowerQuota.getLeft() <= aiPowerQuota.getRight()) {
            log.warn("用户 {} AI算力额度不足", uid);
            return Result.error(
                    String.format(
                            "AI算力不足。用户资源额度为%d，已使用%d",
                            aiPowerQuota.getLeft(), aiPowerQuota.getRight()),
                    null);
        }

        return Result.OK("OK");
    }
}
