using System.Net.Sockets;
using System.Text;
using System.Text.Json;
using NATS.Server.Subscriptions;

namespace NATS.Server.LeafNodes;

/// <summary>
/// Represents a single leaf node connection (inbound or outbound).
/// Handles LEAF handshake, LS+/LS- interest propagation, and LMSG forwarding.
/// The JetStreamDomain property is propagated during handshake for domain-aware routing.
/// Go reference: leafnode.go.
/// </summary>
public sealed class LeafConnection(Socket socket) : IAsyncDisposable
{
    private readonly NetworkStream _stream = new(socket, ownsSocket: true);
    private readonly SemaphoreSlim _writeGate = new(1, 1);
    private readonly CancellationTokenSource _closedCts = new();
    private TimeSpan _connectDelay;
    private string? _remoteCluster;
    private Task? _loopTask;

    public string? RemoteId { get; internal set; }
    public string RemoteEndpoint => socket.RemoteEndPoint?.ToString() ?? Guid.NewGuid().ToString("N");
    public Func<RemoteSubscription, Task>? RemoteSubscriptionReceived { get; set; }
    public Func<LeafMessage, Task>? MessageReceived { get; set; }

    /// <summary>
    /// True when this connection was solicited (outbound dial) rather than accepted inbound.
    /// Go reference: isSolicitedLeafNode.
    /// </summary>
    public bool IsSolicited { get; internal set; }

    /// <summary>
    /// True when this connection is a spoke-side leaf connection.
    /// Go reference: isSpokeLeafNode / isHubLeafNode.
    /// </summary>
    public bool IsSpoke { get; set; }

    /// <summary>
    /// True when this leaf connection is isolated from hub propagation.
    /// Go reference: isIsolatedLeafNode.
    /// </summary>
    public bool Isolated { get; set; }

    /// <summary>
    /// JetStream domain for this leaf connection. When set, the domain is propagated
    /// in the LEAF handshake and included in LMSG frames for domain-aware routing.
    /// Go reference: leafnode.go — jsClusterDomain field in leafInfo.
    /// </summary>
    public string? JetStreamDomain { get; set; }

    /// <summary>
    /// The JetStream domain advertised by the remote side during handshake.
    /// </summary>
    public string? RemoteJetStreamDomain { get; private set; }

    /// <summary>
    /// The account associated with this leaf connection.
    /// Go reference: leafnode.go — client.acc field on leaf connections.
    /// </summary>
    public string? AccountName { get; set; }

    /// <summary>
    /// Subjects this leaf connection is allowed to publish to.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo, client.perm.pub.allow.
    /// </summary>
    public List<string> AllowedPublishSubjects { get; } = [];

    /// <summary>
    /// Subjects this leaf connection is allowed to subscribe to.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo, client.perm.sub.allow.
    /// </summary>
    public List<string> AllowedSubscribeSubjects { get; } = [];

    /// <summary>
    /// Indicates whether permissions have been synchronised to this leaf connection.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo sets leaf permissions.
    /// </summary>
    public bool PermsSynced { get; private set; }

    /// <summary>
    /// Returns the currently configured reconnect delay for this connection.
    /// Go reference: leafnode.go setLeafConnectDelayIfSoliciting.
    /// </summary>
    public TimeSpan GetConnectDelay() => _connectDelay;

    /// <summary>
    /// Sets the allowed publish and subscribe subjects for this connection and marks
    /// permissions as synced. Passing null for either list clears that list.
    /// Go reference: leafnode.go — sendPermsAndAccountInfo.
    /// </summary>
    public void SetPermissions(IEnumerable<string>? publishAllow, IEnumerable<string>? subscribeAllow)
    {
        AllowedPublishSubjects.Clear();
        AllowedSubscribeSubjects.Clear();

        if (publishAllow != null)
            AllowedPublishSubjects.AddRange(publishAllow);

        if (subscribeAllow != null)
            AllowedSubscribeSubjects.AddRange(subscribeAllow);

        PermsSynced = true;
    }

    public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        var handshakeLine = BuildHandshakeLine(serverId);
        await WriteLineAsync(handshakeLine, ct);
        var line = await ReadLineAsync(ct);
        ParseHandshakeResponse(line);
    }

    public async Task PerformInboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        var line = await ReadLineAsync(ct);
        ParseHandshakeResponse(line);
        var handshakeLine = BuildHandshakeLine(serverId);
        await WriteLineAsync(handshakeLine, ct);
    }

    public void StartLoop(CancellationToken ct)
    {
        if (_loopTask != null)
            return;

        var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _closedCts.Token);
        _loopTask = Task.Run(() => ReadLoopAsync(linked.Token), linked.Token);
    }

    public Task WaitUntilClosedAsync(CancellationToken ct)
        => _loopTask?.WaitAsync(ct) ?? Task.CompletedTask;

    public Task SendLsPlusAsync(string account, string subject, string? queue, CancellationToken ct)
        => SendLsPlusAsync(account, subject, queue, queueWeight: 0, ct);

    public Task SendLsPlusAsync(string account, string subject, string? queue, int queueWeight, CancellationToken ct)
    {
        string frame;
        if (queue is { Length: > 0 } && queueWeight > 0)
            frame = $"LS+ {account} {subject} {queue} {queueWeight}";
        else if (queue is { Length: > 0 })
            frame = $"LS+ {account} {subject} {queue}";
        else
            frame = $"LS+ {account} {subject}";

        return WriteLineAsync(frame, ct);
    }

    public Task SendLsMinusAsync(string account, string subject, string? queue, CancellationToken ct)
        => WriteLineAsync(queue is { Length: > 0 } ? $"LS- {account} {subject} {queue}" : $"LS- {account} {subject}", ct);

    /// <summary>
    /// Sends a CONNECT protocol line with JSON payload for solicited leaf links.
    /// Go reference: leafnode.go sendLeafConnect.
    /// </summary>
    public Task SendLeafConnectAsync(LeafConnectInfo connectInfo, CancellationToken ct)
    {
        ArgumentNullException.ThrowIfNull(connectInfo);
        var json = JsonSerializer.Serialize(connectInfo);
        return WriteLineAsync($"CONNECT {json}", ct);
    }

    public async Task SendMessageAsync(string account, string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        var reply = string.IsNullOrEmpty(replyTo) ? "-" : replyTo;
        await _writeGate.WaitAsync(ct);
        try
        {
            var control = Encoding.ASCII.GetBytes($"LMSG {account} {subject} {reply} {payload.Length}\r\n");
            await _stream.WriteAsync(control, ct);
            if (!payload.IsEmpty)
                await _stream.WriteAsync(payload, ct);
            await _stream.WriteAsync("\r\n"u8.ToArray(), ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _closedCts.CancelAsync();
        if (_loopTask != null)
            await _loopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);
        _closedCts.Dispose();
        _writeGate.Dispose();
        await _stream.DisposeAsync();
    }

    private string BuildHandshakeLine(string serverId)
    {
        if (!string.IsNullOrEmpty(JetStreamDomain))
            return $"LEAF {serverId} domain={JetStreamDomain}";

        return $"LEAF {serverId}";
    }

    public bool IsSolicitedLeafNode() => IsSolicited;
    public bool IsSpokeLeafNode() => IsSpoke;
    public bool IsHubLeafNode() => !IsSpoke;
    public bool IsIsolatedLeafNode() => Isolated;
    public string? RemoteCluster() => _remoteCluster;

    /// <summary>
    /// Applies connect delay only when this is a solicited leaf connection.
    /// Go reference: leafnode.go setLeafConnectDelayIfSoliciting.
    /// </summary>
    public void SetLeafConnectDelayIfSoliciting(TimeSpan delay)
    {
        if (IsSolicited)
            _connectDelay = delay;
    }

    /// <summary>
    /// Handles remote ERR protocol for leaf links and applies reconnect delay hints.
    /// Go reference: leafnode.go leafProcessErr.
    /// </summary>
    public void LeafProcessErr(string errStr)
    {
        if (string.IsNullOrWhiteSpace(errStr))
            return;

        if (errStr.Contains("permission", StringComparison.OrdinalIgnoreCase))
        {
            SetLeafConnectDelayIfSoliciting(LeafNodeManager.LeafNodeReconnectAfterPermViolation);
            return;
        }

        if (errStr.Contains("loop", StringComparison.OrdinalIgnoreCase))
        {
            SetLeafConnectDelayIfSoliciting(LeafNodeManager.LeafNodeReconnectDelayAfterLoopDetected);
            return;
        }

        if (errStr.Contains("cluster name", StringComparison.OrdinalIgnoreCase)
            && errStr.Contains("same", StringComparison.OrdinalIgnoreCase))
        {
            SetLeafConnectDelayIfSoliciting(LeafNodeManager.LeafNodeReconnectDelayAfterClusterNameSame);
        }
    }

    /// <summary>
    /// Handles subscription permission violations.
    /// Go reference: leafnode.go leafSubPermViolation.
    /// </summary>
    public void LeafSubPermViolation(string subj) => LeafPermViolation(pub: false, subj);

    /// <summary>
    /// Handles publish/subscribe permission violations.
    /// Go reference: leafnode.go leafPermViolation.
    /// </summary>
    public void LeafPermViolation(bool pub, string subj)
        => SetLeafConnectDelayIfSoliciting(LeafNodeManager.LeafNodeReconnectAfterPermViolation);

    private void ParseHandshakeResponse(string line)
    {
        if (!line.StartsWith("LEAF ", StringComparison.OrdinalIgnoreCase))
            throw new InvalidOperationException("Invalid leaf handshake");

        var rest = line[5..].Trim();
        if (rest.Length == 0)
            throw new InvalidOperationException("Leaf handshake missing id");

        // Parse "serverId [domain=xxx]" format
        var spaceIdx = rest.IndexOf(' ');
        if (spaceIdx > 0)
        {
            RemoteId = rest[..spaceIdx];
            var attrs = rest[(spaceIdx + 1)..];
            foreach (var token in attrs.Split(' ', StringSplitOptions.RemoveEmptyEntries))
            {
                const string domainPrefix = "domain=";
                if (token.StartsWith(domainPrefix, StringComparison.OrdinalIgnoreCase))
                {
                    RemoteJetStreamDomain = token[domainPrefix.Length..].Trim();
                    continue;
                }

                const string clusterPrefix = "cluster=";
                if (token.StartsWith(clusterPrefix, StringComparison.OrdinalIgnoreCase))
                    _remoteCluster = token[clusterPrefix.Length..].Trim();
            }
        }
        else
        {
            RemoteId = rest;
        }
    }

    private async Task ReadLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            string line;
            try
            {
                line = await ReadLineAsync(ct);
            }
            catch
            {
                break;
            }

            if (line.StartsWith("LS+ ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (RemoteSubscriptionReceived != null &&
                    TryParseAccountScopedInterest(parts, out var parsedAccount, out var parsedSubject, out var queue, out var queueWeight))
                {
                    await RemoteSubscriptionReceived(new RemoteSubscription(parsedSubject, queue, RemoteId ?? string.Empty, parsedAccount, QueueWeight: queueWeight));
                }
                continue;
            }

            if (line.StartsWith("LS- ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (RemoteSubscriptionReceived != null &&
                    TryParseAccountScopedInterest(parts, out var parsedAccount, out var parsedSubject, out var queue, out _))
                {
                    await RemoteSubscriptionReceived(RemoteSubscription.Removal(parsedSubject, queue, RemoteId ?? string.Empty, parsedAccount));
                }
                continue;
            }

            if (!line.StartsWith("LMSG ", StringComparison.Ordinal))
                continue;

            var args = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
            if (args.Length < 4)
                continue;

            var account = "$G";
            string subject;
            string replyToken;
            string sizeToken;

            // New format: LMSG <account> <subject> <reply> <size>
            // Legacy format: LMSG <subject> <reply> <size>
            if (args.Length >= 5 && !LooksLikeSubject(args[1]))
            {
                account = args[1];
                subject = args[2];
                replyToken = args[3];
                sizeToken = args[4];
            }
            else
            {
                subject = args[1];
                replyToken = args[2];
                sizeToken = args[3];
            }

            if (!int.TryParse(sizeToken, out var size) || size < 0)
                continue;

            var payload = await ReadPayloadAsync(size, ct);
            if (MessageReceived != null)
                await MessageReceived(new LeafMessage(subject, replyToken == "-" ? null : replyToken, payload, account));
        }
    }

    private async Task<ReadOnlyMemory<byte>> ReadPayloadAsync(int size, CancellationToken ct)
    {
        var payload = new byte[size];
        var offset = 0;
        while (offset < size)
        {
            var read = await _stream.ReadAsync(payload.AsMemory(offset, size - offset), ct);
            if (read == 0)
                throw new IOException("Leaf payload read closed");
            offset += read;
        }

        var trailer = new byte[2];
        _ = await _stream.ReadAsync(trailer, ct);
        return payload;
    }

    private async Task WriteLineAsync(string line, CancellationToken ct)
    {
        await _writeGate.WaitAsync(ct);
        try
        {
            var bytes = Encoding.ASCII.GetBytes($"{line}\r\n");
            await _stream.WriteAsync(bytes, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    private async Task<string> ReadLineAsync(CancellationToken ct)
    {
        var bytes = new List<byte>(64);
        var single = new byte[1];
        while (true)
        {
            var read = await _stream.ReadAsync(single, ct);
            if (read == 0)
                throw new IOException("Leaf closed");
            if (single[0] == (byte)'\n')
                break;
            if (single[0] != (byte)'\r')
                bytes.Add(single[0]);
        }

        return Encoding.ASCII.GetString([.. bytes]);
    }

    private static bool TryParseAccountScopedInterest(string[] parts, out string account, out string subject, out string? queue, out int queueWeight)
    {
        account = "$G";
        subject = string.Empty;
        queue = null;
        queueWeight = 1;

        if (parts.Length < 2)
            return false;

        // New format: LS+ <account> <subject> [queue]
        // Legacy format: LS+ <subject> [queue]
        if (parts.Length >= 3 && !LooksLikeSubject(parts[1]))
        {
            account = parts[1];
            subject = parts[2];
            queue = parts.Length >= 4 ? parts[3] : null;
            if (queue is { Length: > 0 } && parts.Length >= 5)
                queueWeight = ParseQueueWeight(parts[4]);
            return true;
        }

        subject = parts[1];
        queue = parts.Length >= 3 ? parts[2] : null;
        if (queue is { Length: > 0 } && parts.Length >= 4)
            queueWeight = ParseQueueWeight(parts[3]);
        return true;
    }

    private static bool LooksLikeSubject(string token)
        => token.Contains('.', StringComparison.Ordinal)
            || token.Contains('*', StringComparison.Ordinal)
            || token.Contains('>', StringComparison.Ordinal);

    private static int ParseQueueWeight(string token)
        => int.TryParse(token, out var parsed) && parsed > 0 ? parsed : 1;
}

public sealed record LeafMessage(string Subject, string? ReplyTo, ReadOnlyMemory<byte> Payload, string Account = "$G");
