package com.aigeo.util;

import com.aigeo.entity.*;
import com.aigeo.socket.WebSocketAiServer;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import okhttp3.*;
import okio.BufferedSource;
import org.apache.tika.exception.TikaException;
import org.springframework.web.multipart.MultipartFile;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;

/**
 * <b>Dify AI工具类.</b>
 * <p>
 * 文件详情: 如需修改请联系开发人员<br>
 * 版权所有: 深圳市科飞时速网络技术有限公司(0755-88843705)<br>
 * 技术支持: info@21gmail.com<br>
 * 开始时间: 2025/08/19 16:36<br>
 * 最后修改: 2025/08/19 16:36<br>
 *
 * @author : 温志锋
 * @version : 3.0
 */
public class DifyAiUtil {

    //测试知识库ID：8b3c4cdd-01a3-453b-8e12-bc7c0d5268db
    private static final String DIFY_BASE_URL="http://193.112.177.8:8088/v1";
    private static final String API_KEY="app-sxppRrqW8OMd6Ak07iprjn2t";
    private static final String DATASET_KEY="dataset-2nhuEkzhNshtvZ6L6Ut3G5hC";

    private static final String DIFY_CHAT_URL="/chat-messages";
    private static final String DATASET_ENDPOINT="/datasets/";
    private static final String DATASET_DOCUMENT_ENDPOINT="/document/create-by-text";
    private static final String DATASET_UPLOADFILE_ENDPOINT="/document/create-by-file";

    private static final MediaType JSON_TYPE = MediaType.get("application/json; charset=utf-8");

    private static final OkHttpClient client = new OkHttpClient();

    /**
     * 回复
     * @param message 用户消息
     * @param sessionId 会话ID
     */
    public static void reply(String message,String sessionId){
        DifyRequestDTO difyRequestDTO = JSONObject.parseObject(message, DifyRequestDTO.class);
        difyRequestDTO.setInputs(new Object());
        difyRequestDTO.setUser("anhssd");
        difyRequestDTO.setResponse_mode("streaming");

        String jsonBody = JSON.toJSONString(difyRequestDTO);

        // 创建请求体
        RequestBody requestBody = RequestBody.create(
                jsonBody,
                JSON_TYPE
        );

        // 创建请求
        Request request = new Request.Builder()
                .url(DIFY_BASE_URL+DIFY_CHAT_URL)
                .addHeader("Authorization", "Bearer " + API_KEY)
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build();

        // 发送请求并处理流式响应
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                e.printStackTrace();
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (!response.isSuccessful()) {
                    System.err.println("请求失败: " + response.code() + " - " + response.body().string());
                    return;
                }

                processStream(response,sessionId);
            }
        });
    }

    /**
     * 创建知识库
     * @param datasetName 知识库名称
     * @param description 知识库描述
     */
    public static void createDataset(String datasetName, String description) {
        HashMap<String, Object> map = new HashMap<>();
        map.put("name", datasetName);
        map.put("description", description);
        map.put("permission", "all_team_members");

        String jsonBody = JSON.toJSONString(map);

        RequestBody requestBody = RequestBody.create(jsonBody, JSON_TYPE);
        // 创建请求
        Request request = new Request.Builder()
                .url(DIFY_BASE_URL + DATASET_ENDPOINT)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build();

        // 发送请求
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                e.printStackTrace();
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (!response.isSuccessful()) {
                    System.err.println("请求失败: " + response.code() + " - " + response.body().string());
                    return;
                }

                System.out.println("知识库创建成功: " + response.body().string());
            }
        });
    }

    /**
     * 获取知识库列表
     */
    public static void getDatasetList() {
        // 创建请求
        Request request = new Request.Builder()
                .url(DIFY_BASE_URL + DATASET_ENDPOINT)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .build();

        try {
            // 发送请求并获取响应
            Response response = client.newCall(request).execute();

            if (!response.isSuccessful()) {
                System.err.println("请求失败: " + response.code() + " - " + response.body().string());
                return;
            }

            if (response.body() != null) {
                System.out.println("数据集列表: " + response.body().string());
            }
        } catch (IOException e) {
            System.err.println("获取数据集列表失败: " + e.getMessage());
            e.printStackTrace();
        }
    }

    /**
     * 上传文件到知识库
     * @param datasetId 知识库ID
     * @param file  文件
     * @param fileName 文件名
     * @throws IOException IO异常
     */
    public static DocumentUploadRes uploadFileToDataset(String datasetId, MultipartFile file, String fileName) throws IOException, TikaException {
        // 构建URL
        String uploadUrl = DIFY_BASE_URL + DATASET_ENDPOINT + datasetId + DATASET_UPLOADFILE_ENDPOINT;
        DifyUploadFileToDatasetDTO datasetDTO = new DifyUploadFileToDatasetDTO();
        datasetDTO.setIndexing_technique("high_quality");
        DifyUploadFileToDatasetDTO.PrecessRule precessRule = new DifyUploadFileToDatasetDTO.PrecessRule();
        precessRule.setMode("custom");
        DifyUploadFileToDatasetDTO.PrecessRule.Rules rules = new DifyUploadFileToDatasetDTO.PrecessRule.Rules();
        DifyUploadFileToDatasetDTO.PrecessRule.Rules.Segmentation segmentation = new DifyUploadFileToDatasetDTO.PrecessRule.Rules.Segmentation();
        segmentation.setSeparator("###");
        segmentation.setMax_tokens(500);
        DifyUploadFileToDatasetDTO.PrecessRule.Rules.PreProcessingRules processingRules = new DifyUploadFileToDatasetDTO.PrecessRule.Rules.PreProcessingRules();
        processingRules.setId("remove_extra_spaces");
        processingRules.setEnabled("true");
        DifyUploadFileToDatasetDTO.PrecessRule.Rules.PreProcessingRules processingRules1 = new DifyUploadFileToDatasetDTO.PrecessRule.Rules.PreProcessingRules();
        processingRules1.setId("remove_urls_emails");
        processingRules1.setEnabled("true");
        ArrayList<DifyUploadFileToDatasetDTO.PrecessRule.Rules.PreProcessingRules> list = new ArrayList<>();
        list.add(processingRules);
        list.add(processingRules1);
        rules.setPre_processing_rules(list);
        rules.setSegmentation(segmentation);
        precessRule.setRules(rules);
        datasetDTO.setProcess_rule(precessRule);
        String jsonString = JSONObject.toJSONString(datasetDTO);
        //System.out.println("Request JSON: " + jsonString);
        // 关键修改：使用正确的媒体类型 application/json
        RequestBody dataBody = RequestBody.create(
                jsonString,
                MediaType.parse("text/plain")
        );

        RequestBody requestBody = new MultipartBody.Builder()
                .setType(MultipartBody.FORM)
                // 修改点：使用 application/json 类型传递JSON数据
                .addFormDataPart("data", null, dataBody)
                .addFormDataPart("file", fileName,
                        RequestBody.create(file.getBytes(), MediaType.parse("application/octet-stream")))
                .build();
        // 创建请求
        Request request = new Request.Builder()
                .url(uploadUrl)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .post(requestBody)
                .build();

        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                String errorBody = response.body() != null ? response.body().string() : "null";
                System.err.println("上传文件请求失败: " + response.code() + " - " + errorBody);
                DocumentUploadRes documentUploadRes = new DocumentUploadRes();
                documentUploadRes.setCode(response.code());
                return documentUploadRes;
            }
            if (response.body() != null) {
                String responseBody = response.body().string();
                return JSONObject.parseObject(responseBody, DocumentUploadRes.class); // 直接返回结果
            } else {
                DocumentUploadRes documentUploadRes = new DocumentUploadRes();
                documentUploadRes.setCode(400);
                return documentUploadRes;
            }
        }
    }

    /**
     * 上传知识库内容并创建文档至知识库
     * @param datasetId 知识库ID
     * @param documentName 文档名称
     * @param text 文本内容
     */
    public static void createDocumentByText(String datasetId, String documentName, String text) throws IOException {
        // 构建URL
        String url = DIFY_BASE_URL + DATASET_ENDPOINT + datasetId + DATASET_DOCUMENT_ENDPOINT;

        // 构建请求数据
        JSONObject requestData = new JSONObject();
        requestData.put("name", documentName);
        requestData.put("text", text);
        requestData.put("indexing_technique", "high_quality");

        // 构建处理规则
        JSONObject processRule = new JSONObject();
        processRule.put("mode", "automatic");
        requestData.put("process_rule", processRule);

        String jsonBody = requestData.toJSONString();

        // 创建请求体
        RequestBody requestBody = RequestBody.create(jsonBody, JSON_TYPE);

        // 创建请求
        Request request = new Request.Builder()
                .url(url)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .addHeader("Content-Type", "application/json")
                .post(requestBody)
                .build();

        // 发送请求
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                System.err.println("创建文档失败: " + e.getMessage());
                e.printStackTrace();
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (!response.isSuccessful()) {
                    String errorBody = response.body() != null ? response.body().string() : "null";
                    System.err.println("创建文档请求失败: " + response.code() + " - " + errorBody);
                    return;
                }

                if (response.body() != null) {
                    System.out.println("文档创建成功: " + response.body().string());
                }
            }
        });
    }

    /**
     * 获取指定数据集中特定文档的详细信息
     * @param datasetId 数据集ID
     * @param documentId 文档ID
     * @throws IOException IO异常
     */
    public static void getDocumentById(String datasetId, String documentId) throws IOException {
        // 构建URL
        String url = DIFY_BASE_URL + DATASET_ENDPOINT + datasetId + "/documents/" + documentId;

        // 创建请求
        Request request = new Request.Builder()
                .url(url)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .get() // 明确指定GET方法
                .build();

        // 发送请求
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                System.err.println("获取文档失败: " + e.getMessage());
                e.printStackTrace();
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                if (!response.isSuccessful()) {
                    String errorBody = response.body() != null ? response.body().string() : "null";
                    System.err.println("获取文档请求失败: " + response.code() + " - " + errorBody);
                    return;
                }

                if (response.body() != null) {
                    System.out.println("文档详情: " + response.body().string());
                }
            }
        });
    }

    /**
     * 获取指定知识库中的文档列表
     * @param datasetId 数据集ID
     * @param page 页码（从1开始）
     * @param limit 每页返回条数
     * @param keyword 搜索关键词
     * @return 响应内容
     * @throws IOException IO异常
     */
    public static DifyDatasetDocuments getDatasetDocuments(String datasetId, Integer page, Integer limit, String keyword) throws IOException {
        // 构建基础URL
        String url = DIFY_BASE_URL + DATASET_ENDPOINT + datasetId + "/documents";

        // 构建查询参数
        HttpUrl.Builder urlBuilder = HttpUrl.parse(url).newBuilder();

        // 添加可选参数
        if (page != null && page > 0) {
            urlBuilder.addQueryParameter("page", String.valueOf(page));
        }

        if (limit != null && limit > 0) {
            urlBuilder.addQueryParameter("limit", String.valueOf(limit));
        }

        if (keyword != null && !keyword.trim().isEmpty()) {
            urlBuilder.addQueryParameter("keyword", keyword);
        }

        // 构建最终URL
        String finalUrl = urlBuilder.build().toString();

        // 创建请求
        Request request = new Request.Builder()
                .url(finalUrl)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .build();

        // 发送请求并处理响应
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                DifyDatasetDocuments documents = new DifyDatasetDocuments();
                documents.setCode(500);
                return documents;
            }

            if (response.body() != null) {
                DifyDatasetDocuments documents = JSONObject.parseObject(response.body().string(), DifyDatasetDocuments.class);
                documents.setCode(200);
                return documents;
            } else {
                DifyDatasetDocuments documents = new DifyDatasetDocuments();
                documents.setCode(500);
                return documents;
            }
        } catch (IOException e) {
            System.err.println("获取文档列表失败: " + e.getMessage());
            throw e;
        }
    }

    /**
     * 同步方式删除指定数据集中的特定文档
     * @param datasetId 数据集ID
     * @param documentId 文档ID
     * @return 删除操作结果
     * @throws IOException IO异常
     */
    public static boolean deleteDocumentSync(String datasetId, String documentId) throws IOException {
        // 构建URL
        String url = DIFY_BASE_URL + DATASET_ENDPOINT + datasetId + "/documents/" + documentId;

        // 创建DELETE请求
        Request request = new Request.Builder()
                .url(url)
                .addHeader("Authorization", "Bearer " + DATASET_KEY)
                .delete()
                .build();

        // 发送请求并处理响应
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                String errorBody = response.body() != null ? response.body().string() : "null";
                System.err.println("删除文档请求失败: " + response.code() + " - " + errorBody);
                return false;
            }

            System.out.println("文档删除成功: " + response.code());
            return true;
        } catch (IOException e) {
            System.err.println("删除文档失败: " + e.getMessage());
            throw e;
        }
    }


    private static void processStream(Response response,String sessionId) throws IOException {
        try (BufferedSource source = response.body().source()) {
            StringBuilder eventBuffer = new StringBuilder();

            while (!source.exhausted()) {
                String line = source.readUtf8Line();
                if (line == null) break;

                if (line.startsWith("data:")) {
                    eventBuffer.setLength(0);
                    eventBuffer.append(line.substring(5).trim());
                } else if (line.isEmpty() && eventBuffer.length() > 0) {
                    handleEvent(eventBuffer.toString(),sessionId);
                    eventBuffer.setLength(0);
                }
            }
        }
    }

    private static void handleEvent(String eventData,String sessionId) {
        try {
            DifyResponse event = JSON.parseObject(eventData, DifyResponse.class);

            switch (event.getEvent()) {
                case "agent_thought":
                    //System.out.println("[思考] " + event.getThought());
                    break;
                case "agent_message":
                    System.out.print(event.getAnswer());
                    event.setCode(200);
                    WebSocketAiServer.sendMessage(sessionId,event);
                    break;
                case "message_end":
                    event.setCode(203);
                    System.out.println("\n[对话结束]");
                    WebSocketAiServer.sendMessage(sessionId,event);
                    break;

                case "error":
                    event.setCode(500);
                    System.err.println("[错误] " + event.getMessage());
                    WebSocketAiServer.sendMessage(sessionId,event);
                    break;
                default:
                    System.out.println("[事件] " + event.getEvent());
            }
        } catch (Exception e) {
            System.err.println("解析错误: " + eventData);
            e.printStackTrace();
        }
    }
}