package com.ruoyi.cms.controller.app;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.cms.config.ChatClient;
import com.ruoyi.cms.config.ChatConfig;
import com.ruoyi.cms.domain.ai.AiChatHistory;
import com.ruoyi.cms.domain.chat.ChatRequest;
import com.ruoyi.cms.service.AiChatHistoryService;
import com.ruoyi.common.utils.StringUtils;
import com.ruoyi.common.utils.spring.SpringUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * WebSocket AI 聊天处理器
 */
@Slf4j
@Component
@ServerEndpoint("/ws/chat/{userId}")
public class ChatWebSocketHandler {

    // 存储所有连接的会话
    private static final Map<String, Session> SESSIONS = new ConcurrentHashMap<>();

    // 由于 WebSocket 是多例的，需要通过 SpringUtils 获取 Bean
    private ChatClient getChatClient() {
        return SpringUtils.getBean(ChatClient.class);
    }

    private ChatConfig getChatConfig() {
        return SpringUtils.getBean(ChatConfig.class);
    }

    private AiChatHistoryService getAiChatHistoryService() {
        return SpringUtils.getBean(AiChatHistoryService.class);
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        SESSIONS.put(userId + "_" + session.getId(), session);
        log.info("WebSocket 连接建立，userId: {}, sessionId: {}", userId, session.getId());
        sendMessage(session, buildResponse("connected", "连接成功", null));
    }

    @OnClose
    public void onClose(Session session, @PathParam("userId") String userId) {
        SESSIONS.remove(userId + "_" + session.getId());
        log.info("WebSocket 连接关闭，userId: {}, sessionId: {}", userId, session.getId());
    }

    @OnError
    public void onError(Session session, Throwable error, @PathParam("userId") String userId) {
        log.error("WebSocket 发生错误，userId: {}, error: {}", userId, error.getMessage());
        sendMessage(session, buildResponse("error", error.getMessage(), null));
    }


    /**
     * 接收客户端消息
     * 消息格式：
     * {
     *   "action": "chat",  // chat: 聊天, history: 获取历史, detail: 获取详情, guest: 获取建议
     *   "data": "用户输入的内容",
     *   "sessionId": "会话ID",
     *   "dataId": "数据ID",
     *   "fileUrl": ["文件URL列表"]
     * }
     */
    @OnMessage
    public void onMessage(String message, Session session, @PathParam("userId") String userId) {
        log.info("收到消息，userId: {}, message: {}", userId, message);
        try {
            JSONObject json = JSONObject.parseObject(message);
            String action = json.getString("action");

            if (StringUtils.isEmpty(action)) {
                action = "chat";
            }

            switch (action) {
                case "chat":
                    handleChat(json, session, userId);
                    break;
                case "history":
                    handleHistory(json, session);
                    break;
                case "detail":
                    handleDetail(json, session);
                    break;
                case "guest":
                    handleGuest(json, session);
                    break;
                default:
                    sendMessage(session, buildResponse("error", "未知的操作类型: " + action, null));
            }
        } catch (Exception e) {
            log.error("处理消息失败: {}", e.getMessage(), e);
            sendMessage(session, buildResponse("error", "处理消息失败: " + e.getMessage(), null));
        }
    }

    /**
     * 处理聊天请求
     */
    private void handleChat(JSONObject json, Session session, String userId) {
        ChatRequest request = buildChatRequest(json, userId);
        
        JSONObject contentObject = new JSONObject();
        contentObject.put("content", request.getData());
        contentObject.put("role", "user");
        
        JSONArray array = getAiChatHistoryService().getChatHistoryData(request.getSessionId());
        if (array == null || array.isEmpty()) {
            array = new JSONArray();
        }
        array.add(contentObject);
        request.setMessages(array);

        List<String> answerList = new ArrayList<>();
        long[] timeStart = {0};

        getChatClient().sendStreamingChat(request, new ChatClient.StreamCallback() {
            @Override
            public void onData(String chunk) {
                if (timeStart[0] == 0) {
                    timeStart[0] = System.currentTimeMillis();
                }
                String processedChunk = chunk.trim();
                String content = parseChatChunk(processedChunk);
                if (StringUtils.isNotEmpty(content)) {
                    answerList.add(content);
                }
                // 发送流式数据
                sendMessage(session, buildResponse("data", null, processedChunk));
            }

            @Override
            public void onComplete() {
                // 保存聊天记录
                saveChatHistory(request, userId, answerList, timeStart[0]);
                sendMessage(session, buildResponse("complete", "对话完成", null));
            }

            @Override
            public void onError(Throwable e) {
                log.error("聊天请求失败: {}", e.getMessage());
                sendMessage(session, buildResponse("error", e.getMessage(), null));
            }
        });
    }


    /**
     * 处理获取历史记录请求
     */
    private void handleHistory(JSONObject json, Session session) {
        AiChatHistory chatHistory = new AiChatHistory();
        if (json.containsKey("userId")) {
            chatHistory.setUserId(json.getLong("userId"));
        }
        JSONObject result = getAiChatHistoryService().getList(chatHistory);
        sendMessage(session, buildResponse("history", null, result.toJSONString()));
    }

    /**
     * 处理获取聊天详情请求
     */
    private void handleDetail(JSONObject json, Session session) {
        String sessionId = json.getString("sessionId");
        JSONObject result = getAiChatHistoryService().getDetailList(sessionId);
        sendMessage(session, buildResponse("detail", null, result.toJSONString()));
    }

    /**
     * 处理获取建议请求
     */
    private void handleGuest(JSONObject json, Session session) {
        try {
            ChatRequest request = new ChatRequest();
            request.setSessionId(json.getString("sessionId"));
            JSONArray array = getAiChatHistoryService().getChatHistoryData(request.getSessionId());
            request.setMessages(array);
            
            String result = getChatClient().sendChatGuest(request);
            String[] strList = result.split("？");
            List<String> list = new ArrayList<>();
            for (String str : strList) {
                if (StringUtils.isNotEmpty(str)) {
                    str = str + "？";
                    list.add(str);
                }
            }
            sendMessage(session, buildResponse("guest", null, JSONObject.toJSONString(list)));
        } catch (Exception e) {
            log.error("获取建议失败: {}", e.getMessage());
            sendMessage(session, buildResponse("error", "获取建议失败: " + e.getMessage(), null));
        }
    }

    /**
     * 构建 ChatRequest 对象
     */
    private ChatRequest buildChatRequest(JSONObject json, String userId) {
        ChatRequest request = new ChatRequest();
        request.setData(json.getString("data"));
        request.setSessionId(json.getString("sessionId"));
        request.setDataId(json.getString("dataId"));
        
        if (json.containsKey("fileUrl")) {
            request.setFileUrl(json.getJSONArray("fileUrl").toJavaList(String.class));
        }
        
        if (json.containsKey("userId") && json.getLong("userId") != 0) {
            request.setUserId(json.getLong("userId"));
        } else if (StringUtils.isNotEmpty(userId) && !"0".equals(userId)) {
            request.setUserId(Long.parseLong(userId));
        }
        
        return request;
    }

    /**
     * 保存聊天记录
     */
    private void saveChatHistory(ChatRequest request, String userId, List<String> answerList, long timeStart) {
        long timeEnd = System.currentTimeMillis();
        double duration = (timeEnd - timeStart) / 1000.0;
        
        AiChatHistory chatHistory = new AiChatHistory();
        chatHistory.setChatId(request.getSessionId());
        
        Long uid = null;
        if (request.getUserId() != 0) {
            uid = request.getUserId();
        } else if (StringUtils.isNotEmpty(userId) && !"0".equals(userId)) {
            uid = Long.parseLong(userId);
        }
        chatHistory.setUserId(uid);
        chatHistory.setAppId(getChatConfig().getAppId());
        chatHistory.setDataId(request.getDataId());
        chatHistory.setTitle(request.getData());
        chatHistory.setAnswerStringList(answerList);
        chatHistory.setDurationSeconds(duration);
        
        getAiChatHistoryService().saveChatHistory(chatHistory);
    }

    /**
     * 解析流式数据
     */
    private String parseChatChunk(String chunk) {
        try {
            String processed = chunk.trim();
            if (processed.startsWith("data:")) {
                processed = processed.substring("data:".length()).trim();
            }
            if (processed.isEmpty() || "[DONE]".equals(processed)) {
                return null;
            }
            JSONObject json = JSONObject.parseObject(processed);
            String choices = json.getString("choices");
            if (choices != null && !choices.trim().isEmpty()) {
                JSONArray jsonArray = JSONArray.parseArray(choices);
                json = JSONObject.parseObject(jsonArray.getString(0));
                json = json.getJSONObject("delta");
                return json.getString("content");
            }
            return null;
        } catch (Exception e) {
            return null;
        }
    }

    /**
     * 构建响应消息
     */
    private String buildResponse(String type, String message, String data) {
        JSONObject response = new JSONObject();
        response.put("type", type);
        if (message != null) {
            response.put("message", message);
        }
        if (data != null) {
            response.put("data", data);
        }
        return response.toJSONString();
    }

    /**
     * 发送消息
     */
    private void sendMessage(Session session, String message) {
        if (session != null && session.isOpen()) {
            try {
                synchronized (session) {
                    session.getBasicRemote().sendText(message);
                }
            } catch (IOException e) {
                log.error("发送消息失败: {}", e.getMessage());
            }
        }
    }

    /**
     * 向指定用户发送消息
     */
    public static void sendMessageToUser(String userId, String message) {
        SESSIONS.forEach((key, session) -> {
            if (key.startsWith(userId + "_") && session.isOpen()) {
                try {
                    synchronized (session) {
                        session.getBasicRemote().sendText(message);
                    }
                } catch (IOException e) {
                    log.error("发送消息给用户 {} 失败: {}", userId, e.getMessage());
                }
            }
        });
    }
}
