using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using Milvus.Client;

namespace MilvusDemo
{
    /// <summary>
    /// 提供与Milvus向量数据库交互的服务类
    /// </summary>
    public class MilvusService
    {
        private readonly MilvusClient _milvusClient;
        private readonly string _collectionName;
        private readonly int _dim;

        /// <summary>
        /// 初始化MilvusService的新实例
        /// </summary>
        /// <param name="host">Milvus服务器主机名</param>
        /// <param name="port">Milvus服务器端口</param>
        /// <param name="collectionName">集合名称</param>
        /// <param name="dim">向量维度</param>
        public MilvusService(string host, int port, string collectionName, int dim)
        {
            _milvusClient = new MilvusClient(host, port);
            _collectionName = collectionName;
            _dim = dim;
        }

        /// <summary>
        /// 初始化集合，如果存在则删除并重新创建
        /// </summary>
        public async Task InitializeCollectionAsync()
        {
            // 检查集合是否存在，如果存在则删除
            if (await _milvusClient.HasCollectionAsync(_collectionName))
            {
                var collectionToDrop = _milvusClient.GetCollection(_collectionName);
                await collectionToDrop.DropAsync();
                Console.WriteLine($"Collection '{_collectionName}' dropped.");
            }

            // 创建新的集合
            await _milvusClient.CreateCollectionAsync(
                _collectionName,
                new[]
                {
                    FieldSchema.Create<long>("id", isPrimaryKey: true),
                    FieldSchema.CreateVarchar("varchar", 256),
                    FieldSchema.CreateFloatVector("vector", _dim)
                });
            Console.WriteLine($"Collection '{_collectionName}' created.");
        }

        /// <summary>
        /// 创建索引并加载集合
        /// </summary>
        public async Task CreateIndexAndLoadAsync()
        {
            var collection = _milvusClient.GetCollection(_collectionName);
            
            // 创建索引
            await collection.CreateIndexAsync("vector", IndexType.Flat, SimilarityMetricType.L2);
            Console.WriteLine("Index created.");

            // 加载集合
            await collection.LoadAsync();
            Console.WriteLine("Collection loaded.");
        }

        /// <summary>
        /// 插入随机向量数据
        /// </summary>
        /// <param name="count">要插入的向量数量</param>
        public async Task InsertRandomDataAsync(int count)
        {
            var collection = _milvusClient.GetCollection(_collectionName);
            var random = new Random();
            
            // 生成随机向量
            var vectors = new List<ReadOnlyMemory<float>>();
            for (int i = 0; i < count; i++)
            {
                var vector = new float[_dim];
                for (int j = 0; j < _dim; j++)
                {
                    vector[j] = (float)random.NextDouble();
                }
                vectors.Add(vector);
            }

            // 生成ID和字符串字段
            var ids = Enumerable.Range(0, count).Select(i => (long)i).ToList();
            var varchars = ids.Select(i => $"varchar_{i}").ToList();

            // 打包数据
            var insertData = new List<FieldData>
            {
                FieldData.Create("id", ids),
                FieldData.Create("varchar", varchars),
                FieldData.CreateFloatVector("vector", vectors)
            };

            // 执行插入
            await collection.InsertAsync(insertData);
            Console.WriteLine("Data inserted.");
        }

        /// <summary>
        /// 执行向量相似度搜索
        /// </summary>
        /// <param name="topK">返回最相似的结果数量</param>
        /// <returns>搜索结果</returns>
        public async Task<SearchResults> SearchVectorAsync(int topK = 5)
        {
            var collection = _milvusClient.GetCollection(_collectionName);
            var random = new Random();
            
            // 生成随机目标向量
            var targetVector = new float[_dim];
            for (int j = 0; j < _dim; j++)
            {
                targetVector[j] = (float)random.NextDouble();
            }

            // 设置搜索参数
            var searchParameters = new SearchParameters();
            searchParameters.OutputFields.Add("id");
            searchParameters.OutputFields.Add("varchar");

            // 执行搜索
            var searchResults = await collection.SearchAsync(
                "vector", 
                new[] { (ReadOnlyMemory<float>)targetVector },
                SimilarityMetricType.L2,
                topK,
                searchParameters);

            return searchResults;
        }
    }
}