using Microsoft.Extensions.Logging;
using System.Collections.Concurrent;

namespace StellaOps.Microservice;

/// <summary>
/// Tracks in-flight requests and manages their cancellation tokens.
/// </summary>
public sealed class InflightRequestTracker : IDisposable
{
    private readonly ConcurrentDictionary<Guid, InflightRequest> _inflight = new();
    private readonly ILogger<InflightRequestTracker> _logger;
    private bool _disposed;

    /// <summary>
    /// Initializes a new instance of the <see cref="InflightRequestTracker"/> class.
    /// </summary>
    public InflightRequestTracker(ILogger<InflightRequestTracker> logger)
    {
        _logger = logger;
    }

    /// <summary>
    /// Gets the count of in-flight requests.
    /// </summary>
    public int Count => _inflight.Count;

    /// <summary>
    /// Starts tracking a request and returns a cancellation token for it.
    /// </summary>
    /// <param name="correlationId">The correlation ID of the request.</param>
    /// <returns>A cancellation token that will be triggered if the request is cancelled.</returns>
    public CancellationToken Track(Guid correlationId)
    {
        ObjectDisposedException.ThrowIf(_disposed, this);

        var cts = new CancellationTokenSource();
        var request = new InflightRequest(cts);

        if (!_inflight.TryAdd(correlationId, request))
        {
            cts.Dispose();
            throw new InvalidOperationException($"Request {correlationId} is already being tracked");
        }

        _logger.LogDebug("Started tracking request {CorrelationId}", correlationId);
        return cts.Token;
    }

    /// <summary>
    /// Cancels a specific request.
    /// </summary>
    /// <param name="correlationId">The correlation ID of the request to cancel.</param>
    /// <param name="reason">The reason for cancellation.</param>
    /// <returns>True if the request was found and cancelled; otherwise false.</returns>
    public bool Cancel(Guid correlationId, string? reason)
    {
        if (_inflight.TryGetValue(correlationId, out var request))
        {
            try
            {
                request.Cts.Cancel();
                _logger.LogInformation(
                    "Cancelled request {CorrelationId}: {Reason}",
                    correlationId,
                    reason ?? "Unknown");
                return true;
            }
            catch (ObjectDisposedException)
            {
                // CTS was already disposed, request completed
                return false;
            }
        }

        _logger.LogDebug(
            "Cannot cancel request {CorrelationId}: not found (may have already completed)",
            correlationId);
        return false;
    }

    /// <summary>
    /// Marks a request as completed and removes it from tracking.
    /// </summary>
    /// <param name="correlationId">The correlation ID of the completed request.</param>
    public void Complete(Guid correlationId)
    {
        if (_inflight.TryRemove(correlationId, out var request))
        {
            request.Cts.Dispose();
            _logger.LogDebug("Completed request {CorrelationId}", correlationId);
        }
    }

    /// <summary>
    /// Cancels all in-flight requests.
    /// </summary>
    /// <param name="reason">The reason for cancellation.</param>
    public void CancelAll(string reason)
    {
        var count = 0;
        foreach (var kvp in _inflight)
        {
            try
            {
                kvp.Value.Cts.Cancel();
                count++;
            }
            catch (ObjectDisposedException)
            {
                // Already disposed
            }
        }

        _logger.LogInformation("Cancelled {Count} in-flight requests: {Reason}", count, reason);

        // Clear and dispose all
        foreach (var kvp in _inflight)
        {
            if (_inflight.TryRemove(kvp.Key, out var request))
            {
                request.Cts.Dispose();
            }
        }
    }

    /// <inheritdoc />
    public void Dispose()
    {
        if (_disposed) return;
        _disposed = true;

        CancelAll("Disposing tracker");
    }

    private sealed class InflightRequest
    {
        public CancellationTokenSource Cts { get; }

        public InflightRequest(CancellationTokenSource cts)
        {
            Cts = cts;
        }
    }
}