///////////////////////////////////////////////////////////////////////////////
//
// (C) 2024 ICE TEA GROUP LLC - ALL RIGHTS RESERVED
//
// 
//
// ALL INFORMATION CONTAINED HEREIN IS, AND REMAINS
// THE PROPERTY OF ICE TEA GROUP LLC AND ITS SUPPLIERS, IF ANY.
// THE INTELLECTUAL PROPERTY AND TECHNICAL CONCEPTS CONTAINED
// HEREIN ARE PROPRIETARY TO ICE TEA GROUP LLC AND ITS SUPPLIERS
// AND MAY BE COVERED BY U.S. AND FOREIGN PATENTS, PATENT IN PROCESS, AND
// ARE PROTECTED BY TRADE SECRET OR COPYRIGHT LAW.
//
// DISSEMINATION OF THIS INFORMATION OR REPRODUCTION OF THIS MATERIAL
// IS STRICTLY FORBIDDEN UNLESS PRIOR WRITTEN PERMISSION IS OBTAINED
// FROM ICE TEA GROUP LLC.
//
///////////////////////////////////////////////////////////////////////////////

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Linq;
using System.Threading.Tasks;
using Wisej.AI.Embeddings;
using Wisej.AI.Endpoints;

namespace Wisej.AI.Services
{
	/// <summary>
	/// Represents a service that provides functionality to generate embeddings
	/// for given text inputs using a <see cref="SmartEndpoint"/>
	/// </summary>
	[ApiCategory("Services/IEmbeddingGenerationService")]
	public class DefaultEmbeddingGenerationService : IEmbeddingGenerationService
	{
		/// <summary>
		/// Initializes a new instance of the <see cref="DefaultEmbeddingGenerationService"/> class with the specified endpoint and vector size.
		/// </summary>
		/// <param name="endpoint">The endpoint used for embedding generation. If not specified, defaults to a new instance of <see cref="OpenAIEndpoint"/>.</param>
		/// <param name="vectorSize">The size of the embedding vector. Default is 1536, which corresponds to the size for "text-embedding-3-small".</param>
		public DefaultEmbeddingGenerationService(SmartEndpoint endpoint = null, int vectorSize = 1536)
		{
			// 1536 = size for "text-embedding-3-small"

			this.Endpoint = endpoint ?? new OpenAIEndpoint();
		}

		/// <summary>
		/// Gets or sets the default endpoint used for generating embeddings.
		/// </summary>
		[DefaultValue(typeof(OpenAIEndpoint))]
		public virtual SmartEndpoint Endpoint
		{
			get;
			set;
		}

		/// <summary>
		/// Gets or sets the maximum number of text chunks that can be sent as an array to the provider.
		/// </summary>
		[DefaultValue(512)]
		public virtual int MaxArraySize
		{
			get;
			set;
		} = 512; // 1536 = size for "text-embedding-3-small"

		/// <summary>
		/// Gets or sets the dimension of the embedding vectors.
		/// </summary>
		[DefaultValue(1536)]
		public virtual int VectorSize
		{
			get;
			set;
		} = 1536;

		/// <summary>
		/// Asynchronously generates an embedding for a given array of text chunks.
		/// </summary>
		/// <param name="chunks">An array of text strings to be embedded. Must not be null or empty.</param>
		/// <returns>A task representing the asynchronous operation, with a result of type <see cref="Embedding"/>.</returns>
		/// <remarks>
		/// This method provides a way to generate embeddings for text data which can be utilized in various NLP tasks.
		/// Embeddings are numerical representations of text that capture semantic information.
		/// <para>
		/// Applications may use embeddings for:
		/// <list type="bullet">
		///		<item>Similarity comparisons</item>
		///		<item>Semantic searches</item>
		///		<item>Machine learning models</item>
		/// </list>
		/// </para>
		/// <code>
		/// <![CDATA[
		///		var textChunks = new[] { "Hello, world!", "C# is awesome." };
		///		var embeddings = await myEmbeddingService.EmbedAsync(textChunks);
		/// ]]>
		/// </code>
		/// </remarks>
		public async Task<Embedding> EmbedAsync(string[] chunks)
		{
			var endpoint = this.Endpoint;
			if (endpoint == null)
				throw new ArgumentNullException(nameof(endpoint));

			if (chunks == null || chunks.Length == 0)
				return null;

			// single request?
			if (chunks.Length <= this.MaxArraySize)
			{
				return await endpoint.AskEmbeddingsAsync(chunks);
			}
			else
			{
				// parallelize the requests.

				var used = 0;
				var tasks = new List<Task<Embedding>>();
				var totalItems = chunks.Length;
				var maxSize = this.MaxArraySize;

				do
				{
					var inputs = chunks.Skip(used).Take(maxSize).ToArray();
					used += inputs.Length;
					tasks.Add(endpoint.AskEmbeddingsAsync(inputs));

				} while (used < totalItems);

				await Task.WhenAll(tasks.ToArray());

				if (tasks.Count == 1)
					return tasks[0].Result;

				// combine the embeddings.
				var embedding = tasks[0].Result;
				for (var i = 1; i < tasks.Count; i++)
				{
					embedding.Add(tasks[i].Result);
				}

				return embedding;
			}
		}
	}
}
