/*
 *  Copyright (c) 2023-2026, Easy-Agents (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.easyagents.store.milvus;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.easyagents.core.document.Document;
import com.easyagents.core.store.DocumentStore;
import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult;
import com.easyagents.core.util.CollectionUtil;
import com.easyagents.core.util.Maps;
import com.easyagents.core.util.StringUtil;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.exception.MilvusClientException;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.net.URI;
import java.util.*;

/**
 * Milvus vector store based on Milvus Java SDK v2.
 */
public class MilvusVectorStore extends DocumentStore implements AutoCloseable {

    private static final Logger LOG = LoggerFactory.getLogger(MilvusVectorStore.class);
    private static final long LOAD_TIMEOUT_MS = 30_000L;
    private static final long LOAD_POLL_INTERVAL_MS = 200L;

    private static final String FIELD_ID = "id";
    private static final String FIELD_CONTENT = "content";
    private static final String FIELD_METADATA = "metadata";
    private static final String FIELD_VECTOR = "vector";

    private final MilvusClientV2 client;
    private final MilvusVectorStoreConfig config;
    private final String defaultCollectionName;
    private final Set<String> initializedCollections = Collections.synchronizedSet(new HashSet<String>());
    private final Set<String> loadedCollections = Collections.synchronizedSet(new HashSet<String>());

    public MilvusVectorStore(MilvusVectorStoreConfig config) {
        this.config = config;
        this.defaultCollectionName = config.getDefaultCollectionName();
        String uri = normalizeAndValidateUri(config.getUri());
        String dbName = StringUtil.hasText(config.getDatabaseName()) ? config.getDatabaseName().trim() : "default";

        ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder()
            .uri(uri)
            .dbName(dbName);

        if (StringUtil.hasText(config.getToken())) {
            builder.token(config.getToken().trim());
        }

        if (StringUtil.hasText(config.getUsername()) && StringUtil.hasText(config.getPassword())) {
            builder.username(config.getUsername().trim());
            builder.password(config.getPassword().trim());
        }

        ConnectConfig connectConfig = builder.build();
        this.client = new MilvusClientV2(connectConfig);
    }

    private String normalizeAndValidateUri(String uri) {
        if (StringUtil.noText(uri)) {
            throw new IllegalArgumentException("Milvus uri is required. Example: http://127.0.0.1:19530");
        }

        String normalized = uri.trim();
        if (!normalized.contains("://")) {
            normalized = "http://" + normalized;
        }

        URI parsed;
        try {
            parsed = URI.create(normalized);
        } catch (Exception e) {
            throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530", e);
        }

        if (StringUtil.noText(parsed.getHost()) || parsed.getPort() <= 0) {
            throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530");
        }

        return normalized;
    }

    @Override
    public StoreResult doStore(List<Document> documents, StoreOptions options) {
        if (CollectionUtil.noItems(documents)) {
            return StoreResult.success();
        }
        String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
        if (StringUtil.noText(collectionName)) {
            throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
        }

        int dimension = getDimension(documents);
        ensureCollectionExists(collectionName, dimension);

        try {
            InsertReq.InsertReqBuilder<?, ?> builder = InsertReq.builder();
            if (StringUtil.hasText(options.getPartitionName())) {
                builder.partitionName(options.getPartitionName());
            }
            InsertReq insertReq = builder
                .collectionName(collectionName)
                .data(toMilvusDocuments(documents))
                .build();
            client.insert(insertReq);
            return StoreResult.successWithIds(documents);
        } catch (MilvusClientException e) {
            return StoreResult.fail();
        }
    }

    @Override
    public StoreResult doDelete(Collection<?> ids, StoreOptions options) {
        if (CollectionUtil.noItems(ids)) {
            return StoreResult.success();
        }
        String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
        if (StringUtil.noText(collectionName)) {
            throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
        }

        try {
            DeleteReq.DeleteReqBuilder<?, ?> builder = DeleteReq.builder();
            if (StringUtil.hasText(options.getPartitionName())) {
                builder.partitionName(options.getPartitionName());
            }
            DeleteReq deleteReq = builder
                .collectionName(collectionName)
                .ids(MilvusPrimaryKeySupport.normalize(ids))
                .build();
            client.delete(deleteReq);
            return StoreResult.success();
        } catch (Exception e) {
            LOG.error("Milvus delete failed. collection={}, message={}",
                collectionName, e.getMessage(), e);
            return StoreResult.fail(e.getMessage());
        }
    }

    @Override
    public StoreResult doUpdate(List<Document> documents, StoreOptions options) {
        if (CollectionUtil.noItems(documents)) {
            return StoreResult.success();
        }
        String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
        if (StringUtil.noText(collectionName)) {
            throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
        }

        int dimension = getDimension(documents);
        ensureCollectionExists(collectionName, dimension);

        try {
            UpsertReq upsertReq = UpsertReq.builder()
                .collectionName(collectionName)
                .partitionName(options.getPartitionName())
                .data(toMilvusDocuments(documents))
                .build();
            client.upsert(upsertReq);
            return StoreResult.successWithIds(documents);
        } catch (Exception e) {
            return StoreResult.fail();
        }
    }

    @Override
    public List<Document> doSearch(SearchWrapper wrapper, StoreOptions options) {
        String collectionName = options.getCollectionNameOrDefault(defaultCollectionName);
        if (StringUtil.noText(collectionName)) {
            throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
        }
        ensureCollectionLoaded(collectionName);

        if (wrapper.getVector() == null || wrapper.getVector().length == 0) {
            return queryByCondition(wrapper, options, collectionName);
        }
        return searchByVector(wrapper, options, collectionName);
    }

    private List<Document> searchByVector(SearchWrapper wrapper, StoreOptions options, String collectionName) {
        SearchReq searchReq = buildSearchReq(wrapper, options, collectionName);
        try {
            SearchResp resp = client.search(searchReq);
            return parseSearchResults(resp, wrapper.getMinScore());
        } catch (Exception e) {
            if (isCollectionNotLoaded(e)) {
                loadedCollections.remove(collectionName);
                try {
                    ensureCollectionLoaded(collectionName);
                    SearchResp retryResp = client.search(searchReq);
                    return parseSearchResults(retryResp, wrapper.getMinScore());
                } catch (Exception retryException) {
                    LOG.warn("Milvus search retry failed after load. collection={}, message={}", collectionName, retryException.getMessage());
                    return Collections.emptyList();
                }
            }
            LOG.warn("Milvus search failed. collection={}, message={}", collectionName, e.getMessage());
            return Collections.emptyList();
        }
    }

    private List<Document> queryByCondition(SearchWrapper wrapper, StoreOptions options, String collectionName) {
        QueryReq queryReq = buildQueryReq(wrapper, options, collectionName);
        try {
            QueryResp resp = client.query(queryReq);
            return parseQueryResults(resp);
        } catch (Exception e) {
            if (isCollectionNotLoaded(e)) {
                loadedCollections.remove(collectionName);
                try {
                    ensureCollectionLoaded(collectionName);
                    QueryResp retryResp = client.query(queryReq);
                    return parseQueryResults(retryResp);
                } catch (Exception retryException) {
                    LOG.warn("Milvus query retry failed after load. collection={}, message={}", collectionName, retryException.getMessage());
                    return Collections.emptyList();
                }
            }
            LOG.warn("Milvus query failed. collection={}, message={}", collectionName, e.getMessage());
            return Collections.emptyList();
        }
    }

    private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) {
        SearchReq.SearchReqBuilder<?, ?> builder = SearchReq.builder()
            .collectionName(collectionName)
            .consistencyLevel(ConsistencyLevel.STRONG)
            .outputFields(getOutputFields(wrapper))
            .topK(wrapper.getMaxResults())
            .annsField(FIELD_VECTOR)
            .data(Collections.singletonList(toFloatList(wrapper.getVector())))
            .searchParams(Maps.of("ef", 64));

        if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) {
            builder.partitionNames(options.getPartitionNamesOrEmpty());
        }
        String filterExpression = wrapper.toFilterExpression(MilvusExpressionAdaptor.DEFAULT);
        if (StringUtil.hasText(filterExpression)) {
            builder.filter(filterExpression);
        }
        return builder.build();
    }

    private QueryReq buildQueryReq(SearchWrapper wrapper, StoreOptions options, String collectionName) {
        QueryReq.QueryReqBuilder<?, ?> builder = QueryReq.builder()
            .collectionName(collectionName)
            .consistencyLevel(ConsistencyLevel.STRONG)
            .outputFields(getOutputFields(wrapper));

        if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) {
            builder.partitionNames(options.getPartitionNamesOrEmpty());
        }
        String filterExpression = wrapper.toFilterExpression(MilvusExpressionAdaptor.DEFAULT);
        if (StringUtil.hasText(filterExpression)) {
            builder.filter(filterExpression);
        }
        return builder.build();
    }

    private List<Document> parseSearchResults(SearchResp resp, Double minScore) {
        List<List<SearchResp.SearchResult>> results = resp.getSearchResults();
        if (CollectionUtil.noItems(results)) {
            return Collections.emptyList();
        }

        List<Document> documents = new ArrayList<Document>();
        for (List<SearchResp.SearchResult> resultList : results) {
            if (CollectionUtil.noItems(resultList)) {
                continue;
            }
            for (SearchResp.SearchResult result : resultList) {
                Document document = toDocument(result.getEntity());
                if (document == null) {
                    continue;
                }
                document.setId(result.getId());
                Float distance = result.getDistance();
                if (distance != null) {
                    double score = (distance + 1.0d) / 2.0d;
                    document.setScore(score);
                }
                if (minScore == null || document.getScore() == null || document.getScore() >= minScore) {
                    documents.add(document);
                }
            }
        }
        return documents;
    }

    private List<Document> parseQueryResults(QueryResp resp) {
        List<QueryResp.QueryResult> results = resp.getQueryResults();
        if (CollectionUtil.noItems(results)) {
            return Collections.emptyList();
        }
        List<Document> documents = new ArrayList<Document>(results.size());
        for (QueryResp.QueryResult result : results) {
            Document document = toDocument(result.getEntity());
            if (document != null) {
                documents.add(document);
            }
        }
        return documents;
    }

    private Document toDocument(Map<String, Object> entity) {
        if (entity == null || entity.isEmpty()) {
            return null;
        }

        Document document = new Document();
        document.setId(entity.get(FIELD_ID));

        Object vectorObj = entity.get(FIELD_VECTOR);
        if (vectorObj instanceof List) {
            @SuppressWarnings("unchecked")
            List<? extends Number> vectorList = (List<? extends Number>) vectorObj;
            document.setVector(vectorList);
        }

        Object contentObj = entity.get(FIELD_CONTENT);
        if (contentObj != null) {
            document.setContent(contentObj.toString());
        }

        Object metadataObj = entity.get(FIELD_METADATA);
        if (metadataObj instanceof Map) {
            @SuppressWarnings("unchecked")
            Map<String, Object> metadata = (Map<String, Object>) metadataObj;
            document.addMetadata(metadata);
        } else if (metadataObj != null) {
            @SuppressWarnings("unchecked")
            Map<String, Object> metadata = JSON.parseObject(JSON.toJSONString(metadataObj), Map.class);
            document.addMetadata(metadata);
        }

        return document;
    }

    private List<JSONObject> toMilvusDocuments(List<Document> documents) {
        List<JSONObject> rows = new ArrayList<JSONObject>(documents.size());
        for (Document doc : documents) {
            JSONObject row = new JSONObject();
            row.put(FIELD_ID, String.valueOf(doc.getId()));
            row.put(FIELD_CONTENT, doc.getContent());
            row.put(FIELD_VECTOR, toFloatList(doc.getVector()));
            Map<String, Object> metadatas = doc.getMetadataMap();
            row.put(FIELD_METADATA, metadatas == null ? new JSONObject() : new JSONObject(metadatas));
            rows.add(row);
        }
        return rows;
    }

    private List<String> getOutputFields(SearchWrapper wrapper) {
        List<String> outputFields = wrapper.isOutputVector()
            ? new ArrayList<String>(Arrays.asList(FIELD_ID, FIELD_VECTOR, FIELD_CONTENT, FIELD_METADATA))
            : new ArrayList<String>(Arrays.asList(FIELD_ID, FIELD_CONTENT, FIELD_METADATA));

        if (CollectionUtil.hasItems(wrapper.getOutputFields())) {
            outputFields.addAll(wrapper.getOutputFields());
        }
        return outputFields;
    }

    private List<Float> toFloatList(float[] vector) {
        List<Float> values = new ArrayList<Float>(vector.length);
        for (float item : vector) {
            values.add(item);
        }
        return values;
    }

    private int getDimension(List<Document> documents) {
        for (Document document : documents) {
            float[] vector = document.getVector();
            if (vector != null && vector.length > 0) {
                return vector.length;
            }
        }
        if (this.getEmbeddingModel() != null) {
            return this.getEmbeddingModel().dimensions();
        }
        throw new IllegalStateException("Unable to determine vector dimension for Milvus collection.");
    }

    private void ensureCollectionExists(String collectionName, int dimension) {
        if (initializedCollections.contains(collectionName)) {
            return;
        }
        synchronized (initializedCollections) {
            if (initializedCollections.contains(collectionName)) {
                return;
            }
            Boolean exists = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build());
            if (Boolean.TRUE.equals(exists)) {
                initializedCollections.add(collectionName);
                return;
            }
            if (!config.isAutoCreateCollection()) {
                throw new IllegalStateException("Milvus collection not found and autoCreateCollection is disabled: " + collectionName);
            }
            createCollection(collectionName, dimension);
            initializedCollections.add(collectionName);
        }
    }

    private void ensureCollectionLoaded(String collectionName) {
        if (loadedCollections.contains(collectionName)) {
            return;
        }
        synchronized (loadedCollections) {
            if (loadedCollections.contains(collectionName)) {
                return;
            }
            boolean loaded = false;
            try {
                loaded = Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()));
            } catch (Exception e) {
                LOG.warn("Milvus getLoadState failed. collection={}, message={}", collectionName, e.getMessage());
            }

            if (!loaded) {
                client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build());
                waitForCollectionLoaded(collectionName);
            }
            loadedCollections.add(collectionName);
        }
    }

    private void waitForCollectionLoaded(String collectionName) {
        long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS;
        while (System.currentTimeMillis() < deadline) {
            if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) {
                return;
            }
            try {
                Thread.sleep(LOAD_POLL_INTERVAL_MS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new IllegalStateException("Interrupted while loading Milvus collection: " + collectionName, e);
            }
        }
        throw new IllegalStateException("Timeout waiting for Milvus collection loaded: " + collectionName);
    }

    private boolean isCollectionNotLoaded(Exception e) {
        Throwable current = e;
        while (current != null) {
            String message = current.getMessage();
            if (message != null && message.toLowerCase().contains("collection not loaded")) {
                return true;
            }
            current = current.getCause();
        }
        return false;
    }

    private void createCollection(String collectionName, int dimension) {
        List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<CreateCollectionReq.FieldSchema>();
        fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
            .name(FIELD_ID)
            .dataType(DataType.VarChar)
            .maxLength(64)
            .isPrimaryKey(true)
            .autoID(false)
            .build());
        fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
            .name(FIELD_CONTENT)
            .dataType(DataType.VarChar)
            .maxLength(65535)
            .build());
        fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
            .name(FIELD_METADATA)
            .dataType(DataType.JSON)
            .build());
        fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
            .name(FIELD_VECTOR)
            .dataType(DataType.FloatVector)
            .dimension(dimension)
            .build());

        CreateCollectionReq.CollectionSchema collectionSchema = CreateCollectionReq.CollectionSchema.builder()
            .fieldSchemaList(fieldSchemaList)
            .build();

        List<IndexParam> indexParams = new ArrayList<IndexParam>();
        indexParams.add(IndexParam.builder()
            .fieldName(FIELD_VECTOR)
            .indexName(FIELD_VECTOR + "_hnsw_idx")
            .indexType(IndexParam.IndexType.HNSW)
            .metricType(IndexParam.MetricType.COSINE)
            .extraParams(Maps.of("M", 16).set("efConstruction", 200))
            .build());

        CreateCollectionReq createCollectionReq = CreateCollectionReq.builder()
            .collectionName(collectionName)
            .collectionSchema(collectionSchema)
            .primaryFieldName(FIELD_ID)
            .vectorFieldName(FIELD_VECTOR)
            .description("Easy-Agents Milvus Vector Store")
            .indexParams(indexParams)
            .build();
        client.createCollection(createCollectionReq);
        ensureCollectionLoaded(collectionName);
    }

    public MilvusClientV2 getClient() {
        return client;
    }

    public boolean checkAvailable() {
        try {
            return client.hasCollection(HasCollectionReq.builder()
                .collectionName("__milvus_boot_probe__")
                .build()) != null;
        } catch (Exception e) {
            LOG.warn("Milvus availability check failed. message={}", e.getMessage());
            return false;
        }
    }

    @Override
    public void close() {
        try {
            client.close(1L);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            LOG.warn("Interrupted while closing Milvus client. uri={}", config.getUri(), e);
        }
    }
}
