using System.Collections.Concurrent;
using System.Threading;

namespace Gott.Services;

public class StreamingTracker
{
    private readonly ConcurrentDictionary<string, StreamingState> _streaming = new ConcurrentDictionary<string, StreamingState>();

    public CancellationTokenSource Start(string sessionId, string chainId)
    {
        var cts = new CancellationTokenSource();
        _streaming[sessionId] = new StreamingState(chainId, cts);
        return cts;
    }

    public void Stop(string sessionId)
    {
        if (_streaming.TryRemove(sessionId, out var state))
        {
            state.Cts.Dispose();
        }
    }

    /// <summary>
    /// Cancels the running stream for a session. Returns true if a stream was cancelled.
    /// </summary>
    public bool Cancel(string sessionId)
    {
        if (_streaming.TryGetValue(sessionId, out var state))
        {
            try { state.Cts.Cancel(); } catch { }
            return true;
        }
        return false;
    }

    public bool IsStreaming(string sessionId) => _streaming.ContainsKey(sessionId);

    public string GetChainId(string sessionId) => _streaming.TryGetValue(sessionId, out var state) ? state.ChainId : null;

    private record StreamingState(string ChainId, CancellationTokenSource Cts);
}