using System.Collections.Concurrent;
using System.Data;
using Kusto.Cloud.Platform.Data;
using Kusto.Cloud.Platform.Utils;
using Kusto.Data;
using Kusto.Data.Common;
using Kusto.Data.Net.Client;
using Kusto.Ingest;
using Microsoft.Extensions.Logging;
using Newtonsoft.Json;

public class KustoIngestService
{
    private readonly IKustoIngestClient _ingestClient;
    private readonly ICslAdminProvider _adminClient;
    private readonly string _database;
    private readonly Microsoft.Extensions.Logging.ILogger<KustoIngestService> _logger;

    public KustoIngestService(string kustoUri, string database, string appId, string appKey, string tenantId, Microsoft.Extensions.Logging.ILogger<KustoIngestService> logger)
    {
        var kustoCsb = new KustoConnectionStringBuilder(kustoUri)
            .WithAadApplicationKeyAuthentication(appId, appKey, tenantId);

        _ingestClient = KustoIngestFactory.CreateDirectIngestClient(kustoCsb);
        _adminClient = KustoClientFactory.CreateCslAdminProvider(kustoCsb);
        _database = database;
        _logger = logger;
    }

    public async Task SaveToKustoInBatchAsync<T>(T[] dataToWrite, string table)
    {
        foreach (var parition in Program.Partition(dataToWrite, 150))
        {
            var dataTable = parition.ToDataTable();
            await IngestDataAsync(dataTable, table);
        }

        await IngestDataAsync(dataToWrite.ToDataTable(), table);
    }

    public async Task IngestDataAsync(DataTable dataTable, string tableName)
    {
        try
        {
            _logger.LogInformation("Starting ingestion of {RowCount} rows into table {TableName}", dataTable.Rows.Count, tableName);
            var ingestionProperties = new KustoIngestionProperties(_database, tableName);
            await _ingestClient.IngestFromDataReaderAsync(dataTable.CreateDataReader(), ingestionProperties);
            _logger.LogInformation("Successfully ingested {RowCount} rows into table {TableName}", dataTable.Rows.Count, tableName);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error during ingestion into table {TableName}", tableName);
            throw;
        }
    }

    public async Task EnsureTableExistsAsync<T>(string tableName)
    {
        _logger.LogInformation("Ensuring table {TableName} exists.", tableName);

        var schema = await GetDatabaseSchemaAsync();
        if (schema == null)
        {
            _logger.LogWarning("Database schema is null, creating new table {TableName}", tableName);
            await CreateMergeTable<T>(tableName);
        }
        else if (!schema.Tables.TryGetValue(tableName, out var tableSchema))
        {
            _logger.LogInformation("Table {TableName} does not exist. Creating it.", tableName);
            await CreateMergeTable<T>(tableName);
        }
        else
        {
            _logger.LogInformation("Table {TableName} exists. Verifying schema.", tableName);
            var existingColumns = new HashSet<string>(tableSchema.OrderedColumns.Select(c => c.Name));
            var newColumns = typeof(T).GetProperties().Select(p => p.Name).Where(c => !existingColumns.Contains(c));

            if (newColumns.Any())
            {
                _logger.LogWarning("New columns detected in table {TableName}, updating schema.", tableName);
                await CreateMergeTable<T>(tableName);
            }
        }
    }

    private async Task<DatabaseSchema> GetDatabaseSchemaAsync()
    {
        _logger.LogInformation("Fetching database schema for {Database}", _database);
        try
        {
            var command = CslCommandGenerator.GenerateDatabaseSchemaShowAsJsonCommand(_database);
            using var reader = await _adminClient.ExecuteControlCommandAsync(_database, command);
            var schemaJson = reader.ToEnumerable<string>().FirstOrDefault();
            var clusterSchema = JsonConvert.DeserializeObject<ClusterSchema>(schemaJson);
            return clusterSchema?.Databases.GetValueOrDefault(_database);
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to retrieve database schema for {Database}", _database);
            throw;
        }
    }

    private async Task CreateMergeTable<T>(string tableName)
    {
        _logger.LogInformation("Creating or merging table schema for {TableName}", tableName);
        var columns = typeof(T).GetProperties().Select(p => ColumnSchema.FromNameAndCslType(p.Name, ConvertTypeToCslType(p.PropertyType)));
        var command = CslCommandGenerator.GenerateTableCreateMergeCommand(new TableSchema(tableName, columns));
        await _adminClient.ExecuteControlCommandAsync(_database, command);
        _logger.LogInformation("Table {TableName} schema created/merged successfully", tableName);
    }

    private static string ConvertTypeToCslType(Type type) => type switch
    {
        _ when type == typeof(bool) => "bool",
        _ when type == typeof(DateTime) => "datetime",
        _ when type == typeof(Guid) => "guid",
        _ when type == typeof(int) => "int",
        _ when type == typeof(long) => "long",
        _ when type == typeof(double) => "real",
        _ when type == typeof(string) => "string",
        _ when type == typeof(TimeSpan) => "timespan",
        _ when type.IsEnum => "string",
        _ when type == typeof(byte[]) => "string",
        _ => "dynamic"
    };
}

// 数据模型
public class MyData
{
    public int Id { get; set; }
    public string Name { get; set; }
    public DateTime Timestamp { get; set; }
}

// 启动入口
public class Program
{
    public static async Task Main()
    {
        var loggerFactory = LoggerFactory.Create(builder => builder.AddConsole());
        var logger = loggerFactory.CreateLogger<KustoIngestService>();

        var kustoUri = "https://testkustoanduin.eastasia.kusto.windows.net";
        var database = "yourdatabase";
        var appId = "1c4b1053-5aef-45dc-9d42-a69ea84b422d";
        var appKey = "7fm8Q~CxycQfV3qW5v6G71vRGVCuYM2eLfZ0obC5";
        var tenantId = "538bf3b7-8deb-4179-9c63-f09d13f65838";

        var ingestService = new KustoIngestService(
            kustoUri: kustoUri,
            database: database,
            appId: appId,
            appKey: appKey,
            tenantId: tenantId,
            logger: logger
        );

        await ingestService.EnsureTableExistsAsync<MyData>("your-kusto-table");

        Console.WriteLine("Generating records...");
        var data = GenerateTestData(10_0000);

        Console.WriteLine("Ingesting data into Kusto...");
        var kustoBuffer = new KustoBuffer<MyData>(ingestService, "your-kusto-table");
        foreach (var parition in Partition(data, 2000))
        {
            kustoBuffer.Add(parition.ToArray());
        }
        Console.ReadLine();
    }

    private static List<MyData> GenerateTestData(int count)
    {
        var data = new List<MyData>(count);
        var now = DateTime.UtcNow;

        Parallel.For(0, count, i =>
        {
            var item = new MyData
            {
                Id = i + 1,
                Name = $"Item_{i}",
                Timestamp = now.AddSeconds(i)
            };

            lock (data) { data.Add(item); }
        });

        return data;
    }

    public static IEnumerable<List<T>> Partition<T>(IEnumerable<T> source, int maxSize)
    {
        if (source == null)
        {
            throw new ArgumentNullException(nameof(source));
        }
        if (maxSize <= 0)
        {
            throw new ArgumentOutOfRangeException(nameof(maxSize), "MaxSize must larger than 0");
        }

        var buffer = new List<T>();
        foreach (var item in source)
        {
            buffer.Add(item);
            if (buffer.Count == maxSize)
            {
                yield return buffer;
                buffer = new List<T>();
            }
        }
        if (buffer.Count > 0)
        {
            yield return buffer;
        }
    }
}

/// <summary>
/// KustoBuffer provides a non-blocking way to buffer incoming data and batch write them to Kusto.
/// This design is inspired by the proven ArrayDb implementation:
/// 
/// https://github.com/AiursoftWeb/ArrayDb/blob/master/src/Aiursoft.ArrayDb.WriteBuffer/BufferedObjectBucket.cs
/// </summary>
public class KustoBuffer<T>
{
    // Dependency: Kusto service for persisting data
    private readonly KustoIngestService _kustoService;
    private readonly string _tableName;

    // Configurable parameters: maximum sleep time when cold and threshold to stop sleeping
    private readonly int _maxSleepMilliSecondsWhenCold;
    private readonly int _stopSleepingWhenWriteBufferItemsMoreThan;

    // Engine tasks controlling the write process
    private Task _engine = Task.CompletedTask;
    private Task _coolDownEngine = Task.CompletedTask;

    // Locks for protecting buffer swapping and engine status switching
    private readonly ReaderWriterLockSlim _bufferLock = new ReaderWriterLockSlim();
    private readonly object _bufferWriteSwapLock = new object();
    private readonly object _engineStatusSwitchLock = new object();

    // Double buffering: _activeBuffer is used for enqueuing new data,
    // while _secondaryBuffer is used for swapping and persisting data.
    private ConcurrentQueue<T> _activeBuffer = new ConcurrentQueue<T>();
    private ConcurrentQueue<T> _secondaryBuffer = new ConcurrentQueue<T>();

    public bool IsCold => _engine.IsCompleted && _coolDownEngine.IsCompleted;

    public bool IsHot => !IsCold;

    public KustoBuffer(
        KustoIngestService kustoService,
        string tableName,
        int maxSleepMilliSecondsWhenCold = 2000,
        int stopSleepingWhenWriteBufferItemsMoreThan = 1000)
    {
        _kustoService = kustoService ?? throw new ArgumentNullException(nameof(kustoService));
        _tableName = tableName;
        _maxSleepMilliSecondsWhenCold = maxSleepMilliSecondsWhenCold;
        _stopSleepingWhenWriteBufferItemsMoreThan = stopSleepingWhenWriteBufferItemsMoreThan;
    }

    private static int CalculateSleepTime(double maxSleepMilliSecondsWhenCold, double stopSleepingWhenWriteBufferItemsMoreThan, int writeBufferItemsCount)
    {
        if (stopSleepingWhenWriteBufferItemsMoreThan <= 0)
        {
            throw new ArgumentException("B must be a positive number.");
        }

        if (writeBufferItemsCount > stopSleepingWhenWriteBufferItemsMoreThan)
        {
            return 0;
        }

        var y = maxSleepMilliSecondsWhenCold * (1 - (Math.Log(1 + writeBufferItemsCount) / Math.Log(1 + stopSleepingWhenWriteBufferItemsMoreThan)));
        return (int)y;
    }

    /// <summary>
    /// Adds objects to the buffer in a non-blocking manner.
    /// This method wakes up the engine to persist the data.
    /// </summary>
    public void Add(params T[] objs)
    {
        lock (_bufferWriteSwapLock)
        {
            foreach (var obj in objs)
            {
                _activeBuffer.Enqueue(obj);
            }
        }

        if (objs.Length == 0)
        {
            return;
        }

        // Get the engine status in advanced to avoid lock contention.
        if (!IsCold)
        {
            // Most of the cases in a high-frequency environment, the engine is still running.
            return;
        }

        // Engine might be sleeping. Wake it up.
        lock (_engineStatusSwitchLock)
        {
            // Avoid multiple threads to wake up the engine at the same time.
            if (!IsCold)
            {
                return;
            }

            _engine = Task.Run(WriteBuffered);
        }
    }

    /// <summary>
    /// Persists the buffered data in batches to the Kusto database.
    /// This method is guaranteed to be executed by a single thread at a time.
    /// </summary>
    private async Task WriteBuffered()
    {
        _bufferLock.EnterWriteLock();
        try
        {
            ConcurrentQueue<T> bufferToPersist;
            lock (_bufferWriteSwapLock)
            {
                // Swap active and secondary buffers
                bufferToPersist = _activeBuffer;
                _activeBuffer = _secondaryBuffer;
                _secondaryBuffer = new ConcurrentQueue<T>();
            }

            var dataToWrite = bufferToPersist.ToArray();

            // Release the buffer to avoid memory leak.
            bufferToPersist.Clear();

            // If there is data, batch write to Kusto
            if (dataToWrite.Length > 0)
            {
                // Process the buffer to persist
                await _kustoService.SaveToKustoInBatchAsync(dataToWrite, _tableName);
            }
        }
        finally
        {
            _bufferLock.ExitWriteLock();
        }

        // While we are writing, new data may be added to the buffer. If so, we need to write it too.
        if (!_activeBuffer.IsEmpty)
        {
            // Restart the engine to write the new added data.
            // Before engine quits, it wakes up cool down engine to ensure the engine will be restarted.
            // Before cool down quits, it wakes up the engine to ensure the engine will be restarted.
            // So if one of the two tasks is running, the engine will be restarted. And this buffer is in a hot state.
            // In a hot state, you don't have to start the engine again.
            _coolDownEngine = Task.Run(async () =>
            {
                // Slow down a little bit to wait for more data to come.
                // If we persist too fast, and the buffer is almost empty, frequent write will cause data fragmentation.
                // If we persist too slow, and a lot of new data has been added to the buffer, and the engine wasted time in sleeping.
                // So the time to sleep is calculated by the number of items in the buffer.
                var sleepTime = CalculateSleepTime(
                    _maxSleepMilliSecondsWhenCold,
                    _stopSleepingWhenWriteBufferItemsMoreThan,
                    _activeBuffer.Count);
                await Task.Delay(sleepTime);

                // Wake up the engine to write the new added data.
                _engine = Task.Run(WriteBuffered);
            });
        }
    }

    /// <summary>
    /// Synchronizes the buffer, ensuring that all pending data is persisted.
    /// </summary>
    public async Task SyncAsync()
    {
        // Case 1:
        // Engine is working. However, the buffer may still have data that after this phase, the data is still not written.
        // Wait two rounds of engine to finish to ensure all data is written.
        // Cool down engine will ensure restart the engine to write the remaining data.
        // Wait for the engine to finish.

        // Case 2:
        // The engine is not working. In this case, it might be in the cool down phase.
        // The first wait is just await a completed task.
        // Cool down engine will ensure restart the engine to write the remaining data.
        // Then wait for the engine to finish.
        await _engine;
        await _coolDownEngine;
        await _engine;
    }
}