namespace ZB.MOM.NatsNet.Server;

public sealed partial class Account
{
    private static Dictionary<string, object> DefaultJetStreamAccountTiers()
    {
        return new Dictionary<string, object>(StringComparer.Ordinal)
        {
            [string.Empty] = new JetStreamAccountLimits
            {
                MaxMemory = -1,
                MaxStore = -1,
                MaxStreams = -1,
                MaxConsumers = -1,
                MaxAckPending = -1,
                MemoryMaxStreamBytes = -1,
                StoreMaxStreamBytes = -1,
            },
        };
    }

    private static Dictionary<string, JetStreamAccountLimits> ToTypedLimits(Dictionary<string, object> limits)
    {
        var typed = new Dictionary<string, JetStreamAccountLimits>(StringComparer.Ordinal);
        foreach (var (tier, value) in limits)
        {
            if (value is JetStreamAccountLimits v)
                typed[tier] = v;
        }

        return typed;
    }

    private static JetStreamAccountLimits SelectLimits(Dictionary<string, JetStreamAccountLimits> limits)
    {
        if (limits.TryGetValue(string.Empty, out var selected))
            return selected;

        foreach (var (_, value) in limits)
            return value;

        return new JetStreamAccountLimits();
    }

    internal void AssignJetStreamLimits(Dictionary<string, object> limits)
    {
        _mu.EnterWriteLock();
        try
        {
            JetStreamLimits = limits;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal Exception? EnableJetStream(Dictionary<string, object>? limits)
    {
        _mu.EnterReadLock();
        var server = Server as NatsServer;
        _mu.ExitReadLock();

        if (server == null)
            return new InvalidOperationException("jetstream account not registered");
        if (ReferenceEquals(server.SystemAccount(), this))
            return new InvalidOperationException("jetstream can not be enabled on the system account");

        limits ??= DefaultJetStreamAccountTiers();
        if (limits.Count == 0)
            limits = DefaultJetStreamAccountTiers();

        AssignJetStreamLimits(limits);
        var typedLimits = ToTypedLimits(limits);

        var js = server.GetJetStreamState();
        if (js == null)
            return new InvalidOperationException("jetstream not enabled");

        js.Lock.EnterWriteLock();
        try
        {
            if (js.Accounts.TryGetValue(Name, out var existing))
            {
                _mu.EnterWriteLock();
                JetStream = existing;
                _mu.ExitWriteLock();
                return EnableAllJetStreamServiceImportsAndMappings();
            }

            var jsa = new JsAccount
            {
                Js = js,
                Account = this,
                StoreDir = Path.Combine(js.Config.StoreDir, Name),
            };

            foreach (var (tier, tierLimits) in typedLimits)
                jsa.Limits[tier] = tierLimits;

            js.Accounts[Name] = jsa;
            _mu.EnterWriteLock();
            JetStream = jsa;
            _mu.ExitWriteLock();
        }
        finally
        {
            js.Lock.ExitWriteLock();
        }

        return EnableAllJetStreamServiceImportsAndMappings();
    }

    internal (NatsServer? Server, JsAccount? JetStreamAccount, Exception? Error) CheckForJetStream()
    {
        _mu.EnterReadLock();
        var server = Server as NatsServer;
        var jsa = JetStream;
        _mu.ExitReadLock();

        if (server == null || jsa == null)
        {
            var description = JsApiErrors.NewJSNotEnabledForAccountError().Description ?? "jetstream not enabled for account";
            return (null, null, new InvalidOperationException(description));
        }

        return (server, jsa, null);
    }

    internal (NatsServer? Server, JetStream? JetStream, JsAccount? JetStreamAccount) GetJetStreamFromAccount()
    {
        _mu.EnterReadLock();
        var jetStreamAccount = JetStream;
        _mu.ExitReadLock();
        if (jetStreamAccount == null)
            return (null, null, null);

        jetStreamAccount.Lock.EnterReadLock();
        var jetStream = jetStreamAccount.Js as JetStream;
        jetStreamAccount.Lock.ExitReadLock();
        if (jetStream == null)
            return (null, null, null);

        return (jetStream.Server as NatsServer, jetStream, jetStreamAccount);
    }

    internal bool JetStreamIsStreamLeader(string stream)
    {
        var (server, jetStream, _) = GetJetStreamFromAccount();
        if (server == null || jetStream == null)
            return false;

        jetStream.Lock.EnterReadLock();
        try
        {
            return (jetStream.Cluster as JetStreamCluster)?.IsStreamLeader(Name, stream) == true;
        }
        finally
        {
            jetStream.Lock.ExitReadLock();
        }
    }

    internal bool JetStreamIsConsumerLeader(string stream, string consumer)
    {
        var (server, jetStream, _) = GetJetStreamFromAccount();
        if (server == null || jetStream == null)
            return false;

        jetStream.Lock.EnterReadLock();
        try
        {
            return (jetStream.Cluster as JetStreamCluster)?.IsConsumerLeader(Name, stream, consumer) == true;
        }
        finally
        {
            jetStream.Lock.ExitReadLock();
        }
    }

    internal void TrackAPI()
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa is null)
            return;

        jsa.UsageLock.EnterWriteLock();
        try
        {
            jsa.UsageApi++;
            jsa.ApiTotal++;
            jsa.SendClusterUsageUpdate();

            if (jsa.Js is JetStream js)
                Interlocked.Add(ref js.ApiTotal, 1);
        }
        finally
        {
            jsa.UsageLock.ExitWriteLock();
        }
    }

    internal void TrackAPIErr()
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa is null)
            return;

        jsa.UsageLock.EnterWriteLock();
        try
        {
            jsa.UsageApi++;
            jsa.ApiTotal++;
            jsa.UsageErr++;
            jsa.ApiErrors++;
            jsa.SendClusterUsageUpdate();

            if (jsa.Js is JetStream js)
            {
                Interlocked.Add(ref js.ApiTotal, 1);
                Interlocked.Add(ref js.ApiErrors, 1);
            }
        }
        finally
        {
            jsa.UsageLock.ExitWriteLock();
        }
    }

    internal (bool Enabled, bool ShouldError) CheckJetStream()
    {
        _mu.EnterReadLock();
        try
        {
            return (JetStream is not null, _nleafs + _nrleafs == 0);
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal (bool MaxBytesRequired, long MaxStreamBytes) MaxBytesLimits(StreamConfig? cfg)
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa == null)
            return (false, 0);

        jsa.UsageLock.EnterReadLock();
        try
        {
            var selected = SelectLimits(jsa.Limits);
            var maxStreamBytes = cfg?.Storage == StorageType.MemoryStorage
                ? selected.MemoryMaxStreamBytes
                : selected.StoreMaxStreamBytes;
            return (selected.MaxBytesRequired, maxStreamBytes);
        }
        finally
        {
            jsa.UsageLock.ExitReadLock();
        }
    }

    internal int NumStreams()
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa == null)
            return 0;

        jsa.Lock.EnterReadLock();
        try
        {
            return jsa.Streams.Count;
        }
        finally
        {
            jsa.Lock.ExitReadLock();
        }
    }

    internal List<NatsStream> Streams() => FilteredStreams(string.Empty);

    internal List<NatsStream> FilteredStreams(string filter)
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa == null)
            return [];

        jsa.Lock.EnterReadLock();
        try
        {
            var streams = new List<NatsStream>();
            foreach (var stream in jsa.Streams.Values.OfType<NatsStream>())
            {
                if (string.IsNullOrWhiteSpace(filter))
                {
                    streams.Add(stream);
                    continue;
                }

                foreach (var subject in stream.Config.Subjects ?? [])
                {
                    if (Internal.DataStructures.SubscriptionIndex.SubjectsCollide(filter, subject))
                    {
                        streams.Add(stream);
                        break;
                    }
                }
            }

            return streams;
        }
        finally
        {
            jsa.Lock.ExitReadLock();
        }
    }

    internal (NatsStream? Stream, Exception? Error) LookupStream(string name)
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        _mu.ExitReadLock();
        if (jsa == null)
            return (null, new InvalidOperationException("jetstream not enabled for account"));

        jsa.Lock.EnterReadLock();
        try
        {
            if (jsa.Streams.TryGetValue(name, out var stream) && stream is NatsStream ns)
                return (ns, null);

            return (null, new InvalidOperationException("stream not found"));
        }
        finally
        {
            jsa.Lock.ExitReadLock();
        }
    }

    internal Exception? UpdateJetStreamLimits(Dictionary<string, object>? limits)
    {
        _mu.EnterReadLock();
        var server = Server as NatsServer;
        var jsa = JetStream;
        _mu.ExitReadLock();

        if (server == null)
            return new InvalidOperationException("jetstream account not registered");
        if (server.GetJetStreamState() == null)
            return new InvalidOperationException("jetstream not enabled");
        if (jsa == null)
            return new InvalidOperationException("jetstream not enabled for account");

        limits ??= DefaultJetStreamAccountTiers();
        if (limits.Count == 0)
            limits = DefaultJetStreamAccountTiers();
        AssignJetStreamLimits(limits);

        var typed = ToTypedLimits(limits);
        jsa.UsageLock.EnterWriteLock();
        try
        {
            jsa.Limits.Clear();
            foreach (var (tier, tierLimits) in typed)
                jsa.Limits[tier] = tierLimits;
        }
        finally
        {
            jsa.UsageLock.ExitWriteLock();
        }

        return null;
    }

    public JetStreamAccountStats JetStreamUsage()
    {
        _mu.EnterReadLock();
        var jsa = JetStream;
        var accountName = Name;
        var configuredLimits = JetStreamLimits;
        _mu.ExitReadLock();

        var stats = new JetStreamAccountStats();
        if (jsa == null)
            return stats;

        var (js, _) = jsa.JetStreamAndClustered();
        if (js == null)
            return stats;

        jsa.UsageLock.EnterReadLock();
        try
        {
            long mem = 0;
            long store = 0;
            foreach (var usage in jsa.Usage.Values)
            {
                mem += usage.Total.Mem;
                store += usage.Total.Store;
            }

            stats.Memory = (ulong)Math.Max(0, mem);
            stats.Store = (ulong)Math.Max(0, store);
            stats.Domain = js.Config.Domain;
            stats.Api = new JetStreamApiStats
            {
                Level = JetStreamVersioning.JsApiLevel,
                Total = jsa.ApiTotal,
                Errors = jsa.ApiErrors,
            };

            if (jsa.Limits.TryGetValue(string.Empty, out var defaultTier))
            {
                stats.Limits = defaultTier;
            }
            else
            {
                stats.Tiers = new Dictionary<string, JetStreamTier>(StringComparer.Ordinal);
                foreach (var (tier, usage) in jsa.Usage)
                {
                    jsa.Limits.TryGetValue(tier, out var tierLimits);
                    stats.Tiers[tier] = new JetStreamTier
                    {
                        Memory = (ulong)Math.Max(0, usage.Total.Mem),
                        Store = (ulong)Math.Max(0, usage.Total.Store),
                        Limits = tierLimits ?? new JetStreamAccountLimits(),
                    };
                }

                if (configuredLimits != null)
                {
                    foreach (var (tier, value) in configuredLimits)
                    {
                        if (stats.Tiers.ContainsKey(tier))
                            continue;
                        if (value is not JetStreamAccountLimits lim)
                            continue;
                        stats.Tiers[tier] = new JetStreamTier { Limits = lim };
                    }
                }
            }
        }
        finally
        {
            jsa.UsageLock.ExitReadLock();
        }

        var allStreams = Streams();
        stats.Streams = allStreams.Count;
        foreach (var stream in allStreams)
            stats.Consumers += stream.State().Consumers;

        if (stats.Tiers != null)
        {
            foreach (var stream in allStreams)
            {
                var tier = JetStreamEngine.TierName(stream.Config.Replicas);
                if (!stats.Tiers.TryGetValue(tier, out var u))
                    u = new JetStreamTier();
                u.Streams++;
                u.Consumers += stream.State().Consumers;
                stats.Tiers[tier] = u;
            }
        }

        if (stats.Tiers == null || stats.Tiers.Count == 0)
        {
            var (rmem, rstore) = jsa.ReservedStorage(string.Empty);
            stats.ReservedMemory = rmem;
            stats.ReservedStore = rstore;
        }
        else
        {
            foreach (var tier in stats.Tiers.Keys.ToArray())
            {
                var tierStats = stats.Tiers[tier];
                (tierStats.ReservedMemory, tierStats.ReservedStore) = jsa.ReservedStorage(tier);
                stats.Tiers[tier] = tierStats;
            }
        }

        _ = accountName;
        return stats;
    }

    internal Exception? DisableJetStream() => RemoveJetStream();

    internal Exception? RemoveJetStream()
    {
        _mu.EnterWriteLock();
        var server = Server as NatsServer;
        var jsa = JetStream;
        JetStream = null;
        _mu.ExitWriteLock();

        if (server == null)
            return new InvalidOperationException("jetstream account not registered");
        var js = server.GetJetStream();
        if (js == null)
            return new InvalidOperationException("jetstream not enabled for account");

        return js.DisableJetStream(jsa);
    }

    internal bool JetStreamConfigured()
    {
        _mu.EnterReadLock();
        try
        {
            return JetStreamLimits != null && JetStreamLimits.Count > 0;
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal bool JetStreamEnabled()
    {
        _mu.EnterReadLock();
        try
        {
            return JetStream != null;
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }

    internal Exception? EnableAllJetStreamServiceImportsAndMappings()
    {
        _mu.EnterReadLock();
        var server = Server as NatsServer;
        _mu.ExitReadLock();

        if (server == null)
            return new InvalidOperationException("jetstream account not registered");

        var systemAccount = server.SystemAccount();
        var destinationName = systemAccount?.Name ?? string.Empty;

        if (systemAccount != null && !ServiceImportExists(destinationName, JsApiSubjects.JsAllApi))
        {
            var err = AddServiceImport(systemAccount, JsApiSubjects.JsAllApi, JsApiSubjects.JsAllApi);
            if (err != null)
                return new InvalidOperationException($"error setting up jetstream service imports for account: {err.Message}", err);
        }

        var domain = server.GetOpts().JetStreamDomain;
        if (!string.IsNullOrWhiteSpace(domain))
        {
            var mappings = new Dictionary<string, string>(StringComparer.Ordinal)
            {
                [$"$JS.{domain}.API.>"] = JsApiSubjects.JsAllApi,
                [$"$JS.{domain}.API.INFO"] = JsApiSubjects.JsApiAccountInfo,
            };

            _mu.EnterReadLock();
            try
            {
                foreach (var mapping in _mappings)
                    mappings.Remove(mapping.Source);
            }
            finally
            {
                _mu.ExitReadLock();
            }

            foreach (var (src, dest) in mappings)
            {
                var err = AddMapping(src, dest);
                if (err != null)
                    server.Errorf("Error adding JetStream domain mapping: {0}", err.Message);
            }
        }

        return null;
    }

    internal Exception? EnableJetStreamInfoServiceImportOnly()
    {
        if (ServiceImportShadowed(JsApiSubjects.JsApiAccountInfo))
            return null;

        return EnableAllJetStreamServiceImportsAndMappings();
    }

    internal Exception? JsNonClusteredStreamLimitsCheck(StreamConfig config)
    {
        var (server, jsa, err) = CheckForJetStream();
        if (err is not null || server is null || jsa is null)
            return err ?? new InvalidOperationException("jetstream not enabled for account");

        var selected = jsa.SelectLimits(config.Replicas);
        if (!selected.Found)
            return new InvalidOperationException(JsApiErrors.NewJSNoLimitsError().Description ?? "jetstream limits not configured");

        var reservation = jsa.TieredReservation(selected.Tier, config);
        var js = server.GetJetStream();
        if (js is null)
            return new InvalidOperationException("jetstream not enabled");

        return js.CheckAccountLimits(selected.Limits, config, reservation);
    }

    internal (JetStreamAccountLimits? Limits, string Tier, JsAccount? JsAccount, JsApiError? Error) SelectLimits(int replicas)
    {
        _mu.EnterReadLock();
        try
        {
            var jsa = JetStream;
            if (jsa == null)
                return (null, string.Empty, null, JsApiErrors.NewJSNotEnabledForAccountError());

            var (selected, tier, found) = jsa.SelectLimits(replicas);
            if (!found)
                return (null, string.Empty, jsa, JsApiErrors.NewJSNoLimitsError());

            return (selected, tier, jsa, null);
        }
        finally
        {
            _mu.ExitReadLock();
        }
    }
}