package org.dromara.chat.controller;

import cn.dev33.satoken.annotation.SaCheckLogin;
import cn.hutool.core.util.StrUtil;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.dromara.ai.domain.dto.AiChatRequest;
import org.dromara.ai.domain.dto.AiChatResponse;
import org.dromara.ai.service.AiChatService;
import org.dromara.chat.domain.bo.ChatMessageBo;
import org.dromara.chat.domain.vo.ChatSessionVo;
import org.dromara.chat.service.ChatMessageService;
import org.dromara.chat.service.ChatSessionService;
import org.dromara.common.core.utils.StringUtils;
import org.dromara.common.json.utils.JsonUtils;
import org.dromara.common.satoken.utils.LoginHelper;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.math.BigDecimal;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

/**
 * 开放聊天接口
 *
 * @author pengles
 */
@Slf4j
@SaCheckLogin
@RestController
@RequiredArgsConstructor
@RequestMapping("/chat")
public class ChatOpenApiController {

    private final AiChatService aiChatService;
    private final ChatSessionService chatSessionService;
    private final ChatMessageService chatMessageService;

    @PostMapping(value = "/send", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter send(@RequestBody Map<String, Object> body) {
        SseEmitter emitter = new SseEmitter(TimeUnit.MINUTES.toMillis(10));
        Long userId;
        try {
            userId = LoginHelper.getUserId();
        } catch (Exception e) {
            emitter.completeWithError(new RuntimeException("未登录"));
            return emitter;
        }

        try {
            AiChatRequest req = toRequest(body);
            req.setUserId(String.valueOf(userId));

            // 获取或创建会话
            String firstMsgContent = "";
            if (req.getMessages() != null && !req.getMessages().isEmpty()) {
                firstMsgContent = req.getMessages().get(req.getMessages().size() - 1).getContent();
            }
            Long sessionId = getOrCreateSession(userId, req.getSessionId(), firstMsgContent);
            req.setSessionId(String.valueOf(sessionId));

            // 如果是新会话，尝试通知前端sessionId
            if (StrUtil.isBlank(req.getSessionId())) {
                try {
                    emitter.send(SseEmitter.event().name("session").data("{\"sessionId\":\"" + sessionId + "\"}"));
                } catch (IOException e) {
                    log.error("发送会话ID失败", e);
                }
            }

            // 保存用户消息
            saveUserMessage(sessionId, userId, firstMsgContent, req.getModel());

            final StringBuilder fullText = new StringBuilder();
            final Long finalSessionId = sessionId;
            final String modelName = req.getModel();

            aiChatService.chatStream(req, new AiChatService.StreamCallback() {
                @Override
                public void onChunk(String chunk) {
                    String c = normalizeChunk(chunk);
                    if ("[DONE]".equals(c)) {
                        // ignore
                    } else {
                        fullText.append(c);
                        try {
                            emitter.send(toOpenAiDelta(c, false));
                        } catch (IOException e) {
                            log.debug("发送SSE失败", e);
                        }
                    }
                }

                @Override
                public void onComplete(AiChatResponse response) {
                    try {
                        emitter.send(toOpenAiDelta("", true));
                        emitter.send("[DONE]");
                        emitter.complete();

                        // 保存AI响应消息
                        saveAiMessage(finalSessionId, userId, fullText.toString(), modelName);
                    } catch (IOException e) {
                        log.error("发送SSE完成标记失败", e);
                        emitter.completeWithError(e);
                    }
                }

                @Override
                public void onError(Throwable error) {
                    log.error("AI聊天出错", error);
                    try {
                        emitter.send("{\"error\":\"" + error.getMessage().replace("\"", "\\\"") + "\"}");
                    } catch (IOException ignored) {
                    }
                    emitter.completeWithError(error);
                }
            });

        } catch (Exception e) {
            log.error("处理聊天请求失败", e);
            try {
                emitter.send("{\"error\":\"" + e.getMessage().replace("\"", "\\\"") + "\"}");
            } catch (IOException ignored) {
            }
            emitter.completeWithError(e);
        }
        return emitter;
    }

    private Long getOrCreateSession(Long userId, String sessionIdStr, String firstMsgContent) {
        if (StrUtil.isNotBlank(sessionIdStr)) {
            Long sessionId = Long.parseLong(sessionIdStr);
            ChatSessionVo session = chatSessionService.queryById(sessionId);
            if (session == null) {
                throw new RuntimeException("会话不存在");
            }
            if (!session.getUserId().equals(userId)) {
                throw new RuntimeException("无权访问该会话");
            }
            return sessionId;
        } else {
            String title = "新对话";
            if (StrUtil.isNotBlank(firstMsgContent)) {
                title = StringUtils.substring(firstMsgContent, 0, 20);
            }
            ChatSessionVo session = chatSessionService.createNewSession(userId, title);
            return session.getId();
        }
    }

    private void saveUserMessage(Long sessionId, Long userId, String content, String modelName) {
        if (StrUtil.isNotBlank(content)) {
            ChatMessageBo userMsg = new ChatMessageBo();
            userMsg.setSessionId(sessionId);
            userMsg.setUserId(userId);
            userMsg.setRole("user");
            userMsg.setContent(content);
            userMsg.setModelName(modelName);
            userMsg.setCreateTime(new Date());
            chatMessageService.insertByBo(userMsg);
        }
    }

    private void saveAiMessage(Long sessionId, Long userId, String content, String modelName) {
        if (StrUtil.isNotBlank(content)) {
            ChatMessageBo aiMsg = new ChatMessageBo();
            aiMsg.setSessionId(sessionId);
            aiMsg.setUserId(userId);
            aiMsg.setRole("assistant");
            aiMsg.setContent(content);
            aiMsg.setModelName(modelName);
            aiMsg.setCreateTime(new Date());
            aiMsg.setTotalTokens(0L);
            aiMsg.setDeductCost(BigDecimal.ZERO);
            chatMessageService.insertByBo(aiMsg);
        }
    }

    private AiChatRequest toRequest(Map<String, Object> body) {
        AiChatRequest req = new AiChatRequest();
        req.setModel(asString(body.get("model")));
        req.setTemperature(asDouble(body.get("temperature")));
        req.setTopP(asDouble(body.get("top_p")));
        req.setPresencePenalty(asDouble(body.get("presence_penalty")));
        req.setFrequencyPenalty(asDouble(body.get("frequency_penalty")));
        req.setMaxTokens(asInt(body.get("max_tokens")));
        req.setStream(Boolean.TRUE);
        req.setSessionId(asString(body.get("sessionId")));
        Object msgs = body.get("messages");
        if (msgs instanceof List<?> list) {
            java.util.ArrayList<AiChatRequest.AiMessage> arr = new java.util.ArrayList<>();
            for (Object o : list) {
                if (o instanceof Map<?, ?> m) {
                    AiChatRequest.AiMessage am = new AiChatRequest.AiMessage();
                    am.setRole(asString(m.get("role")));
                    am.setContent(asString(m.get("content")));
                    arr.add(am);
                }
            }
            req.setMessages(arr);
        }
        return req;
    }

    private String asString(Object o) {
        return o == null ? null : String.valueOf(o);
    }

    private Double asDouble(Object o) {
        try {
            return o == null ? null : Double.valueOf(String.valueOf(o));
        } catch (Exception e) {
            return null;
        }
    }

    private Integer asInt(Object o) {
        try {
            return o == null ? null : Integer.valueOf(String.valueOf(o));
        } catch (Exception e) {
            return null;
        }
    }

    private String toOpenAiDelta(String content, boolean finish) {
        String c = content == null ? "" : content;
        String esc = c.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t");
        String finishReason = finish ? "\"stop\"" : "null";
        return "{\"choices\":[{\"delta\":{\"content\":\"" + esc + "\"},\"finish_reason\":" + finishReason + "}]}";
    }

    private String normalizeChunk(String chunk) {
        if (chunk == null) return "";
        String s = chunk.trim();
        if ("[DONE]".equals(s)) return "[DONE]";
        try {
            Map<?, ?> m = JsonUtils.parseObject(s, Map.class);
            if (m != null) {
                Object choices = m.get("choices");
                if (choices instanceof List<?> list && !list.isEmpty()) {
                    Object first = list.get(0);
                    if (first instanceof Map<?, ?> fm) {
                        Object delta = fm.get("delta");
                        if (delta instanceof Map<?, ?> dm) {
                            Object c = dm.get("content");
                            if (c == null) c = dm.get("reasoning_content");
                            if (c != null) return String.valueOf(c);
                        }
                        Object text = fm.get("text");
                        if (text != null) return String.valueOf(text);
                    }
                }
            }
        } catch (Exception ignore) {
        }
        return s;
    }
}
