package com.cy.guoyan.admin.module.ai.util;

import com.alibaba.fastjson.JSONObject;
import org.apache.http.client.methods.*;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.StringEntity;
import org.apache.http.entity.mime.HttpMultipartMode;
import org.apache.http.entity.mime.MultipartEntityBuilder;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import javax.annotation.PostConstruct;
import java.io.File;
import java.io.IOException;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Map;

/**
 * Dify AI HTTP 客户端工具类
 */
@Component
public class DifyClient  {

    private static final String AUTH_HEADER       = "Authorization";
    private static final String CONTENT_TYPE_JSON = "application/json";
    private static final String BEARER            = "Bearer ";

    @Value("${dify.baseUrl}")
    private String baseUrl;    // 从 application.properties 或 yml 中注入

    @Value("${dify.datasetsKey}")
    private String apiKey;     // 从 application.properties 或 yml 中注入

    private CloseableHttpClient httpClient;


    @PostConstruct
    public void init() {
        // 去掉末尾斜杠
        if (baseUrl.endsWith("/")) {
            baseUrl = baseUrl.substring(0, baseUrl.length() - 1);
        }
        this.httpClient = HttpClients.createDefault();
    }

    /**
     * 上传文档（文件方式）
     */
    public JSONObject createDocumentByFile(File file, String name, Map<String, Object> requestData) throws IOException {
        String url = baseUrl + "/v1/datasets/" + requestData.get("datasetId") + "/document/create-by-file";

        MultipartEntityBuilder builder = MultipartEntityBuilder.create()
                .setCharset(StandardCharsets.UTF_8)
                // 关键：启用 RFC6532，允许头部使用原生 UTF-8
                .setMode(HttpMultipartMode.RFC6532)
                .addBinaryBody("file", file, ContentType.APPLICATION_OCTET_STREAM, name)
                .addTextBody("data", JSONObject.toJSONString(requestData), ContentType.APPLICATION_JSON.withCharset(StandardCharsets.UTF_8));

        HttpPost post = new HttpPost(url);
        post.setHeader(AUTH_HEADER, BEARER + apiKey);
        post.setEntity(builder.build());

        return executeJson(post, 200);
    }


    /**
     * 获取文档嵌入（索引）状态
     */
    public JSONObject getDocumentIndexingStatus(String datasetId, String batch) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s/indexing-status", baseUrl, datasetId, batch);
        HttpGet get = new HttpGet(url);
        get.setHeader(AUTH_HEADER, BEARER + apiKey);
        return executeJson(get, 200);
    }

    /**
     * 删除整个文档
     */
    public void deleteDocument(String datasetId, String documentId) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s", baseUrl, datasetId, documentId);
        HttpDelete delete = new HttpDelete(url);
        delete.setHeader(AUTH_HEADER, BEARER + apiKey);
        executeStatus(delete, 200);
    }

    /**
     * 新增分段
     */
    public JSONObject addDocumentSegments(String datasetId, String documentId, List<Map<String, Object>> segments) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s/segments", baseUrl, datasetId, documentId);
        HttpPost post = new HttpPost(url);
        post.setHeader(AUTH_HEADER, BEARER + apiKey);
        post.setHeader("Content-Type", CONTENT_TYPE_JSON);

        JSONObject body = new JSONObject();
        body.put("segments", segments);
        post.setEntity(new StringEntity(body.toJSONString(), StandardCharsets.UTF_8));

        return executeJson(post, 200);
    }

    /**
     * 查询知识库中文档列表（分页 + 关键字可选）
     */
    public JSONObject listDocuments(String datasetId, String keyword, Integer page, Integer limit) throws IOException {
        StringBuilder sb = new StringBuilder(baseUrl)
                .append("/v1/datasets/").append(datasetId).append("/documents?");
        if (keyword != null) sb.append("keyword=").append(URLEncoder.encode(keyword, "UTF-8")).append("&");
        if (page    != null) sb.append("page=").append(page).append("&");
        if (limit   != null) sb.append("limit=").append(limit).append("&");
        String url = sb.toString().replaceAll("&$", "");

        HttpGet get = new HttpGet(url);
        get.setHeader(AUTH_HEADER, BEARER + apiKey);
        return executeJson(get, 200);
    }

    /**
     * 查询文档分段（支持 keyword/status 过滤）
     */
    public JSONObject listDocumentSegments(String datasetId, String documentId, String keyword, String status) throws IOException {
        StringBuilder sb = new StringBuilder(baseUrl)
                .append("/v1/datasets/").append(datasetId)
                .append("/documents/").append(documentId).append("/segments?");
        if (keyword != null) sb.append("keyword=").append(URLEncoder.encode(keyword, "UTF-8")).append("&");
        if (status  != null) sb.append("status=").append(status).append("&");
        String url = sb.toString().replaceAll("&$", "");

        HttpGet get = new HttpGet(url);
        get.setHeader(AUTH_HEADER, BEARER + apiKey);
        return executeJson(get, 200);
    }

    /**
     * 删除指定分段
     */
    public void deleteDocumentSegment(String datasetId, String documentId, String segmentId) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s/segments/%s",
                baseUrl, datasetId, documentId, segmentId);
        HttpDelete delete = new HttpDelete(url);
        delete.setHeader(AUTH_HEADER, BEARER + apiKey);
        executeStatus(delete, 200);
    }

    /**
     * 获取上传文件信息（预览 & 下载地址）
     */
    public JSONObject getUploadedFileInfo(String datasetId, String documentId) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s/upload-file",
                baseUrl, datasetId, documentId);
        HttpGet get = new HttpGet(url);
        get.setHeader(AUTH_HEADER, BEARER + apiKey);
        return executeJson(get, 200);
    }

    /**
     * 更新分段
     */
    public JSONObject updateDocumentSegment(String datasetId,
                                            String documentId,
                                            String segmentId,
                                            Map<String, Object> segmentData) throws IOException {
        String url = String.format("%s/v1/datasets/%s/documents/%s/segments/%s",
                baseUrl, datasetId, documentId, segmentId);
        HttpPost post = new HttpPost(url);
        post.setHeader(AUTH_HEADER, BEARER + apiKey);
        post.setHeader("Content-Type", CONTENT_TYPE_JSON);

        JSONObject body = new JSONObject();
        body.put("segment", segmentData);
        post.setEntity(new StringEntity(body.toJSONString(), StandardCharsets.UTF_8));

        return executeJson(post, 200);
    }

    /**
     * 检索知识库
     */
    public JSONObject retrieveFromDataset(String datasetId, String query, JSONObject retrievalModel) throws IOException {
        String url = baseUrl + "/v1/datasets/" + datasetId + "/retrieve";
        HttpPost post = new HttpPost(url);
        post.setHeader(AUTH_HEADER, BEARER + apiKey);
        post.setHeader("Content-Type", CONTENT_TYPE_JSON);

        JSONObject body = new JSONObject();
        body.put("query", query);
        if (retrievalModel != null) {
            body.put("retrieval_model", retrievalModel);
        }
        post.setEntity(new StringEntity(body.toJSONString(), StandardCharsets.UTF_8));

        return executeJson(post, 200);
    }

    // —— 私有工具方法 —— //

    /** 执行并返回 JSON，检查期望状态码 */
    private JSONObject executeJson(HttpRequestBase request, int expectStatus) throws IOException {
        try (CloseableHttpResponse resp = httpClient.execute(request)) {
            int code = resp.getStatusLine().getStatusCode();
            String body = EntityUtils.toString(resp.getEntity(), StandardCharsets.UTF_8);
            if (code != expectStatus) {
                throw new IOException("请求失败，状态码：" + code + "，内容：" + body);
            }
            return JSONObject.parseObject(body);
        }
    }

    /** 执行并仅校验状态码，常用于 204 No Content */
    private void executeStatus(HttpRequestBase request, int expectStatus) throws IOException {
        try (CloseableHttpResponse resp = httpClient.execute(request)) {
            int code = resp.getStatusLine().getStatusCode();
            if (code != expectStatus) {
                String body = EntityUtils.toString(resp.getEntity(), StandardCharsets.UTF_8);
                throw new IOException("请求失败，状态码：" + code + "，内容：" + body);
            }
        }
    }

    /**
     * 关闭 httpClient，应用退出时调用
     */
    public void close() throws IOException {
        httpClient.close();
    }
}