using Microsoft.Extensions.AI;
using OpenAI;

namespace AnchorCli;

/// <summary>
/// Handles streaming responses from the chat client, including token usage capture.
/// </summary>
internal sealed class ResponseStreamer
{
    private readonly ChatSession _session;

    public ResponseStreamer(ChatSession session)
    {
        _session = session;
    }

    /// <summary>
    /// Streams a response from the session and captures token usage.
    /// Returns an async enumerable that yields text chunks as they arrive.
    /// </summary>
    public async IAsyncEnumerable<string> StreamAsync(
        [System.Runtime.CompilerServices.EnumeratorCancellation] CancellationToken cancellationToken)
    {
        await using var stream = _session
            .GetStreamingResponseAsync(cancellationToken)
            .GetAsyncEnumerator(cancellationToken);

        int respIn = 0, respOut = 0;

        void CaptureUsage(ChatResponseUpdate update)
        {
            if (update.RawRepresentation is OpenAI.Chat.StreamingChatCompletionUpdate raw
                && raw.Usage != null)
            {
                respIn = raw.Usage.InputTokenCount;
                respOut += raw.Usage.OutputTokenCount;
            }
        }

        // Stream all chunks
        while (await stream.MoveNextAsync())
        {
            cancellationToken.ThrowIfCancellationRequested();
            CaptureUsage(stream.Current);
            var text = stream.Current.Text;
            if (!string.IsNullOrEmpty(text))
            {
                yield return text;
            }
        }

        // Store final usage stats
        LastInputTokens = respIn;
        LastOutputTokens = respOut;
    }

    public int LastInputTokens { get; private set; }
    public int LastOutputTokens { get; private set; }
}