namespace DbaMessageInserter;

public enum DeflateBlockType
{
    /// <summary>
    /// The encoder evaluates Stored, Fixed, and Dynamic candidates at the
    /// current bit position and writes the candidate with the smallest exact
    /// encoded bit count.
    /// </summary>
    Auto,

    Stored,
    Fixed,
    Dynamic,
    Npncb
}

public sealed record DeflateBlockSpec(DeflateBlockType Type, int Offset, int Length)
{
    public static DeflateBlockSpec CreateNpncb() =>
        new DeflateBlockSpec(DeflateBlockType.Npncb, 0, 0);
}

public static class DeflateBlockPlanner
{
    /// <summary>
    /// Splits the input into data regions. The physical DEFLATE type of every
    /// region is selected later by RawDeflateEncoder after exact Stored, Fixed,
    /// and data-dependent Dynamic costs have been calculated.
    /// </summary>
    public static List<DeflateBlockSpec> BuildAdaptivePlan(
        int inputLength,
        int dataBlockSize,
        int npncbCountAfterEachDataBlock,
        bool includeNpncbAfterFinalDataBlock)
    {
        return BuildPlan(
            inputLength,
            dataBlockSize,
            new[] { DeflateBlockType.Auto },
            npncbCountAfterEachDataBlock,
            includeNpncbAfterFinalDataBlock);
    }

    public static List<DeflateBlockSpec> BuildPatternPlan(
        int inputLength,
        int dataBlockSize,
        IReadOnlyList<DeflateBlockType> dataBlockPattern,
        bool insertNpncbBetweenDataBlocks)
    {
        return BuildPatternPlan(
            inputLength,
            dataBlockSize,
            dataBlockPattern,
            insertNpncbBetweenDataBlocks ? 1 : 0,
            false);
    }

    /// <summary>
    /// Cycles through an explicitly requested data-block pattern and inserts a
    /// train of NPNCB blocks after each data block. This method remains
    /// available for controlled experiments; normal operation uses
    /// BuildAdaptivePlan().
    /// </summary>
    public static List<DeflateBlockSpec> BuildPatternPlan(
        int inputLength,
        int dataBlockSize,
        IReadOnlyList<DeflateBlockType> dataBlockPattern,
        int npncbCountAfterEachDataBlock,
        bool includeNpncbAfterFinalDataBlock)
    {
        return BuildPlan(
            inputLength,
            dataBlockSize,
            dataBlockPattern,
            npncbCountAfterEachDataBlock,
            includeNpncbAfterFinalDataBlock);
    }

    public static void Validate(IReadOnlyList<DeflateBlockSpec> plan, int inputLength)
    {
        if (plan is null) throw new ArgumentNullException(nameof(plan));
        if (inputLength < 0) throw new ArgumentOutOfRangeException(nameof(inputLength));
        if (plan.Count == 0) throw new ArgumentException("The block plan cannot be empty.", nameof(plan));

        int expectedOffset = 0;
        for (int i = 0; i < plan.Count; i++)
        {
            DeflateBlockSpec block = plan[i];
            if (block.Type == DeflateBlockType.Npncb)
            {
                if (block.Offset != 0 || block.Length != 0)
                    throw new ArgumentException($"NPNCB {i} must have Offset=0 and Length=0.", nameof(plan));
                continue;
            }

            if (block.Offset != expectedOffset)
                throw new ArgumentException($"Block {i} starts at {block.Offset}; expected {expectedOffset}.", nameof(plan));
            if (block.Length < 0 || block.Offset < 0 || block.Offset > inputLength - block.Length)
                throw new ArgumentException($"Block {i} is outside the input range.", nameof(plan));
            if (block.Type == DeflateBlockType.Stored && block.Length > ushort.MaxValue)
                throw new ArgumentException($"Stored block {i} exceeds 65535 bytes.", nameof(plan));
            expectedOffset += block.Length;
        }

        if (expectedOffset != inputLength)
            throw new ArgumentException($"Plan covers {expectedOffset} bytes; input has {inputLength}.", nameof(plan));
    }

    private static List<DeflateBlockSpec> BuildPlan(
        int inputLength,
        int dataBlockSize,
        IReadOnlyList<DeflateBlockType> dataBlockPattern,
        int npncbCountAfterEachDataBlock,
        bool includeNpncbAfterFinalDataBlock)
    {
        ValidateBuildArguments(
            inputLength,
            dataBlockSize,
            dataBlockPattern,
            npncbCountAfterEachDataBlock);

        var plan = new List<DeflateBlockSpec>();

        if (inputLength == 0)
        {
            // A Fixed EOB-only block avoids confusing an empty data block with
            // an NPNCB. The optional NPNCB train still provides DBA capacity.
            plan.Add(new DeflateBlockSpec(DeflateBlockType.Fixed, 0, 0));
            if (includeNpncbAfterFinalDataBlock)
                AddNpncbTrain(plan, npncbCountAfterEachDataBlock);
            return plan;
        }

        int offset = 0;
        int patternIndex = 0;

        while (offset < inputLength)
        {
            int length = Math.Min(dataBlockSize, inputLength - offset);
            DeflateBlockType type = dataBlockPattern[patternIndex % dataBlockPattern.Count];
            plan.Add(new DeflateBlockSpec(type, offset, length));
            offset += length;
            patternIndex++;
            bool hasAnotherDataBlock = offset < inputLength;

            // ---------------------------------------------------------------
            // NPNCB insertion point. This is intentionally unchanged.
            // Comment out this if block to disable NPNCB insertion.
            // ---------------------------------------------------------------
            if (npncbCountAfterEachDataBlock > 0 &&
                (hasAnotherDataBlock || includeNpncbAfterFinalDataBlock))
            {
                AddNpncbTrain(plan, npncbCountAfterEachDataBlock);
            }
            // ---------------------------------------------------------------
        }

        return plan;
    }

    private static void ValidateBuildArguments(
        int inputLength,
        int dataBlockSize,
        IReadOnlyList<DeflateBlockType> dataBlockPattern,
        int npncbCountAfterEachDataBlock)
    {
        if (inputLength < 0) throw new ArgumentOutOfRangeException(nameof(inputLength));
        if (dataBlockSize <= 0 || dataBlockSize > ushort.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(dataBlockSize));
        if (dataBlockPattern is null) throw new ArgumentNullException(nameof(dataBlockPattern));
        if (dataBlockPattern.Count == 0)
            throw new ArgumentException("A data-block pattern is required.", nameof(dataBlockPattern));
        if (npncbCountAfterEachDataBlock < 0 || npncbCountAfterEachDataBlock > 100_000)
            throw new ArgumentOutOfRangeException(nameof(npncbCountAfterEachDataBlock));

        foreach (DeflateBlockType type in dataBlockPattern)
        {
            if (type == DeflateBlockType.Npncb)
                throw new ArgumentException("NPNCB is inserted separately.", nameof(dataBlockPattern));
        }
    }

    private static void AddNpncbTrain(ICollection<DeflateBlockSpec> plan, int count)
    {
        for (int i = 0; i < count; i++)
            plan.Add(DeflateBlockSpec.CreateNpncb());
    }
}

public sealed class DeflatePlanBuilder
{
    private readonly int _inputLength;
    private readonly List<DeflateBlockSpec> _blocks = new();
    private int _offset;

    public DeflatePlanBuilder(int inputLength)
    {
        if (inputLength < 0) throw new ArgumentOutOfRangeException(nameof(inputLength));
        _inputLength = inputLength;
    }

    public int CurrentInputOffset => _offset;
    public int RemainingInputBytes => _inputLength - _offset;

    public DeflatePlanBuilder AddDataBlock(DeflateBlockType type, int length)
    {
        if (type == DeflateBlockType.Npncb)
            throw new ArgumentException("Use AddNpncb().", nameof(type));
        if (length <= 0 || length > RemainingInputBytes)
            throw new ArgumentOutOfRangeException(nameof(length));
        if (type == DeflateBlockType.Stored && length > ushort.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(length));

        _blocks.Add(new DeflateBlockSpec(type, _offset, length));
        _offset += length;
        return this;
    }

    public DeflatePlanBuilder AddAdaptiveDataBlock(int length)
    {
        return AddDataBlock(DeflateBlockType.Auto, length);
    }

    public DeflatePlanBuilder AddNpncb()
    {
        _blocks.Add(DeflateBlockSpec.CreateNpncb());
        return this;
    }

    public DeflatePlanBuilder AddNpncb(int count)
    {
        if (count < 0) throw new ArgumentOutOfRangeException(nameof(count));
        for (int i = 0; i < count; i++) AddNpncb();
        return this;
    }

    public DeflatePlanBuilder AddRemaining(
        DeflateBlockType type,
        int maximumBlockSize = 65_535)
    {
        if (type == DeflateBlockType.Npncb)
            throw new ArgumentException("Use AddNpncb().", nameof(type));
        if (maximumBlockSize <= 0 || maximumBlockSize > ushort.MaxValue)
            throw new ArgumentOutOfRangeException(nameof(maximumBlockSize));

        while (RemainingInputBytes > 0)
            AddDataBlock(type, Math.Min(maximumBlockSize, RemainingInputBytes));
        return this;
    }

    public DeflatePlanBuilder AddRemainingAdaptive(int maximumBlockSize = 65_535)
    {
        return AddRemaining(DeflateBlockType.Auto, maximumBlockSize);
    }

    public List<DeflateBlockSpec> Build()
    {
        if (_blocks.Count == 0 && _inputLength == 0)
            _blocks.Add(new DeflateBlockSpec(DeflateBlockType.Fixed, 0, 0));

        DeflateBlockPlanner.Validate(_blocks, _inputLength);
        return new List<DeflateBlockSpec>(_blocks);
    }
}
