using System.Collections.Concurrent;
using System.Net;
using System.Net.Sockets;
using NATS.Server.Auth;

namespace NATS.Server.Mqtt;

public sealed class MqttListener(
    string host,
    int port,
    string? requiredUsername = null,
    string? requiredPassword = null) : IAsyncDisposable
{
    private readonly string _host = host;
    private int _port = port;
    private readonly string? _requiredUsername = requiredUsername;
    private readonly string? _requiredPassword = requiredPassword;
    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 TcpListener? _listener;
    private Task? _acceptLoop;
    private readonly CancellationTokenSource _cts = new();

    public int Port => _port;

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

    internal bool TryAuthenticate(string? username, string? password)
        => AuthService.ValidateMqttCredentials(_requiredUsername, _requiredPassword, username, password);

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

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

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

    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();
        _cts.Dispose();
    }

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

            var connection = new MqttConnection(client, this);
            _connections[connection] = 0;
            _ = Task.Run(async () =>
            {
                try
                {
                    await connection.RunAsync(ct);
                }
                finally
                {
                    await connection.DisposeAsync();
                }
            }, ct);
        }
    }

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

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