using System.Diagnostics;
using Microsoft.Extensions.Logging;
using Serilog.Context;
using ZB.MOM.WW.CBDD.Bson;
using ZB.MOM.WW.CBDD.Core.Collections;
using ZB.MOM.WW.CBDD.Core.Compression;
using ZB.MOM.WW.CBDD.Core.Storage;

namespace ZB.MOM.WW.CBDD.Tests.Benchmark;

internal static class DatabaseSizeBenchmark
{
    private static readonly int[] TargetCounts = [10_000, 1_000_000, 10_000_000];
    private static readonly Scenario[] Scenarios =
    [
        new(
            "Uncompressed",
            CompressionOptions.Default),
        new(
            "Compressed-BrotliFast",
            new CompressionOptions
            {
                EnableCompression = true,
                MinSizeBytes = 256,
                MinSavingsPercent = 0,
                Codec = CompressionCodec.Brotli,
                Level = System.IO.Compression.CompressionLevel.Fastest
            })
    ];

    private const int BatchSize = 50_000;
    private const int ProgressInterval = 1_000_000;

    public static void Run(ILogger logger)
    {
        var results = new List<SizeResult>(TargetCounts.Length * Scenarios.Length);

        logger.LogInformation("=== CBDD Database Size Benchmark ===");
        logger.LogInformation("Targets: {Targets}", string.Join(", ", TargetCounts.Select(x => x.ToString("N0"))));
        logger.LogInformation("Scenarios: {Scenarios}", string.Join(", ", Scenarios.Select(x => x.Name)));
        logger.LogInformation("Batch size: {BatchSize:N0}", BatchSize);

        foreach (var targetCount in TargetCounts)
        {
            foreach (var scenario in Scenarios)
            {
                var dbPath = Path.Combine(Path.GetTempPath(), $"cbdd_size_{scenario.Name}_{targetCount}_{Guid.NewGuid():N}.db");
                var walPath = Path.ChangeExtension(dbPath, ".wal");
                using var _ = LogContext.PushProperty("TargetCount", targetCount);
                using var __ = LogContext.PushProperty("Scenario", scenario.Name);

                logger.LogInformation("Starting scenario {Scenario} for target {TargetCount:N0} docs", scenario.Name, targetCount);

                var insertStopwatch = Stopwatch.StartNew();
                CompressionStats compressionStats = default;
                CompactionStats compactionStats;
                long preCompactDbBytes;
                long preCompactWalBytes;
                long postCompactDbBytes;
                long postCompactWalBytes;

                using (var storage = new StorageEngine(dbPath, PageFileConfig.Default, scenario.CompressionOptions))
                using (var transactionHolder = new BenchmarkTransactionHolder(storage))
                {
                    var collection = new DocumentCollection<SizeBenchmarkDocument>(
                        storage,
                        transactionHolder,
                        new SizeBenchmarkDocumentMapper());

                    var inserted = 0;
                    while (inserted < targetCount)
                    {
                        var currentBatchSize = Math.Min(BatchSize, targetCount - inserted);
                        var documents = new SizeBenchmarkDocument[currentBatchSize];
                        var baseValue = inserted;

                        for (var i = 0; i < currentBatchSize; i++)
                        {
                            documents[i] = CreateDocument(baseValue + i);
                        }

                        collection.InsertBulk(documents);
                        transactionHolder.CommitAndReset();

                        inserted += currentBatchSize;
                        if (inserted == targetCount || inserted % ProgressInterval == 0)
                        {
                            logger.LogInformation("Inserted {Inserted:N0}/{TargetCount:N0}", inserted, targetCount);
                        }
                    }

                    insertStopwatch.Stop();
                    preCompactDbBytes = File.Exists(dbPath) ? new FileInfo(dbPath).Length : 0;
                    preCompactWalBytes = File.Exists(walPath) ? new FileInfo(walPath).Length : 0;

                    compactionStats = storage.Compact(new CompactionOptions
                    {
                        EnableTailTruncation = true,
                        DefragmentSlottedPages = true,
                        NormalizeFreeList = true
                    });

                    postCompactDbBytes = File.Exists(dbPath) ? new FileInfo(dbPath).Length : 0;
                    postCompactWalBytes = File.Exists(walPath) ? new FileInfo(walPath).Length : 0;
                    compressionStats = storage.GetCompressionStats();
                }

                var result = new SizeResult(
                    scenario.Name,
                    targetCount,
                    insertStopwatch.Elapsed,
                    preCompactDbBytes,
                    preCompactWalBytes,
                    postCompactDbBytes,
                    postCompactWalBytes,
                    compactionStats,
                    compressionStats);
                results.Add(result);

                logger.LogInformation(
                    "Completed {Scenario} {TargetCount:N0} docs in {Elapsed}. pre={PreTotal}, post={PostTotal}, shrink={Shrink}, compRatio={CompRatio}",
                    scenario.Name,
                    targetCount,
                    insertStopwatch.Elapsed,
                    FormatBytes(result.PreCompactTotalBytes),
                    FormatBytes(result.PostCompactTotalBytes),
                    FormatBytes(result.ShrinkBytes),
                    result.CompressionRatioText);

                TryDelete(dbPath);
                TryDelete(walPath);
            }
        }

        logger.LogInformation("=== Size Benchmark Summary ===");
        foreach (var result in results.OrderBy(x => x.TargetCount).ThenBy(x => x.Scenario))
        {
            logger.LogInformation(
                "{Scenario,-22} | {Count,12:N0} docs | insert={Elapsed,12} | pre={Pre,12} | post={Post,12} | shrink={Shrink,12} | compact={CompactBytes,12} | ratio={Ratio}",
                result.Scenario,
                result.TargetCount,
                result.InsertElapsed,
                FormatBytes(result.PreCompactTotalBytes),
                FormatBytes(result.PostCompactTotalBytes),
                FormatBytes(result.ShrinkBytes),
                FormatBytes(result.CompactionStats.ReclaimedFileBytes),
                result.CompressionRatioText);
        }
    }

    private static SizeBenchmarkDocument CreateDocument(int value)
    {
        return new SizeBenchmarkDocument
        {
            Id = ObjectId.NewObjectId(),
            Value = value,
            Name = $"doc-{value:D8}"
        };
    }

    private static void TryDelete(string path)
    {
        if (File.Exists(path))
        {
            File.Delete(path);
        }
    }

    private static string FormatBytes(long bytes)
    {
        string[] units = ["B", "KB", "MB", "GB", "TB"];
        double size = bytes;
        var unitIndex = 0;

        while (size >= 1024 && unitIndex < units.Length - 1)
        {
            size /= 1024;
            unitIndex++;
        }

        return $"{size:N2} {units[unitIndex]}";
    }

    private sealed record Scenario(string Name, CompressionOptions CompressionOptions);

    private sealed record SizeResult(
        string Scenario,
        int TargetCount,
        TimeSpan InsertElapsed,
        long PreCompactDbBytes,
        long PreCompactWalBytes,
        long PostCompactDbBytes,
        long PostCompactWalBytes,
        CompactionStats CompactionStats,
        CompressionStats CompressionStats)
    {
        public long PreCompactTotalBytes => PreCompactDbBytes + PreCompactWalBytes;
        public long PostCompactTotalBytes => PostCompactDbBytes + PostCompactWalBytes;
        public long ShrinkBytes => PreCompactTotalBytes - PostCompactTotalBytes;

        public string CompressionRatioText =>
            CompressionStats.BytesAfterCompression > 0
                ? $"{(double)CompressionStats.BytesBeforeCompression / CompressionStats.BytesAfterCompression:N2}x"
                : "n/a";
    }

    private sealed class SizeBenchmarkDocument
    {
        public ObjectId Id { get; set; }
        public int Value { get; set; }
        public string Name { get; set; } = string.Empty;
    }

    private sealed class SizeBenchmarkDocumentMapper : ObjectIdMapperBase<SizeBenchmarkDocument>
    {
        public override string CollectionName => "size_documents";

        public override ObjectId GetId(SizeBenchmarkDocument entity) => entity.Id;

        public override void SetId(SizeBenchmarkDocument entity, ObjectId id) => entity.Id = id;

        public override int Serialize(SizeBenchmarkDocument entity, BsonSpanWriter writer)
        {
            var sizePos = writer.BeginDocument();
            writer.WriteObjectId("_id", entity.Id);
            writer.WriteInt32("value", entity.Value);
            writer.WriteString("name", entity.Name);
            writer.EndDocument(sizePos);
            return writer.Position;
        }

        public override SizeBenchmarkDocument Deserialize(BsonSpanReader reader)
        {
            var document = new SizeBenchmarkDocument();

            reader.ReadDocumentSize();
            while (reader.Remaining > 0)
            {
                var bsonType = reader.ReadBsonType();
                if (bsonType == BsonType.EndOfDocument)
                {
                    break;
                }

                var name = reader.ReadElementHeader();
                switch (name)
                {
                    case "_id":
                        document.Id = reader.ReadObjectId();
                        break;
                    case "value":
                        document.Value = reader.ReadInt32();
                        break;
                    case "name":
                        document.Name = reader.ReadString();
                        break;
                    default:
                        reader.SkipValue(bsonType);
                        break;
                }
            }

            return document;
        }
    }
}
