namespace ZB.MOM.NatsNet.Server;

internal sealed class JetStreamEngine(JetStream state)
{
    private readonly JetStream _state = state;
    private static readonly TimeSpan MinUsageUpdateWindow = TimeSpan.FromMilliseconds(250);

    internal void SetStarted()
    {
        _state.Lock.EnterWriteLock();
        try
        {
            _state.Started = DateTime.UtcNow;
        }
        finally
        {
            _state.Lock.ExitWriteLock();
        }
    }

    internal bool IsEnabled() => Interlocked.CompareExchange(ref _state.Disabled, 0, 0) == 0;

    internal void SetJetStreamStandAlone(bool isStandAlone)
    {
        _state.Lock.EnterWriteLock();
        try
        {
            _state.StandAlone = isStandAlone;
        }
        finally
        {
            _state.Lock.ExitWriteLock();
        }
    }

    internal bool IsShuttingDown()
    {
        _state.Lock.EnterReadLock();
        try
        {
            return _state.ShuttingDown;
        }
        finally
        {
            _state.Lock.ExitReadLock();
        }
    }

    internal Exception? DisableJetStream(JsAccount? account)
    {
        if (account?.Account is not Account a)
            return new InvalidOperationException("jetstream not enabled for account");

        _state.Lock.EnterWriteLock();
        try
        {
            _state.Accounts.Remove(a.Name);
        }
        finally
        {
            _state.Lock.ExitWriteLock();
        }

        account.Delete();
        return null;
    }

    internal bool WouldExceedLimits(StorageType storageType, int size)
    {
        var total = storageType == StorageType.MemoryStorage
            ? Interlocked.Read(ref _state.MemUsed)
            : Interlocked.Read(ref _state.StoreUsed);
        var max = storageType == StorageType.MemoryStorage
            ? _state.Config.MaxMemory
            : _state.Config.MaxStore;
        return total + size > max;
    }

    internal bool LimitsExceeded(StorageType storageType) => WouldExceedLimits(storageType, 0);

    internal static string TierName(int replicas) => $"R{(replicas <= 0 ? 1 : replicas)}";

    internal static bool IsSameTier(StreamConfig cfgA, StreamConfig cfgB) =>
        cfgA.Replicas == cfgB.Replicas;

    internal static Dictionary<string, JetStreamAccountLimits> DiffCheckedLimits(
        Dictionary<string, JetStreamAccountLimits> a,
        Dictionary<string, JetStreamAccountLimits> b)
    {
        var diff = new Dictionary<string, JetStreamAccountLimits>(StringComparer.Ordinal);

        foreach (var (tier, oldLimit) in a)
        {
            b.TryGetValue(tier, out var newLimit);
            newLimit ??= new JetStreamAccountLimits();
            diff[tier] = new JetStreamAccountLimits
            {
                MaxMemory = newLimit.MaxMemory - oldLimit.MaxMemory,
                MaxStore = newLimit.MaxStore - oldLimit.MaxStore,
            };
        }

        foreach (var (tier, newLimit) in b)
        {
            if (a.ContainsKey(tier))
                continue;

            diff[tier] = new JetStreamAccountLimits
            {
                MaxMemory = newLimit.MaxMemory,
                MaxStore = newLimit.MaxStore,
            };
        }

        return diff;
    }

    internal static (ulong Mem, ulong Store) ReservedStorage(
        Dictionary<string, object?> streamAssignments,
        string tier)
    {
        ulong mem = 0;
        ulong store = 0;

        foreach (var assignment in streamAssignments.Values.OfType<StreamAssignmentView>())
        {
            var cfg = assignment.Config;
            if (!string.IsNullOrEmpty(tier) && !string.Equals(tier, TierName(cfg.Replicas), StringComparison.Ordinal))
                continue;
            if (cfg.MaxBytes <= 0)
                continue;

            if (cfg.Storage == StorageType.FileStorage)
                store += (ulong)cfg.MaxBytes;
            else if (cfg.Storage == StorageType.MemoryStorage)
                mem += (ulong)cfg.MaxBytes;
        }

        return (mem, store);
    }

    internal static bool ShouldSendUsageUpdate(DateTime lastUpdateUtc) =>
        DateTime.UtcNow - lastUpdateUtc >= MinUsageUpdateWindow;
}

internal sealed class StreamAssignmentView
{
    public required StreamConfig Config { get; init; }
}