// Copyright 2020-2026 The NATS Authors
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Adapted from server/mqtt.go in the NATS server Go source.

using System.Collections.Concurrent;
using System.Text;
using System.Text.Json;
using System.Threading.Channels;

// All JetStream API types (StreamInfo, ConsumerConfig, JSPubAckResponse, etc.)
// live in the parent namespace ZB.MOM.NatsNet.Server and are accessible here.

namespace ZB.MOM.NatsNet.Server.Mqtt;

// ============================================================================
// Helper types
// ============================================================================

/// <summary>
/// A pending JS API request message to be sent via the JSA send queue.
/// Mirrors Go <c>mqttJSPubMsg</c> struct in server/mqtt.go.
/// </summary>
internal sealed class MqttJsPubMsg
{
    public string  Subject { get; set; } = string.Empty;
    public string  Reply   { get; set; } = string.Empty;
    public int     Hdr     { get; set; }
    public byte[]? Msg     { get; set; }
}

/// <summary>
/// A response received from a JS API request.
/// Mirrors Go <c>mqttJSAResponse</c> struct in server/mqtt.go.
/// </summary>
internal sealed class MqttJsaResponse
{
    public string  Reply { get; set; } = string.Empty;
    public object? Value { get; set; }
}

/// <summary>
/// Deletion request for a retained message (JSON payload).
/// Mirrors Go <c>mqttRetMsgDel</c> struct in server/mqtt.go.
/// </summary>
internal sealed class MqttRetMsgDel
{
    public string Subject { get; set; } = string.Empty;
    public ulong  Seq     { get; set; }
}

// ============================================================================
// MqttJsa — JetStream API bridge (Sub-batch A: features 2269-2290)
// ============================================================================

/// <summary>
/// Per-account MQTT JetStream API bridge.
/// Mirrors Go <c>mqttJSA</c> struct in server/mqtt.go.
/// Sends requests to JS API subjects and waits for replies via a reply subject.
/// </summary>
internal sealed class MqttJsa
{
    private readonly Lock _mu = new();

    // ------------------------------------------------------------------
    // Identity / routing
    // ------------------------------------------------------------------

    /// <summary>Node identifier (server name hash used in reply subjects).</summary>
    public string Id { get; set; } = string.Empty;

    /// <summary>Internal client connection used to send JS API messages.</summary>
    public ClientConnection? Client { get; set; }

    /// <summary>Reply subject prefix: <c>$MQTT.JSA.{id}.</c></summary>
    public string Rplyr { get; set; } = string.Empty;

    /// <summary>Send queue for pending JS API messages.</summary>
    public Channel<MqttJsPubMsg> SendQ { get; } = Channel.CreateUnbounded<MqttJsPubMsg>(
        new UnboundedChannelOptions { SingleReader = true });

    /// <summary>
    /// Map of reply subjects to their response channels.
    /// Mirrors Go <c>replies sync.Map</c> in mqttJSA.
    /// </summary>
    public ConcurrentDictionary<string, Channel<MqttJsaResponse>> Replies { get; } = new();

    /// <summary>Monotonically increasing counter used for unique reply IDs.</summary>
    private long _nuid;

    /// <summary>Cancellation token (mirrors Go <c>quitCh chan struct{}</c>).</summary>
    public CancellationToken QuitCt { get; set; }

    /// <summary>JS domain (may be empty).</summary>
    public string Domain { get; set; } = string.Empty;

    /// <summary>Whether <see cref="Domain"/> was explicitly set (even to empty).</summary>
    public bool DomainSet { get; set; }

    /// <summary>Timeout for JS API requests (mirrors Go <c>timeout time.Duration</c>).</summary>
    public TimeSpan Timeout { get; set; } = TimeSpan.FromSeconds(5);

    // ------------------------------------------------------------------
    // Feature 2270: prefixDomain
    // ------------------------------------------------------------------

    /// <summary>
    /// Rewrites a JS API subject to include the configured domain.
    /// Mirrors Go <c>(*mqttJSA).prefixDomain()</c>.
    /// </summary>
    public string PrefixDomain(string subject)
    {
        if (string.IsNullOrEmpty(Domain))
            return subject;

        const string jsApiPrefix = "$JS.API.";
        if (subject.StartsWith(jsApiPrefix, StringComparison.Ordinal))
        {
            var sub = subject[jsApiPrefix.Length..];
            return $"$JS.{Domain}.API.{sub}";
        }
        return subject;
    }

    // ------------------------------------------------------------------
    // Feature 2269: newRequest
    // ------------------------------------------------------------------

    /// <summary>
    /// Creates a new JS API request (single message, no client-ID hash).
    /// Mirrors Go <c>(*mqttJSA).newRequest()</c>.
    /// </summary>
    public async Task<object?> NewRequestAsync(
        string kind, string subject, int hdr, byte[]? msg,
        CancellationToken ct = default)
    {
        var responses = await NewRequestExMultiAsync(
            kind, subject, string.Empty, [hdr], [msg ?? []], ct);
        if (responses.Length != 1)
            throw new InvalidOperationException(
                $"unreachable: invalid number of responses ({responses.Length})");
        return responses[0]?.Value;
    }

    // ------------------------------------------------------------------
    // Feature 2271: newRequestEx
    // ------------------------------------------------------------------

    /// <summary>
    /// Extended request with an optional client-ID hash embedded in the reply subject.
    /// Mirrors Go <c>(*mqttJSA).newRequestEx()</c>.
    /// </summary>
    public async Task<object?> NewRequestExAsync(
        string kind, string subject, string cidHash, int hdr, byte[]? msg,
        CancellationToken ct = default)
    {
        var responses = await NewRequestExMultiAsync(
            kind, subject, cidHash, [hdr], [msg ?? []], ct);
        if (responses.Length != 1)
            throw new InvalidOperationException(
                $"unreachable: invalid number of responses ({responses.Length})");
        return responses[0]?.Value;
    }

    // ------------------------------------------------------------------
    // Feature 2272: newRequestExMulti
    // ------------------------------------------------------------------

    /// <summary>
    /// Sends multiple messages on the same subject and waits for all responses.
    /// Returns a sparse array (same length as <paramref name="msgs"/>)
    /// where timed-out entries remain <c>null</c>.
    /// Mirrors Go <c>(*mqttJSA).newRequestExMulti()</c>.
    /// </summary>
    public async Task<MqttJsaResponse?[]> NewRequestExMultiAsync(
        string kind, string subject, string cidHash,
        int[] hdrs, byte[][] msgs,
        CancellationToken ct = default)
    {
        if (hdrs.Length != msgs.Length)
            throw new ArgumentException(
                $"mismatched hdrs ({hdrs.Length}) and msgs ({msgs.Length}) counts");

        var responseCh = Channel.CreateBounded<MqttJsaResponse>(
            new BoundedChannelOptions(msgs.Length) { FullMode = BoundedChannelFullMode.Wait });
        var replyToIndex = new Dictionary<string, int>(msgs.Length);
        var responses    = new MqttJsaResponse?[msgs.Length];

        for (int i = 0; i < msgs.Length; i++)
        {
            string uid;
            string replyBase;
            lock (_mu)
            {
                uid       = Interlocked.Increment(ref _nuid).ToString();
                replyBase = Rplyr;
            }

            var sb = new StringBuilder(replyBase);
            sb.Append(kind);
            sb.Append('.');
            if (!string.IsNullOrEmpty(cidHash))
            {
                sb.Append(cidHash);
                sb.Append('.');
            }
            sb.Append(uid);
            string reply = sb.ToString();

            Replies[reply] = responseCh;

            string finalSubject = PrefixDomain(subject);
            await SendQ.Writer.WriteAsync(new MqttJsPubMsg
            {
                Subject = finalSubject,
                Reply   = reply,
                Hdr     = hdrs[i],
                Msg     = msgs[i],
            }, ct);

            replyToIndex[reply] = i;
        }

        int count = 0;
        using var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, QuitCt);
        linked.CancelAfter(Timeout);

        try
        {
            while (count < msgs.Length)
            {
                var r = await responseCh.Reader.ReadAsync(linked.Token);
                if (replyToIndex.TryGetValue(r.Reply, out int idx))
                {
                    responses[idx] = r;
                    count++;
                }
            }
        }
        catch (OperationCanceledException) when (QuitCt.IsCancellationRequested)
        {
            throw new InvalidOperationException("server not running");
        }
        catch (OperationCanceledException)
        {
            // Timeout — clean up pending reply registrations.
            foreach (var key in replyToIndex.Keys)
                Replies.TryRemove(key, out _);
        }

        return responses;
    }

    // ------------------------------------------------------------------
    // Feature 2273: sendAck
    // ------------------------------------------------------------------

    /// <summary>
    /// Sends a JS acknowledgement to the given ack subject (empty payload).
    /// Mirrors Go <c>(*mqttJSA).sendAck()</c>.
    /// </summary>
    public void SendAck(string ackSubject) => SendMsg(ackSubject, null);

    // ------------------------------------------------------------------
    // Feature 2274: sendMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Queues a message for fire-and-forget delivery (no JS reply).
    /// Mirrors Go <c>(*mqttJSA).sendMsg()</c>.
    /// </summary>
    public void SendMsg(string subj, byte[]? msg)
    {
        if (string.IsNullOrEmpty(subj))
            return;
        // hdr = -1 signals the send loop to skip adding client-info header.
        SendQ.Writer.TryWrite(new MqttJsPubMsg { Subject = subj, Msg = msg, Hdr = -1 });
    }

    // ------------------------------------------------------------------
    // Feature 2275: createEphemeralConsumer
    // ------------------------------------------------------------------

    /// <summary>
    /// Creates an ephemeral JetStream consumer for MQTT delivery.
    /// Mirrors Go <c>(*mqttJSA).createEphemeralConsumer()</c>.
    /// </summary>
    public async Task<JsApiConsumerCreateResponse?> CreateEphemeralConsumerAsync(
        CreateConsumerRequest cfg, CancellationToken ct = default)
    {
        var cfgBytes = JsonSerializer.SerializeToUtf8Bytes(cfg);
        string subj  = string.Format(JsApiSubjects.JsApiConsumerCreateT, cfg.Stream);
        var resp     = await NewRequestAsync(MqttTopics.JsaConsumerCreate, subj, 0, cfgBytes, ct);
        return resp as JsApiConsumerCreateResponse;
    }

    // ------------------------------------------------------------------
    // Feature 2276: createDurableConsumer
    // ------------------------------------------------------------------

    /// <summary>
    /// Creates a durable JetStream consumer for MQTT QoS delivery.
    /// Mirrors Go <c>(*mqttJSA).createDurableConsumer()</c>.
    /// </summary>
    public async Task<JsApiConsumerCreateResponse?> CreateDurableConsumerAsync(
        CreateConsumerRequest cfg, CancellationToken ct = default)
    {
        var cfgBytes = JsonSerializer.SerializeToUtf8Bytes(cfg);
        string durable = cfg.Config?.Durable ?? string.Empty;
        string subj    = string.Format(JsApiSubjects.JsApiDurableCreateT, cfg.Stream, durable);
        var resp       = await NewRequestAsync(MqttTopics.JsaConsumerCreate, subj, 0, cfgBytes, ct);
        return resp as JsApiConsumerCreateResponse;
    }

    // ------------------------------------------------------------------
    // Feature 2277: deleteConsumer
    // ------------------------------------------------------------------

    /// <summary>
    /// Deletes a JetStream consumer. If <paramref name="noWait"/> is <c>true</c>,
    /// the request is fire-and-forget.
    /// Mirrors Go <c>(*mqttJSA).deleteConsumer()</c>.
    /// </summary>
    public async Task<JsApiConsumerDeleteResponse?> DeleteConsumerAsync(
        string streamName, string consName, bool noWait,
        CancellationToken ct = default)
    {
        string subj = string.Format(JsApiSubjects.JsApiConsumerDeleteT, streamName, consName);
        if (noWait)
        {
            SendMsg(subj, null);
            return null;
        }
        var resp = await NewRequestAsync(MqttTopics.JsaConsumerDel, subj, 0, null, ct);
        return resp as JsApiConsumerDeleteResponse;
    }

    // ------------------------------------------------------------------
    // Feature 2278: createStream
    // ------------------------------------------------------------------

    /// <summary>
    /// Creates a JetStream stream.
    /// Returns (<see cref="StreamInfo"/>, didCreate).
    /// Mirrors Go <c>(*mqttJSA).createStream()</c>.
    /// </summary>
    public async Task<(StreamInfo? Info, bool DidCreate)> CreateStreamAsync(
        StreamConfig cfg, CancellationToken ct = default)
    {
        var cfgBytes = JsonSerializer.SerializeToUtf8Bytes(cfg);
        string subj  = string.Format(JsApiSubjects.JsApiStreamCreateT, cfg.Name);
        var resp     = await NewRequestAsync(MqttTopics.JsaStreamCreate, subj, 0, cfgBytes, ct);
        if (resp is JsApiStreamCreateResponse scr)
        {
            var si = scr.Config != null ? new StreamInfo { Config = scr.Config, State = scr.State } : null;
            return (si, scr.DidCreate);
        }
        return (null, false);
    }

    // ------------------------------------------------------------------
    // Feature 2279: updateStream
    // ------------------------------------------------------------------

    /// <summary>
    /// Updates an existing JetStream stream.
    /// Mirrors Go <c>(*mqttJSA).updateStream()</c>.
    /// </summary>
    public async Task<StreamInfo?> UpdateStreamAsync(
        StreamConfig cfg, CancellationToken ct = default)
    {
        var cfgBytes = JsonSerializer.SerializeToUtf8Bytes(cfg);
        string subj  = string.Format(JsApiSubjects.JsApiStreamUpdateT, cfg.Name);
        var resp     = await NewRequestAsync(MqttTopics.JsaStreamUpdate, subj, 0, cfgBytes, ct);
        return (resp as JsApiStreamUpdateResponse)?.Info;
    }

    // ------------------------------------------------------------------
    // Feature 2280: lookupStream
    // ------------------------------------------------------------------

    /// <summary>
    /// Looks up a JetStream stream by name.
    /// Mirrors Go <c>(*mqttJSA).lookupStream()</c>.
    /// </summary>
    public async Task<StreamInfo?> LookupStreamAsync(
        string name, CancellationToken ct = default)
    {
        string subj = string.Format(JsApiSubjects.JsApiStreamInfoT, name);
        var resp    = await NewRequestAsync(MqttTopics.JsaStreamLookup, subj, 0, null, ct);
        return (resp as JsApiStreamInfoResponse)?.Info;
    }

    // ------------------------------------------------------------------
    // Feature 2281: deleteStream
    // ------------------------------------------------------------------

    /// <summary>
    /// Deletes a JetStream stream by name.
    /// Mirrors Go <c>(*mqttJSA).deleteStream()</c>.
    /// </summary>
    public async Task<bool> DeleteStreamAsync(
        string name, CancellationToken ct = default)
    {
        string subj = string.Format(JsApiSubjects.JsApiStreamDeleteT, name);
        var resp    = await NewRequestAsync(MqttTopics.JsaStreamDel, subj, 0, null, ct);
        return (resp as JsApiStreamDeleteResponse)?.Success ?? false;
    }

    // ------------------------------------------------------------------
    // Feature 2282: loadLastMsgFor
    // ------------------------------------------------------------------

    /// <summary>
    /// Loads the last message stored for a subject in a stream.
    /// Mirrors Go <c>(*mqttJSA).loadLastMsgFor()</c>.
    /// </summary>
    public async Task<StoredMsg?> LoadLastMsgForAsync(
        string streamName, string subject, CancellationToken ct = default)
    {
        var req      = new JsApiMsgGetRequest { LastFor = subject };
        var reqBytes = JsonSerializer.SerializeToUtf8Bytes(req);
        string subj  = string.Format(JsApiSubjects.JsApiMsgGetT, streamName);
        var resp     = await NewRequestAsync(MqttTopics.JsaMsgLoad, subj, 0, reqBytes, ct);
        return (resp as JsApiMsgGetResponse)?.Message;
    }

    // ------------------------------------------------------------------
    // Feature 2283: loadLastMsgForMulti
    // ------------------------------------------------------------------

    /// <summary>
    /// Loads the last message for each of the supplied subjects.
    /// Returns a sparse array in the same order as <paramref name="subjects"/>.
    /// Mirrors Go <c>(*mqttJSA).loadLastMsgForMulti()</c>.
    /// </summary>
    public async Task<JsApiMsgGetResponse?[]> LoadLastMsgForMultiAsync(
        string streamName, string[] subjects, CancellationToken ct = default)
    {
        var marshaled   = new byte[subjects.Length][];
        var headerBytes = new int[subjects.Length];
        for (int i = 0; i < subjects.Length; i++)
        {
            var req      = new JsApiMsgGetRequest { LastFor = subjects[i] };
            marshaled[i] = JsonSerializer.SerializeToUtf8Bytes(req);
        }

        string subj = string.Format(JsApiSubjects.JsApiMsgGetT, streamName);
        var all     = await NewRequestExMultiAsync(
            MqttTopics.JsaMsgLoad, subj, string.Empty, headerBytes, marshaled, ct);

        var responses = new JsApiMsgGetResponse?[all.Length];
        for (int i = 0; i < all.Length; i++)
            responses[i] = all[i]?.Value as JsApiMsgGetResponse;
        return responses;
    }

    // ------------------------------------------------------------------
    // Feature 2284: loadNextMsgFor
    // ------------------------------------------------------------------

    /// <summary>
    /// Loads the next message after a given subject sequence in a stream.
    /// Mirrors Go <c>(*mqttJSA).loadNextMsgFor()</c>.
    /// </summary>
    public async Task<StoredMsg?> LoadNextMsgForAsync(
        string streamName, string subject, CancellationToken ct = default)
    {
        var req      = new JsApiMsgGetRequest { NextFor = subject };
        var reqBytes = JsonSerializer.SerializeToUtf8Bytes(req);
        string subj  = string.Format(JsApiSubjects.JsApiMsgGetT, streamName);
        var resp     = await NewRequestAsync(MqttTopics.JsaMsgLoad, subj, 0, reqBytes, ct);
        return (resp as JsApiMsgGetResponse)?.Message;
    }

    // ------------------------------------------------------------------
    // Feature 2285: loadMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Loads a specific message by sequence from a stream.
    /// Mirrors Go <c>(*mqttJSA).loadMsg()</c>.
    /// </summary>
    public async Task<StoredMsg?> LoadMsgAsync(
        string streamName, ulong seq, CancellationToken ct = default)
    {
        var req      = new JsApiMsgGetRequest { Seq = seq };
        var reqBytes = JsonSerializer.SerializeToUtf8Bytes(req);
        string subj  = string.Format(JsApiSubjects.JsApiMsgGetT, streamName);
        var resp     = await NewRequestAsync(MqttTopics.JsaMsgLoad, subj, 0, reqBytes, ct);
        return (resp as JsApiMsgGetResponse)?.Message;
    }

    // ------------------------------------------------------------------
    // Feature 2286: storeMsgNoWait
    // ------------------------------------------------------------------

    /// <summary>
    /// Queues a message for JetStream storage without waiting for the pub-ack.
    /// Mirrors Go <c>(*mqttJSA).storeMsgNoWait()</c>.
    /// </summary>
    public void StoreMsgNoWait(string subject, int hdrLen, byte[]? msg)
    {
        SendQ.Writer.TryWrite(new MqttJsPubMsg
        {
            Subject = subject,
            Msg     = msg,
            Hdr     = hdrLen,
        });
    }

    // ------------------------------------------------------------------
    // Feature 2287: storeMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Stores a message in JetStream and waits for the pub-ack.
    /// Mirrors Go <c>(*mqttJSA).storeMsg()</c>.
    /// </summary>
    public async Task<JSPubAckResponse?> StoreMsgAsync(
        string subject, int headers, byte[]? msg, CancellationToken ct = default)
    {
        var resp = await NewRequestAsync(MqttTopics.JsaMsgStore, subject, headers, msg, ct);
        return resp as JSPubAckResponse;
    }

    // ------------------------------------------------------------------
    // Feature 2288: storeSessionMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Stores a session-state message in the MQTT sessions JetStream stream.
    /// Mirrors Go <c>(*mqttJSA).storeSessionMsg()</c>.
    /// </summary>
    public async Task<JSPubAckResponse?> StoreSessionMsgAsync(
        string domainTk, string cidHash, int hdr, byte[]? msg, CancellationToken ct = default)
    {
        string subject = $"{domainTk}{MqttTopics.SessStreamSubjectPrefix}{cidHash}";
        var resp = await NewRequestExAsync(
            MqttTopics.JsaSessPersist, subject, cidHash, hdr, msg, ct);
        return resp as JSPubAckResponse;
    }

    // ------------------------------------------------------------------
    // Feature 2289: loadSessionMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Loads the most recent session-state message for a client ID hash.
    /// Mirrors Go <c>(*mqttJSA).loadSessionMsg()</c>.
    /// </summary>
    public Task<StoredMsg?> LoadSessionMsgAsync(
        string domainTk, string cidHash, CancellationToken ct = default)
    {
        string streamSubject = $"{domainTk}{MqttTopics.SessStreamSubjectPrefix}{cidHash}";
        return LoadLastMsgForAsync(MqttTopics.SessStreamName, streamSubject, ct);
    }

    // ------------------------------------------------------------------
    // Feature 2290: deleteMsg
    // ------------------------------------------------------------------

    /// <summary>
    /// Deletes a message by sequence from a JetStream stream.
    /// If <paramref name="wait"/> is <c>false</c>, the request is fire-and-forget.
    /// Mirrors Go <c>(*mqttJSA).deleteMsg()</c>.
    /// </summary>
    public async Task DeleteMsgAsync(
        string stream, ulong seq, bool wait, CancellationToken ct = default)
    {
        var dreq     = new JsApiMsgDeleteRequest { Seq = seq, NoErase = true };
        var reqBytes = JsonSerializer.SerializeToUtf8Bytes(dreq);
        string subj  = PrefixDomain(string.Format(JsApiSubjects.JsApiMsgDeleteT, stream));

        if (!wait)
        {
            SendQ.Writer.TryWrite(new MqttJsPubMsg { Subject = subj, Msg = reqBytes });
            return;
        }

        await NewRequestAsync(MqttTopics.JsaMsgDelete, subj, 0, reqBytes, ct);
    }
}
