using System.Net.Sockets;
using System.Text;

namespace NATS.Server.Mqtt;

public sealed class MqttConnection(TcpClient client, MqttListener listener) : IAsyncDisposable
{
    private readonly TcpClient _client = client;
    private readonly NetworkStream _stream = client.GetStream();
    private readonly MqttListener _listener = listener;
    private readonly MqttProtocolParser _parser = new();
    private readonly SemaphoreSlim _writeGate = new(1, 1);
    private string _clientId = string.Empty;
    private bool _cleanSession = true;
    private TimeSpan _idleTimeout = Timeout.InfiniteTimeSpan;

    public async Task RunAsync(CancellationToken ct)
    {
        while (!ct.IsCancellationRequested)
        {
            string line;
            try
            {
                line = await ReadLineAsync(ct, _idleTimeout);
            }
            catch
            {
                break;
            }

            var packet = _parser.ParseLine(line);
            switch (packet.Type)
            {
                case MqttPacketType.Connect:
                    if (!_listener.TryAuthenticate(packet.Username, packet.Password))
                    {
                        await WriteLineAsync("ERR mqtt auth failed", ct);
                        return;
                    }

                    _clientId = packet.ClientId;
                    _cleanSession = packet.CleanSession;
                    _idleTimeout = _listener.ResolveKeepAliveTimeout(packet.KeepAliveSeconds);
                    var pending = _listener.OpenSession(_clientId, _cleanSession);
                    await WriteLineAsync("CONNACK", ct);
                    foreach (var redelivery in pending)
                        await WriteLineAsync($"REDLIVER {redelivery.PacketId} {redelivery.Topic} {redelivery.Payload}", ct);
                    break;
                case MqttPacketType.Subscribe:
                    _listener.RegisterSubscription(this, packet.Topic);
                    await WriteLineAsync($"SUBACK {packet.Topic}", ct);
                    break;
                case MqttPacketType.Publish:
                    await _listener.PublishAsync(packet.Topic, packet.Payload, this, ct);
                    break;
                case MqttPacketType.PublishQos1:
                    _listener.RecordPendingPublish(_clientId, packet.PacketId, packet.Topic, packet.Payload);
                    await WriteLineAsync($"PUBACK {packet.PacketId}", ct);
                    await _listener.PublishAsync(packet.Topic, packet.Payload, this, ct);
                    break;
                case MqttPacketType.Ack:
                    _listener.AckPendingPublish(_clientId, packet.PacketId);
                    break;
            }
        }
    }

    public Task SendMessageAsync(string topic, string payload, CancellationToken ct)
        => WriteLineAsync($"MSG {topic} {payload}", ct);

    public async ValueTask DisposeAsync()
    {
        _listener.Unregister(this);
        _writeGate.Dispose();
        await _stream.DisposeAsync();
        _client.Dispose();
    }

    private async Task WriteLineAsync(string line, CancellationToken ct)
    {
        await _writeGate.WaitAsync(ct);
        try
        {
            var bytes = Encoding.UTF8.GetBytes(line + "\n");
            await _stream.WriteAsync(bytes, ct);
            await _stream.FlushAsync(ct);
        }
        finally
        {
            _writeGate.Release();
        }
    }

    private async Task<string> ReadLineAsync(CancellationToken ct, TimeSpan idleTimeout)
    {
        CancellationToken token = ct;
        CancellationTokenSource? timeoutCts = null;
        if (idleTimeout != Timeout.InfiniteTimeSpan && idleTimeout > TimeSpan.Zero)
        {
            timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
            timeoutCts.CancelAfter(idleTimeout);
            token = timeoutCts.Token;
        }

        var bytes = new List<byte>(64);
        var single = new byte[1];
        try
        {
            while (true)
            {
                var read = await _stream.ReadAsync(single.AsMemory(0, 1), token);
                if (read == 0)
                    throw new IOException("mqtt closed");
                if (single[0] == (byte)'\n')
                    break;
                if (single[0] != (byte)'\r')
                    bytes.Add(single[0]);
            }
        }
        catch (OperationCanceledException) when (timeoutCts != null && !ct.IsCancellationRequested)
        {
            throw new IOException("mqtt keepalive timeout");
        }
        finally
        {
            timeoutCts?.Dispose();
        }

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