using System.Security.Cryptography;
using System.Text;
using ZB.MOM.NatsNet.Server.Internal.DataStructures;

namespace ZB.MOM.NatsNet.Server;

public sealed partial class NatsServer
{
    private const string JetStreamStoreDir = "jetstream";
    private const long JetStreamMaxMemDefault = 1024L * 1024L * 256L;

    public Exception? EnableJetStream(JetStreamConfig? config)
    {
        if (JetStreamEnabled())
            return new InvalidOperationException("jetstream already enabled");

        Noticef("Starting JetStream");

        if (config == null || config.MaxMemory <= 0 || config.MaxStore <= 0)
        {
            config = new JetStreamConfig
            {
                StoreDir = string.IsNullOrWhiteSpace(GetOpts().StoreDir)
                    ? Path.Combine(Path.GetTempPath(), JetStreamStoreDir)
                    : Path.Combine(GetOpts().StoreDir, JetStreamStoreDir),
                MaxMemory = GetOpts().JetStreamMaxMemory > 0 ? GetOpts().JetStreamMaxMemory : 1,
                MaxStore = GetOpts().JetStreamMaxStore > 0 ? GetOpts().JetStreamMaxStore : 1,
                SyncInterval = GetOpts().SyncInterval,
                SyncAlways = GetOpts().SyncAlways,
                Domain = GetOpts().JetStreamDomain,
            };
        }
        else if (!string.IsNullOrWhiteSpace(config.StoreDir))
        {
            config.StoreDir = Path.Combine(config.StoreDir, JetStreamStoreDir);
        }

        if (string.IsNullOrWhiteSpace(config.StoreDir))
        {
            config.StoreDir = Path.Combine(Path.GetTempPath(), JetStreamStoreDir);
            Warnf("Temporary storage directory used, data could be lost on system reboot");
        }

        var err = CheckStoreDir(config);
        if (err != null)
            return err;

        return EnableJetStreamInternal(config);
    }

    private KeyGen? JsKeyGen(string jsKey, string info)
    {
        if (string.IsNullOrEmpty(jsKey))
            return null;

        return context =>
        {
            using var hmac = new HMACSHA256(Encoding.UTF8.GetBytes(jsKey));
            hmac.TransformBlock(Encoding.UTF8.GetBytes(info), 0, info.Length, null, 0);
            hmac.TransformFinalBlock(context, 0, context.Length);
            return hmac.Hash ?? [];
        };
    }

    internal (byte[]? Plain, bool UsedFallback, Exception? Error) DecryptMeta(
        StoreCipher storeCipher,
        byte[] encryptedKey,
        byte[] encryptedBuffer,
        string accountName,
        string context)
    {
        if (encryptedKey.Length == 0)
            return (null, false, new InvalidOperationException("encryption key missing"));

        var ciphers = storeCipher == StoreCipher.Aes
            ? new[] { StoreCipher.Aes, StoreCipher.ChaCha }
            : new[] { StoreCipher.ChaCha, StoreCipher.Aes };

        var candidates = new List<(KeyGen Prf, StoreCipher Cipher)>();
        var opts = GetOpts();

        var prf = JsKeyGen(opts.JetStreamKey, accountName);
        if (prf == null)
            return (null, false, new InvalidOperationException("jetstream encryption key is not configured"));

        foreach (var cipher in ciphers)
            candidates.Add((prf, cipher));

        var oldPrf = JsKeyGen(opts.JetStreamOldKey, accountName);
        if (oldPrf != null)
        {
            foreach (var cipher in ciphers)
                candidates.Add((oldPrf, cipher));
        }

        for (var i = 0; i < candidates.Count; i++)
        {
            try
            {
                var rb = candidates[i].Prf(Encoding.UTF8.GetBytes(context));
                var kek = JetStreamFileStore.GenEncryptionKey(candidates[i].Cipher, rb);
                var ns = kek.NonceSize;
                if (encryptedKey.Length < ns || encryptedBuffer.Length < ns)
                    continue;

                var seed = kek.Open(encryptedKey.AsSpan(0, ns), encryptedKey.AsSpan(ns));
                var aek = JetStreamFileStore.GenEncryptionKey(candidates[i].Cipher, seed);
                var plain = aek.Open(encryptedBuffer.AsSpan(0, ns), encryptedBuffer.AsSpan(ns));
                return (plain, i > 0, null);
            }
            catch
            {
                // Try the next candidate.
            }
        }

        return (null, false, new InvalidOperationException("unable to recover encrypted metadata"));
    }

    internal Exception? CheckStoreDir(JetStreamConfig cfg)
    {
        if (string.IsNullOrWhiteSpace(cfg.StoreDir))
            return new InvalidOperationException("jetstream store directory is required");

        try
        {
            Directory.CreateDirectory(cfg.StoreDir);
            return null;
        }
        catch (Exception ex)
        {
            return ex;
        }
    }

    internal Exception? InitJetStreamEncryption()
    {
        var opts = GetOpts();

        if (!string.IsNullOrEmpty(opts.JetStreamKey) && !string.IsNullOrEmpty(opts.JetStreamTpm.KeysFile))
            return new InvalidOperationException("JetStream encryption key may not be used with TPM options");

        return null;
    }

    private Exception? EnableJetStreamInternal(JetStreamConfig cfg)
    {
        var encryptionErr = InitJetStreamEncryption();
        if (encryptionErr != null)
            return encryptionErr;

        try
        {
            Directory.CreateDirectory(cfg.StoreDir);
        }
        catch (Exception ex)
        {
            return ex;
        }

        var js = new JetStream
        {
            Server = this,
            Config = cfg,
            Started = DateTime.UtcNow,
            StandAlone = true,
        };

        _mu.EnterWriteLock();
        try
        {
            _jetStream = js;
            _info.JetStream = true;
            _info.Domain = cfg.Domain;
        }
        finally
        {
            _mu.ExitWriteLock();
        }

        var err = EnableJetStreamAccounts();
        if (err != null)
        {
            _mu.EnterWriteLock();
            try
            {
                _jetStream = null;
                _info.JetStream = false;
            }
            finally
            {
                _mu.ExitWriteLock();
            }
        }

        return err;
    }

    internal bool CanExtendOtherDomain()
    {
        var opts = GetOpts();
        var sysAcc = SystemAccount()?.GetName();
        if (string.IsNullOrEmpty(sysAcc))
            return false;

        foreach (var remote in opts.LeafNode.Remotes)
        {
            if (!string.Equals(remote.LocalAccount, sysAcc, StringComparison.Ordinal))
                continue;

            foreach (var denyImport in remote.DenyImports)
            {
                if (SubscriptionIndex.SubjectIsSubsetMatch(denyImport, JsApiSubjects.JsAllApi))
                    return false;
            }

            return true;
        }

        return false;
    }

    internal void UpdateJetStreamInfoStatus(bool enabled)
    {
        _mu.EnterWriteLock();
        try
        {
            _info.JetStream = enabled;
        }
        finally
        {
            _mu.ExitWriteLock();
        }
    }

    internal Exception? RestartJetStream()
    {
        var opts = GetOpts();
        var cfg = new JetStreamConfig
        {
            StoreDir = opts.StoreDir,
            SyncInterval = opts.SyncInterval,
            SyncAlways = opts.SyncAlways,
            MaxMemory = opts.JetStreamMaxMemory,
            MaxStore = opts.JetStreamMaxStore,
            Domain = opts.JetStreamDomain,
            Strict = !opts.NoJetStreamStrict,
        };

        Noticef("Restarting JetStream");
        var err = EnableJetStream(cfg);
        if (err != null)
        {
            Warnf("Can't start JetStream: {0}", err.Message);
            _ = DisableJetStream();
            return err;
        }

        UpdateJetStreamInfoStatus(true);
        return null;
    }

    internal void CheckJetStreamExports()
    {
        if (SystemAccount() != null)
            SetupJetStreamExports();
    }

    internal void SetupJetStreamExports()
    {
        var err = SetJetStreamExportSubs();
        if (err != null)
            Warnf("Error setting up jetstream service exports: {0}", err.Message);
    }

    internal bool JetStreamOOSPending()
    {
        var js = _jetStream;
        if (js == null)
            return false;

        js.Lock.EnterWriteLock();
        try
        {
            var wasPending = js.Oos;
            js.Oos = true;
            return wasPending;
        }
        finally
        {
            js.Lock.ExitWriteLock();
        }
    }

    internal void SetJetStreamDisabled()
    {
        var js = _jetStream;
        if (js != null)
            Interlocked.Exchange(ref js.Disabled, 1);
    }

    internal void HandleOutOfSpace(NatsStream? stream)
    {
        if (!JetStreamEnabled() || JetStreamOOSPending())
            return;

        if (stream != null)
            Errorf("JetStream out of resources for stream {0}, will be DISABLED", stream.Config.Name);
        else
            Errorf("JetStream out of resources, will be DISABLED");

        _ = Task.Run(() => DisableJetStream());
    }

    public Exception? DisableJetStream()
    {
        if (!JetStreamEnabled())
            return null;

        SetJetStreamDisabled();
        UpdateJetStreamInfoStatus(false);

        _mu.EnterWriteLock();
        try
        {
            _jetStream = null;
        }
        finally
        {
            _mu.ExitWriteLock();
        }

        ShutdownJetStream();
        ShutdownRaftNodes();
        return null;
    }

    private Exception? EnableJetStreamAccounts()
    {
        if (GlobalAccountOnly())
        {
            var gacc = GlobalAccount();
            if (gacc == null)
                return new InvalidOperationException("global account not found");

            gacc.JetStreamLimits ??= new Dictionary<string, object>(StringComparer.Ordinal)
            {
                [string.Empty] = new JetStreamAccountLimits
                {
                    MaxMemory = -1,
                    MaxStore = -1,
                    MaxStreams = -1,
                    MaxConsumers = -1,
                    MaxAckPending = -1,
                    MemoryMaxStreamBytes = -1,
                    StoreMaxStreamBytes = -1,
                },
            };

            return ConfigJetStream(gacc);
        }

        return ConfigAllJetStreamAccounts();
    }

    internal Exception? ConfigJetStream(Account? acc)
    {
        if (acc == null)
            return null;

        var jsLimits = acc.JetStreamLimits;
        if (jsLimits != null)
            return acc.EnableAllJetStreamServiceImportsAndMappings();

        if (!ReferenceEquals(acc, SystemAccount()))
        {
            acc.JetStream = null;
            return acc.EnableJetStreamInfoServiceImportOnly();
        }

        return null;
    }

    internal Exception? ConfigAllJetStreamAccounts()
    {
        CheckJetStreamExports();
        if (_jetStream == null)
            return null;

        foreach (var acc in _accounts.Values)
        {
            var err = ConfigJetStream(acc);
            if (err != null)
                return err;
        }

        var storeDir = _jetStream.Config.StoreDir;
        if (!Directory.Exists(storeDir))
            return null;

        foreach (var directory in Directory.EnumerateDirectories(storeDir))
        {
            var accountName = Path.GetFileName(directory);
            if (string.IsNullOrWhiteSpace(accountName) || _accounts.ContainsKey(accountName))
                continue;

            var (resolved, _) = LookupAccount(accountName);
            if (resolved == null)
                continue;

            var err = ConfigJetStream(resolved);
            if (err != null)
                return err;
        }

        return null;
    }

    public bool JetStreamEnabled()
    {
        var js = _jetStream;
        return js != null && Interlocked.CompareExchange(ref js.Disabled, 0, 0) == 0;
    }

    public bool JetStreamEnabledForDomain()
    {
        if (JetStreamEnabled())
            return true;

        foreach (var value in _nodeToInfo.Values)
        {
            if (value is NodeInfo { Js: true })
                return true;
        }

        return false;
    }

    public JetStreamConfig? JetStreamConfig()
    {
        var js = _jetStream;
        if (js == null)
            return null;

        var cfg = js.Config;
        return new JetStreamConfig
        {
            MaxMemory = cfg.MaxMemory,
            MaxStore = cfg.MaxStore,
            StoreDir = cfg.StoreDir,
            SyncInterval = cfg.SyncInterval,
            SyncAlways = cfg.SyncAlways,
            Domain = cfg.Domain,
            CompressOK = cfg.CompressOK,
            UniqueTag = cfg.UniqueTag,
            Strict = cfg.Strict,
        };
    }

    public string StoreDir()
    {
        var js = _jetStream;
        return js == null ? string.Empty : js.Config.StoreDir;
    }

    public int JetStreamNumAccounts()
    {
        var js = _jetStream;
        if (js == null)
            return 0;

        js.Lock.EnterReadLock();
        try
        {
            return js.Accounts.Count;
        }
        finally
        {
            js.Lock.ExitReadLock();
        }
    }

    public (long MemReserved, long StoreReserved, Exception? Error) JetStreamReservedResources()
    {
        var js = _jetStream;
        if (js == null)
            return (-1, -1, new InvalidOperationException("jetstream not enabled"));

        return (
            Interlocked.Read(ref js.MemReserved),
            Interlocked.Read(ref js.StoreReserved),
            null);
    }

    internal JetStreamConfig DynJetStreamConfig(string storeDir, long maxStore, long maxMem)
    {
        var cfg = new JetStreamConfig();
        if (!string.IsNullOrWhiteSpace(storeDir))
        {
            cfg.StoreDir = Path.Combine(storeDir, JetStreamStoreDir);
        }
        else
        {
            cfg.StoreDir = Path.Combine(Path.GetTempPath(), "nats", JetStreamStoreDir);
            Warnf("Temporary storage directory used, data could be lost on system reboot");
        }

        var opts = GetOpts();
        cfg.Strict = !opts.NoJetStreamStrict;
        cfg.SyncInterval = opts.SyncInterval;
        cfg.SyncAlways = opts.SyncAlways;

        cfg.MaxStore = opts.MaxStoreSet && maxStore >= 0
            ? maxStore
            : DiskAvailability.DiskAvailable(cfg.StoreDir);

        if (opts.MaxMemSet && maxMem >= 0)
        {
            cfg.MaxMemory = maxMem;
        }
        else
        {
            var totalAvailable = GC.GetGCMemoryInfo().TotalAvailableMemoryBytes;
            cfg.MaxMemory = totalAvailable > 0 && totalAvailable < long.MaxValue
                ? totalAvailable / 4 * 3
                : JetStreamMaxMemDefault;
        }

        return cfg;
    }

    internal void ResourcesExceededError(StorageType storeType)
    {
        var didAlert = false;
        lock (_resourceErrorLock)
        {
            var now = DateTime.UtcNow;
            if (now - _resourceErrorLastUtc > TimeSpan.FromSeconds(10))
            {
                var storeName = storeType switch
                {
                    StorageType.MemoryStorage => "memory",
                    StorageType.FileStorage => "file",
                    _ => storeType.ToString().ToLowerInvariant(),
                };

                Errorf("JetStream {0} resource limits exceeded for server", storeName);
                _resourceErrorLastUtc = now;
                didAlert = true;
            }
        }

        if (!didAlert)
            return;

        var js = GetJetStreamState();
        if (js?.Cluster is JetStreamCluster { Meta: not null } cluster)
            cluster.Meta.StepDown();
    }

    internal void HandleWritePermissionError()
    {
        if (!JetStreamEnabled())
            return;

        Errorf("File system permission denied while writing, disabling JetStream");
        _ = Task.Run(() => DisableJetStream());
    }

    internal JetStreamEngine? GetJetStream() =>
        _jetStream == null ? null : new JetStreamEngine(_jetStream);

    internal JetStream? GetJetStreamState() => _jetStream;
}