// Copyright (c) Microsoft. All rights reserved.

using System.ClientModel;
using System.Diagnostics.CodeAnalysis;
using Microsoft.Agents.AI.OpenAI;
using Microsoft.Shared.DiagnosticIds;
using Microsoft.Shared.Diagnostics;
using OpenAI.Chat;
using OpenAI.Responses;

namespace Microsoft.Agents.AI;

/// <summary>
/// Provides extension methods for <see cref="AIAgent"/> to simplify interaction with OpenAI chat messages
/// and return native OpenAI <see cref="ChatCompletion"/> responses.
/// </summary>
/// <remarks>
/// These extensions bridge the gap between the Microsoft Extensions AI framework and the OpenAI SDK,
/// allowing developers to work with native OpenAI types while leveraging the AI Agent framework.
/// The methods handle the conversion between OpenAI chat message types and Microsoft Extensions AI types,
/// and return OpenAI <see cref="ChatCompletion"/> objects directly from the agent's <see cref="AgentResponse"/>.
/// </remarks>
[Experimental(DiagnosticIds.Experiments.AIOpenAIResponses)]
public static class AIAgentWithOpenAIExtensions
{
    /// <summary>
    /// Runs the AI agent with a collection of OpenAI chat messages and returns the response as a native OpenAI <see cref="ChatCompletion"/>.
    /// </summary>
    /// <param name="agent">The AI agent to run.</param>
    /// <param name="messages">The collection of OpenAI chat messages to send to the agent.</param>
    /// <param name="session">The conversation session to continue with this invocation. If not provided, creates a new session. The session will be mutated with the provided messages and agent response.</param>
    /// <param name="options">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>A <see cref="Task{ChatCompletion}"/> representing the asynchronous operation that returns a native OpenAI <see cref="ChatCompletion"/> response.</returns>
    /// <exception cref="ArgumentNullException">Thrown when <paramref name="agent"/> or <paramref name="messages"/> is <see langword="null"/>.</exception>
    /// <exception cref="ArgumentException">Thrown when the agent's response cannot be converted to a <see cref="ChatCompletion"/>, typically when the underlying representation is not an OpenAI response.</exception>
    /// <exception cref="NotSupportedException">Thrown when any message in <paramref name="messages"/> has a type that is not supported by the message conversion method.</exception>
    /// <remarks>
    /// This method converts the OpenAI chat messages to the Microsoft Extensions AI format using the appropriate conversion method,
    /// runs the agent with the converted message collection, and then extracts the native OpenAI <see cref="ChatCompletion"/> from the response using <see cref="AgentResponseExtensions.AsOpenAIChatCompletion"/>.
    /// </remarks>
    public static async Task<ChatCompletion> RunAsync(this AIAgent agent, IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
    {
        Throw.IfNull(agent);
        Throw.IfNull(messages);

        var response = await agent.RunAsync([.. messages.AsChatMessages()], session, options, cancellationToken).ConfigureAwait(false);

        return response.AsOpenAIChatCompletion();
    }

    /// <summary>
    /// Runs the AI agent with a single OpenAI chat message and returns the response as collection of native OpenAI <see cref="StreamingChatCompletionUpdate"/>.
    /// </summary>
    /// <param name="agent">The AI agent to run.</param>
    /// <param name="messages">The collection of OpenAI chat messages to send to the agent.</param>
    /// <param name="session">The conversation session to continue with this invocation. If not provided, creates a new session. The session will be mutated with the provided message and agent response.</param>
    /// <param name="options">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>A <see cref="Task{ChatCompletion}"/> representing the asynchronous operation that returns a native OpenAI <see cref="ChatCompletion"/> response.</returns>
    /// <exception cref="ArgumentNullException">Thrown when <paramref name="agent"/> or <paramref name="messages"/> is <see langword="null"/>.</exception>
    /// <exception cref="ArgumentException">Thrown when the agent's response cannot be converted to a <see cref="ChatCompletion"/>, typically when the underlying representation is not an OpenAI response.</exception>
    /// <exception cref="NotSupportedException">Thrown when the <paramref name="messages"/> type is not supported by the message conversion method.</exception>
    /// <remarks>
    /// This method converts the OpenAI chat messages to the Microsoft Extensions AI format using the appropriate conversion method,
    /// runs the agent, and then extracts the native OpenAI <see cref="ChatCompletion"/> from the response using <see cref="AgentResponseExtensions.AsOpenAIChatCompletion"/>.
    /// </remarks>
    public static AsyncCollectionResult<StreamingChatCompletionUpdate> RunStreamingAsync(this AIAgent agent, IEnumerable<ChatMessage> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
    {
        Throw.IfNull(agent);
        Throw.IfNull(messages);

        IAsyncEnumerable<AgentResponseUpdate> response = agent.RunStreamingAsync([.. messages.AsChatMessages()], session, options, cancellationToken);

        return new AsyncStreamingChatCompletionUpdateCollectionResult(response);
    }

    /// <summary>
    /// Runs the AI agent with a collection of OpenAI response items and returns the response as a native OpenAI <see cref="ResponseResult"/>.
    /// </summary>
    /// <param name="agent">The AI agent to run.</param>
    /// <param name="messages">The collection of OpenAI response items to send to the agent.</param>
    /// <param name="session">The conversation session to continue with this invocation. If not provided, creates a new session. The session will be mutated with the provided messages and agent response.</param>
    /// <param name="options">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>A <see cref="Task{ResponseResult}"/> representing the asynchronous operation that returns a native OpenAI <see cref="ResponseResult"/> response.</returns>
    /// <exception cref="ArgumentNullException">Thrown when <paramref name="agent"/> or <paramref name="messages"/> is <see langword="null"/>.</exception>
    /// <exception cref="ArgumentException">Thrown when the agent's response cannot be converted to an <see cref="ResponseResult"/>, typically when the underlying representation is not an OpenAI response.</exception>
    /// <exception cref="NotSupportedException">Thrown when any message in <paramref name="messages"/> has a type that is not supported by the message conversion method.</exception>
    /// <remarks>
    /// This method converts the OpenAI response items to the Microsoft Extensions AI format using the appropriate conversion method,
    /// runs the agent with the converted message collection, and then extracts the native OpenAI <see cref="ResponseResult"/> from the response using <see cref="AgentResponseExtensions.AsOpenAIResponse"/>.
    /// </remarks>
    public static async Task<ResponseResult> RunAsync(this AIAgent agent, IEnumerable<ResponseItem> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
    {
        Throw.IfNull(agent);
        Throw.IfNull(messages);

        var response = await agent.RunAsync(messages.AsChatMessages(), session, options, cancellationToken).ConfigureAwait(false);

        return response.AsOpenAIResponse();
    }

    /// <summary>
    /// Runs the AI agent in streaming mode with a collection of OpenAI response items and returns the response as a collection of native OpenAI <see cref="StreamingResponseUpdate"/>.
    /// </summary>
    /// <param name="agent">The AI agent to run.</param>
    /// <param name="messages">The collection of OpenAI response items to send to the agent.</param>
    /// <param name="session">The conversation session to continue with this invocation. If not provided, creates a new session. The session will be mutated with the provided messages and agent response updates.</param>
    /// <param name="options">Optional parameters for agent invocation.</param>
    /// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests. The default is <see cref="CancellationToken.None"/>.</param>
    /// <returns>An <see cref="AsyncCollectionResult{StreamingResponseUpdate}"/> representing the asynchronous enumerable that yields native OpenAI <see cref="StreamingResponseUpdate"/> instances as they are streamed.</returns>
    /// <exception cref="ArgumentNullException">Thrown when <paramref name="agent"/> or <paramref name="messages"/> is <see langword="null"/>.</exception>
    /// <exception cref="ArgumentException">Thrown when the agent's response cannot be converted to <see cref="StreamingResponseUpdate"/> instances, typically when the underlying representation is not an OpenAI response.</exception>
    /// <exception cref="NotSupportedException">Thrown when any message in <paramref name="messages"/> has a type that is not supported by the message conversion method.</exception>
    /// <remarks>
    /// This method converts the OpenAI response items to the Microsoft Extensions AI format using the appropriate conversion method,
    /// runs the agent in streaming mode, and then yields native OpenAI <see cref="StreamingResponseUpdate"/> instances as they are produced.
    /// The method attempts to extract <see cref="StreamingResponseUpdate"/> from the underlying response representation. If a raw update is not available,
    /// it is skipped because the OpenAI library does not currently expose model factory methods for creating such instances.
    /// </remarks>
    public static AsyncCollectionResult<StreamingResponseUpdate> RunStreamingAsync(this AIAgent agent, IEnumerable<ResponseItem> messages, AgentSession? session = null, AgentRunOptions? options = null, CancellationToken cancellationToken = default)
    {
        Throw.IfNull(agent);
        Throw.IfNull(messages);

        IAsyncEnumerable<AgentResponseUpdate> response = agent.RunStreamingAsync([.. messages.AsChatMessages()], session, options, cancellationToken);

        return new AsyncStreamingResponseUpdateCollectionResult(response);
    }
}
