///////////////////////////////////////////////////////////////////////////////
//
// (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.IO;
using System.Linq;
using System.Net.Http;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Wisej.AI.Embeddings;
using Wisej.AI.Helpers;

namespace Wisej.AI.Services
{
	/// <summary>
	/// Represents a service for generating embeddings using the Hugging Face API.
	/// </summary>
	/// <remarks>
	/// This class implements the <see cref="IEmbeddingGenerationService"/> interface and provides methods to generate embeddings from text chunks using the Hugging Face API.
	/// </remarks>
	[ApiCategory("Services/IEmbeddingGenerationService")]
	public class HuggingFaceEmbeddingGenerationService : IEmbeddingGenerationService
	{
		/// <summary>
		/// Initializes a new instance of the <see cref="HuggingFaceEmbeddingGenerationService"/> class with an optional URL.
		/// </summary>
		/// <param name="url">The base URL of the Hugging Face API. Default is <c>null</c>.</param>
		public HuggingFaceEmbeddingGenerationService(string url = null)
		{
			this.URL = url;

			_client = ServiceHelper.GetService<IHttpClientService>();
		}

		/// <summary>
		/// Gets or sets the base URL of the Hugging Face API.
		/// </summary>
		[DefaultValue(null)]
		public virtual string URL
		{
			get;
			set;
		}

		/// <summary>
		/// Gets or sets the maximum size of the array that can be processed in a single request.
		/// </summary>
		[DefaultValue(100)]
		public virtual int MaxArraySize
		{
			get;
			set;
		} = 100;

		/// <summary>
		/// Gets or sets the size of the vector for the embeddings. This property is not used.
		/// </summary>
		[DefaultValue(0)]
		public virtual int VectorSize
		{
			get;
			set;
		}

		/// <summary>
		/// Gets the HTTP client service used for making requests.
		/// </summary>
		internal virtual IHttpClientService Client
		{
			get => _client;
		}
		private IHttpClientService _client;

		/// <summary>
		/// Asynchronously generates embeddings for the provided text chunks.
		/// </summary>
		/// <param name="chunks">An array of text chunks to generate embeddings for.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains the generated <see cref="Embedding"/>.</returns>
		/// <remarks>
		/// This method handles chunking of input data and parallelizes requests if necessary. It combines results from multiple requests if the input exceeds the maximum array size.
		/// <code>
		/// <![CDATA[
		/// var service = new HuggingFaceEmbeddingGenerationService("http://api.example.com");
		/// var embedding = await service.EmbedAsync(new string[] { "text1", "text2" });
		/// ]]>
		/// </code>
		/// </remarks>
		public async Task<Embedding> EmbedAsync(string[] chunks)
		{
			if (chunks == null || chunks.Length == 0)
				return null;

			// single request?
			if (chunks.Length <= this.MaxArraySize)
			{
				return await PostEmbeddingRequest(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(PostEmbeddingRequest(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;
			}
		}

		private async Task<Embedding> PostEmbeddingRequest(string[] inputs)
		{
			for (var retry = 0; retry < 3; retry++)
			{
				try
				{
					var client = this.Client;
					var payload = (string)JSON.Stringify(new { inputs = inputs });
					var content = new StringContent(payload, Encoding.UTF8, "application/json");
					var response = await client.PostAsync($"{this.URL}/embed", content);
					response.EnsureSuccessStatusCode();
					var stream = await response.Content.ReadAsStreamAsync();
					var reader = new StreamReader(stream);
					var json = reader.ReadToEnd();
					var embeddings = (float[][])JSON.Parse(json);

					var vectors = new float[inputs.Length][];
					for (var i = 0; i < vectors.Length && i < embeddings.Length; i++)
					{
						vectors[i] = embeddings[i];
					}

					return new Embedding(inputs, vectors, "");
				}
				catch
				{
					// TODO: Improve it, HF sometimes returns all nulls.
					Thread.Sleep(1000);
				}
			}

			throw new Exception($"HF embedding at {this.URL} failed to respond.");
		}

		/// <summary>
		/// Not used.
		/// </summary>
		/// <remarks>
		/// This property is not implemented and will throw a <see cref="NotSupportedException"/> if accessed.
		/// </remarks>
		[Browsable(false)]
		SmartEndpoint IEmbeddingGenerationService.Endpoint
		{
			get => null;
			set { throw new NotSupportedException(); }
		}
	}
}