using System.Collections.Concurrent;
using System.Net.Sockets;
using System.Text;
using NATS.Server.Subscriptions;

namespace NATS.Server.Gateways;

public sealed class GatewayConnection(Socket socket) : IAsyncDisposable
{
    private readonly NetworkStream _stream = new(socket, ownsSocket: true);
    private readonly SemaphoreSlim _writeGate = new(1, 1);
    private readonly CancellationTokenSource _closedCts = new();
    private readonly GatewayInterestTracker _interestTracker = new();
    private readonly ConcurrentDictionary<string, HashSet<string>> _accountSubscriptions = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, HashSet<string>> _queueSubscriptions = new(StringComparer.Ordinal);
    private Task? _loopTask;

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

    /// <summary>
    /// Per-connection interest mode tracker.
    /// Go: gateway.go:100-150 — each outbound gateway connection maintains its own interest state.
    /// </summary>
    public GatewayInterestTracker InterestTracker => _interestTracker;

    /// <summary>
    /// Adds a subject to the account-specific subscription set for this gateway connection.
    /// Go: gateway.go — per-account subscription routing state on outbound connections.
    /// </summary>
    public void AddAccountSubscription(string account, string subject)
    {
        var subs = _accountSubscriptions.GetOrAdd(account, _ => new HashSet<string>(StringComparer.Ordinal));
        lock (subs) subs.Add(subject);
    }

    /// <summary>
    /// Removes a subject from the account-specific subscription set for this gateway connection.
    /// </summary>
    public void RemoveAccountSubscription(string account, string subject)
    {
        if (_accountSubscriptions.TryGetValue(account, out var subs))
            lock (subs) subs.Remove(subject);
    }

    /// <summary>
    /// Returns a snapshot of all subjects tracked for the given account on this connection.
    /// </summary>
    public IReadOnlySet<string> GetAccountSubscriptions(string account)
    {
        if (_accountSubscriptions.TryGetValue(account, out var subs))
            lock (subs) return new HashSet<string>(subs, StringComparer.Ordinal);
        return new HashSet<string>(StringComparer.Ordinal);
    }

    /// <summary>
    /// Returns the number of subjects tracked for the given account. Returns 0 for unknown accounts.
    /// </summary>
    public int AccountSubscriptionCount(string account)
    {
        if (_accountSubscriptions.TryGetValue(account, out var subs))
            lock (subs) return subs.Count;
        return 0;
    }

    /// <summary>
    /// Registers a queue group subscription for propagation to this gateway.
    /// Go reference: gateway.go — sendQueueSubsToGateway.
    /// </summary>
    public void AddQueueSubscription(string subject, string queueGroup)
    {
        var groups = _queueSubscriptions.GetOrAdd(subject, _ => new HashSet<string>(StringComparer.Ordinal));
        lock (groups) groups.Add(queueGroup);
    }

    /// <summary>
    /// Removes a queue group subscription from this gateway connection's tracking state.
    /// Go reference: gateway.go — sendQueueSubsToGateway (removal path).
    /// </summary>
    public void RemoveQueueSubscription(string subject, string queueGroup)
    {
        if (_queueSubscriptions.TryGetValue(subject, out var groups))
            lock (groups) groups.Remove(queueGroup);
    }

    /// <summary>
    /// Returns a snapshot of all queue group names registered for the given subject.
    /// </summary>
    public IReadOnlySet<string> GetQueueGroups(string subject)
    {
        if (_queueSubscriptions.TryGetValue(subject, out var groups))
            lock (groups) return new HashSet<string>(groups, StringComparer.Ordinal);
        return new HashSet<string>(StringComparer.Ordinal);
    }

    /// <summary>
    /// Returns the number of distinct subjects that have at least one queue group registered.
    /// </summary>
    public int QueueSubscriptionCount => _queueSubscriptions.Count;

    /// <summary>
    /// Returns true if the given subject/queueGroup pair is currently registered on this gateway connection.
    /// </summary>
    public bool HasQueueSubscription(string subject, string queueGroup)
    {
        if (!_queueSubscriptions.TryGetValue(subject, out var groups))
            return false;
        lock (groups) return groups.Contains(queueGroup);
    }

    public async Task PerformOutboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        await WriteLineAsync($"GATEWAY {serverId}", ct);
        var line = await ReadLineAsync(ct);
        RemoteId = ParseHandshake(line);
    }

    public async Task PerformInboundHandshakeAsync(string serverId, CancellationToken ct)
    {
        var line = await ReadLineAsync(ct);
        RemoteId = ParseHandshake(line);
        await WriteLineAsync($"GATEWAY {serverId}", 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 SendAPlusAsync(string account, string subject, string? queue, CancellationToken ct)
        => WriteLineAsync(queue is { Length: > 0 } ? $"A+ {account} {subject} {queue}" : $"A+ {account} {subject}", ct);

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

    public async Task SendMessageAsync(string account, string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        // Go: gateway.go:2900 (shouldForwardMsg) — check interest tracker before sending
        if (!_interestTracker.ShouldForward(account, subject))
            return;

        var reply = string.IsNullOrEmpty(replyTo) ? "-" : replyTo;
        await _writeGate.WaitAsync(ct);
        try
        {
            var control = Encoding.ASCII.GetBytes($"GMSG {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 async Task ReadLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            string line;
            try
            {
                line = await ReadLineAsync(ct);
            }
            catch
            {
                break;
            }

            if (line.StartsWith("A+ ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (TryParseAccountScopedInterest(parts, out var parsedAccount, out var parsedSubject, out var queue))
                {
                    // Go: gateway.go:1540 — track positive interest on A+
                    _interestTracker.TrackInterest(parsedAccount, parsedSubject);
                    if (RemoteSubscriptionReceived != null)
                        await RemoteSubscriptionReceived(new RemoteSubscription(parsedSubject, queue, RemoteId ?? string.Empty, parsedAccount));
                }
                continue;
            }

            if (line.StartsWith("A- ", StringComparison.Ordinal))
            {
                var parts = line.Split(' ', StringSplitOptions.RemoveEmptyEntries);
                if (TryParseAccountScopedInterest(parts, out var parsedAccount, out var parsedSubject, out var queue))
                {
                    // Go: gateway.go:1560 — track no-interest on A-, may trigger mode switch
                    _interestTracker.TrackNoInterest(parsedAccount, parsedSubject);
                    if (RemoteSubscriptionReceived != null)
                        await RemoteSubscriptionReceived(RemoteSubscription.Removal(parsedSubject, queue, RemoteId ?? string.Empty, parsedAccount));
                }
                continue;
            }

            if (!line.StartsWith("GMSG ", 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: GMSG <account> <subject> <reply> <size>
            // Legacy format: GMSG <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 GatewayMessage(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("Gateway 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("Gateway closed");
            if (single[0] == (byte)'\n')
                break;
            if (single[0] != (byte)'\r')
                bytes.Add(single[0]);
        }

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

    private static string ParseHandshake(string line)
    {
        if (!line.StartsWith("GATEWAY ", StringComparison.OrdinalIgnoreCase))
            throw new InvalidOperationException("Invalid gateway handshake");

        var id = line[8..].Trim();
        if (id.Length == 0)
            throw new InvalidOperationException("Gateway handshake missing id");
        return id;
    }

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

        if (parts.Length < 2)
            return false;

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

        subject = parts[1];
        queue = parts.Length >= 3 ? parts[2] : null;
        return true;
    }

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

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