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.Gateways;

/// <summary>
/// Lifecycle states for a registered gateway connection.
/// Go reference: server/gateway.go gwConnState / gwReplyMapping.
/// </summary>
public enum GatewayConnectionState
{
    Connecting,
    Connected,
    Disconnected,
    Draining,
}

/// <summary>
/// Tracks the registration and runtime metrics for a single named gateway connection.
/// Go reference: server/gateway.go srvGateway / outboundGateway structs.
/// </summary>
public sealed class GatewayRegistration
{
    internal long _messagesSent;
    internal long _messagesReceived;

    public required string Name { get; init; }
    public GatewayConnectionState State { get; set; } = GatewayConnectionState.Connecting;
    public DateTime ConnectedAtUtc { get; set; }
    public DateTime? DisconnectedAtUtc { get; set; }
    public string? RemoteAddress { get; set; }
    public long MessagesSent { get => Interlocked.Read(ref _messagesSent); set => Interlocked.Exchange(ref _messagesSent, value); }
    public long MessagesReceived { get => Interlocked.Read(ref _messagesReceived); set => Interlocked.Exchange(ref _messagesReceived, value); }
}

/// <summary>
/// Controls retry timing for outbound gateway reconnections using exponential backoff with jitter.
/// Go reference: server/gateway.go solicitGateway / reconnectGateway delay logic.
/// </summary>
public sealed class GatewayReconnectPolicy
{
    public TimeSpan InitialDelay { get; init; } = TimeSpan.FromSeconds(1);
    public TimeSpan MaxDelay { get; init; } = TimeSpan.FromSeconds(30);
    public double JitterFactor { get; init; } = 0.2;
    public int MaxAttempts { get; init; } = int.MaxValue; // 0 = unlimited

    public TimeSpan CalculateDelay(int attempt)
    {
        var baseDelay = InitialDelay.TotalMilliseconds * Math.Pow(2, Math.Min(attempt, 10));
        var capped = Math.Min(baseDelay, MaxDelay.TotalMilliseconds);
        return TimeSpan.FromMilliseconds(capped);
    }

    public TimeSpan CalculateDelayWithJitter(int attempt)
    {
        var delay = CalculateDelay(attempt);
        var jitter = Random.Shared.NextDouble() * JitterFactor * delay.TotalMilliseconds;
        return TimeSpan.FromMilliseconds(delay.TotalMilliseconds + jitter);
    }
}

public sealed class GatewayManager : IAsyncDisposable
{
    public const string GatewayTlsInsecureWarning =
        "Gateway TLS insecure configuration is enabled; verify certificates and hostname validation for production.";

    private readonly GatewayOptions _options;
    private readonly ServerStats _stats;
    private readonly string _serverId;
    private readonly Action<RemoteSubscription> _remoteSubSink;
    private readonly Action<GatewayMessage> _messageSink;
    private readonly ILogger<GatewayManager> _logger;
    private readonly ConcurrentDictionary<string, GatewayConnection> _connections = new(StringComparer.Ordinal);
    private readonly HashSet<string> _discoveredGateways = new(StringComparer.OrdinalIgnoreCase);
    private readonly ConcurrentDictionary<string, int> _reconnectAttempts = new(StringComparer.OrdinalIgnoreCase);
    private readonly ConcurrentDictionary<string, GatewayRegistration> _registrations = new(StringComparer.OrdinalIgnoreCase);
    private long _forwardedJetStreamClusterMessages;

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

    public string ListenEndpoint => $"{_options.Host}:{_options.Port}";
    public long ForwardedJetStreamClusterMessages => Interlocked.Read(ref _forwardedJetStreamClusterMessages);

    internal static bool ShouldForwardInterestOnly(SubList subList, string account, string subject)
        => subList.HasRemoteInterest(account, subject);

    public GatewayManager(
        GatewayOptions options,
        ServerStats stats,
        string serverId,
        Action<RemoteSubscription> remoteSubSink,
        Action<GatewayMessage> messageSink,
        ILogger<GatewayManager> logger)
    {
        _options = options;
        _stats = stats;
        _serverId = serverId;
        _remoteSubSink = remoteSubSink;
        _messageSink = messageSink;
        _logger = logger;
    }

    /// <summary>
    /// Validates gateway options for required fields and basic endpoint correctness.
    /// Go reference: validateGatewayOptions.
    /// </summary>
    public static bool ValidateGatewayOptions(GatewayOptions? options, out string? error)
    {
        if (options is null)
        {
            error = "Gateway options are required.";
            return false;
        }

        if (string.IsNullOrWhiteSpace(options.Name))
        {
            error = "Gateway name is required.";
            return false;
        }

        if (options.Port is < 0 or > 65535)
        {
            error = "Gateway port must be in range 0-65535.";
            return false;
        }

        foreach (var remote in options.Remotes)
        {
            if (string.IsNullOrWhiteSpace(remote))
            {
                error = "Gateway remote entries cannot be empty.";
                return false;
            }
        }

        error = null;
        return true;
    }

    /// <summary>
    /// Gateway clusters auto-discovered via INFO gossip.
    /// Go reference: server/gateway.go processImplicitGateway.
    /// </summary>
    public IReadOnlyCollection<string> DiscoveredGateways
    {
        get { lock (_discoveredGateways) return _discoveredGateways.ToList(); }
    }

    /// <summary>
    /// Processes a gateway info message from a peer, discovering new gateway clusters.
    /// Go reference: server/gateway.go:800-850 (processImplicitGateway).
    /// </summary>
    public void ProcessImplicitGateway(GatewayInfo gwInfo)
    {
        ArgumentNullException.ThrowIfNull(gwInfo);

        lock (_discoveredGateways)
        {
            _discoveredGateways.Add(gwInfo.Name);
        }
    }

    /// <summary>
    /// Returns the current reconnect attempt count for a named gateway.
    /// Returns 0 if no reconnect attempt has been recorded yet.
    /// Go reference: server/gateway.go reconnectGateway attempt tracking.
    /// </summary>
    public int GetReconnectAttempts(string gatewayName)
        => _reconnectAttempts.TryGetValue(gatewayName, out var n) ? n : 0;

    /// <summary>
    /// Resets the reconnect attempt counter for a named gateway (called on successful connection).
    /// Go reference: server/gateway.go solicitGateway successful connect path.
    /// </summary>
    public void ResetReconnectAttempts(string gatewayName)
        => _reconnectAttempts.TryRemove(gatewayName, out _);

    /// <summary>
    /// Performs a single reconnect cycle for a named gateway using exponential backoff.
    /// Increments the attempt counter, waits the backoff delay, then attempts to connect.
    /// Go reference: server/gateway.go reconnectGateway / solicitGateway.
    /// </summary>
    public async Task ReconnectGatewayAsync(string gatewayName, CancellationToken ct)
    {
        var attempt = _reconnectAttempts.AddOrUpdate(gatewayName, 1, (_, n) => n + 1);
        var policy = new GatewayReconnectPolicy();
        var delay = policy.CalculateDelayWithJitter(attempt - 1);

        _logger.LogDebug("Gateway reconnect attempt {Attempt} for {Gateway}, delay={Delay}ms",
            attempt, gatewayName, delay.TotalMilliseconds);

        await Task.Delay(delay, ct);
    }

    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));

        // Collect outbound endpoints from both the legacy Remotes list and the
        // config-parsed RemoteGateways list (populated by ConfigProcessor).
        var endpoints = new HashSet<string>(_options.Remotes, StringComparer.OrdinalIgnoreCase);
        foreach (var rgw in _options.RemoteGateways)
        {
            foreach (var uri in rgw.GetUrls())
                endpoints.Add($"{uri.Host}:{uri.Port}");
        }

        foreach (var remote in endpoints)
            _ = Task.Run(() => ConnectWithRetryAsync(remote, _cts.Token));

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

    public async Task ForwardMessageAsync(string account, string subject, string? replyTo, ReadOnlyMemory<byte> payload, CancellationToken ct)
    {
        foreach (var connection in _connections.Values)
            await connection.SendMessageAsync(account, subject, replyTo, payload, ct);
    }

    public async Task ForwardJetStreamClusterMessageAsync(GatewayMessage message, CancellationToken ct)
    {
        Interlocked.Increment(ref _forwardedJetStreamClusterMessages);
        await ForwardMessageAsync(message.Account, message.Subject, message.ReplyTo, message.Payload, ct);
    }

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

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

    /// <summary>
    /// Sends a batch of account-scoped subscriptions to a named gateway connection, recording
    /// them in that connection's per-account subscription set.
    /// Go: gateway.go — account-specific subscription propagation on outbound routes.
    /// </summary>
    public void SendAccountSubscriptions(string gatewayName, string account, IEnumerable<string> subjects)
    {
        if (!_connections.TryGetValue(gatewayName, out var conn)) return;
        foreach (var subject in subjects)
            conn.AddAccountSubscription(account, subject);
    }

    /// <summary>
    /// Returns a snapshot of all subjects tracked for the given account on the named gateway connection.
    /// Returns an empty set when the connection is not found.
    /// </summary>
    public IReadOnlySet<string> GetAccountSubscriptions(string gatewayName, string account)
    {
        if (!_connections.TryGetValue(gatewayName, out var conn))
            return new HashSet<string>(StringComparer.Ordinal);
        return conn.GetAccountSubscriptions(account);
    }

    // ── Gateway connection registry (Gap 11.7) ────────────────────────────────

    /// <summary>
    /// Registers a new gateway by name, starting in the Connecting state.
    /// Go reference: server/gateway.go solicitGateway creates outbound entry before dialling.
    /// </summary>
    public void RegisterGateway(string name, string? remoteAddress = null)
    {
        var reg = new GatewayRegistration
        {
            Name = name,
            State = GatewayConnectionState.Connecting,
            RemoteAddress = remoteAddress,
        };
        _registrations[name] = reg;
    }

    /// <summary>
    /// Updates the connection state of a registered gateway.
    /// Setting Connected stamps ConnectedAtUtc; setting Disconnected stamps DisconnectedAtUtc.
    /// Go reference: server/gateway.go gwConnState transitions.
    /// </summary>
    public void UpdateState(string name, GatewayConnectionState state)
    {
        if (!_registrations.TryGetValue(name, out var reg)) return;
        reg.State = state;
        if (state == GatewayConnectionState.Connected)
            reg.ConnectedAtUtc = DateTime.UtcNow;
        else if (state == GatewayConnectionState.Disconnected)
            reg.DisconnectedAtUtc = DateTime.UtcNow;
    }

    /// <summary>
    /// Returns the registration for the named gateway, or null if not registered.
    /// Go reference: server/gateway.go server.getOutboundGatewayConnection.
    /// </summary>
    public GatewayRegistration? GetRegistration(string name)
        => _registrations.TryGetValue(name, out var reg) ? reg : null;

    /// <summary>
    /// Returns a snapshot of all current gateway registrations.
    /// </summary>
    public IReadOnlyCollection<GatewayRegistration> GetAllRegistrations()
        => [.. _registrations.Values];

    /// <summary>
    /// Removes the named gateway registration.
    /// Go reference: server/gateway.go outboundGateway teardown.
    /// </summary>
    public void UnregisterGateway(string name)
        => _registrations.TryRemove(name, out _);

    /// <summary>
    /// Returns the number of gateways currently in the Connected state.
    /// Go reference: server/gateway.go numOutboundGatewayConnections.
    /// </summary>
    public int GetConnectedGatewayCount()
        => _registrations.Values.Count(r => r.State == GatewayConnectionState.Connected);

    /// <summary>
    /// Returns the number of active outbound gateway connections.
    /// Go reference: server/gateway.go NumOutboundGateways / numOutboundGateways.
    /// </summary>
    public int NumOutboundGateways()
        => _connections.Values.Count(c => c.IsOutbound);

    /// <summary>
    /// Returns the number of active inbound gateway connections.
    /// Go reference: server/gateway.go numInboundGateways.
    /// </summary>
    public int NumInboundGateways()
        => _connections.Values.Count(c => !c.IsOutbound);

    /// <summary>
    /// Returns true if an inbound gateway connection exists for the given remote server id.
    /// Go reference: server/gateway.go srvGateway.hasInbound.
    /// </summary>
    public bool HasInbound(string remoteServerId)
        => _connections.Values.Any(c => !c.IsOutbound && string.Equals(c.RemoteId, remoteServerId, StringComparison.Ordinal));

    /// <summary>
    /// Returns the first outbound gateway connection for the given remote server id, or null.
    /// Go reference: server/gateway.go getOutboundGatewayConnection.
    /// </summary>
    public GatewayConnection? GetOutboundGatewayConnection(string remoteServerId)
        => _connections.Values.FirstOrDefault(c => c.IsOutbound && string.Equals(c.RemoteId, remoteServerId, StringComparison.Ordinal));

    /// <summary>
    /// Returns all outbound gateway connections.
    /// Go reference: server/gateway.go getOutboundGatewayConnections.
    /// </summary>
    public IReadOnlyList<GatewayConnection> GetOutboundGatewayConnections()
        => _connections.Values.Where(c => c.IsOutbound).ToList();

    /// <summary>
    /// Returns all inbound gateway connections.
    /// Go reference: server/gateway.go getInboundGatewayConnections.
    /// </summary>
    public IReadOnlyList<GatewayConnection> GetInboundGatewayConnections()
        => _connections.Values.Where(c => !c.IsOutbound).ToList();

    /// <summary>
    /// Atomically increments the messages-sent counter for the named gateway.
    /// Go reference: server/gateway.go outboundGateway.msgs.
    /// </summary>
    public void IncrementMessagesSent(string name)
    {
        if (_registrations.TryGetValue(name, out var reg))
            Interlocked.Increment(ref reg._messagesSent);
    }

    /// <summary>
    /// Atomically increments the messages-received counter for the named gateway.
    /// Go reference: server/gateway.go inboundGateway.msgs.
    /// </summary>
    public void IncrementMessagesReceived(string name)
    {
        if (_registrations.TryGetValue(name, out var reg))
            Interlocked.Increment(ref reg._messagesReceived);
    }

    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.Gateways, 0);
        _cts.Dispose();
        _cts = null;
        _logger.LogDebug("Gateway manager stopped");
    }

    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 GatewayConnection(socket);
        try
        {
            await connection.PerformInboundHandshakeAsync(_serverId, ct);
            Register(connection);
        }
        catch
        {
            await connection.DisposeAsync();
        }
    }

    private async Task ConnectWithRetryAsync(string remote, CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            try
            {
                var endPoint = ParseEndpoint(remote);
                var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
                await socket.ConnectAsync(endPoint.Address, endPoint.Port, ct);
                var connection = new GatewayConnection(socket) { IsOutbound = true };
                await connection.PerformOutboundHandshakeAsync(_serverId, ct);
                Register(connection);
                return;
            }
            catch (OperationCanceledException)
            {
                return;
            }
            catch (Exception ex)
            {
                _logger.LogDebug(ex, "Gateway connect retry for {Remote}", remote);
            }

            try
            {
                await Task.Delay(250, ct);
            }
            catch (OperationCanceledException)
            {
                return;
            }
        }
    }

    private void Register(GatewayConnection 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 =>
        {
            _messageSink(msg);
            return Task.CompletedTask;
        };
        connection.StartLoop(_cts!.Token);
        Interlocked.Increment(ref _stats.Gateways);
        _ = Task.Run(() => WatchConnectionAsync(key, connection, _cts!.Token));
    }

    private async Task WatchConnectionAsync(string key, GatewayConnection connection, CancellationToken ct)
    {
        try
        {
            await connection.WaitUntilClosedAsync(ct);
        }
        catch
        {
        }
        finally
        {
            if (_connections.TryRemove(key, out _))
                Interlocked.Decrement(ref _stats.Gateways);
            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]));
    }
}
