using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using Microsoft.Extensions.Logging;
using NATS.Server.Configuration;
using NATS.Server.Subscriptions;

namespace NATS.Server.LeafNodes;

/// <summary>
/// Manages leaf node connections — both inbound (accepted) and outbound (solicited).
/// Outbound connections use exponential backoff retry: 1s, 2s, 4s, ..., capped at 60s.
/// Subject filtering via DenyExports (hub→leaf) and DenyImports (leaf→hub) is applied
/// to both message forwarding and subscription propagation.
/// Go reference: leafnode.go.
/// </summary>
public sealed class LeafNodeManager : IAsyncDisposable
{
    private readonly LeafNodeOptions _options;
    private readonly ServerStats _stats;
    private readonly string _serverId;
    private readonly Action<RemoteSubscription> _remoteSubSink;
    private readonly Action<LeafMessage> _messageSink;
    private readonly ILogger<LeafNodeManager> _logger;
    private readonly ConcurrentDictionary<string, LeafConnection> _connections = new(StringComparer.Ordinal);
    private readonly LeafHubSpokeMapper _subjectFilter;

    private CancellationTokenSource? _cts;
    private Socket? _listener;
    private Task? _acceptLoopTask;

    /// <summary>
    /// Initial retry delay for solicited connections (1 second).
    /// Go reference: leafnode.go — DEFAULT_LEAF_NODE_RECONNECT constant.
    /// </summary>
    internal static readonly TimeSpan InitialRetryDelay = TimeSpan.FromSeconds(1);

    /// <summary>
    /// Maximum retry delay for solicited connections (60 seconds).
    /// </summary>
    internal static readonly TimeSpan MaxRetryDelay = TimeSpan.FromSeconds(60);

    public string ListenEndpoint => $"{_options.Host}:{_options.Port}";

    public LeafNodeManager(
        LeafNodeOptions options,
        ServerStats stats,
        string serverId,
        Action<RemoteSubscription> remoteSubSink,
        Action<LeafMessage> messageSink,
        ILogger<LeafNodeManager> logger)
    {
        _options = options;
        _stats = stats;
        _serverId = serverId;
        _remoteSubSink = remoteSubSink;
        _messageSink = messageSink;
        _logger = logger;
        _subjectFilter = new LeafHubSpokeMapper(
            new Dictionary<string, string>(),
            options.DenyExports,
            options.DenyImports,
            options.ExportSubjects,
            options.ImportSubjects);
    }

    public Task StartAsync(CancellationToken ct)
    {
        _cts = CancellationTokenSource.CreateLinkedTokenSource(ct);
        _listener = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        _listener.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.ReuseAddress, true);
        _listener.Bind(new IPEndPoint(IPAddress.Parse(_options.Host), _options.Port));
        _listener.Listen(128);

        if (_options.Port == 0)
            _options.Port = ((IPEndPoint)_listener.LocalEndPoint!).Port;

        _acceptLoopTask = Task.Run(() => AcceptLoopAsync(_cts.Token));
        foreach (var remote in _options.Remotes.Distinct(StringComparer.OrdinalIgnoreCase))
            _ = Task.Run(() => ConnectSolicitedWithRetryAsync(remote, _options.JetStreamDomain, _cts.Token));

        _logger.LogDebug("Leaf manager started (listen={Host}:{Port})", _options.Host, _options.Port);
        return Task.CompletedTask;
    }

    /// <summary>
    /// Establishes a single solicited (outbound) leaf connection to the specified URL.
    /// Performs socket connection and LEAF handshake. If a JetStream domain is specified,
    /// it is propagated during the handshake.
    /// Go reference: leafnode.go — connectSolicited.
    /// </summary>
    public async Task<LeafConnection> ConnectSolicitedAsync(string url, string? account, CancellationToken ct)
    {
        var endPoint = ParseEndpoint(url);
        var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
        try
        {
            await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
            var connection = new LeafConnection(socket)
            {
                JetStreamDomain = _options.JetStreamDomain,
            };
            await connection.PerformOutboundHandshakeAsync(_serverId, ct);
            Register(connection);
            _logger.LogDebug("Solicited leaf connection established to {Url} (account={Account})", url, account ?? "$G");
            return connection;
        }
        catch
        {
            socket.Dispose();
            throw;
        }
    }

    public async Task ForwardMessageAsync(string account, string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        // Apply subject filtering: outbound direction is hub→leaf (DenyExports).
        // The subject may be loop-marked ($LDS.{serverId}.{realSubject}), so we
        // strip the marker before checking the filter against the logical subject.
        // Go reference: leafnode.go:475-478 (DenyExports → Publish deny list).
        var filterSubject = LeafLoopDetector.TryUnmark(subject, out var unmarked) ? unmarked : subject;
        if (!_subjectFilter.IsSubjectAllowed(filterSubject, LeafMapDirection.Outbound))
        {
            _logger.LogDebug("Leaf outbound message denied for subject {Subject} (DenyExports)", filterSubject);
            return;
        }

        foreach (var connection in _connections.Values)
            await connection.SendMessageAsync(account, subject, replyTo, payload, ct);
    }

    public void PropagateLocalSubscription(string account, string subject, string? queue)
    {
        // Subscription propagation is also subject to export filtering:
        // we don't propagate subscriptions for subjects that are denied.
        if (!_subjectFilter.IsSubjectAllowed(subject, LeafMapDirection.Outbound))
        {
            _logger.LogDebug("Leaf subscription propagation denied for subject {Subject} (DenyExports)", subject);
            return;
        }

        foreach (var connection in _connections.Values)
            _ = connection.SendLsPlusAsync(account, subject, queue, _cts?.Token ?? CancellationToken.None);
    }

    public void PropagateLocalUnsubscription(string account, string subject, string? queue)
    {
        foreach (var connection in _connections.Values)
            _ = connection.SendLsMinusAsync(account, subject, queue, _cts?.Token ?? CancellationToken.None);
    }

    public async ValueTask DisposeAsync()
    {
        if (_cts == null)
            return;

        await _cts.CancelAsync();
        _listener?.Dispose();
        if (_acceptLoopTask != null)
            await _acceptLoopTask.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

        foreach (var connection in _connections.Values)
            await connection.DisposeAsync();
        _connections.Clear();
        Interlocked.Exchange(ref _stats.Leafs, 0);
        _cts.Dispose();
        _cts = null;
        _logger.LogDebug("Leaf manager stopped");
    }

    /// <summary>
    /// Computes the next backoff delay using exponential backoff with a cap.
    /// Delay sequence: 1s, 2s, 4s, 8s, 16s, 32s, 60s, 60s, ...
    /// </summary>
    internal static TimeSpan ComputeBackoff(int attempt)
    {
        if (attempt < 0) attempt = 0;
        var seconds = Math.Min(InitialRetryDelay.TotalSeconds * Math.Pow(2, attempt), MaxRetryDelay.TotalSeconds);
        return TimeSpan.FromSeconds(seconds);
    }

    private async Task AcceptLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            Socket socket;
            try
            {
                socket = await _listener!.AcceptAsync(ct);
            }
            catch
            {
                break;
            }

            _ = Task.Run(() => HandleInboundAsync(socket, ct), ct);
        }
    }

    private async Task HandleInboundAsync(Socket socket, CancellationToken ct)
    {
        var connection = new LeafConnection(socket)
        {
            JetStreamDomain = _options.JetStreamDomain,
        };
        try
        {
            await connection.PerformInboundHandshakeAsync(_serverId, ct);
            Register(connection);
        }
        catch
        {
            await connection.DisposeAsync();
        }
    }

    private async Task ConnectSolicitedWithRetryAsync(string remote, string? jetStreamDomain, CancellationToken ct)
    {
        var attempt = 0;
        while (!ct.IsCancellationRequested)
        {
            try
            {
                var endPoint = ParseEndpoint(remote);
                var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                try
                {
                    await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
                    var connection = new LeafConnection(socket)
                    {
                        JetStreamDomain = jetStreamDomain,
                    };
                    await connection.PerformOutboundHandshakeAsync(_serverId, ct);
                    Register(connection);
                    _logger.LogDebug("Solicited leaf connection established to {Remote}", remote);
                    return;
                }
                catch
                {
                    socket.Dispose();
                    throw;
                }
            }
            catch (OperationCanceledException)
            {
                return;
            }
            catch (Exception ex)
            {
                _logger.LogDebug(ex, "Leaf connect retry for {Remote} (attempt {Attempt})", remote, attempt);
            }

            var delay = ComputeBackoff(attempt);
            attempt++;
            try
            {
                await Task.Delay(delay, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
        }
    }

    private void Register(LeafConnection connection)
    {
        var key = $"{connection.RemoteId}:{connection.RemoteEndpoint}:{Guid.NewGuid():N}";
        if (!_connections.TryAdd(key, connection))
        {
            _ = connection.DisposeAsync();
            return;
        }

        connection.RemoteSubscriptionReceived = sub =>
        {
            _remoteSubSink(sub);
            return Task.CompletedTask;
        };
        connection.MessageReceived = msg =>
        {
            // Apply inbound filtering: DenyImports restricts leaf→hub messages.
            // The subject may be loop-marked ($LDS.{serverId}.{realSubject}), so we
            // strip the marker before checking the filter against the logical subject.
            // Go reference: leafnode.go:480-481 (DenyImports → Subscribe deny list).
            var filterSubject = LeafLoopDetector.TryUnmark(msg.Subject, out var unmarked)
                ? unmarked
                : msg.Subject;
            if (!_subjectFilter.IsSubjectAllowed(filterSubject, LeafMapDirection.Inbound))
            {
                _logger.LogDebug("Leaf inbound message denied for subject {Subject} (DenyImports)", filterSubject);
                return Task.CompletedTask;
            }

            _messageSink(msg);
            return Task.CompletedTask;
        };
        connection.StartLoop(_cts!.Token);
        Interlocked.Increment(ref _stats.Leafs);
        _ = Task.Run(() => WatchConnectionAsync(key, connection, _cts!.Token));
    }

    private async Task WatchConnectionAsync(string key, LeafConnection connection, CancellationToken ct)
    {
        try
        {
            await connection.WaitUntilClosedAsync(ct);
        }
        catch
        {
        }
        finally
        {
            if (_connections.TryRemove(key, out _))
                Interlocked.Decrement(ref _stats.Leafs);
            await connection.DisposeAsync();
        }
    }

    private static IPEndPoint ParseEndpoint(string endpoint)
    {
        var parts = endpoint.Split(':', 2, StringSplitOptions.TrimEntries | StringSplitOptions.RemoveEmptyEntries);
        if (parts.Length != 2)
            throw new FormatException($"Invalid endpoint: {endpoint}");

        return new IPEndPoint(IPAddress.Parse(parts[0]), int.Parse(parts[1]));
    }
}
