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<string> 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));
        ChatClient chatClient = azureClient.GetChatClient(deploymentName);

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

        await foreach (StreamingChatCompletionUpdate update in response)
        {
            foreach (ChatMessageContentPart updatePart in update.ContentUpdate)
            {
                yield return updatePart.Text;
            }
        }
    }
}