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.client.MilvusClient;
import io.milvus.grpc.*;
import io.milvus.param.*;
import io.milvus.param.collection.*;
import io.milvus.param.dml.DeleteParam;
import io.milvus.param.highlevel.collection.ListCollectionsParam;
import io.milvus.param.highlevel.collection.response.ListCollectionsResponse;
import io.milvus.param.highlevel.dml.QuerySimpleParam;
import io.milvus.param.highlevel.dml.response.QueryResponse;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.QueryResultsWrapper;
import org.springframework.stereotype.Service;

import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.TimeZone;
import java.util.stream.Collectors;

@Service
public class MilvusServiceImpl implements IMilvusService {

    /**
     * 新建知识库Collection（含Schema、Field、Index）
     */
    @Override
    public R<RpcStatus> createCollection(MilvusClient milvusClient, String collectionName, String description, int dimension) {
        FieldType idField = FieldType.newBuilder()
                .withName("id")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(true)
                .build();

        FieldType fileNameField = FieldType.newBuilder()
                .withName("file_name")
                .withDataType(DataType.VarChar)
                .withMaxLength(256)
                .build();

        FieldType fieldTypeField = FieldType.newBuilder()
                .withName("file_type")
                .withDataType(DataType.VarChar)
                .withMaxLength(256)
                .build();

        FieldType contentField = FieldType.newBuilder()
                .withName("content")
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build();

        FieldType vectorField = FieldType.newBuilder()
                .withName("embedding")
                .withDataType(DataType.FloatVector)
                .withDimension(dimension)
                .build();

        CreateCollectionParam createCollectionParam = CreateCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .withDescription(description)
                .addFieldType(idField)
                .addFieldType(fileNameField)
                .addFieldType(fieldTypeField)
                .addFieldType(contentField)
                .addFieldType(vectorField)
                .build();

        R<RpcStatus> createResult = milvusClient.createCollection(createCollectionParam);

        // 创建索引
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName(collectionName)
                .withFieldName("embedding")
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.COSINE)
                .withExtraParam("{\"nlist\": 64}")
                .build();

        milvusClient.createIndex(createIndexParam);
        return createResult;
    }

    /**
     * 查询知识库Collection
     */
    public JSONArray getCollectionNames(MilvusClient milvusClient) {
        JSONArray jsonArray = new JSONArray();
        ListCollectionsParam listParam = ListCollectionsParam.newBuilder().build();
        ListCollectionsResponse listResponse = milvusClient.listCollections(listParam).getData();
        if (listResponse != null) {
            List<String> collectionNames = milvusClient.listCollections(listParam).getData().collectionNames;
            for (String collectionName : collectionNames) {
                JSONObject jsonObject = new JSONObject();
                DescribeCollectionParam describeParam = DescribeCollectionParam.newBuilder()
                                                        .withCollectionName(collectionName)
                                                        .build();
                DescribeCollectionResponse describeResponse = milvusClient.describeCollection(describeParam).getData();
                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.getCreatedUtcTimestamp());
                jsonObject.put("createdTime", beijingTime);
                jsonObject.put("description", describeResponse.getSchema().getDescription());
                jsonArray.add(jsonObject);
            }
        }
        return jsonArray;
    }

    /**
     * 修改知识库Collection
     */
    public void collectionRename(MilvusClient milvusClient, String collectionName, String newCollectionName) {
        RenameCollectionParam renameCollectionReq = RenameCollectionParam.newBuilder()
                .withOldCollectionName(collectionName)
                .withNewCollectionName(newCollectionName)
                .build();

        milvusClient.renameCollection(renameCollectionReq);
    }

    /**
     * 删除知识库Collection
     */
    public void deleteCollectionName(MilvusClient milvusClient, String collectionName) {
        DropCollectionParam param = DropCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        milvusClient.dropCollection(param);
    }

    /**
     * 查询知识库文件
     */
    public List<String> listDocument(MilvusClient milvusClient, String collectionName, String fileType) {
        List<String> documentList = new ArrayList<>();
        loadCollectionName(milvusClient,collectionName);
        QuerySimpleParam queryParam = QuerySimpleParam.newBuilder()        
                .withCollectionName(collectionName)
                .withFilter("id > 0")
                .withOutputFields(Arrays.asList("file_type", "file_name"))
                .withLimit(16384L)
                .build();
        if (fileType != null && !fileType.equals(""))
            queryParam = QuerySimpleParam.newBuilder()
                .withCollectionName(collectionName)
                .withFilter(String.format("file_type == \"%s\"", fileType))
                .withOutputFields(Arrays.asList("file_type", "file_name"))
                .withLimit(16384L)
                .build();
        QueryResponse queryResults = milvusClient.query(queryParam).getData();
        List<QueryResultsWrapper.RowRecord> rowRecordList = new ArrayList<>();
        if (queryResults != null) {
            rowRecordList = queryResults.getRowRecords();
            for (QueryResultsWrapper.RowRecord rowRecord : rowRecordList) {
                documentList.add(rowRecord.get("file_name").toString());
            }
        }
        releaseCollectionName(milvusClient, collectionName);
        return documentList.stream().distinct().collect(Collectors.toList());
    }

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

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

    /**
     * 加载知识库Collection
     */
    public void loadCollectionName(MilvusClient milvusClient, String collectionName) {
        LoadCollectionParam param = LoadCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        milvusClient.loadCollection(param);
    }

    /**
     * 释放知识库Collection
     */
    public void releaseCollectionName(MilvusClient milvusClient, String collectionName) {
        ReleaseCollectionParam param = ReleaseCollectionParam.newBuilder()
                .withCollectionName(collectionName)
                .build();
        milvusClient.releaseCollection(param);
    }

}