MeterVault/src/Infrastructure/Ingestion/HomeAssistantWebSocketWorker.cs

using System.Collections.Concurrent;
using System.Net.WebSockets;
using System.Text;
using System.Text.Json;
using MeterVault.Core.Domain;
using MeterVault.Infrastructure.Persistence;
using Microsoft.EntityFrameworkCore;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Hosting;
using Microsoft.Extensions.Logging;

namespace MeterVault.Infrastructure.Ingestion;

/// <summary>
/// Real-time Home Assistant ingestion via the WebSocket API (SDD §6.2 push path). For each enabled
/// HA endpoint whose config sets <c>UseWebSocket</c>, holds a persistent connection that authenticates,
/// subscribes to <c>state_changed</c> events, and ingests changes for the endpoint's configured
/// entities as they happen. Reconnects with capped backoff. Endpoints without <c>UseWebSocket</c> stay
/// on the REST poll worker (<see cref="HomeAssistantWorker"/>) — each endpoint is served by exactly one.
/// Tokens are resolved from environment variables named in the endpoint config, never stored plaintext.
/// </summary>
public sealed class HomeAssistantWebSocketWorker(
    IServiceScopeFactory scopeFactory,
    ILogger<HomeAssistantWebSocketWorker> logger) : BackgroundService
{
    private static readonly TimeSpan SuperviseInterval = TimeSpan.FromSeconds(15);
    private static readonly TimeSpan InitialBackoff = TimeSpan.FromSeconds(2);
    private static readonly TimeSpan MaxBackoff = TimeSpan.FromSeconds(60);
    private static readonly TimeSpan EntityMapTtl = TimeSpan.FromSeconds(30);

    private readonly IServiceScopeFactory _scopeFactory = scopeFactory;
    private readonly ILogger<HomeAssistantWebSocketWorker> _logger = logger;
    private readonly ConcurrentDictionary<int, Task> _connections = new();

    protected override async Task ExecuteAsync(CancellationToken stoppingToken)
    {
        using var timer = new PeriodicTimer(SuperviseInterval);
        do
        {
            try
            {
                await SuperviseAsync(stoppingToken).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
            {
                break;
            }
            catch (Exception ex)
            {
                _logger.LogWarning(ex, "Home Assistant WebSocket supervisor tick failed; will retry");
            }
        }
        while (await timer.WaitForNextTickAsync(stoppingToken).ConfigureAwait(false));
    }

    /// <summary>Starts a connection loop for each WebSocket-enabled endpoint that isn't already running.</summary>
    private async Task SuperviseAsync(CancellationToken stoppingToken)
    {
        List<int> wsEndpointIds;
        await using (var scope = _scopeFactory.CreateAsyncScope())
        {
            var db = scope.ServiceProvider.GetRequiredService<MeterVaultDbContext>();
            var endpoints = await db.IngestionEndpoints
                .Where(e => e.IsEnabled && e.Type == EndpointType.HomeAssistant)
                .ToListAsync(stoppingToken).ConfigureAwait(false);
            wsEndpointIds = endpoints
                .Where(e => HaEndpointConfig.Parse(e.Config).UseWebSocket)
                .Select(e => e.Id)
                .ToList();
        }

        // Reap finished connection loops (endpoint disabled/removed, or terminal failure).
        foreach (var id in _connections.Where(kv => kv.Value.IsCompleted).Select(kv => kv.Key).ToList())
        {
            _connections.TryRemove(id, out _);
        }

        foreach (var endpointId in wsEndpointIds)
        {
            // Fire-and-forget: the loop stores its own Task in _connections and self-terminates when
            // the endpoint is disabled/removed; the supervisor reaps completed entries above.
            _ = _connections.GetOrAdd(endpointId, id => Task.Run(() => RunConnectionAsync(id, stoppingToken), stoppingToken));
        }
    }

    /// <summary>Connect → listen → reconnect loop for one endpoint, until it is disabled or the app stops.</summary>
    private async Task RunConnectionAsync(int endpointId, CancellationToken stoppingToken)
    {
        var backoff = InitialBackoff;
        while (!stoppingToken.IsCancellationRequested)
        {
            HaEndpointConfig config;
            string? token;
            await using (var scope = _scopeFactory.CreateAsyncScope())
            {
                var db = scope.ServiceProvider.GetRequiredService<MeterVaultDbContext>();
                var endpoint = await db.IngestionEndpoints
                    .FirstOrDefaultAsync(e => e.Id == endpointId, stoppingToken).ConfigureAwait(false);
                if (endpoint is null || !endpoint.IsEnabled)
                {
                    return; // gone/disabled — stop; the supervisor will restart it if it comes back.
                }

                config = HaEndpointConfig.Parse(endpoint.Config);
                if (!config.UseWebSocket)
                {
                    return; // switched to poll mode.
                }

                token = config.ResolveToken();
            }

            if (string.IsNullOrWhiteSpace(config.BaseUrl) || string.IsNullOrWhiteSpace(token))
            {
                _logger.LogWarning("HA WebSocket endpoint {EndpointId} missing base URL or token; retrying", endpointId);
            }
            else
            {
                try
                {
                    await ListenAsync(endpointId, config.BaseUrl!, token!, stoppingToken).ConfigureAwait(false);
                    backoff = InitialBackoff; // clean close → reset backoff.
                }
                catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
                {
                    return;
                }
                catch (Exception ex)
                {
                    _logger.LogWarning(ex, "HA WebSocket connection to endpoint {EndpointId} dropped; reconnecting", endpointId);
                    await UpdateEndpointStatusAsync(endpointId, "disconnected", CancellationToken.None).ConfigureAwait(false);
                }
            }

            try
            {
                await Task.Delay(backoff, stoppingToken).ConfigureAwait(false);
            }
            catch (OperationCanceledException) when (stoppingToken.IsCancellationRequested)
            {
                return;
            }

            backoff = TimeSpan.FromSeconds(Math.Min(MaxBackoff.TotalSeconds, backoff.TotalSeconds * 2));
        }
    }

    private async Task ListenAsync(int endpointId, string baseUrl, string token, CancellationToken stoppingToken)
    {
        using var ws = new ClientWebSocket();
        await ws.ConnectAsync(HaWebSocketProtocol.WebSocketUri(baseUrl), stoppingToken).ConfigureAwait(false);

        // Handshake: auth_required → auth → auth_ok.
        using (var required = await ReceiveJsonAsync(ws, stoppingToken).ConfigureAwait(false))
        {
            if (!HaWebSocketProtocol.IsAuthRequired(required.RootElement))
            {
                // Some setups may not send auth_required first; proceed to auth regardless.
                _logger.LogDebug("HA WebSocket did not send auth_required first (type={Type})", HaWebSocketProtocol.MessageType(required.RootElement));
            }
        }

        await SendAsync(ws, HaWebSocketProtocol.AuthMessage(token), stoppingToken).ConfigureAwait(false);
        using (var authResult = await ReceiveJsonAsync(ws, stoppingToken).ConfigureAwait(false))
        {
            if (!HaWebSocketProtocol.IsAuthOk(authResult.RootElement))
            {
                throw new InvalidOperationException(
                    $"Home Assistant WebSocket auth failed (type={HaWebSocketProtocol.MessageType(authResult.RootElement)}).");
            }
        }

        await SendAsync(ws, HaWebSocketProtocol.SubscribeStateChanged(1), stoppingToken).ConfigureAwait(false);
        await UpdateEndpointStatusAsync(endpointId, "connected (ws)", stoppingToken).ConfigureAwait(false);
        _logger.LogInformation("HA WebSocket connected for endpoint {EndpointId}", endpointId);

        var entityMap = await LoadEntityMapAsync(endpointId, stoppingToken).ConfigureAwait(false);
        var mapLoadedAt = DateTimeOffset.UtcNow;

        while (!stoppingToken.IsCancellationRequested && ws.State == WebSocketState.Open)
        {
            using var doc = await ReceiveJsonAsync(ws, stoppingToken).ConfigureAwait(false);
            if (!HaWebSocketProtocol.TryReadStateChanged(doc.RootElement, out var entityId, out var newState))
            {
                continue;
            }

            if (DateTimeOffset.UtcNow - mapLoadedAt > EntityMapTtl)
            {
                entityMap = await LoadEntityMapAsync(endpointId, stoppingToken).ConfigureAwait(false);
                mapLoadedAt = DateTimeOffset.UtcNow;
            }

            if (!entityMap.TryGetValue(entityId, out var sources))
            {
                continue;
            }

            foreach (var (sourceId, attribute) in sources)
            {
                if (HaStateClient.ParseStateElement(newState, attribute) is { } state)
                {
                    await IngestAsync(sourceId, state, stoppingToken).ConfigureAwait(false);
                }
            }
        }
    }

    /// <summary>Maps each configured entity id to the sources (and optional attribute) that read it.</summary>
    private async Task<Dictionary<string, List<(int SourceId, string? Attribute)>>> LoadEntityMapAsync(
        int endpointId, CancellationToken cancellationToken)
    {
        await using var scope = _scopeFactory.CreateAsyncScope();
        var db = scope.ServiceProvider.GetRequiredService<MeterVaultDbContext>();
        var sources = await db.MeterSources.AsNoTracking()
            .Where(s => s.IsEnabled && s.SourceType == SourceType.HomeAssistant && s.EndpointId == endpointId)
            .ToListAsync(cancellationToken).ConfigureAwait(false);

        var map = new Dictionary<string, List<(int, string?)>>(StringComparer.OrdinalIgnoreCase);
        foreach (var source in sources)
        {
            var config = SourceConfig.Parse(source.Config);
            if (string.IsNullOrWhiteSpace(config.EntityId))
            {
                continue;
            }

            if (!map.TryGetValue(config.EntityId, out var list))
            {
                map[config.EntityId] = list = [];
            }

            list.Add((source.Id, config.Attribute));
        }

        return map;
    }

    private async Task IngestAsync(int sourceId, HaState state, CancellationToken cancellationToken)
    {
        await using var scope = _scopeFactory.CreateAsyncScope();
        var ingestion = scope.ServiceProvider.GetRequiredService<IngestionService>();
        await ingestion.IngestAsync(sourceId, state.Time, state.Value, cancellationToken).ConfigureAwait(false);
    }

    private async Task UpdateEndpointStatusAsync(int endpointId, string status, CancellationToken cancellationToken)
    {
        try
        {
            await using var scope = _scopeFactory.CreateAsyncScope();
            var db = scope.ServiceProvider.GetRequiredService<MeterVaultDbContext>();
            var endpoint = await db.IngestionEndpoints
                .FirstOrDefaultAsync(e => e.Id == endpointId, cancellationToken).ConfigureAwait(false);
            if (endpoint is not null)
            {
                endpoint.LastStatus = status;
                endpoint.LastSeenAt = DateTimeOffset.UtcNow;
                await db.SaveChangesAsync(cancellationToken).ConfigureAwait(false);
            }
        }
        catch (Exception ex)
        {
            _logger.LogDebug(ex, "Failed to update HA endpoint {EndpointId} status", endpointId);
        }
    }

    private static Task SendAsync(ClientWebSocket ws, string json, CancellationToken cancellationToken) =>
        ws.SendAsync(Encoding.UTF8.GetBytes(json), WebSocketMessageType.Text, endOfMessage: true, cancellationToken);

    /// <summary>Reads one (possibly fragmented) text message and parses it as JSON.</summary>
    private static async Task<JsonDocument> ReceiveJsonAsync(ClientWebSocket ws, CancellationToken cancellationToken)
    {
        using var buffer = new MemoryStream();
        var chunk = new byte[8192];
        WebSocketReceiveResult result;
        do
        {
            result = await ws.ReceiveAsync(chunk, cancellationToken).ConfigureAwait(false);
            if (result.MessageType == WebSocketMessageType.Close)
            {
                throw new InvalidOperationException(
                    $"Home Assistant closed the WebSocket ({result.CloseStatus}: {result.CloseStatusDescription}).");
            }

            buffer.Write(chunk, 0, result.Count);
        }
        while (!result.EndOfMessage);

        buffer.Position = 0;
        return await JsonDocument.ParseAsync(buffer, cancellationToken: cancellationToken).ConfigureAwait(false);
    }
}