Cleanuparr/code/backend/Cleanuparr.Infrastructure/Health/HealthCheckService.cs

using Cleanuparr.Domain.Enums;
using Cleanuparr.Infrastructure.Features.Arr.Interfaces;
using Cleanuparr.Infrastructure.Features.DownloadClient;
using Cleanuparr.Persistence;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;

namespace Cleanuparr.Infrastructure.Health;

/// <summary>
/// Service for checking the health of download clients and arr instances
/// </summary>
public class HealthCheckService : IHealthCheckService
{
    private readonly ILogger<HealthCheckService> _logger;
    private readonly Dictionary<Guid, HealthStatus> _healthStatuses = new();
    private readonly Dictionary<Guid, ArrHealthStatus> _arrHealthStatuses = new();
    private readonly IServiceScopeFactory _scopeFactory;
    private readonly object _lockObject = new();

    /// <summary>
    /// Occurs when a client's health status changes
    /// </summary>
    public event EventHandler<ClientHealthChangedEventArgs>? ClientHealthChanged;

    /// <summary>
    /// Occurs when a client is no longer covered by the health sweep
    /// </summary>
    public event EventHandler<ClientHealthRemovedEventArgs>? ClientHealthRemoved;

    /// <summary>
    /// Occurs when an arr instance is no longer covered by the health sweep
    /// </summary>
    public event EventHandler<ArrInstanceHealthRemovedEventArgs>? ArrInstanceHealthRemoved;

    public HealthCheckService(
        ILogger<HealthCheckService> logger,
        IServiceScopeFactory scopeFactory
    )
    {
        _logger = logger;
        _scopeFactory = scopeFactory;
    }

    /// <inheritdoc />
    public async Task<HealthStatus> CheckClientHealthAsync(Guid clientId)
    {
        _logger.LogDebug("Checking health for client {ClientId}", clientId);

        try
        {
            await using var scope = _scopeFactory.CreateAsyncScope();
            await using var dataContext = scope.ServiceProvider.GetRequiredService<DataContext>();
            
            // Get the client configuration
            var downloadClientConfig = await dataContext.DownloadClients
                .Where(x => x.Id == clientId)
                .FirstOrDefaultAsync();
            
            if (downloadClientConfig is null)
            {
                _logger.LogWarning("Client {ClientId} not found in configuration", clientId);
                var notFoundStatus = new HealthStatus
                {
                    ClientId = clientId,
                    IsHealthy = false,
                    LastChecked = DateTimeOffset.UtcNow,
                    ErrorMessage = "Client not found in configuration"
                };
                
                UpdateHealthStatus(notFoundStatus);
                return notFoundStatus;
            }

            // Get the client instance
            var downloadServiceFactory = scope.ServiceProvider.GetRequiredService<IDownloadServiceFactory>();
            var client = downloadServiceFactory.GetDownloadService(downloadClientConfig);
            
            // Execute the health check
            var healthResult = await client.HealthCheckAsync();
            
            // Create health status object
            var status = new HealthStatus
            {
                ClientId = clientId,
                ClientName = downloadClientConfig.Name,
                ClientTypeName = downloadClientConfig.TypeName,
                IsHealthy = healthResult.IsHealthy,
                LastChecked = DateTimeOffset.UtcNow,
                ErrorMessage = healthResult.ErrorMessage,
                ResponseTime = healthResult.ResponseTime
            };
            
            UpdateHealthStatus(status);
            return status;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error performing health check for client {ClientId}", clientId);
            
            var status = new HealthStatus
            {
                ClientId = clientId,
                IsHealthy = false,
                LastChecked = DateTimeOffset.UtcNow,
                ErrorMessage = $"Error: {ex.Message}"
            };
            
            UpdateHealthStatus(status);
            return status;
        }
    }

    /// <inheritdoc />
    public async Task<IDictionary<Guid, HealthStatus>> CheckAllClientsHealthAsync()
    {
        _logger.LogDebug("Checking health for all enabled clients");

        // Captured before anything is awaited: a later check replaces the instance and survives the prune.
        Dictionary<Guid, HealthStatus> knownAtStart = SnapshotClients();

        try
        {
            await using var scope = _scopeFactory.CreateAsyncScope();
            await using var dataContext = scope.ServiceProvider.GetRequiredService<DataContext>();
            
            // Get all enabled client configurations
            var enabledClients = (await dataContext.DownloadClients
                    .Where(x => x.Enabled)
                    .ToListAsync())
                .Where(x => !EnumSentinel.IsUnknown(x.TypeName) && !EnumSentinel.IsUnknown(x.Type))
                .ToList();
            var results = new Dictionary<Guid, HealthStatus>();
            
            // Check health of each enabled client
            foreach (var clientConfig in enabledClients)
            {
                var status = await CheckClientHealthAsync(clientConfig.Id);
                results[clientConfig.Id] = status;
            }
            
            PruneClientHealthStatuses(knownAtStart);
            
            return results;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error checking health for all clients");
            return new Dictionary<Guid, HealthStatus>();
        }
    }

    /// <inheritdoc />
    public HealthStatus? GetClientHealth(Guid clientId)
    {
        lock (_lockObject)
        {
            return _healthStatuses.TryGetValue(clientId, out var status) ? status : null;
        }
    }

    /// <inheritdoc />
    public IDictionary<Guid, HealthStatus> GetAllClientHealth()
    {
        lock (_lockObject)
        {
            return new Dictionary<Guid, HealthStatus>(_healthStatuses);
        }
    }

    /// <inheritdoc />
    public async Task<ArrHealthStatus> CheckArrInstanceHealthAsync(Guid instanceId)
    {
        _logger.LogDebug("Checking health for arr instance {InstanceId}", instanceId);

        try
        {
            await using var scope = _scopeFactory.CreateAsyncScope();
            await using var dataContext = scope.ServiceProvider.GetRequiredService<DataContext>();

            // Get the arr instance with its config (needed for InstanceType)
            // Load config with instances first, then find in memory (SQLite doesn't support APPLY)
            var config = await dataContext.ArrConfigs
                .Include(x => x.Instances)
                .FirstOrDefaultAsync(c => c.Instances.Any(i => i.Id == instanceId));

            var arrInstance = config is null ? null : new
            {
                Instance = config.Instances.First(i => i.Id == instanceId),
                Config = config
            };

            if (arrInstance is null)
            {
                _logger.LogWarning("Arr instance {InstanceId} not found in configuration", instanceId);
                var notFoundStatus = new ArrHealthStatus
                {
                    InstanceId = instanceId,
                    IsHealthy = false,
                    LastChecked = DateTimeOffset.UtcNow,
                    ErrorMessage = "Arr instance not found in configuration"
                };

                UpdateArrHealthStatus(notFoundStatus);
                return notFoundStatus;
            }

            // Get the arr client and execute health check
            var healthChecker = scope.ServiceProvider.GetRequiredService<IInstanceHealthChecker>();
            await healthChecker.CheckAsync(arrInstance.Config.Type, arrInstance.Instance);

            var status = new ArrHealthStatus
            {
                InstanceId = instanceId,
                InstanceName = arrInstance.Instance.Name,
                InstanceType = arrInstance.Config.Type,
                IsHealthy = true,
                LastChecked = DateTimeOffset.UtcNow
            };

            UpdateArrHealthStatus(status);
            return status;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error performing health check for arr instance {InstanceId}", instanceId);

            var status = new ArrHealthStatus
            {
                InstanceId = instanceId,
                IsHealthy = false,
                LastChecked = DateTimeOffset.UtcNow,
                ErrorMessage = $"Error: {ex.Message}"
            };

            UpdateArrHealthStatus(status);
            return status;
        }
    }

    /// <inheritdoc />
    public async Task<IDictionary<Guid, ArrHealthStatus>> CheckAllArrInstancesHealthAsync()
    {
        _logger.LogDebug("Checking health for all enabled arr instances");

        // Captured before anything is awaited: a later check replaces the instance and survives the prune.
        Dictionary<Guid, ArrHealthStatus> knownAtStart = SnapshotArrInstances();

        try
        {
            await using var scope = _scopeFactory.CreateAsyncScope();
            await using var dataContext = scope.ServiceProvider.GetRequiredService<DataContext>();

            // Get all enabled arr instances across all configs
            // Load configs with instances first, then flatten in memory (SQLite doesn't support APPLY)
            var configs = (await dataContext.ArrConfigs
                    .Include(x => x.Instances)
                    .ToListAsync())
                .Where(x => !EnumSentinel.IsUnknown(x.Type))
                .ToList();

            var enabledInstances = configs
                .SelectMany(c => c.Instances
                    .Where(i => i.Enabled)
                    .Select(i => new { Instance = i, Config = c }))
                .ToList();

            var results = new Dictionary<Guid, ArrHealthStatus>();
            var healthChecker = scope.ServiceProvider.GetRequiredService<IInstanceHealthChecker>();

            foreach (var entry in enabledInstances)
            {
                try
                {
                    await healthChecker.CheckAsync(entry.Config.Type, entry.Instance);

                    var status = new ArrHealthStatus
                    {
                        InstanceId = entry.Instance.Id,
                        InstanceName = entry.Instance.Name,
                        InstanceType = entry.Config.Type,
                        IsHealthy = true,
                        LastChecked = DateTimeOffset.UtcNow
                    };

                    UpdateArrHealthStatus(status);
                    results[entry.Instance.Id] = status;
                }
                catch (Exception ex)
                {
                    _logger.LogError(ex, "Error performing health check for arr instance {InstanceId} ({InstanceName})",
                        entry.Instance.Id, entry.Instance.Name);

                    var status = new ArrHealthStatus
                    {
                        InstanceId = entry.Instance.Id,
                        InstanceName = entry.Instance.Name,
                        InstanceType = entry.Config.Type,
                        IsHealthy = false,
                        LastChecked = DateTimeOffset.UtcNow,
                        ErrorMessage = $"Error: {ex.Message}"
                    };

                    UpdateArrHealthStatus(status);
                    results[entry.Instance.Id] = status;
                }
            }

            PruneArrHealthStatuses(knownAtStart);

            return results;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Error checking health for all arr instances");
            return new Dictionary<Guid, ArrHealthStatus>();
        }
    }

    /// <inheritdoc />
    public ArrHealthStatus? GetArrInstanceHealth(Guid instanceId)
    {
        lock (_lockObject)
        {
            return _arrHealthStatuses.TryGetValue(instanceId, out var status) ? status : null;
        }
    }

    /// <inheritdoc />
    public IDictionary<Guid, ArrHealthStatus> GetAllArrInstanceHealth()
    {
        lock (_lockObject)
        {
            return new Dictionary<Guid, ArrHealthStatus>(_arrHealthStatuses);
        }
    }

    private Dictionary<Guid, HealthStatus> SnapshotClients()
    {
        lock (_lockObject)
        {
            return new Dictionary<Guid, HealthStatus>(_healthStatuses);
        }
    }

    private Dictionary<Guid, ArrHealthStatus> SnapshotArrInstances()
    {
        lock (_lockObject)
        {
            return new Dictionary<Guid, ArrHealthStatus>(_arrHealthStatuses);
        }
    }

    /// <summary>
    /// Drops cached clients that have been invalidated.
    /// </summary>
    private void PruneClientHealthStatuses(Dictionary<Guid, HealthStatus> knownAtStart)
    {
        List<Guid> removed = [];

        lock (_lockObject)
        {
            foreach ((Guid clientId, HealthStatus snapshot) in knownAtStart)
            {
                if (!_healthStatuses.TryGetValue(clientId, out HealthStatus? cached)
                    || !ReferenceEquals(cached, snapshot))
                {
                    continue;
                }

                _healthStatuses.Remove(clientId);
                removed.Add(clientId);
            }
        }

        foreach (Guid clientId in removed)
        {
            _logger.LogDebug("Client {ClientId} dropped from the health cache", clientId);
            ClientHealthRemoved?.Invoke(this, new ClientHealthRemovedEventArgs(clientId));
        }
    }

    /// <summary>
    /// Drops cached arr instances that have been invalidated.
    /// </summary>
    private void PruneArrHealthStatuses(Dictionary<Guid, ArrHealthStatus> knownAtStart)
    {
        List<Guid> removed = [];

        lock (_lockObject)
        {
            foreach ((Guid instanceId, ArrHealthStatus snapshot) in knownAtStart)
            {
                if (!_arrHealthStatuses.TryGetValue(instanceId, out ArrHealthStatus? cached)
                    || !ReferenceEquals(cached, snapshot))
                {
                    continue;
                }

                _arrHealthStatuses.Remove(instanceId);
                removed.Add(instanceId);
            }
        }

        foreach (Guid instanceId in removed)
        {
            _logger.LogDebug("Arr instance {InstanceId} dropped from the health cache", instanceId);
            ArrInstanceHealthRemoved?.Invoke(this, new ArrInstanceHealthRemovedEventArgs(instanceId));
        }
    }

    private void UpdateArrHealthStatus(ArrHealthStatus newStatus)
    {
        ArrHealthStatus? previousStatus;

        lock (_lockObject)
        {
            _arrHealthStatuses.TryGetValue(newStatus.InstanceId, out previousStatus);
            _arrHealthStatuses[newStatus.InstanceId] = newStatus;
        }

        bool isStateChange = previousStatus == null ||
                             previousStatus.IsHealthy != newStatus.IsHealthy;

        if (isStateChange)
        {
            _logger.LogInformation(
                "Arr instance {InstanceId} ({InstanceName}) health changed: {Status}",
                newStatus.InstanceId,
                newStatus.InstanceName,
                newStatus.IsHealthy ? "Healthy" : "Unhealthy");
        }
    }

    private void UpdateHealthStatus(HealthStatus newStatus)
    {
        HealthStatus? previousStatus;
        
        lock (_lockObject)
        {
            // Get previous status for comparison
            _healthStatuses.TryGetValue(newStatus.ClientId, out previousStatus);
            
            // Update status
            _healthStatuses[newStatus.ClientId] = newStatus;
        }
        
        // Determine if there's a significant change
        bool isStateChange = previousStatus == null || 
                             previousStatus.IsHealthy != newStatus.IsHealthy;

        // Raise event if there's a significant change
        if (isStateChange)
        {
            _logger.LogInformation(
                "Client {ClientId} health changed: {Status}", 
                newStatus.ClientId, 
                newStatus.IsHealthy ? "Healthy" : "Unhealthy");
            
            OnClientHealthChanged(new ClientHealthChangedEventArgs(
                newStatus.ClientId, 
                newStatus, 
                previousStatus));
        }
    }
    
    private void OnClientHealthChanged(ClientHealthChangedEventArgs e)
    {
        ClientHealthChanged?.Invoke(this, e);
    }
}