JiShe.CollectBus/src/JiShe.CollectBus.IoTDBProvider/Provider/IoTDBProvider.cs

using Apache.IoTDB;
using Apache.IoTDB.DataStructure;
using JiShe.CollectBus.Common.Extensions;
using JiShe.CollectBus.Common.Helpers;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Threading.Tasks;

namespace JiShe.CollectBus.IoTDBProvider
{
    /// <summary>
    /// IoTDB数据源
    /// </summary>
    public class IoTDBProvider : IIoTDBProvider, IDisposable
    {
        private readonly IoTDBOptions _options;
        private readonly TableSessionPool _sessionPool;
        private static readonly ConcurrentDictionary<Type, DeviceMetadata> _metadataCache = new();
        private readonly ILogger<IoTDBProvider> _logger;

        public IoTDBProvider(IOptions<IoTDBOptions> options, ILogger<IoTDBProvider> logger)
        {
            _options = options.Value;

            _sessionPool = new TableSessionPool.Builder()
            .SetNodeUrls(_options.ClusterList)
            .SetUsername(_options.UserName)
            .SetPassword(_options.Password)
            .SetFetchSize(_options.PoolSize)
            .Build();

            _sessionPool.Open(false).Wait();
            _logger = logger;
        }

        /// <summary>
        /// 插入数据
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="entity"></param>
        /// <returns></returns>
        public async Task InsertAsync<T>(T entity, int buildTabletMode) where T : IoTEntity
        {
            var metadata = GetMetadata<T>();

            var tablet = BuildTablet(new[] { entity }, metadata, buildTabletMode);
            await _sessionPool.InsertAsync(tablet);
        }

        /// <summary>
        /// 批量插入数据
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="entities"></param>
        /// <returns></returns>
        public async Task BatchInsertAsync<T>(IEnumerable<T> entities, int buildTabletMode) where T : IoTEntity
        {
            var metadata = GetMetadata<T>();

            var batchSize = 1000;
            var batches = entities.Chunk(batchSize);

            foreach (var batch in batches)
            {
                var tablet = BuildTablet(batch, metadata, buildTabletMode);
                await _sessionPool.InsertAsync(tablet);
            }
        }


        /// <summary>
        /// 删除数据
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="options"></param>
        /// <returns></returns>
        public async Task<object> DeleteAsync<T>(QueryOptions options) where T : IoTEntity
        {
            var query = BuildDeleteSQL<T>(options);
            var sessionDataSet = await _sessionPool.ExecuteQueryStatementAsync(query);

            if (!sessionDataSet.HasNext())
            {
                _logger.LogWarning($"{typeof(T).Name} 删除数据时，没有返回受影响记录数量。");
                return 0;
            }

            //获取唯一结果行
            var row = sessionDataSet.Next();
            return row.Values[0];
        }




        /// <summary>
        /// 查询数据
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="options"></param>
        /// <returns></returns>
        public async Task<PagedResult<T>> QueryAsync<T>(QueryOptions options) where T : IoTEntity, new()
        {
            var query = BuildQuerySQL<T>(options);
            var sessionDataSet = await _sessionPool.ExecuteQueryStatementAsync(query);

            var result = new PagedResult<T>
            {
                TotalCount = await GetTotalCount<T>(options),
                Items = ParseResults<T>(sessionDataSet, options.PageSize)
            };

            return result;
        }

        /// <summary>
        /// 构建表模型
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="entities">表实体</param>
        /// <param name="metadata">设备元数据</param>
        /// <param name="buildTabletMode">构建表模型方式，1 根据实体《T》直接显示指定Tag，2根据实体《T》的名称指定表名</param>
        /// <returns></returns>
        private Tablet BuildTablet<T>(IEnumerable<T> entities, DeviceMetadata metadata, int buildTabletMode) where T : IoTEntity
        {
            var timestamps = new List<long>();
            var values = new List<List<object>>();

            foreach (var entity in entities)
            {
                timestamps.Add(entity.Timestamps);
                var rowValues = new List<object>();
                foreach (var measurement in metadata.ColumnNames)
                {
                    var value = typeof(T).GetProperty(measurement)?.GetValue(entity);
                    if (value == null)
                    {
                        throw new Exception($"{nameof(BuildTablet)} 构建表模型{typeof(T).Name}时，属性{measurement}值为空，不符合IoTDB设计标准，请赋值以后重新处理。");
                    }
                    rowValues.Add(value);
                }
                values.Add(rowValues);
            }

            if (buildTabletMode == 1)
            {
                return new Tablet(
                    DevicePathBuilder.GetDeviceId(entities.First()),
                    metadata.ColumnNames,
                    metadata.ColumnCategories,
                    metadata.DataTypes,
                    values,
                    timestamps
                );
            }
            else if (buildTabletMode == 2)
            {
                return new Tablet(
                    DevicePathBuilder.GetTableName<T>(),
                    metadata.ColumnNames,
                    metadata.ColumnCategories,
                    metadata.DataTypes,
                    values,
                    timestamps
                );
            }
            else
            {
                throw new Exception($"{nameof(BuildTablet)} 构建表模型{typeof(T).Name}时，buildTabletMode参数值不正确，请赋值以后重新处理。");
            }
        }

        /// <summary>
        /// 构建查询语句
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="options"></param>
        /// <returns></returns>
        private string BuildQuerySQL<T>(QueryOptions options) where T : IoTEntity
        {
            var metadata = GetMetadata<T>();
            var sb = new StringBuilder("SELECT ");
            sb.AppendJoin(", ", metadata.ColumnNames);
            sb.Append($" FROM {options.TableNameOrTagName}");

            if (options.Conditions.Any())
            {
                sb.Append(" WHERE ");
                sb.AppendJoin(" AND ", options.Conditions.Select(TranslateCondition));
            }

            sb.Append($" LIMIT {options.PageSize} OFFSET {options.Page * options.PageSize}");
            return sb.ToString();
        }

        /// <summary>
        /// 构建删除语句
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="options"></param>
        /// <returns></returns>
        private string BuildDeleteSQL<T>(QueryOptions options) where T : IoTEntity
        {
            var metadata = GetMetadata<T>();
            var sb = new StringBuilder("DELETE ");

            sb.Append($" FROM {options.TableNameOrTagName}");

            sb.AppendJoin(", ", metadata.ColumnNames);

            if (options.Conditions.Any())
            {
                sb.Append(" WHERE ");
                sb.AppendJoin(" AND ", options.Conditions.Select(TranslateCondition));
            }

            return sb.ToString();
        }

        /// <summary>
        /// 将查询条件转换为SQL语句
        /// </summary>
        /// <param name="condition"></param>
        /// <returns></returns>
        /// <exception cref="NotSupportedException"></exception>
        private string TranslateCondition(QueryCondition condition)
        {
            return condition.Operator switch
            {
                ">" => $"{condition.Field} > {condition.Value}",
                "<" => $"{condition.Field} < {condition.Value}",
                "=" => $"{condition.Field} = '{condition.Value}'",
                _ => throw new NotSupportedException($"Operator {condition.Operator} not supported")
            };
        }

        /// <summary>
        /// 获取查询条件的总数量
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="options"></param>
        /// <returns></returns>
        private async Task<int> GetTotalCount<T>(QueryOptions options) where T : IoTEntity
        {
            var countQuery = $"SELECT COUNT(*) FROM {options.TableNameOrTagName}";
            if (options.Conditions.Any())
            {
                countQuery += " WHERE " + string.Join(" AND ", options.Conditions.Select(TranslateCondition));
            }

            var result = await _sessionPool.ExecuteQueryStatementAsync(countQuery);
            return result.HasNext() ? Convert.ToInt32(result.Next().Values[0]) : 0;
        }

        /// <summary>
        /// 解析查询结果
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <param name="dataSet"></param>
        /// <param name="pageSize"></param>
        /// <returns></returns>
        private IEnumerable<T> ParseResults<T>(SessionDataSet dataSet, int pageSize) where T : IoTEntity, new()
        {
            var results = new List<T>();
            var metadata = GetMetadata<T>();

            var properties = typeof(T).GetProperties();

            while (dataSet.HasNext() && results.Count < pageSize)
            {
                var record = dataSet.Next();
                var entity = new T
                {
                    Timestamps = record.Timestamps
                };


                foreach (var measurement in metadata.ColumnNames)
                {
                    var value = record.Values;

                    var prop = properties.FirstOrDefault(p =>
                        p.Name.Equals(measurement, StringComparison.OrdinalIgnoreCase));
                    if (prop != null)
                    {
                        typeof(T).GetProperty(measurement)?.SetValue(entity, value);
                    }

                }

                results.Add(entity);
            }
            return results;
        }

        /// <summary>
        /// 释放资源
        /// </summary>
        public void Dispose()
        {
            _sessionPool?.Close().Wait();
        }

        /// <summary>
        /// 获取设备元数据
        /// </summary>
        /// <typeparam name="T"></typeparam>
        /// <returns></returns>
        private DeviceMetadata GetMetadata<T>() where T : IoTEntity
        {
            return _metadataCache.GetOrAdd(typeof(T), type =>
            {
                var columns = CollectColumnMetadata(type);
                var metadata = BuildDeviceMetadata(columns);
                return metadata;
            });

            //return _metadataCache.GetOrAdd(typeof(T), type =>
            //{
            //    var metadata = new DeviceMetadata();
            //    List<Tuple<string, ColumnCategory, TSDataType>> columns = new();
            //    foreach (var prop in type.GetProperties())
            //    {
            //        //标签列
            //        var attrTAG = prop.GetCustomAttribute<TAGColumnAttribute>();
            //        if (attrTAG != null)
            //        {
            //            columns.Add(Tuple.Create(prop.PropertyType.Name, ColumnCategory.TAG, GetDataTypeFromStr(prop.PropertyType.Name)));
            //            continue;
            //        }

            //        //属性列
            //        var attrATTRIBUTE = prop.GetCustomAttribute<ATTRIBUTEColumnAttribute>();
            //        if (attrATTRIBUTE != null)
            //        {
            //            columns.Add(Tuple.Create(prop.PropertyType.Name, ColumnCategory.ATTRIBUTE, GetDataTypeFromStr(prop.PropertyType.Name)));
            //            continue;
            //        }

            //        //数据列
            //        var attrFIELD = prop.GetCustomAttribute<FIELDColumnAttribute>();
            //        if (attrFIELD != null)
            //        {
            //            columns.Add(Tuple.Create(prop.PropertyType.Name, ColumnCategory.FIELD, GetDataTypeFromStr(prop.PropertyType.Name)));
            //        }
            //    }
            //    var columnCategories = EnumExtensions.ToEnumDictionary<ColumnCategory>();
            //    foreach (var item in columnCategories)
            //    {
            //        if (item.Value == ColumnCategory.TAG)
            //        {
            //            metadata.ColumnNames.AddRange(columns.Where(d => d.Item2 == ColumnCategory.FIELD).Select(d => d.Item1).ToList());
            //            metadata.ColumnCategories.AddRange(columns.Where(d => d.Item2 == ColumnCategory.TAG).Select(d => d.Item2).ToList());
            //            metadata.DataTypes.AddRange(columns.Where(d => d.Item2 == ColumnCategory.TAG).Select(d => d.Item3).ToList());
            //        }
            //        else if (item.Value == ColumnCategory.ATTRIBUTE)
            //        {
            //            metadata.ColumnNames.AddRange(columns.Where(d => d.Item2 == ColumnCategory.ATTRIBUTE).Select(d => d.Item1).ToList());
            //            metadata.ColumnCategories.AddRange(columns.Where(d => d.Item2 == ColumnCategory.ATTRIBUTE).Select(d => d.Item2).ToList());
            //            metadata.DataTypes.AddRange(columns.Where(d => d.Item2 == ColumnCategory.ATTRIBUTE).Select(d => d.Item3).ToList());
            //        }
            //        else if (item.Value == ColumnCategory.FIELD)
            //        {
            //            metadata.ColumnNames.AddRange(columns.Where(d => d.Item2 == ColumnCategory.FIELD).Select(d => d.Item1).ToList());
            //            metadata.ColumnCategories.AddRange(columns.Where(d => d.Item2 == ColumnCategory.FIELD).Select(d => d.Item2).ToList());
            //            metadata.DataTypes.AddRange(columns.Where(d => d.Item2 == ColumnCategory.FIELD).Select(d => d.Item3).ToList());
            //        }
            //    }

            //    return metadata;
            //});
        }

        /// <summary>
        /// 获取设备元数据的列
        /// </summary>
        /// <param name="type"></param>
        /// <returns></returns>
        private List<ColumnInfo> CollectColumnMetadata(Type type)
        {
            var columns = new List<ColumnInfo>();

            foreach (var prop in type.GetProperties())
            {
                //按优先级顺序检查属性，避免重复反射
                ColumnInfo? column = prop.GetCustomAttribute<TAGColumnAttribute>() is not null ? new ColumnInfo(
                    name: prop.Name,  //使用属性名
                    category: ColumnCategory.TAG,
                    dataType: GetDataTypeFromTypeName(prop.PropertyType.Name)
                ) : prop.GetCustomAttribute<ATTRIBUTEColumnAttribute>() is not null ? new ColumnInfo(
                    prop.Name,
                    ColumnCategory.ATTRIBUTE,
                    GetDataTypeFromTypeName(prop.PropertyType.Name)
                ) : prop.GetCustomAttribute<FIELDColumnAttribute>() is not null ? new ColumnInfo(
                    prop.Name,
                    ColumnCategory.FIELD,
                    GetDataTypeFromTypeName(prop.PropertyType.Name)
                ) : null;

                if (column.HasValue)
                {
                    columns.Add(column.Value);
                }
            }
            return columns;
        }

        /// <summary>
        /// 构建设备元数据
        /// </summary>
        /// <param name="columns"></param>
        /// <returns></returns>
        private DeviceMetadata BuildDeviceMetadata(List<ColumnInfo> columns)
        {
            var metadata = new DeviceMetadata();

            //按业务逻辑顺序处理（TAG -> FIELD -> ATTRIBUTE）
            var groupedColumns = columns
                .GroupBy(c => c.Category)
                .ToDictionary(g => g.Key, g => g.ToList());

            ProcessCategory(groupedColumns, ColumnCategory.TAG, metadata);
            ProcessCategory(groupedColumns, ColumnCategory.ATTRIBUTE, metadata);
            ProcessCategory(groupedColumns, ColumnCategory.FIELD, metadata);

            return metadata;
        }

        /// <summary>
        /// 处理不同列类型的逻辑
        /// </summary>
        /// <param name="groupedColumns"></param>
        /// <param name="category"></param>
        /// <param name="metadata"></param>
        private void ProcessCategory(IReadOnlyDictionary<ColumnCategory, List<ColumnInfo>> groupedColumns, ColumnCategory category, DeviceMetadata metadata)
        {
            if (groupedColumns.TryGetValue(category, out var cols))
            {
                metadata.ColumnNames.AddRange(cols.Select(c => c.Name));
                metadata.ColumnCategories.AddRange(cols.Select(c => c.Category));
                metadata.DataTypes.AddRange(cols.Select(c => c.DataType));
            }
        }

        /// <summary>
        /// 数据列结构
        /// </summary>
        private readonly struct ColumnInfo
        {
            public string Name { get; }
            public ColumnCategory Category { get; }
            public TSDataType DataType { get; }

            public ColumnInfo(string name, ColumnCategory category, TSDataType dataType)
            {
                Name = name;
                Category = category;
                DataType = dataType;
            }
        }

        /// <summary>
        /// 根据类型名称获取对应的 IoTDB 数据类型
        /// </summary>
        /// <param name="typeName">类型名称（不区分大小写）</param>
        /// <returns>对应的 TSDataType，默认返回 TSDataType.STRING</returns>
        private TSDataType GetDataTypeFromTypeName(string typeName)
        {
            if (string.IsNullOrWhiteSpace(typeName))
                return TSDataType.STRING; 

            return DataTypeMap.TryGetValue(typeName.Trim(), out var dataType)
                ? dataType
                : TSDataType.STRING;
        }

        /// <summary>
        /// 根据类型名称获取 IoTDB 数据类型
        /// </summary>
        private readonly IReadOnlyDictionary<string, TSDataType> DataTypeMap =
            new Dictionary<string, TSDataType>(StringComparer.OrdinalIgnoreCase)
            {
                ["BOOLEAN"] = TSDataType.BOOLEAN,
                ["INT32"] = TSDataType.INT32,
                ["INT64"] = TSDataType.INT64,
                ["FLOAT"] = TSDataType.FLOAT,
                ["DOUBLE"] = TSDataType.DOUBLE,
                ["TEXT"] = TSDataType.TEXT,
                ["NULLTYPE"] = TSDataType.NONE,
                ["TIMESTAMP"] = TSDataType.TIMESTAMP,
                ["DATE"] = TSDataType.DATE,
                ["BLOB"] = TSDataType.BLOB,
                ["STRING"] = TSDataType.STRING
            };
    }
}