// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
// Code generated by Microsoft (R) AutoRest Code Generator.
package com.azure.ai.openai;

import com.azure.ai.openai.implementation.CompletionsUtils;
import com.azure.ai.openai.implementation.NonAzureOpenAIClientImpl;
import com.azure.ai.openai.implementation.OpenAIClientImpl;
import com.azure.ai.openai.implementation.OpenAIServerSentEvents;
import com.azure.ai.openai.models.ChatCompletions;
import com.azure.ai.openai.models.ChatCompletionsOptions;
import com.azure.ai.openai.models.Completions;
import com.azure.ai.openai.models.CompletionsOptions;
import com.azure.ai.openai.models.Embeddings;
import com.azure.ai.openai.models.EmbeddingsOptions;
import com.azure.ai.openai.models.ImageGenerationOptions;
import com.azure.ai.openai.models.ImageOperationResponse;
import com.azure.ai.openai.models.ImageResponse;
import com.azure.core.annotation.Generated;
import com.azure.core.annotation.ReturnType;
import com.azure.core.annotation.ServiceClient;
import com.azure.core.annotation.ServiceMethod;
import com.azure.core.exception.ClientAuthenticationException;
import com.azure.core.exception.HttpResponseException;
import com.azure.core.exception.ResourceModifiedException;
import com.azure.core.exception.ResourceNotFoundException;
import com.azure.core.http.rest.RequestOptions;
import com.azure.core.http.rest.Response;
import com.azure.core.util.BinaryData;
import com.azure.core.util.IterableStream;
import com.azure.core.util.logging.ClientLogger;
import com.azure.core.util.polling.SyncPoller;
import reactor.core.Exceptions;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;

/** Initializes a new instance of the synchronous OpenAIClient type. */
@ServiceClient(builder = OpenAIClientBuilder.class)
public final class OpenAIClient {
    private static final ClientLogger LOGGER = new ClientLogger(OpenAIClient.class);

    @Generated private final OpenAIClientImpl serviceClient;

    private final NonAzureOpenAIClientImpl openAIServiceClient;

    /**
     * Initializes an instance of OpenAIClient class.
     *
     * @param serviceClient the service client implementation.
     */
    OpenAIClient(OpenAIClientImpl serviceClient) {
        this.serviceClient = serviceClient;
        openAIServiceClient = null;
    }

    /**
     * Initializes an instance of OpenAIClient class for NonAzure Implementation.
     *
     * @param serviceClient the service client implementation.
     */
    OpenAIClient(NonAzureOpenAIClientImpl serviceClient) {
        this.serviceClient = null;
        openAIServiceClient = serviceClient;
    }

    /**
     * Return the embeddings for a given prompt.
     *
     * <p><strong>Request Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     user: String (Optional)
     *     model: String (Optional)
     *     input (Required): [
     *         String (Required)
     *     ]
     * }
     * }</pre>
     *
     * <p><strong>Response Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     data (Required): [
     *          (Required){
     *             embedding (Required): [
     *                 double (Required)
     *             ]
     *             index: int (Required)
     *         }
     *     ]
     *     usage (Required): {
     *         prompt_tokens: int (Required)
     *         total_tokens: int (Required)
     *     }
     * }
     * }</pre>
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param embeddingsOptions The configuration information for an embeddings request. Embeddings measure the
     *     relatedness of text strings and are commonly used for search, clustering, recommendations, and other similar
     *     scenarios.
     * @param requestOptions The options to configure the HTTP request before HTTP client sends it.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @return representation of the response data from an embeddings request. Embeddings measure the relatedness of
     *     text strings and are commonly used for search, clustering, recommendations, and other similar scenarios along
     *     with {@link Response}.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Response<BinaryData> getEmbeddingsWithResponse(
            String deploymentOrModelName, BinaryData embeddingsOptions, RequestOptions requestOptions) {
        return openAIServiceClient != null
                ? openAIServiceClient.getEmbeddingsWithResponse(
                        deploymentOrModelName, embeddingsOptions, requestOptions)
                : serviceClient.getEmbeddingsWithResponse(deploymentOrModelName, embeddingsOptions, requestOptions);
    }

    /**
     * Gets completions for the provided input prompts. Completions support a wide variety of tasks and generate text
     * that continues from or "completes" provided prompt data.
     *
     * <p><strong>Request Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     prompt (Required): [
     *         String (Required)
     *     ]
     *     max_tokens: Integer (Optional)
     *     temperature: Double (Optional)
     *     top_p: Double (Optional)
     *     logit_bias (Optional): {
     *         String: int (Optional)
     *     }
     *     user: String (Optional)
     *     n: Integer (Optional)
     *     logprobs: Integer (Optional)
     *     echo: Boolean (Optional)
     *     stop (Optional): [
     *         String (Optional)
     *     ]
     *     presence_penalty: Double (Optional)
     *     frequency_penalty: Double (Optional)
     *     best_of: Integer (Optional)
     *     stream: Boolean (Optional)
     *     model: String (Optional)
     * }
     * }</pre>
     *
     * <p><strong>Response Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     id: String (Required)
     *     created: int (Required)
     *     choices (Required): [
     *          (Required){
     *             text: String (Required)
     *             index: int (Required)
     *             logprobs (Required): {
     *                 tokens (Required): [
     *                     String (Required)
     *                 ]
     *                 token_logprobs (Required): [
     *                     double (Required)
     *                 ]
     *                 top_logprobs (Required): [
     *                      (Required){
     *                         String: double (Required)
     *                     }
     *                 ]
     *                 text_offset (Required): [
     *                     int (Required)
     *                 ]
     *             }
     *             finish_reason: String(stopped/tokenLimitReached/contentFiltered) (Required)
     *         }
     *     ]
     *     usage (Required): {
     *         completion_tokens: int (Required)
     *         prompt_tokens: int (Required)
     *         total_tokens: int (Required)
     *     }
     * }
     * }</pre>
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param completionsOptions The configuration information for a completions request. Completions support a wide
     *     variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @param requestOptions The options to configure the HTTP request before HTTP client sends it.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @return completions for the provided input prompts. Completions support a wide variety of tasks and generate text
     *     that continues from or "completes" provided prompt data along with {@link Response}.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Response<BinaryData> getCompletionsWithResponse(
            String deploymentOrModelName, BinaryData completionsOptions, RequestOptions requestOptions) {
        return openAIServiceClient != null
                ? openAIServiceClient.getCompletionsWithResponse(
                        deploymentOrModelName, completionsOptions, requestOptions)
                : serviceClient.getCompletionsWithResponse(deploymentOrModelName, completionsOptions, requestOptions);
    }

    /**
     * Gets chat completions for the provided chat messages. Completions support a wide variety of tasks and generate
     * text that continues from or "completes" provided prompt data.
     *
     * <p><strong>Request Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     messages (Required): [
     *          (Required){
     *             role: String(system/assistant/user) (Required)
     *             content: String (Optional)
     *         }
     *     ]
     *     max_tokens: Integer (Optional)
     *     temperature: Double (Optional)
     *     top_p: Double (Optional)
     *     logit_bias (Optional): {
     *         String: int (Optional)
     *     }
     *     user: String (Optional)
     *     n: Integer (Optional)
     *     stop (Optional): [
     *         String (Optional)
     *     ]
     *     presence_penalty: Double (Optional)
     *     frequency_penalty: Double (Optional)
     *     stream: Boolean (Optional)
     *     model: String (Optional)
     * }
     * }</pre>
     *
     * <p><strong>Response Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     id: String (Required)
     *     created: int (Required)
     *     choices (Required): [
     *          (Required){
     *             message (Optional): {
     *                 role: String(system/assistant/user) (Required)
     *                 content: String (Optional)
     *             }
     *             index: int (Required)
     *             finish_reason: String(stopped/tokenLimitReached/contentFiltered) (Required)
     *             delta (Optional): {
     *                 role: String(system/assistant/user) (Optional)
     *                 content: String (Optional)
     *             }
     *         }
     *     ]
     *     usage (Required): {
     *         completion_tokens: int (Required)
     *         prompt_tokens: int (Required)
     *         total_tokens: int (Required)
     *     }
     * }
     * }</pre>
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param chatCompletionsOptions The configuration information for a chat completions request. Completions support a
     *     wide variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @param requestOptions The options to configure the HTTP request before HTTP client sends it.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @return chat completions for the provided chat messages. Completions support a wide variety of tasks and generate
     *     text that continues from or "completes" provided prompt data along with {@link Response}.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Response<BinaryData> getChatCompletionsWithResponse(
            String deploymentOrModelName, BinaryData chatCompletionsOptions, RequestOptions requestOptions) {
        return openAIServiceClient != null
                ? openAIServiceClient.getChatCompletionsWithResponse(
                        deploymentOrModelName, chatCompletionsOptions, requestOptions)
                : serviceClient.getChatCompletionsWithResponse(
                        deploymentOrModelName, chatCompletionsOptions, requestOptions);
    }

    /**
     * Return the embeddings for a given prompt.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param embeddingsOptions The configuration information for an embeddings request. Embeddings measure the
     *     relatedness of text strings and are commonly used for search, clustering, recommendations, and other similar
     *     scenarios.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return representation of the response data from an embeddings request. Embeddings measure the relatedness of
     *     text strings and are commonly used for search, clustering, recommendations, and other similar scenarios.
     */
    @Generated
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Embeddings getEmbeddings(String deploymentOrModelName, EmbeddingsOptions embeddingsOptions) {
        // Generated convenience method for getEmbeddingsWithResponse
        RequestOptions requestOptions = new RequestOptions();
        return getEmbeddingsWithResponse(
                        deploymentOrModelName, BinaryData.fromObject(embeddingsOptions), requestOptions)
                .getValue()
                .toObject(Embeddings.class);
    }

    /**
     * Gets completions for the provided input prompts. Completions support a wide variety of tasks and generate text
     * that continues from or "completes" provided prompt data.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param completionsOptions The configuration information for a completions request. Completions support a wide
     *     variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return completions for the provided input prompts. Completions support a wide variety of tasks and generate text
     *     that continues from or "completes" provided prompt data.
     */
    @Generated
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Completions getCompletions(String deploymentOrModelName, CompletionsOptions completionsOptions) {
        // Generated convenience method for getCompletionsWithResponse
        RequestOptions requestOptions = new RequestOptions();
        return getCompletionsWithResponse(
                        deploymentOrModelName, BinaryData.fromObject(completionsOptions), requestOptions)
                .getValue()
                .toObject(Completions.class);
    }

    /**
     * Gets completions for the provided input prompt. Completions support a wide variety of tasks and generate text
     * that continues from or "completes" provided prompt data.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param prompt The prompt to generate completion text from.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return completions for the provided input prompts. Completions support a wide variety of tasks and generate text
     *     that continues from or "completes" provided prompt data.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public Completions getCompletions(String deploymentOrModelName, String prompt) {
        return getCompletions(deploymentOrModelName, CompletionsUtils.defaultCompletionsOptions(prompt));
    }

    /**
     * Gets completions as a stream for the provided input prompts. Completions support a wide variety of tasks and
     * generate text that continues from or "completes" provided prompt data.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param completionsOptions The configuration information for a completions request. Completions support a wide
     *     variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return an {@link IterableStream} of completions for the provided input prompts. Completions support a wide
     *     variety of tasks and generate text that continues from or "completes" provided prompt data.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public IterableStream<Completions> getCompletionsStream(
            String deploymentOrModelName, CompletionsOptions completionsOptions) {
        completionsOptions.setStream(true);
        RequestOptions requestOptions = new RequestOptions();
        Flux<ByteBuffer> responseStream =
                getCompletionsWithResponse(
                                deploymentOrModelName, BinaryData.fromObject(completionsOptions), requestOptions)
                        .getValue()
                        .toFluxByteBuffer();
        OpenAIServerSentEvents<Completions> completionsStream =
                new OpenAIServerSentEvents<>(responseStream, Completions.class);
        return new IterableStream<>(completionsStream.getEvents());
    }

    /**
     * Gets chat completions for the provided chat messages. Completions support a wide variety of tasks and generate
     * text that continues from or "completes" provided prompt data.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param chatCompletionsOptions The configuration information for a chat completions request. Completions support a
     *     wide variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return chat completions for the provided chat messages. Completions support a wide variety of tasks and generate
     *     text that continues from or "completes" provided prompt data.
     */
    @Generated
    @ServiceMethod(returns = ReturnType.SINGLE)
    public ChatCompletions getChatCompletions(
            String deploymentOrModelName, ChatCompletionsOptions chatCompletionsOptions) {
        // Generated convenience method for getChatCompletionsWithResponse
        RequestOptions requestOptions = new RequestOptions();
        return getChatCompletionsWithResponse(
                        deploymentOrModelName, BinaryData.fromObject(chatCompletionsOptions), requestOptions)
                .getValue()
                .toObject(ChatCompletions.class);
    }

    /**
     * Gets chat completions for the provided chat messages. Chat completions support a wide variety of tasks and
     * generate text that continues from or "completes" provided prompt data.
     *
     * @param deploymentOrModelName Specifies either the model deployment name (when using Azure OpenAI) or model name
     *     (when using non-Azure OpenAI) to use for this request.
     * @param chatCompletionsOptions The configuration information for a chat completions request. Completions support a
     *     wide variety of tasks and generate text that continues from or "completes" provided prompt data.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return chat completions stream for the provided chat messages. Completions support a wide variety of tasks and
     *     generate text that continues from or "completes" provided prompt data.
     */
    @ServiceMethod(returns = ReturnType.COLLECTION)
    public IterableStream<ChatCompletions> getChatCompletionsStream(
            String deploymentOrModelName, ChatCompletionsOptions chatCompletionsOptions) {
        chatCompletionsOptions.setStream(true);
        RequestOptions requestOptions = new RequestOptions();
        Flux<ByteBuffer> responseStream =
                getChatCompletionsWithResponse(
                                deploymentOrModelName, BinaryData.fromObject(chatCompletionsOptions), requestOptions)
                        .getValue()
                        .toFluxByteBuffer();
        OpenAIServerSentEvents<ChatCompletions> chatCompletionsStream =
                new OpenAIServerSentEvents<>(responseStream, ChatCompletions.class);
        return new IterableStream<ChatCompletions>(chatCompletionsStream.getEvents());
    }

    /**
     * Starts the generation of a batch of images from a text caption.
     *
     * @param imageGenerationOptions Represents the request data used to generate images.
     * @throws IllegalArgumentException thrown if parameters fail the validation.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @throws RuntimeException all other wrapped checked exceptions if the request fails to be sent.
     * @return the {@link ImageResponse} for the image generation result.
     */
    @ServiceMethod(returns = ReturnType.SINGLE)
    public ImageResponse getImages(ImageGenerationOptions imageGenerationOptions) {
        RequestOptions requestOptions = new RequestOptions();
        BinaryData imageGenerationOptionsBinaryData = BinaryData.fromObject(imageGenerationOptions);

        if (openAIServiceClient != null) {
            return openAIServiceClient
                .generateImageWithResponse(imageGenerationOptionsBinaryData, requestOptions)
                .getValue()
                .toObject(ImageResponse.class);
        } else {
            // TODO: Currently, we use async client block() to avoid a unknown LRO status "notRunning" which Azure Core will
            //       fix the issue in August release and we will reuse the method
            //       "SyncPoller<BinaryData, BinaryData> beginBeginAzureBatchImageGeneration()" after.
            try {
                return this.serviceClient.beginBeginAzureBatchImageGenerationAsync(imageGenerationOptionsBinaryData,
                        requestOptions)
                    .last()
                    .flatMap(it -> it.getFinalResult())
                    .map(it -> it.toObject(ImageOperationResponse.class).getResult()).block();
            } catch (Exception e) {
                Throwable unwrapped = Exceptions.unwrap(e);
                if (unwrapped instanceof RuntimeException) {
                    throw LOGGER.logExceptionAsError((RuntimeException) unwrapped);
                } else if (unwrapped instanceof IOException) {
                    throw LOGGER.logExceptionAsError(new UncheckedIOException((IOException) unwrapped));
                } else {
                    throw LOGGER.logExceptionAsError(new RuntimeException(unwrapped));
                }
            }
        }
    }

    /**
     * Starts the generation of a batch of images from a text caption.
     *
     * <p><strong>Request Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     prompt: String (Required)
     *     n: Integer (Optional)
     *     size: String(256x256/512x512/1024x1024) (Optional)
     *     response_format: String(url/b64_json) (Optional)
     *     user: String (Optional)
     * }
     * }</pre>
     *
     * <p><strong>Response Body Schema</strong>
     *
     * <pre>{@code
     * {
     *     id: String (Required)
     *     status: String (Required)
     *     error (Optional): {
     *         code: String (Required)
     *         message: String (Required)
     *         target: String (Optional)
     *         details (Optional): [
     *             (recursive schema, see above)
     *         ]
     *         innererror (Optional): {
     *             code: String (Optional)
     *             innererror (Optional): (recursive schema, see innererror above)
     *         }
     *     }
     * }
     * }</pre>
     *
     * @param imageGenerationOptions Represents the request data used to generate images.
     * @param requestOptions The options to configure the HTTP request before HTTP client sends it.
     * @throws HttpResponseException thrown if the request is rejected by server.
     * @throws ClientAuthenticationException thrown if the request is rejected by server on status code 401.
     * @throws ResourceNotFoundException thrown if the request is rejected by server on status code 404.
     * @throws ResourceModifiedException thrown if the request is rejected by server on status code 409.
     * @return the {@link SyncPoller} for polling of status details for long running operations.
     */
    @Generated
    @ServiceMethod(returns = ReturnType.LONG_RUNNING_OPERATION)
    SyncPoller<BinaryData, BinaryData> beginBeginAzureBatchImageGeneration(
            BinaryData imageGenerationOptions, RequestOptions requestOptions) {
        return this.serviceClient.beginBeginAzureBatchImageGeneration(imageGenerationOptions, requestOptions);
    }
}
