/*
 *  Copyright (c) 2023-2025, Agents-Flex (fuhai999@gmail.com).
 *  <p>
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  <p>
 *  http://www.apache.org/licenses/LICENSE-2.0
 *  <p>
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package com.agentsflex.store.qdrant;

import com.agentsflex.core.document.Document;
import com.agentsflex.core.store.DocumentStore;
import com.agentsflex.core.store.SearchWrapper;
import com.agentsflex.core.store.StoreOptions;
import com.agentsflex.core.store.StoreResult;
import com.agentsflex.core.util.CollectionUtil;
import com.agentsflex.core.util.StringUtil;
import com.agentsflex.core.util.VectorUtil;
import io.grpc.Grpc;
import io.grpc.TlsChannelCredentials;
import io.qdrant.client.QdrantClient;
import io.grpc.ManagedChannel;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections;
import io.qdrant.client.grpc.JsonWithInt;
import io.qdrant.client.grpc.Points;
import io.qdrant.client.grpc.Points.Filter;
import io.qdrant.client.grpc.Points.PointId;
import io.qdrant.client.grpc.Points.PointStruct;
import io.qdrant.client.grpc.Points.QueryPoints;
import io.qdrant.client.grpc.Points.ScoredPoint;

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

import static io.qdrant.client.ConditionFactory.matchKeyword;
import static io.qdrant.client.PointIdFactory.id;
import static io.qdrant.client.QueryFactory.nearest;
import static io.qdrant.client.ValueFactory.value;
import static io.qdrant.client.VectorsFactory.vectors;
import static io.qdrant.client.WithPayloadSelectorFactory.enable;

public class QdrantVectorStore extends DocumentStore {

    private final QdrantVectorStoreConfig config;
    private final QdrantClient            client;
    private final String                  defaultCollectionName;
    private       boolean                 isCreateCollection = false;

    public QdrantVectorStore(QdrantVectorStoreConfig config) throws IOException {
        this.config = config;
        this.defaultCollectionName = config.getDefaultCollectionName();
        String uri = config.getUri();
        int port = 6334;
        QdrantGrpcClient.Builder builder;
        if (StringUtil.hasText(config.getCaPath())) {
            ManagedChannel channel = Grpc.newChannelBuilder(
                uri,
                TlsChannelCredentials.newBuilder().trustManager(new File(config.getCaPath())).build()
            ).build();
            builder = QdrantGrpcClient.newBuilder(channel, true);
        } else {
            if (uri.contains(":")) {
                uri = uri.split(":")[0];
                port = Integer.parseInt(uri.split(":")[1]);
            }
            builder = QdrantGrpcClient.newBuilder(uri, port, false);
        }
        if (StringUtil.hasText(config.getApiKey())) {
            builder.withApiKey(config.getApiKey());
        }
        this.client = new QdrantClient(builder.build());
    }

    @Override
    public StoreResult storeInternal(List<Document> documents, StoreOptions options) {
        List<PointStruct> points = new ArrayList<>();
        int size = 1024;
        for (Document doc : documents) {
            size = doc.getVector().length;
            Map<String, JsonWithInt.Value> payload = new HashMap<>();
            payload.put("content", value(doc.getContent()));
            points.add(PointStruct.newBuilder()
                .setId(id(Long.parseLong(doc.getId().toString())))
                .setVectors(vectors(VectorUtil.toFloatArray(doc.getVector())))
                .putAllPayload(payload)
                .build());
        }
        try {
            String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
            if (config.isAutoCreateCollection() && !isCreateCollection) {
                Boolean exists = client.collectionExistsAsync(collectionName).get();
                if (!exists) {
                    client.createCollectionAsync(collectionName, Collections.VectorParams.newBuilder()
                            .setDistance(Collections.Distance.Cosine)
                            .setSize(size)
                            .build())
                        .get();
                }
            } else {
                isCreateCollection = true;
            }
            if (CollectionUtil.hasItems(points)) {
                client.upsertAsync(collectionName, points).get();
            }
            return StoreResult.successWithIds(documents);
        } catch (Exception e) {
            return StoreResult.fail();
        }
    }

    @Override
    public StoreResult deleteInternal(Collection<?> ids, StoreOptions options) {
        try {
            String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
            List<PointId> pointIds = ids.stream()
                .map(id -> id((Long) id))
                .collect(Collectors.toList());
            client.deleteAsync(collectionName, pointIds).get();
            return StoreResult.success();
        } catch (Exception e) {
            return StoreResult.fail();
        }
    }

    @Override
    public StoreResult updateInternal(List<Document> documents, StoreOptions options) {
        try {
            List<PointStruct> points = new ArrayList<>();
            for (Document doc : documents) {
                Map<String, JsonWithInt.Value> payload = new HashMap<>();
                payload.put("content", value(doc.getContent()));
                points.add(PointStruct.newBuilder()
                    .setId(id(Long.parseLong(doc.getId().toString())))
                    .setVectors(vectors(VectorUtil.toFloatArray(doc.getVector())))
                    .putAllPayload(payload)
                    .build());
            }
            String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
            if (CollectionUtil.hasItems(points)) {
                client.upsertAsync(collectionName, points).get();
            }
            return StoreResult.successWithIds(documents);
        } catch (Exception e) {
            return StoreResult.fail();
        }
    }

    @Override
    public List<Document> searchInternal(SearchWrapper wrapper, StoreOptions options) {
        List<Document> documents = new ArrayList<>();
        try {
            String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
            QueryPoints.Builder query = QueryPoints.newBuilder()
                .setCollectionName(collectionName)
                .setLimit(wrapper.getMaxResults())
                .setWithVectors(Points.WithVectorsSelector.newBuilder().setEnable(true).build())
                .setWithPayload(enable(true));
            if (wrapper.getVector() != null) {
                query.setQuery(nearest(VectorUtil.toFloatArray(wrapper.getVector())));
            }
            if (StringUtil.hasText(wrapper.getText())) {
                query.setFilter(Filter.newBuilder().addMust(matchKeyword("content", wrapper.getText())));
            }
            List<ScoredPoint> data = client.queryAsync(query.build()).get();
            for (ScoredPoint point : data) {
                Document doc = new Document();
                doc.setId(point.getId().getNum());
                doc.setVector(VectorUtil.convertToVector(point.getVectors().getVector().getDataList()));
                doc.setContent(point.getPayloadMap().get("content").getStringValue());
                documents.add(doc);
            }
            return documents;
        } catch (Exception e) {
            return documents;
        }
    }

    public QdrantClient getClient() {
        return client;
    }
}
