LumaTunnel

using System.Collections.Concurrent;
using LumaTunnel.Server.Core.Persistence;

namespace LumaTunnel.Server.Core.Security;

public sealed class DeviceRegistry(AtomicJsonStore<DeviceDatabase> store) : IDisposable
{
    private readonly ConcurrentDictionary<string, DeviceRecord> _devices = new(StringComparer.Ordinal);
    private readonly AtomicJsonStore<DeviceDatabase> _store = store;
    private readonly SemaphoreSlim _mutationGate = new(1, 1);

    public event Action<string>? DeviceRevoked;

    public IReadOnlyCollection<DeviceRecord> Devices => _devices.Values.OrderBy(static x => x.DeviceName).ToArray();

    public async Task InitializeAsync(CancellationToken cancellationToken = default)
    {
        var database = await _store.LoadOrCreateAsync(static () => new DeviceDatabase(), cancellationToken).ConfigureAwait(false);
        foreach (var device in database.Devices)
            _devices[device.DeviceId] = device;
    }

    public async Task ReloadAsync(CancellationToken cancellationToken = default)
    {
        List<string> revoked = [];
        await _mutationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            var database = await _store.LoadOrCreateAsync(static () => new DeviceDatabase(), cancellationToken).ConfigureAwait(false);
            var incoming = database.Devices.ToDictionary(static item => item.DeviceId, StringComparer.Ordinal);
            revoked.AddRange(_devices.Values.Where(current => current.Enabled && (!incoming.TryGetValue(current.DeviceId, out var replacement) || !replacement.Enabled)).Select(static item => item.DeviceId));
            _devices.Clear();
            foreach (var device in incoming.Values)
                _devices[device.DeviceId] = device;
        }
        finally
        {
            _mutationGate.Release();
        }

        foreach (var deviceId in revoked)
            DeviceRevoked?.Invoke(deviceId);
    }

    public async Task<(DeviceRecord Device, string Token)> EnrollAsync(string deviceId, string deviceName, string clientVersion, CancellationToken cancellationToken = default)
    {
        ArgumentException.ThrowIfNullOrWhiteSpace(deviceId);
        ArgumentException.ThrowIfNullOrWhiteSpace(deviceName);

        var token = TokenUtilities.CreateDeviceToken();
        var record = new DeviceRecord
        {
            DeviceId = deviceId.Trim(),
            DeviceName = deviceName.Trim(),
            ClientVersion = clientVersion.Trim(),
            TokenHash = TokenUtilities.Hash(token)
        };

        await _mutationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            _devices[record.DeviceId] = record;
            await SaveCoreAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _mutationGate.Release();
        }

        return (record, token);
    }

    public bool Validate(string deviceId, string token, out DeviceRecord? record)
    {
        if (_devices.TryGetValue(deviceId, out var candidate)
            && candidate.Enabled
            && TokenUtilities.FixedTimeHashEquals(token, candidate.TokenHash))
        {
            record = candidate;
            return true;
        }

        record = null;
        return false;
    }

    public async Task<bool> RevokeAsync(string deviceId, CancellationToken cancellationToken = default)
    {
        await _mutationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (!_devices.TryGetValue(deviceId, out var existing))
                return false;
            _devices[deviceId] = existing with { Enabled = false };
            await SaveCoreAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _mutationGate.Release();
        }

        DeviceRevoked?.Invoke(deviceId);
        return true;
    }

    public async Task<bool> RenameAsync(string deviceId, string name, CancellationToken cancellationToken = default)
    {
        await _mutationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            if (!_devices.TryGetValue(deviceId, out var existing))
                return false;
            _devices[deviceId] = existing with { DeviceName = name.Trim() };
            await SaveCoreAsync(cancellationToken).ConfigureAwait(false);
            return true;
        }
        finally
        {
            _mutationGate.Release();
        }
    }

    public async Task AddTrafficAsync(string deviceId, long uploadBytes, long downloadBytes, CancellationToken cancellationToken = default)
    {
        if (!_devices.TryGetValue(deviceId, out var existing))
            return;

        _devices[deviceId] = existing with
        {
            LastSeenAtUtc = DateTimeOffset.UtcNow,
            UploadBytes = existing.UploadBytes + Math.Max(0, uploadBytes),
            DownloadBytes = existing.DownloadBytes + Math.Max(0, downloadBytes)
        };

        // Traffic persistence is intentionally batched by the runtime flush timer.
        await Task.CompletedTask.ConfigureAwait(false);
    }

    public async Task SaveAsync(CancellationToken cancellationToken = default)
    {
        await _mutationGate.WaitAsync(cancellationToken).ConfigureAwait(false);
        try
        {
            await SaveCoreAsync(cancellationToken).ConfigureAwait(false);
        }
        finally
        {
            _mutationGate.Release();
        }
    }

    private Task SaveCoreAsync(CancellationToken cancellationToken) =>
        _store.SaveAsync(new DeviceDatabase { Devices = _devices.Values.OrderBy(static x => x.DeviceId).ToList() }, cancellationToken);

    public void Dispose()
    {
        _mutationGate.Dispose();
        _store.Dispose();
        GC.SuppressFinalize(this);
    }
}