package com.cmvr.llm.util;

import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.*;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * SSE 流式对话客户端（纯 OkHttp 实现，无额外依赖）
 * 包含会话管理、消息管理、流式/阻塞对话
 */
public class SseChatClient {

    private static final String DEFAULT_BASE_URL = "https://aiagentplatform.cmft.com/api/proxy/api/v1";
    private static final MediaType JSON = MediaType.parse("application/json; charset=utf-8");
    private static final ObjectMapper objectMapper = new ObjectMapper();

    private final String baseUrl;
    private final String apiKey;
    private final OkHttpClient httpClient;

    // ==================== 数据模型 ====================

    /**
     * 会话信息
     */
    public static class Conversation {
        private String conversationId;
        private String name;
        private long createdAt;

        public Conversation(String conversationId, String name, long createdAt) {
            this.conversationId = conversationId;
            this.name = name;
            this.createdAt = createdAt;
        }

        public String getConversationId() { return conversationId; }
        public String getName() { return name; }
        public long getCreatedAt() { return createdAt; }

        @Override
        public String toString() {
            return String.format("Conversation{id='%s', name='%s'}", conversationId, name);
        }
    }

    /**
     * 消息信息
     */
    public static class ChatMessage {
        private String messageId;
        private String queryId;
        private String content;
        private String role;
        private long createdAt;

        public ChatMessage(String messageId, String queryId, String content, String role, long createdAt) {
            this.messageId = messageId;
            this.queryId = queryId;
            this.content = content;
            this.role = role;
            this.createdAt = createdAt;
        }

        public String getMessageId() { return messageId; }
        public String getQueryId() { return queryId; }
        public String getContent() { return content; }
        public String getRole() { return role; }
        public long getCreatedAt() { return createdAt; }
    }

    /**
     * 流式消息回调接口
     */
    public interface StreamCallback {
        /**
         * 收到消息块时触发（实现打字机效果）
         * @param chunk 文本块
         * @param messageId 消息ID
         */
        void onMessage(String chunk, String messageId);

        /**
         * 流式传输完成
         * @param fullContent 完整内容
         * @param messageId 最终消息ID
         */
        default void onComplete(String fullContent, String messageId) {}

        /**
         * 发生错误
         * @param error 错误信息
         */
        default void onError(Throwable error) {}
    }

    // ==================== 构造函数 ====================

    public SseChatClient(String apiKey) {
        this(DEFAULT_BASE_URL, apiKey);
    }

    public SseChatClient(String baseUrl, String apiKey) {
        this.baseUrl = baseUrl;
        this.apiKey = apiKey;
        this.httpClient = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS)
                .readTimeout(120, TimeUnit.SECONDS)
                .writeTimeout(30, TimeUnit.SECONDS)
                .build();
    }

    // ==================== 会话管理 ====================

    /**
     * 创建新会话
     * @param userId 用户ID
     * @param name 会话名称（可选）
     * @return 会话对象
     */
    public Conversation createConversation(String userId, String name) throws IOException {
        Map<String, Object> body = new HashMap<>();
        body.put("AppKey", apiKey);
        body.put("UserID", userId);
        if (name != null && !name.isEmpty()) {
            body.put("Name", name);
        }

        String jsonBody = objectMapper.writeValueAsString(body);

        Request request = new Request.Builder()
                .url(baseUrl + "/create_conversation")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("创建会话失败: " + response.code() + " - " + response.body().string());
            }

            String respJson = response.body().string();
            Map<String, Object> respMap = objectMapper.readValue(respJson, Map.class);

            // 修正：先取 Conversation 对象，再取 AppConversationID
            Map<String, Object> conversation = (Map<String, Object>) respMap.get("Conversation");
            if (conversation == null) {
                throw new IOException("返回数据格式错误，缺少 Conversation 字段");
            }

            String conversationId = (String) conversation.get("AppConversationID");
            if (conversationId == null || conversationId.isEmpty()) {
                throw new IOException("返回数据格式错误，缺少 AppConversationID");
            }
            String convName = (String) conversation.getOrDefault("ConversationName", "新会话");
            long createdAt = ((Number) conversation.getOrDefault("CreatedAt", System.currentTimeMillis())).longValue();

            return new Conversation(conversationId, convName, createdAt);
        }
    }

    /**
     * 创建会话（自动生成名称）
     */
    public Conversation createConversation(String userId) throws IOException {
        return createConversation(userId, null);
    }

    /**
     * 获取会话列表
     */
    public List<Conversation> getConversationList(String userId, Integer limit, Integer offset) throws IOException {
        Map<String, Object> body = new HashMap<>();
        body.put("AppKey", apiKey);
        body.put("UserID", userId);
        if (limit != null) body.put("Limit", limit);
        if (offset != null) body.put("Offset", offset);

        String jsonBody = objectMapper.writeValueAsString(body);

        Request request = new Request.Builder()
                .url(baseUrl + "/get_conversation_list")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("获取会话列表失败: " + response.code());
            }

            String respJson = response.body().string();
            Map<String, Object> respMap = objectMapper.readValue(respJson, Map.class);

            List<Conversation> conversations = new ArrayList<>();
            List<Map<String, Object>> list = (List<Map<String, Object>>) respMap.get("ConversationList");

            if (list != null) {
                for (Map<String, Object> item : list) {
                    conversations.add(new Conversation(
                            (String) item.get("ConversationID"),
                            (String) item.getOrDefault("Name", "未命名"),
                            ((Number) item.getOrDefault("CreatedAt", 0L)).longValue()
                    ));
                }
            }

            return conversations;
        }
    }

    /**
     * 删除会话
     */
    public void deleteConversation(String conversationId, String userId) throws IOException {
        Map<String, Object> body = new HashMap<>();
        body.put("AppKey", apiKey);
        body.put("ConversationID", conversationId);
        body.put("UserID", userId);

        String jsonBody = objectMapper.writeValueAsString(body);

        Request request = new Request.Builder()
                .url(baseUrl + "/delete_conversation")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("删除会话失败: " + response.code());
            }
        }
    }

    // ==================== 消息管理 ====================

    /**
     * 获取消息详情
     */
    public ChatMessage getMessageInfo(String messageId, String userId) throws IOException {
        Map<String, Object> body = new HashMap<>();
        body.put("AppKey", apiKey);
        body.put("MessageID", messageId);
        body.put("UserID", userId);

        String jsonBody = objectMapper.writeValueAsString(body);

        Request request = new Request.Builder()
                .url(baseUrl + "/get_message_info")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("获取消息失败: " + response.code());
            }

            String respJson = response.body().string();
            Map<String, Object> respMap = objectMapper.readValue(respJson, Map.class);
            Map<String, Object> msgInfo = (Map<String, Object>) respMap.get("MessageInfo");

            if (msgInfo == null) {
                throw new IOException("消息不存在");
            }

            return new ChatMessage(
                    (String) msgInfo.get("MessageID"),
                    (String) msgInfo.get("QueryID"),
                    (String) msgInfo.get("Content"),
                    (String) msgInfo.get("Role"),
                    ((Number) msgInfo.getOrDefault("CreatedAt", 0L)).longValue()
            );
        }
    }

    /**
     * 删除消息
     */
    public void deleteMessage(String messageId, String queryId, String userId) throws IOException {
        Map<String, Object> body = new HashMap<>();
        body.put("AppKey", apiKey);
        body.put("MessageID", messageId);
        body.put("UserID", userId);
        if (queryId != null && !queryId.isEmpty()) {
            body.put("QueryID", queryId);
        }

        String jsonBody = objectMapper.writeValueAsString(body);

        Request request = new Request.Builder()
                .url(baseUrl + "/delete_message")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("删除消息失败: " + response.code());
            }
        }
    }

    // ==================== 对话功能 ====================

    /**
     * 流式对话（自动创建会话）
     * @return 创建的会话ID
     */
    public String chatStreamingAuto(String query, String userId, StreamCallback callback) throws IOException {
        Conversation conv = createConversation(userId, "自动创建-" + System.currentTimeMillis());
        String conversationId = conv.getConversationId();

        System.out.println("自动创建会话: " + conversationId);

        chatStreaming(query, conversationId, userId, null, callback);

        return conversationId;
    }

    /**
     * 流式对话（使用已有会话）
     */
    public void chatStreaming(String query, String conversationId, String userId, StreamCallback callback) {
        chatStreaming(query, conversationId, userId, null, callback);
    }

    /**
     * 流式对话（支持文件上传）
     */
    public void chatStreaming(String query, String conversationId, String userId,
                              List<Map<String, Object>> files, StreamCallback callback) {
        try {
            Map<String, Object> requestBody = buildChatRequestBody(query, conversationId, userId, "streaming", files);
            String jsonBody = objectMapper.writeValueAsString(requestBody);

            Request request = new Request.Builder()
                    .url(baseUrl + "/chat_query")
                    .header("Apikey", apiKey)
                    .header("Content-Type", "application/json")
                    .header("Accept", "text/event-stream")
                    .post(RequestBody.create(jsonBody, JSON))
                    .build();

            executeSseRequest(request, callback);

        } catch (Exception e) {
            callback.onError(e);
        }
    }

    /**
     * 阻塞模式对话
     */
    public String chatBlocking(String query, String conversationId, String userId) throws IOException {
        return chatBlocking(query, conversationId, userId, null);
    }

    /**
     * 阻塞模式对话（支持文件）
     */
    public String chatBlocking(String query, String conversationId, String userId,
                               List<Map<String, Object>> files) throws IOException {
        Map<String, Object> requestBody = buildChatRequestBody(query, conversationId, userId, "blocking", files);
        String jsonBody = objectMapper.writeValueAsString(requestBody);

        Request request = new Request.Builder()
                .url(baseUrl + "/chat_query")
                .header("Apikey", apiKey)
                .header("Content-Type", "application/json")
                .post(RequestBody.create(jsonBody, JSON))
                .build();

        try (Response response = httpClient.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("请求失败: " + response.code() + " - " + response.body().string());
            }
            return response.body().string();
        }
    }

    /**
     * 重新生成回复（流式）
     */
    public void regenerateStreaming(String conversationId, String messageId,
                                    String userId, StreamCallback callback) {
        try {
            Map<String, Object> body = new HashMap<>();
            body.put("AppKey", apiKey);
            body.put("AppConversationID", conversationId);
            body.put("MessageID", messageId);
            body.put("UserID", userId);

            String jsonBody = objectMapper.writeValueAsString(body);

            Request request = new Request.Builder()
                    .url(baseUrl + "/query_again")
                    .header("Apikey", apiKey)
                    .header("Content-Type", "application/json")
                    .header("Accept", "text/event-stream")
                    .post(RequestBody.create(jsonBody, JSON))
                    .build();

            executeSseRequest(request, callback);

        } catch (Exception e) {
            callback.onError(e);
        }
    }

    // ==================== 私有辅助方法 ====================

    private Map<String, Object> buildChatRequestBody(String query, String conversationId,
                                                     String userId, String mode,
                                                     List<Map<String, Object>> files) {
        Map<String, Object> body = new HashMap<>();
        body.put("Query", query);
        body.put("AppConversationID", conversationId);
        body.put("AppKey", apiKey);
        body.put("ResponseMode", mode);
        body.put("UserID", userId);
        body.put("PubAgentJump", true);

        if (files != null && !files.isEmpty()) {
            Map<String, Object> queryExtends = new HashMap<>();
            queryExtends.put("Files", files);
            body.put("QueryExtends", queryExtends);
        }

        return body;
    }

    /**
     * 执行 SSE 请求（核心方法）
     */
    /**
     * 执行 SSE 请求（修正版）
     * 处理实际格式：event:text + data:data: {...}
     */
    private void executeSseRequest(Request request, StreamCallback callback) throws IOException {
        Call call = httpClient.newCall(request);
        Response response = call.execute();

        if (!response.isSuccessful()) {
            String errorBody = response.body() != null ? response.body().string() : "无错误信息";
            throw new IOException("SSE请求失败: " + response.code() + " - " + errorBody);
        }

        StringBuilder fullContent = new StringBuilder();
        String finalMessageId = "";
        String currentEvent = "";

        try (BufferedReader reader = new BufferedReader(
                new InputStreamReader(response.body().byteStream(), StandardCharsets.UTF_8))) {

            String line;
            while ((line = reader.readLine()) != null) {
                // 空行表示一个事件结束
                if (line.isEmpty()) {
                    currentEvent = "";
                    continue;
                }

                // 解析 event 行
                if (line.startsWith("event:")) {
                    currentEvent = line.substring(6).trim();
                    continue;
                }

                // 解析 data 行（注意：实际是 data:data: {...} 格式）
                if (line.startsWith("data:")) {
                    String data = line.substring(5).trim();

                    // 处理双重 data: 前缀（data:data: {...}）
                    if (data.startsWith("data:")) {
                        data = data.substring(5).trim();
                    }

                    // 结束标记
                    if ("[DONE]".equals(data)) {
                        callback.onComplete(fullContent.toString(), finalMessageId);
                        return;
                    }

                    try {
                        Map<String, Object> jsonData = objectMapper.readValue(data, Map.class);

                        // 提取消息ID
                        String messageId = extractString(jsonData, "id", "task_id", "message_id");
                        if (!messageId.isEmpty()) {
                            finalMessageId = messageId;
                        }

                        // 根据事件类型处理
                        String event = (String) jsonData.getOrDefault("event", currentEvent);

                        if ("message".equals(event)) {
                            // 提取回答内容（在 answer 字段）
                            String answer = extractString(jsonData, "answer", "content", "chunk", "text");
                            if (!answer.isEmpty()) {
                                fullContent.append(answer);
                                callback.onMessage(answer, finalMessageId);
                            }
                        } else if ("message_start".equals(event)) {
                            // 消息开始，可忽略或记录
                            System.out.println("消息开始，ID: " + finalMessageId);
                        } else if ("message_output_start".equals(event)) {
                            // 输出开始，可忽略或记录
                            System.out.println("输出开始");
                        } else if ("message_end".equals(event) || "end".equals(event)) {
                            // 消息结束
                            callback.onComplete(fullContent.toString(), finalMessageId);
                            return;
                        }
                        // 其他事件类型可扩展...

                    } catch (Exception e) {
                        // JSON解析失败，打印调试
                        System.out.println("解析失败，原始数据: " + data);
                    }
                }
            }

            // 流正常结束（没有明确结束标记）
            callback.onComplete(fullContent.toString(), finalMessageId);

        } catch (IOException e) {
            callback.onError(e);
            throw e;
        }
    }

    /**
     * 从 Map 中提取字符串（尝试多个 key）
     */
    private String extractString(Map<String, Object> map, String... keys) {
        for (String key : keys) {
            Object value = map.get(key);
            if (value != null) {
                return value.toString();
            }
        }
        return "";
    }

    /**
     * 关闭客户端
     */
    public void close() {
        httpClient.dispatcher().executorService().shutdown();
        httpClient.connectionPool().evictAll();
    }

    // ==================== 使用示例 ====================

    public static void main(String[] args) {
        SseChatClient client = new SseChatClient("d6gl6n6llh49on40fqug");

        try {
            String userId = "user_" + System.currentTimeMillis();

            // ========== 示例1：标准流程 ==========
            System.out.println("=== 步骤1：创建会话 ===");
            Conversation conv = client.createConversation(userId, "测试会话");
            System.out.println("会话创建成功: " + conv.getConversationId());

            System.out.println("\n=== 步骤2：流式对话 ===");
            final StringBuilder contentBuilder = new StringBuilder();
            final String[] lastMessageId = {""};

            client.chatStreaming(
                    "给我讲一个关于人工智能的短故事",
                    conv.getConversationId(),
                    userId,
                    new StreamCallback() {
                        @Override
                        public void onMessage(String chunk, String messageId) {
                            System.out.print(chunk);
                            System.out.flush();
                            contentBuilder.append(chunk);
                            if (!messageId.isEmpty()) {
                                lastMessageId[0] = messageId;
                            }
                        }

                        @Override
                        public void onComplete(String fullContent, String messageId) {
                            System.out.println("\n\n[对话完成] 消息ID: " + messageId);
                        }

                        @Override
                        public void onError(Throwable error) {
                            System.err.println("\n[错误] " + error.getMessage());
                        }
                    }
            );
            String string = contentBuilder.toString();
//            Thread.sleep(15000);
//            // 等待流式传输完成
//            Thread.sleep(15000);
//            String string = contentBuilder.toString();

            // ========== 示例3：获取会话列表 ==========
            System.out.println("\n=== 步骤4：获取会话列表 ===");
            List<Conversation> conversations = client.getConversationList(userId, 10, 0);
            System.out.println("共有 " + conversations.size() + " 个会话:");
            for (Conversation c : conversations) {
                System.out.println(" - " + c);
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            client.close();
        }
    }
}