using System.Net.WebSockets;

namespace MeetingAssistant.Transcription;

public interface IFunAsrBackendReadinessProbe
{
    Task WaitUntilReadyAsync(Uri endpoint, TimeSpan timeout, CancellationToken cancellationToken);
}

public sealed class FunAsrWebSocketBackendReadinessProbe : IFunAsrBackendReadinessProbe
{
    private static readonly TimeSpan RetryDelay = TimeSpan.FromSeconds(1);

    public async Task WaitUntilReadyAsync(
        Uri endpoint,
        TimeSpan timeout,
        CancellationToken cancellationToken)
    {
        if (timeout <= TimeSpan.Zero)
        {
            return;
        }

        using var timeoutSource = new CancellationTokenSource(timeout);
        using var linkedSource = CancellationTokenSource.CreateLinkedTokenSource(
            cancellationToken,
            timeoutSource.Token);

        while (!linkedSource.IsCancellationRequested)
        {
            try
            {
                using var client = new ClientWebSocket();
                await client.ConnectAsync(endpoint, linkedSource.Token).ConfigureAwait(false);
                await CloseQuietlyAsync(client, linkedSource.Token).ConfigureAwait(false);
                return;
            }
            catch (OperationCanceledException)
            {
                break;
            }
            catch (Exception exception) when (IsTransientReadinessFailure(exception))
            {
                try
                {
                    await Task.Delay(RetryDelay, linkedSource.Token).ConfigureAwait(false);
                }
                catch (OperationCanceledException)
                {
                    break;
                }
            }
        }

        throw new TimeoutException(
            $"Timed out waiting {timeout} for managed FunASR backend WebSocket readiness at {endpoint}.");
    }

    private static async Task CloseQuietlyAsync(ClientWebSocket client, CancellationToken cancellationToken)
    {
        if (client.State != WebSocketState.Open)
        {
            return;
        }

        try
        {
            await client.CloseOutputAsync(
                WebSocketCloseStatus.NormalClosure,
                "readiness probe complete",
                cancellationToken).ConfigureAwait(false);
        }
        catch (WebSocketException)
        {
        }
    }

    private static bool IsTransientReadinessFailure(Exception exception)
    {
        return exception is WebSocketException
            or HttpRequestException
            or IOException;
    }
}