// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.


namespace TypeAgent.AIClient;

public class OpenAITextEmbeddingModel : ModelApi, ITextEmbeddingModel
{
    int _dimensions;

    public OpenAITextEmbeddingModel()
        : this(ModelApiSettings.FromEnv(ModelType.Embedding))
    {
    }

    public OpenAITextEmbeddingModel(ModelApiSettings settings, int dimensions = 0, int maxBatchSize = 2048)
        : this(settings, null, dimensions, maxBatchSize)
    {
    }

    public OpenAITextEmbeddingModel(
        ModelApiSettings settings,
        HttpClient? client,
        int dimensions = 0,
        int maxBatchSize = 2048
    )
        : base(settings, client)
    {
        _dimensions = dimensions;
        MaxBatchSize = maxBatchSize;
    }

    public int MaxBatchSize { get; }

    public async Task<float[]> GenerateAsync(string text, CancellationToken cancellationToken)
    {
        Response response = await GetResponseAsync([text], cancellationToken).ConfigureAwait(false);
        return response.data[0].embedding;
    }

    public async Task<IList<float[]>> GenerateAsync(IList<string> texts, CancellationToken cancellationToken)
    {
        Response response = await GetResponseAsync(texts, cancellationToken).ConfigureAwait(false);
        return response.data.Map((m) => m.embedding);
    }

    private async Task<Response> GetResponseAsync(IList<string> texts, CancellationToken cancellationToken)
    {
        ArgumentVerify.ThrowIfNullOrEmpty(texts, nameof(texts));

        Request request = CreateRequest(texts);
        string? apiToken = Settings.ApiTokenProvider is not null
                    ? await Settings.ApiTokenProvider.GetAccessTokenAsync(cancellationToken).ConfigureAwait(false)
                    : null;

        Response response = await Client.GetJsonResponseAsync<Request, Response>(
            Settings.Endpoint,
            request,
            apiToken,
            Settings.Retry,
            cancellationToken
        ).ConfigureAwait(false);

        response.ThrowIfInvalid();

        return response;
    }

    private Request CreateRequest(IList<string> input)
    {
        var request = new Request
        {
            input = input
        };
        if (_dimensions > 0)
        {
            request.dimensions = _dimensions;
        }
        if (!string.IsNullOrEmpty(Settings.ModelName))
        {
            request.model = Settings.ModelName;
        }
        return request;
    }

    private struct Request
    {
        public string? model { get; set; }
        public int? dimensions { get; set; }
        public IList<string> input { get; set; }
    }

    private struct Response
    {
        public EmbeddingData[] data { get; set; }

        public void ThrowIfInvalid()
        {
            if (data is null || data.Length == 0)
            {
                throw new AIClientException(AIClientException.ErrorCode.InvalidEmbeddingResponse);
            }
        }
    };

    private struct EmbeddingData
    {
        public float[] embedding { get; set; }
    }
}
