using dotnet_etcd;
using Etcdserverpb;
using Grpc.Core;

namespace DysonNetwork.Shared.Etcd;

public class EtcdService(string connectionString) : IEtcdService
{
    private readonly EtcdClient _etcdClient = new(connectionString);
    private long _leaseId;
    private string? _serviceKey;
    private readonly CancellationTokenSource _cts = new();

    public async Task RegisterServiceAsync(string serviceName, string serviceAddress, int ttl = 15)
    {
        _serviceKey = $"/services/{serviceName}/{Guid.NewGuid()}";
        var leaseGrantResponse = await _etcdClient.LeaseGrantAsync(new LeaseGrantRequest { TTL = ttl });
        _leaseId = leaseGrantResponse.ID;

        await _etcdClient.PutAsync(new PutRequest
        {
            Key = Google.Protobuf.ByteString.CopyFromUtf8(_serviceKey),
            Value = Google.Protobuf.ByteString.CopyFromUtf8(serviceAddress),
            Lease = _leaseId
        });

        _ = Task.Run(async () =>
        {
            while (!_cts.Token.IsCancellationRequested)
            {
                try
                {
                    await _etcdClient.LeaseKeepAlive(new LeaseKeepAliveRequest { ID = _leaseId },
                        _ => { }, _cts.Token);
                    await Task.Delay(TimeSpan.FromSeconds(ttl / 3), _cts.Token);
                }
                catch (RpcException)
                {
                    // Ignored
                }
            }
        }, _cts.Token);
    }

    public async Task UnregisterServiceAsync()
    {
        if (!string.IsNullOrEmpty(_serviceKey))
        {
            await _etcdClient.DeleteRangeAsync(_serviceKey);
        }
    }

    public async Task<List<string>> DiscoverServicesAsync(string serviceName)
    {
        var prefix = $"/services/{serviceName}/";
        var rangeResponse = await _etcdClient.GetRangeAsync(prefix);
        return rangeResponse.Kvs.Select(kv => kv.Value.ToStringUtf8()).ToList();
    }

    public void Dispose()
    {
        _cts.Cancel();
        if (_leaseId != 0)
        {
            _etcdClient.LeaseRevoke(new LeaseRevokeRequest { ID = _leaseId });
        }

        _etcdClient.Dispose();
    }
}