///////////////////////////////////////////////////////////////////////////////
//
// (C) 2025 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.ComponentModel;
using System.Diagnostics;
using System.Linq;
using System.Threading.Tasks;
using Wisej.AI.Endpoints;
using Wisej.AI.Helpers;

namespace Wisej.AI.Services
{
	/// <summary>
	/// Provides a default implementation of the <see cref="IRerankingService"/> for reranking documents
	/// using a specified endpoint to an AI provider. It uses LLM models to rerank and optionally optimize the
	/// submitted inputs.
	/// </summary>
	[ApiCategory("Services/IRerankingService")]
	public class DefaultRerankingService : IRerankingService
	{
		/// <summary>
		/// Initializes a new instance of the <see cref="DefaultRerankingService"/> class.
		/// </summary>
		/// <param name="endpoint">The endpoint to be used by the service.
		/// If not provided, a default <see cref="OpenAIEndpoint"/> is used.</param>
		public DefaultRerankingService(SmartEndpoint endpoint = null)
		{
			this.Endpoint = endpoint ?? new OpenAIEndpoint();
			this.Prompt = new SmartPrompt("[DefaultRerankingService.Prompt]");
		}

		/// <summary>
		/// Gets or sets the prompt used to perform the reranking task.
		/// </summary>
		public SmartPrompt Prompt
		{
			get;
			set;
		}

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

		/// <summary>
		/// Gets or sets the maximum number of documents to return in the reranked response.
		/// If set to 0 it will return all the documents in the input.
		/// </summary>
		[DefaultValue(10)]
		public virtual int TopN
		{
			get;
			set;
		} = 10;

		/// <summary>
		/// Asynchronously reranks the provided chunks based on the specified query.
		/// </summary>
		/// <param name="query">The query used to rerank the chunks.</param>
		/// <param name="chunks">The array of text inputs to be reranked.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains an array of reranked text inputs.</returns>
		public virtual async Task<string[]> RerankAsync(string query, string[] chunks)
		{
			if (String.IsNullOrEmpty(query))
				return chunks;

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

			var topN = this.TopN == 0 ? chunks.Length : this.TopN;

			this.Prompt.Parameters.Add("topN", topN);
			this.Prompt.Parameters.Add("query", query);

			var payload =
				Enumerable.Range(0, chunks.Length).Select(i => new
				{
					index = i,
					document = chunks[i]
				});

			try
			{
				var result = await this.Prompt.AskAsync(this.Endpoint, JSON.Stringify(payload));
				var json = RegexHelper.GetJSON(result);
				var data = JSON.Parse(json);

				if (data is Array array)
				{
					if (array.Length == 0)
					{
						chunks = [];
					}
					else
					{
						var ranking = ((int[])data);
						chunks = ranking.Take(topN).Select(i => chunks[i]).ToArray();
					}
				}

				return chunks;
			}
			catch (Exception ex)
			{
				LogHelper.Log(TraceLevel.Error, ex);

				return chunks;
			}
		}
	}
}