package com.ruoyi.llm.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.common.utils.milvus.MilvusConnectionPool;
import com.ruoyi.llm.service.IMilvusService;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.*;
import io.milvus.v2.service.collection.response.DescribeCollectionResp;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import io.milvus.v2.service.vector.request.DeleteReq;
import io.milvus.v2.service.vector.request.QueryReq;
import io.milvus.v2.service.vector.response.QueryResp;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.text.SimpleDateFormat;
import java.util.*;
import java.util.stream.Collectors;

@Service
public class MilvusServiceImpl implements IMilvusService {

    @Autowired
    private MilvusConnectionPool milvusConnectionPool;

    /**
     * 新建知识库Collection（含Schema、Field、Index）
     */
    @Override
    public void createCollection(String collectionName, String description, int dimension) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            CreateCollectionReq.CollectionSchema schema = MilvusClientV2.CreateSchema();

            schema.addField(AddFieldReq.builder()
                    .fieldName("id")
                    .dataType(DataType.Int64)
                    .isPrimaryKey(true)
                    .autoID(true)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("file_name")
                    .dataType(DataType.VarChar)
                    .maxLength(256)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("title")
                    .dataType(DataType.VarChar)
                    .maxLength(256)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("file_type")
                    .dataType(DataType.VarChar)
                    .maxLength(10)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("content")
                    .dataType(DataType.VarChar)
                    .maxLength(65535)
                    .build());

            schema.addField(AddFieldReq.builder()
                    .fieldName("embedding")
                    .dataType(DataType.FloatVector)
                    .dimension(dimension)
                    .build());

            // 创建索引
            Map<String, Object> extraParams = new HashMap<>();
            extraParams.put("nlist", 64);
            IndexParam indexParam = IndexParam.builder()
                    .fieldName("embedding")
                    .indexType(IndexParam.IndexType.IVF_FLAT)
                    .metricType(IndexParam.MetricType.COSINE)
                    .extraParams(extraParams)
                    .build();

            CreateCollectionReq createCollectionParam = CreateCollectionReq.builder()
                    .collectionName(collectionName)
                    .description(description)
                    .collectionSchema(schema)
                    .indexParam(indexParam)
                    .build();

            milvusClient.createCollection(createCollectionParam);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 查询知识库Collection
     */
    @Override
    public JSONArray getCollectionNames() {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            JSONArray jsonArray = new JSONArray();
            ListCollectionsResp listResponse = milvusClient.listCollections();
            if (listResponse != null) {
                List<String> collectionNames = listResponse.getCollectionNames();
                for (String collectionName : collectionNames) {
                    JSONObject jsonObject = new JSONObject();
                    DescribeCollectionReq request = DescribeCollectionReq.builder()
                            .collectionName(collectionName)
                            .build();
                    DescribeCollectionResp describeResponse = milvusClient.describeCollection(request);
                    jsonObject.put("collectionId", describeResponse.getCollectionID());
                    jsonObject.put("collectionName", collectionName);
                    SimpleDateFormat beijingFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
                    beijingFormat.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
                    String beijingTime = beijingFormat.format(describeResponse.getCreateUtcTime());
                    jsonObject.put("createdTime", beijingTime);
                    jsonObject.put("description", describeResponse.getDescription());
                    jsonArray.add(jsonObject);
                }
            }
            return jsonArray;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 根据名称查询知识库Collection
     * 返回指定名称的集合信息
     */
    @Override
    public JSONArray listKnowLedgeByCollectionName(String collectionName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            JSONArray jsonArray = new JSONArray();
            ListCollectionsResp listResponse = milvusClient.listCollections();

            if (listResponse != null) {
                List<String> allCollectionNames = listResponse.getCollectionNames();
                for (String name : allCollectionNames) {
                    JSONObject jsonObject = new JSONObject();
                    DescribeCollectionReq request = DescribeCollectionReq.builder()
                            .collectionName(name)
                            .build();
                    DescribeCollectionResp describeResponse = milvusClient.describeCollection(request);
                    if (describeResponse.getDescription().contains(collectionName)) {
                        jsonObject.put("collectionId", describeResponse.getCollectionID());
                        jsonObject.put("collectionName", name);
                        SimpleDateFormat beijingFormat = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
                        beijingFormat.setTimeZone(TimeZone.getTimeZone("Asia/Shanghai"));
                        String beijingTime = beijingFormat.format(describeResponse.getCreateUtcTime());
                        jsonObject.put("createdTime", beijingTime);
                        jsonObject.put("description", describeResponse.getDescription());
                        jsonArray.add(jsonObject);
                    }
                }
            }
            return jsonArray;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 修改知识库Collection
     */
    @Override
    public void collectionRename(String collectionName, String newCollectionName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            RenameCollectionReq renameCollectionReq = RenameCollectionReq.builder()
                    .collectionName(collectionName)
                    .newCollectionName(newCollectionName)
                    .build();

            milvusClient.renameCollection(renameCollectionReq);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 删除知识库Collection
     */
    @Override
    public void deleteCollectionName(String collectionName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            DropCollectionReq dropCollectionReq = DropCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.dropCollection(dropCollectionReq);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 查询知识库文件
     */
    @Override
    public List<String> listDocument(String collectionName, String fileType) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            List<String> documentList = new ArrayList<>();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            QueryReq queryParam = QueryReq.builder()
                    .collectionName(collectionName)
                    .filter("id > 0")
                    .outputFields(Arrays.asList("file_type", "file_name"))
                    .build();
            if (fileType != null && !fileType.equals(""))
                queryParam = QueryReq.builder()
                        .collectionName(collectionName)
                        .filter(String.format("file_type == \"%s\"", fileType))
                        .outputFields(Arrays.asList("file_type", "file_name"))
                        .build();
            QueryResp queryResp = milvusClient.query(queryParam);
            List<QueryResp.QueryResult> rowRecordList;
            if (queryResp != null) {
                rowRecordList = queryResp.getQueryResults();
                for (QueryResp.QueryResult rowRecord : rowRecordList) {
                    documentList.add(rowRecord.getEntity().get("file_name").toString());
                }
            }
            ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.releaseCollection(releaseCollectionReq);
            return documentList.stream().distinct().collect(Collectors.toList());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 删除知识库文件
     */
    @Override
    public void removeDocument(String collectionName, String fileName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            DeleteReq deleteReq = DeleteReq.builder()
                    .collectionName(collectionName)
                    .filter(String.format("file_name == \"%s\"", fileName))
                    .build();
            milvusClient.delete(deleteReq);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 删除知识库所有文件
     */
    @Override
    public void removeAllDocument(String collectionName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            DeleteReq deleteReq = DeleteReq.builder()
                    .collectionName(collectionName)
                    .filter("id > 0")
                    .build();
            milvusClient.delete(deleteReq);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 列出所有的title
     */
    @Override
    public List<String> listTiles(String collectionName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            List<String> titleList = new ArrayList<>();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            QueryReq queryParam = QueryReq.builder()
                    .collectionName(collectionName)
                    .filter("id > 0")
                    .outputFields(Arrays.asList("title"))
                    .build();
            QueryResp queryResp = milvusClient.query(queryParam);
            List<QueryResp.QueryResult> rowRecordList;
            if (queryResp != null) {
                rowRecordList = queryResp.getQueryResults();
                for (QueryResp.QueryResult rowRecord : rowRecordList) {
                    Object title = rowRecord.getEntity().get("title");
                    if (title != null) {
                        titleList.add(title.toString());
                    }
                }
            }
            ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.releaseCollection(releaseCollectionReq);
            return titleList.stream().distinct().collect(Collectors.toList());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 根据文件名获取标题列表
     */
    @Override
    public List<String> listTilesByFile(String collectionName, String fileName) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            List<String> titleList = new ArrayList<>();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            QueryReq queryParam = QueryReq.builder()
                    .collectionName(collectionName)
                    .filter(String.format("id > 0 && file_name == \"%s\"", fileName))
                    .outputFields(Arrays.asList("title"))
                    .build();
            QueryResp queryResp = milvusClient.query(queryParam);
            List<QueryResp.QueryResult> rowRecordList;
            if (queryResp != null) {
                rowRecordList = queryResp.getQueryResults();
                for (QueryResp.QueryResult rowRecord : rowRecordList) {
                    Object title = rowRecord.getEntity().get("title");
                    if (title != null) {
                        titleList.add(title.toString());
                    }
                }
            }
            ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.releaseCollection(releaseCollectionReq);
            return titleList.stream().distinct().collect(Collectors.toList());
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }

    /**
     * 根据title名查询相关content
     */
    @Override
    public JSONArray listByTitle(String collectionName, String title) {
        MilvusClientV2 milvusClient = null;
        try {
            milvusClient = milvusConnectionPool.getClient();
            JSONArray resultArray = new JSONArray();
            LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.loadCollection(loadCollectionReq);
            QueryReq queryParam = QueryReq.builder()
                    .collectionName(collectionName)
                    .filter(String.format("title == \"%s\"", title))
                    .outputFields(Arrays.asList("content", "file_name"))
                    .build();
            QueryResp queryResp = milvusClient.query(queryParam);
            List<QueryResp.QueryResult> rowRecordList;
            if (queryResp != null) {
                rowRecordList = queryResp.getQueryResults();
                for (QueryResp.QueryResult rowRecord : rowRecordList) {
                    JSONObject item = new JSONObject();
                    Object content = rowRecord.getEntity().get("content");
                    Object fileName = rowRecord.getEntity().get("file_name");
                    if (content != null) {
                        item.put("content", content.toString());
                    }
                    if (fileName != null) {
                        item.put("file_name", fileName.toString());
                    }
                    resultArray.add(item);
                }
            }
            ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                    .collectionName(collectionName)
                    .build();
            milvusClient.releaseCollection(releaseCollectionReq);
            return resultArray;
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new RuntimeException("获取Milvus连接失败", e);
        } finally {
            if (milvusClient != null) {
                milvusConnectionPool.returnClient(milvusClient);
            }
        }
    }
}