package com.cmc.agent.service;

import com.google.gson.Gson;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;

import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.MutationResult;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.R;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;
import okhttp3.*;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.*;
import java.util.stream.Collectors;

public class LangChainMilvusService {
    private static final String LLM_SERVICE_URL = "http://localhost:8000/generate";
    private final MilvusServiceClient milvusClient;
    private final String collectionName;
    private  EmbeddingModel embeddingModel;

    /**
     * 连接milvus知识库，加入langchain自带emdding
     */
    public LangChainMilvusService(String host, int port, String collectionName, EmbeddingModel embeddingModel) {
        this.milvusClient = new MilvusServiceClient(
                ConnectParam.newBuilder()
                        .withHost(host)
                        .withPort(port)
                        .build()
        );
        this.collectionName = collectionName;
        this.embeddingModel = embeddingModel;
    }

    /**
     * 新建知识库Collection（含Schema、Field、Index）
     */
    public void createCollection(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 contentField = FieldType.newBuilder()
                .withName("content")
                .withDataType(DataType.VarChar)
                .withMaxLength(65535)
                .build();

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

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

        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);
    }

    public void insertLangchainEmbeddingDocument(File file) throws IOException {
        // 加载文档
        InputStream fileInputStream = new FileInputStream(file);
        Document document = new ApachePdfBoxDocumentParser().parse(fileInputStream);
        DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(150,10);;
        List<TextSegment> segments = splitter.split(document);

        // 提取文本和生成嵌入
        List<String> texts = new ArrayList<>();
        List<List<Float>> embeddings = new ArrayList<>();

        for (TextSegment segment : segments) {
            String text = segment.text();
            if (text.trim().isEmpty())
                continue;
            texts.add(text);
            embeddings.add(embeddingModel.embed(text).content().vectorAsList());
        }

        // 准备插入数据
        List<InsertParam.Field> fields = new ArrayList<>();
        fields.add(new InsertParam.Field("file_name", Arrays.asList(file.getName())));
        fields.add(new InsertParam.Field("content", Arrays.asList(texts)));
        fields.add(new InsertParam.Field("embedding", Arrays.asList(embeddings)));

        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName(collectionName)
                .withFields(fields)
                .build();

        // 执行插入
        R<MutationResult> insertResult = milvusClient.insert(insertParam);

        if (insertResult.getStatus() != R.Status.Success.getCode()) {
            throw new RuntimeException("Failed to insert document: " + insertResult.getMessage());
        }
    }

    // 从Milvus检索相关文档
    public List<String> retrieveFromMilvus(String query, int topK) throws IOException {
        List<Float> queryVector = embeddingModel.embed(query).content().vectorAsList();

        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName(collectionName)
                .withVectors(queryVector)
                .withTopK(topK)
                .withOutFields(Arrays.asList("content"))
                .build();

        R<SearchResults> response = milvusClient.search(searchParam);
        SearchResultsWrapper wrapper = new SearchResultsWrapper(response.getData().getResults());

        return wrapper.getRowRecords(0).stream()
                .map(record -> (String) record.get("content"))
                .collect(Collectors.toList());
    }

    // 调用LLM生成回答
    public String generateAnswer(String question, List<String> contexts) throws IOException {
        String prompt = buildPrompt(question, contexts);
        Gson gson = new Gson();
        Map<String, String> hashMap = new HashMap<>();
        hashMap.put("prompt", prompt);
        hashMap.put("max_tokens", "512");
        RequestBody body = RequestBody.create(
                MediaType.parse("application/json"),
                new Gson().toJson(hashMap));

        Request request = new Request.Builder()
                .url(LLM_SERVICE_URL)
                .post(body)
                .build();

        try (Response response = new OkHttpClient().newCall(request).execute()) {
            return gson.fromJson(response.body().string(), Map.class).get("text").toString();
        }
    }

    private String buildPrompt(String question, List<String> contexts) {
        StringBuilder sb = new StringBuilder();
        sb.append("根据以下上下文回答问题：\n\n");
        for (int i = 0; i < contexts.size(); i++) {
            sb.append("上下文").append(i+1).append(": ").append(contexts.get(i)).append("\n\n");
        }
        sb.append("问题: ").append(question).append("\n回答: ");
        return sb.toString();
    }
    public void close() {
        milvusClient.close();
    }
}