using System.Collections.Concurrent;
using System.Net;
using System.Net.Security;
using System.Net.Sockets;
using System.Security.Cryptography.X509Certificates;
using NATS.Server.Auth;
using NATS.Server.Auth.Jwt;
using NATS.Server.Protocol;
using NATS.Server.Tls;

namespace NATS.Server.Mqtt;

public sealed class MqttListener : IAsyncDisposable
{
    private readonly string _host;
    private int _port;
    private readonly string? _requiredUsername;
    private readonly string? _requiredPassword;
    private readonly AuthService? _authService;
    private readonly MqttOptions? _mqttOptions;
    private readonly SslServerAuthenticationOptions? _sslOptions;
    private readonly ConcurrentDictionary<MqttConnection, byte> _connections = new();
    private readonly ConcurrentDictionary<string, ConcurrentDictionary<MqttConnection, byte>> _subscriptions = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, MqttSessionState> _sessions = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, MqttConnection> _clientIdMap = new(StringComparer.Ordinal);
    private readonly ConcurrentDictionary<string, string> _retainedMessages = new(StringComparer.Ordinal);
    private readonly IMessageRouter? _router;
    private readonly ConcurrentDictionary<ulong, MqttNatsClientAdapter> _mqttAdapters = new();
    private MqttStreamInitializer? _streamInitializer;
    private MqttConsumerManager? _mqttConsumerManager;
    private TcpListener? _listener;
    private Task? _acceptLoop;
    private readonly CancellationTokenSource _cts = new();

    /// <summary>
    /// When false, connections use the legacy text-line protocol for backward compatibility
    /// with existing tests. Default is true (binary MQTT 3.1.1).
    /// TODO: Remove after test migration — deadline: Phase 3 completion.
    /// </summary>
    internal bool UseBinaryProtocol { get; set; } = true;

    public int Port => _port;

    /// <summary>
    /// Simple constructor for tests using static username/password auth (no TLS).
    /// </summary>
    public MqttListener(
        string host,
        int port,
        string? requiredUsername = null,
        string? requiredPassword = null)
    {
        _host = host;
        _port = port;
        _requiredUsername = requiredUsername;
        _requiredPassword = requiredPassword;
    }

    /// <summary>
    /// Full constructor for production use with AuthService, TLS, and optional JetStream support.
    /// </summary>
    public MqttListener(
        string host,
        int port,
        AuthService? authService,
        MqttOptions mqttOptions,
        MqttStreamInitializer? streamInitializer = null,
        MqttConsumerManager? mqttConsumerManager = null,
        IMessageRouter? router = null)
    {
        _host = host;
        _port = port;
        _authService = authService;
        _mqttOptions = mqttOptions;
        _requiredUsername = mqttOptions.Username;
        _requiredPassword = mqttOptions.Password;
        _streamInitializer = streamInitializer;
        _mqttConsumerManager = mqttConsumerManager;
        _router = router;

        // Build TLS options if configured
        if (mqttOptions.HasTls)
        {
            _sslOptions = BuildMqttSslOptions(mqttOptions);
        }
    }

    /// <summary>
    /// The MQTT stream initializer for JetStream persistence, or null if JetStream is not enabled.
    /// </summary>
    internal MqttStreamInitializer? StreamInitializer => _streamInitializer;

    /// <summary>
    /// The MQTT consumer manager for QoS 1/2 JetStream consumers, or null if JetStream is not enabled.
    /// </summary>
    internal MqttConsumerManager? ConsumerManager => _mqttConsumerManager;

    /// <summary>
    /// The message router for cross-protocol delivery (MQTT→NATS SubList routing).
    /// Null when running in test-only mode without NatsServer.
    /// </summary>
    internal IMessageRouter? Router => _router;

    /// <summary>
    /// Delegate to allocate a server-unique client ID for MQTT adapters.
    /// Set by NatsServer after construction.
    /// </summary>
    internal Func<ulong>? AllocateClientId { get; set; }

    /// <summary>
    /// Delegate to resolve an Account by name for MQTT adapters.
    /// Set by NatsServer after construction.
    /// </summary>
    internal Func<string?, Auth.Account?>? ResolveAccount { get; set; }

    internal void RegisterMqttAdapter(MqttNatsClientAdapter adapter)
        => _mqttAdapters[adapter.Id] = adapter;

    internal void UnregisterMqttAdapter(MqttNatsClientAdapter adapter)
        => _mqttAdapters.TryRemove(adapter.Id, out _);

    internal IEnumerable<MqttNatsClientAdapter> GetMqttAdapters()
        => _mqttAdapters.Values;

    /// <summary>
    /// Looks up a specific pending publish by client ID and packet ID.
    /// Used by QoS 2 PUBREL to retrieve the stored message for delivery.
    /// </summary>
    internal MqttPendingPublish? GetPendingPublish(string clientId, int packetId)
    {
        if (string.IsNullOrWhiteSpace(clientId) || packetId <= 0)
            return null;

        if (_sessions.TryGetValue(clientId, out var session)
            && session.Pending.TryGetValue(packetId, out var pending))
            return pending;

        return null;
    }

    public Task StartAsync(CancellationToken ct)
    {
        var linked = CancellationTokenSource.CreateLinkedTokenSource(ct, _cts.Token);
        var ip = string.IsNullOrWhiteSpace(_host) || _host == "0.0.0.0"
            ? IPAddress.Any
            : IPAddress.Parse(_host);
        _listener = new TcpListener(ip, _port);
        _listener.Start();
        _port = ((IPEndPoint)_listener.LocalEndpoint).Port;
        _acceptLoop = Task.Run(() => AcceptLoopAsync(linked.Token), linked.Token);
        return Task.CompletedTask;
    }

    internal void RegisterSubscription(MqttConnection connection, string topic)
    {
        var set = _subscriptions.GetOrAdd(topic, static _ => new ConcurrentDictionary<MqttConnection, byte>());
        set[connection] = 0;
    }

    internal void UnregisterSubscription(MqttConnection connection, string topic)
    {
        if (_subscriptions.TryGetValue(topic, out var set))
            set.TryRemove(connection, out _);
    }

    internal async Task PublishAsync(string topic, string payload, MqttConnection sender, CancellationToken ct)
    {
        if (!_subscriptions.TryGetValue(topic, out var subscribers))
            return;

        foreach (var subscriber in subscribers.Keys)
        {
            if (subscriber == sender)
                continue;

            await subscriber.SendMessageAsync(topic, payload, ct);
        }
    }

    internal IReadOnlyList<MqttPendingPublish> OpenSession(string clientId, bool cleanSession)
    {
        if (string.IsNullOrWhiteSpace(clientId))
            return [];

        if (cleanSession)
        {
            _sessions.TryRemove(clientId, out _);
            return [];
        }

        var session = _sessions.GetOrAdd(clientId, static _ => new MqttSessionState());
        return session.Pending.Values
            .OrderBy(static p => p.PacketId)
            .ToArray();
    }

    internal void RecordPendingPublish(string clientId, int packetId, string topic, string payload)
    {
        if (string.IsNullOrWhiteSpace(clientId) || packetId <= 0)
            return;

        var session = _sessions.GetOrAdd(clientId, static _ => new MqttSessionState());
        session.Pending[packetId] = new MqttPendingPublish(packetId, topic, payload);
    }

    internal void AckPendingPublish(string clientId, int packetId)
    {
        if (string.IsNullOrWhiteSpace(clientId) || packetId <= 0)
            return;

        if (_sessions.TryGetValue(clientId, out var session))
            session.Pending.TryRemove(packetId, out _);
    }

    /// <summary>
    /// Authenticates MQTT CONNECT credentials. Uses AuthService when available,
    /// falls back to static username/password validation.
    /// </summary>
    internal AuthResult? AuthenticateMqtt(string? username, string? password, X509Certificate2? clientCert = null)
    {
        if (_authService != null)
        {
            var context = new ClientAuthContext
            {
                Opts = new ClientOptions
                {
                    Username = username,
                    Password = password,
                },
                Nonce = [],
                ClientCertificate = clientCert,
                ConnectionType = JwtConnectionTypes.Mqtt,
            };

            if (!_authService.IsAuthRequired)
                return new AuthResult { Identity = username ?? string.Empty };

            var result = _authService.Authenticate(context);
            return result;
        }

        // Fallback: static credential check
        if (AuthService.ValidateMqttCredentials(_requiredUsername, _requiredPassword, username, password))
            return new AuthResult { Identity = username ?? string.Empty };

        return null;
    }

    /// <summary>
    /// Backward-compatible simple auth check for text-protocol mode.
    /// </summary>
    internal bool TryAuthenticate(string? username, string? password)
    {
        return AuthenticateMqtt(username, password) != null;
    }

    internal TimeSpan ResolveKeepAliveTimeout(int keepAliveSeconds)
    {
        if (keepAliveSeconds <= 0)
            return Timeout.InfiniteTimeSpan;

        return TimeSpan.FromSeconds(Math.Max(keepAliveSeconds * 1.5, 1));
    }

    /// <summary>
    /// Disconnects an existing connection with the same client-id (takeover).
    /// Go reference: mqtt.go mqttHandleConnect ~line 850 duplicate client handling.
    /// </summary>
    internal void TakeoverExistingConnection(string clientId, MqttConnection newConnection)
    {
        if (_clientIdMap.TryGetValue(clientId, out var existing) && existing != newConnection)
        {
            existing.ForceClose();
            _connections.TryRemove(existing, out _);
        }

        _clientIdMap[clientId] = newConnection;
    }

    /// <summary>
    /// Stores or deletes a retained message. Null payload = tombstone (delete).
    /// </summary>
    internal void SetRetainedMessage(string topic, string? payload)
    {
        if (payload == null)
            _retainedMessages.TryRemove(topic, out _);
        else
            _retainedMessages[topic] = payload;
    }

    /// <summary>
    /// Gets the retained message for a topic, or null if none.
    /// </summary>
    internal string? GetRetainedMessage(string topic)
    {
        _retainedMessages.TryGetValue(topic, out var payload);
        return payload;
    }

    internal void Unregister(MqttConnection connection)
    {
        _connections.TryRemove(connection, out _);
        foreach (var set in _subscriptions.Values)
            set.TryRemove(connection, out _);

        // Remove from client-id map if this connection is the current one
        var clientId = connection.ClientId;
        if (!string.IsNullOrEmpty(clientId))
        {
            _clientIdMap.TryRemove(new KeyValuePair<string, MqttConnection>(clientId, connection));
        }
    }

    public async ValueTask DisposeAsync()
    {
        await _cts.CancelAsync();
        if (_listener != null)
            _listener.Stop();
        if (_acceptLoop != null)
            await _acceptLoop.ConfigureAwait(ConfigureAwaitOptions.SuppressThrowing);

        foreach (var connection in _connections.Keys)
            await connection.DisposeAsync();

        _connections.Clear();
        _subscriptions.Clear();
        _sessions.Clear();
        _clientIdMap.Clear();
        _retainedMessages.Clear();
        _mqttAdapters.Clear();
        _cts.Dispose();
    }

    private async Task AcceptLoopAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            TcpClient client;
            try
            {
                client = await _listener!.AcceptTcpClientAsync(ct);
            }
            catch
            {
                break;
            }

            _ = Task.Run(async () =>
            {
                Stream stream = client.GetStream();
                X509Certificate2? clientCert = null;

                // TLS wrapping for MQTT (TLS-first, no INFO negotiation)
                if (_sslOptions != null)
                {
                    try
                    {
                        var sslStream = new SslStream(stream, leaveInnerStreamOpen: false);
                        using var handshakeCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
                        handshakeCts.CancelAfter(TimeSpan.FromSeconds(
                            _mqttOptions?.TlsTimeout ?? 2.0));
                        await sslStream.AuthenticateAsServerAsync(_sslOptions, handshakeCts.Token);

                        clientCert = sslStream.RemoteCertificate as X509Certificate2;

                        // Validate pinned certs
                        if (_mqttOptions?.TlsPinnedCerts != null && clientCert != null)
                        {
                            if (!TlsHelper.MatchesPinnedCert(clientCert, _mqttOptions.TlsPinnedCerts))
                            {
                                sslStream.Dispose();
                                client.Dispose();
                                return;
                            }
                        }

                        stream = sslStream;
                    }
                    catch
                    {
                        client.Dispose();
                        return;
                    }
                }

                // Lazily initialize MQTT JetStream streams on first connection
                _streamInitializer?.EnsureStreams();

                var connection = new MqttConnection(stream, this, UseBinaryProtocol, clientCert);
                _connections[connection] = 0;

                try
                {
                    await connection.RunAsync(ct);
                }
                finally
                {
                    await connection.DisposeAsync();
                }
            }, ct);
        }
    }

    private static SslServerAuthenticationOptions BuildMqttSslOptions(MqttOptions mqttOptions)
    {
        var cert = TlsHelper.LoadCertificate(mqttOptions.TlsCert!, mqttOptions.TlsKey);
        var authOpts = new SslServerAuthenticationOptions
        {
            ServerCertificate = cert,
            ClientCertificateRequired = mqttOptions.TlsVerify,
        };

        if (mqttOptions.TlsVerify && mqttOptions.TlsCaCert != null)
        {
            var caCerts = TlsHelper.LoadCaCertificates(mqttOptions.TlsCaCert);
            authOpts.RemoteCertificateValidationCallback = (_, cert, chain, errors) =>
            {
                if (cert == null) return false;
                using var chain2 = new X509Chain();
                chain2.ChainPolicy.TrustMode = X509ChainTrustMode.CustomRootTrust;
                foreach (var ca in caCerts)
                    chain2.ChainPolicy.CustomTrustStore.Add(ca);
                chain2.ChainPolicy.RevocationMode = X509RevocationMode.NoCheck;
                var cert2 = cert as X509Certificate2 ?? X509CertificateLoader.LoadCertificate(cert.GetRawCertData());
                return chain2.Build(cert2);
            };
        }

        return authOpts;
    }

    private sealed class MqttSessionState
    {
        public ConcurrentDictionary<int, MqttPendingPublish> Pending { get; } = new();
    }
}

public sealed record MqttPendingPublish(int PacketId, string Topic, string Payload);
