package com.ruoyi.web.llm.service.impl;

import com.alibaba.fastjson2.JSONArray;
import com.alibaba.fastjson2.JSONObject;
import com.ruoyi.web.llm.service.IMilvusService;
import io.milvus.param.dml.DeleteParam;
import io.milvus.v2.client.ConnectConfig;
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.QueryReq;
import io.milvus.v2.service.vector.response.QueryResp;
import org.springframework.stereotype.Service;

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

@Service
public class MilvusServiceImpl implements IMilvusService {

    private static final MilvusClientV2 milvusClient = new MilvusClientV2(
            ConnectConfig.builder()
                    .uri("http://192.168.28.188:19530")
                    .build());

    /**
     * 新建知识库Collection（含Schema、Field、Index）
     */
    @Override
    public void createCollection(String collectionName, String description, int dimension) {
        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("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(5)
                .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);
    }

    /**
     * 查询知识库Collection
     */
    public JSONArray getCollectionNames() {
        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;
    }

    /**
     * 修改知识库Collection
     */
    public void collectionRename(String collectionName, String newCollectionName) {
        RenameCollectionReq renameCollectionReq = RenameCollectionReq.builder()
                .collectionName(collectionName)
                .newCollectionName(newCollectionName)
                .build();

        milvusClient.renameCollection(renameCollectionReq);
    }

    /**
     * 删除知识库Collection
     */
    public void deleteCollectionName(String collectionName) {
        DropCollectionReq dropCollectionReq = DropCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        milvusClient.dropCollection(dropCollectionReq);
    }

    /**
     * 查询知识库文件
     */
    public List<String> listDocument(String collectionName, String fileType) {
        List<String> documentList = new ArrayList<>();
        loadCollectionName(collectionName);
        QueryReq queryParam = QueryReq.builder()
                .collectionName(collectionName)
                .filter("id > 0")
                .outputFields(Arrays.asList("file_type", "file_name"))
                .limit(16384L)
                .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"))
                .limit(16384L)
                .build();
        QueryResp queryResp = milvusClient.query(queryParam);
        List<QueryResp.QueryResult> rowRecordList = new ArrayList<>();
        if (queryResp != null) {
            rowRecordList = queryResp.getQueryResults();
            for (QueryResp.QueryResult rowRecord : rowRecordList) {
                documentList.add(rowRecord.getEntity().get("file_name").toString());
            }
        }
        releaseCollectionName(collectionName);
        return documentList.stream().distinct().collect(Collectors.toList());
    }

    /**
     * 删除知识库文件
     */
    public void removeDocument(String collectionName, String fileName) {
        loadCollectionName(collectionName);
        DeleteParam deleteParam = DeleteParam.newBuilder()
                .withCollectionName(collectionName)
                .withExpr(String.format("file_name == \"%s\"", fileName))
                .build();
//        milvusClient.delete(deleteParam);
    }

    /**
     * 删除知识库所有文件
     */
    public void removeAllDocument(String collectionName) {
        loadCollectionName(collectionName);
        DeleteParam deleteParam = DeleteParam.newBuilder()
                .withCollectionName(collectionName)
                .withExpr("id > 0")
                .build();
//        milvusClient.delete(deleteParam);
    }

    /**
     * 加载知识库Collection
     */
    public void loadCollectionName(String collectionName) {
        LoadCollectionReq loadCollectionReq = LoadCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        milvusClient.loadCollection(loadCollectionReq);
    }

    /**
     * 释放知识库Collection
     */
    public void releaseCollectionName(String collectionName) {
        ReleaseCollectionReq releaseCollectionReq = ReleaseCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        milvusClient.releaseCollection(releaseCollectionReq);
    }

}