using System.Runtime.CompilerServices;
using Azure;
using Azure.AI.OpenAI;
using OpenAI.Chat;
using Yi.Framework.AiHub.Domain.Shared.Dtos;

namespace Yi.Framework.AiHub.Domain.AiChat.Impl;

public class AzureChatService : IChatService
{
    public AzureChatService()
    {
    }

    public async IAsyncEnumerable<CompleteChatResponse> CompleteChatAsync(AiModelDescribe aiModelDescribe,
        List<ChatMessage> messages,
        [EnumeratorCancellation] CancellationToken cancellationToken)
    {
        var endpoint = new Uri(aiModelDescribe.Endpoint);

        var deploymentName = aiModelDescribe.ModelId;
        var apiKey = aiModelDescribe.ApiKey;

        AzureOpenAIClient azureClient = new(
            endpoint,
            new AzureKeyCredential(apiKey), new AzureOpenAIClientOptions()
            {
                NetworkTimeout = TimeSpan.FromSeconds(600),
            });

        ChatClient chatClient = azureClient.GetChatClient(deploymentName);

        var response = chatClient.CompleteChatStreamingAsync(messages, new ChatCompletionOptions()
        {
            //   MaxOutputTokenCount = 2048
        }, cancellationToken: cancellationToken);

        await foreach (StreamingChatCompletionUpdate update in response)
        {
            var result = new CompleteChatResponse();
            var isFinish = update.Usage?.OutputTokenCount is not null;
            if (isFinish)
            {
                result.IsFinish = true;
                result.TokenUsage = new TokenUsage
                {
                    OutputTokenCount = update.Usage.OutputTokenCount,
                    InputTokenCount = update.Usage.InputTokenCount,
                    TotalTokenCount = update.Usage.TotalTokenCount
                };
            }

            foreach (ChatMessageContentPart updatePart in update.ContentUpdate)
            {
                result.Content = updatePart.Text;
                yield return result;
            }

            if (isFinish)
            {
                yield return result;
            }
        }
    }
}