using System.Collections.Concurrent;
using System.Threading.Channels;
using Hc900Crawler.Core.Application.DTOs;

namespace Hc900Crawler.Core.Application.Interfaces;

public interface IPidGraphEventBroadcaster
{
    Task NotifyAsync(string taskId, AnalysisStatus status);
    IAsyncEnumerable<AnalysisStatus> SubscribeAsync(string taskId, CancellationToken ct);
}

public class PidGraphEventBroadcaster : IPidGraphEventBroadcaster
{
    private readonly ConcurrentDictionary<string, List<Channel<AnalysisStatus>>> _subscriptions = new();
    private readonly ILogger<PidGraphEventBroadcaster> _logger;

    public PidGraphEventBroadcaster(ILogger<PidGraphEventBroadcaster> logger)
    {
        _logger = logger;
    }

    public async Task NotifyAsync(string taskId, AnalysisStatus status)
    {
        if (_subscriptions.TryGetValue(taskId, out var channels))
        {
            lock (channels)
            {
                // 이미 종료된 채널 제거 및 전송
                channels.RemoveAll(c => c.Writer.TryComplete()); 
                
                foreach (var channel in channels)
                {
                    channel.Writer.TryWrite(status);
                }
            }
        }
        await Task.CompletedTask;
    }

    public async IAsyncEnumerable<AnalysisStatus> SubscribeAsync(string taskId, [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken ct)
    {
        var channel = Channel.CreateUnbounded<AnalysisStatus>();
        
        var channels = _subscriptions.GetOrAdd(taskId, _ => new List<Channel<AnalysisStatus>>());
        lock (channels)
        {
            channels.Add(channel);
        }

        _logger.LogInformation("New SSE subscription for taskId: {TaskId}", taskId);

        bool hasMore = true;
        while (hasMore)
        {
            try
            {
                hasMore = await channel.Reader.WaitToReadAsync(ct);
            }
            catch (OperationCanceledException)
            {
                _logger.LogInformation("SSE subscription cancelled for taskId: {TaskId}", taskId);
                break;
            }

            while (channel.Reader.TryRead(out var status))
            {
                yield return status;
            }
        }

        lock (channels)
        {
            channels.Remove(channel);
            if (channels.Count == 0)
            {
                _subscriptions.TryRemove(taskId, out _);
            }
        }
    }
}
