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

namespace Wisej.AI.Services
{
	/// <summary>
	/// Represents a service for storing and querying embeddings using Azure AI Search.
	/// </summary>
	/// <remarks>
	/// This class provides methods to store, retrieve, and query embedded documents in a collection.
	/// It utilizes Azure AI Search for vector-based queries and document management.
	/// </remarks>
	[ApiCategory("Services/IEmbeddingStorageService")]
	public class AzureAISearchEmbeddingStorageService : IEmbeddingStorageService
	{
		// metadata keys used internally
		private const string EMBEDDING_MODEL = "embeddingModel";

		/// <summary>
		/// Initializes a new instance of the <see cref="AzureAISearchEmbeddingStorageService"/> class.
		/// </summary>
		/// <param name="url">The base URL for the Azure AI Search service. Default is null.</param>
		public AzureAISearchEmbeddingStorageService(string url = null)
		{
			this.URL = url;
		}

		/// <summary>
		/// Gets or sets the base URL for the Azure AI Search service.
		/// </summary>
		[DefaultValue(null)]
		public virtual string URL
		{
			get;
			set;
		}

		/// <summary>
		/// Gets or sets the API version for the Azure AI Search service.
		/// </summary>
		[DefaultValue("2024-07-01")]
		public virtual string ApiVersion
		{
			get;
			set;
		} = "2024-07-01";

		/// <summary>
		/// Gets or sets the API key for authenticating with the Azure AI Search service.
		/// </summary>
		/// <remarks>
		/// Setting this property will reset the internal HTTP client.
		/// </remarks>
		[DefaultValue(null)]
		public virtual string ApiKey
		{
			get;
			set;
		}

		/// <summary>
		/// Gets the HTTP client service used for making requests.
		/// </summary>
		/// <remarks>
		/// The client is initialized with the API key if it is not already initialized.
		/// </remarks>
		internal virtual IHttpClientService Client
		{
			get
			{
				if (_client == null)
				{
					_client = ServiceHelper.GetService<IHttpClientService>();

					var apiKey = GetApiKey();
					if (!String.IsNullOrEmpty(apiKey))
						_client.SetDefaultHeader("api-key", apiKey);
				}

				return _client;
			}
		}
		private IHttpClientService _client;

		/// <summary>
		/// Retrieves the API key for the service.
		/// </summary>
		/// <returns>The API key as a string.</returns>
		protected virtual string GetApiKey()
		{
			return
				!String.IsNullOrEmpty(this.ApiKey)
					? this.ApiKey
					: ApiKeys.GetApiKey(GetType().Name, "EmbeddingStorageService");
		}

		/// <summary>
		/// Checks if a document exists in the specified collection.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="documentName">The name of the document.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains a boolean indicating whether the document exists.</returns>
		public virtual async Task<bool> ExistsAsync(
			string collectionName,
			string documentName)
		{
			return (await ReadDocumentAsync(collectionName, documentName, false)) != null;
		}

		/// <summary>
		/// Queries the collection for documents similar to the provided query vector.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="query">The query vector.</param>
		/// <param name="topN">The number of top results to return.</param>
		/// <param name="minSimilarity">The minimum similarity threshold.</param>
		/// <param name="filter">An optional filter predicate for the documents. Default is null.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains an array of <see cref="EmbeddedDocument"/> objects.</returns>
		public virtual async Task<EmbeddedDocument[]> QueryAsync(
			string collectionName,
			float[] query,
			int topN,
			float minSimilarity,
			Predicate<EmbeddedDocument> filter = null)
		{
			var list = new List<EmbeddedDocument>();

			collectionName = FixCollectionName(collectionName);

			// if we have a filter, list and filter all documents first
			var filteredNames =
					filter == null
					? null
					: (await RetrieveAsync(collectionName, false, filter)).Select(d => d.Name).ToArray();

			// nothing qualified?
			if (filteredNames?.Length == 0)
				return list.ToArray();

			var result = await PostAsync(
				"docs/search",
				new
				{
					top = topN,
					select = "documentName,chunk",
					filter = CreateDocumentNameFilter(collectionName, filteredNames),

					vectorQueries = new object[] {
						new {
							vector = query,
							k = topN,
							fields = "vector",
							kind = "vector",
							exhaustive = true
						}
					}
				});

			// now we have the topN chunks, need to assemble them back into documents

			var matches = (dynamic[])result.value;
			if (matches.Length > 0)
			{
				var chunks = new List<(string name, string chunk, float similarity)>(matches.Length);

				for (var i = 0; i < matches.Length; i++)
				{
					var similarity = ScoreToSimilarity((float)matches[i]["@search.score"]);
					if (similarity >= minSimilarity)
					{
						var chunk = (string)matches[i].chunk;
						var documentName = (string)matches[i].documentName;

						chunks.Add((documentName, chunk, similarity));
					}
				}

				foreach (var g in chunks.GroupBy(c => c.name))
				{
					var document =
						await ReadDocumentAsync(collectionName, g.Key, false);

					document.SetMatches(
						new Matches(
							g.Select(g => g.chunk),
							g.Select(g => g.similarity)));

					list.Add(document);
				}
			}

			return list.ToArray();
		}

		/// <summary>
		/// Queries a specific document in the collection for similarity to the provided query vector.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="documentName">The name of the document.</param>
		/// <param name="query">The query vector.</param>
		/// <param name="topN">The number of top results to return.</param>
		/// <param name="minSimilarity">The minimum similarity threshold.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="EmbeddedDocument"/> object.</returns>
		/// <exception cref="ArgumentNullException">Thrown when the document is null.</exception>
		public virtual async Task<EmbeddedDocument> QueryAsync(
			string collectionName,
			string documentName,
			float[] query,
			int topN,
			float minSimilarity)
		{
			if (documentName == null)
				throw new ArgumentNullException(nameof(documentName));

			collectionName = FixCollectionName(collectionName);

			var document = await ReadDocumentAsync(collectionName, documentName, false);
			if (document == null)
				return null;

			if (query == null)
			{
				var result = await PostAsync(
					"docs/search",
					new
					{
						top = topN,
						select = "documentName,chunk",
						filter = CreateDocumentNameFilter(collectionName, documentName)
					});

				var matches = (dynamic[])result.value;
				if (matches.Length > 0)
				{
					var chunks = matches.Select(m => (string)m.chunk);
					return
						document.SetMatches(
							new Matches(
								chunks,
								Enumerable.Repeat(1f, chunks.Count())));
				}
			}
			else
			{
				var result = await PostAsync(
					"docs/search",
					new
					{
						top = topN,
						select = "documentName,chunk",
						filter = CreateDocumentNameFilter(collectionName, documentName),

						vectorQueries = new object[] {
							new {
								vector = query,
								k = topN,
								fields = "vector",
								kind = "vector",
								exhaustive = true
							}
						}
					});


				var matches = (dynamic[])result.value;
				if (matches.Length > 0)
				{
					var chunks = new List<(string chunk, float similarity)>();

					for (var i = 0; i < matches.Length; i++)
					{
						var match = matches[i];
						var similarity = ScoreToSimilarity((float)match["@search.score"]);
						if (similarity < minSimilarity)
							break;

						chunks.Add(((string)match.chunk, similarity));
					}

					if (chunks.Count > 0)
					{
						document.SetMatches(
							new Matches(
								chunks.Select(c => c.chunk),
								chunks.Select(c => c.similarity)));
					}
				}
			}

			return document;
		}

		/// <summary>
		/// Removes documents from the collection based on a filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="filter">An optional filter predicate for the documents. Default is null.</param>
		/// <returns>A task that represents the asynchronous operation.</returns>
		public virtual async Task RemoveAsync(
			string collectionName,
			Predicate<EmbeddedDocument> filter = null)
		{

			collectionName = FixCollectionName(collectionName);

			// retrieve all the document ids
			var result = await PostAsync(
					"docs/search",
					new
					{
						top = 999999,
						select = "id",
						filter = CreateCollectionNameFilter(collectionName, null)
					});

			var matches = (dynamic[])result.value;
			var idsToDelete = matches.Select(m => (string)m.id);

			var payload = new DynamicObject();
			payload["@search.action"] = "delete";

			var deleteTasks = new List<Task>();
			foreach (var id in idsToDelete)
			{
				payload["id"] = id;
				deleteTasks.Add(PostAsync("docs/index", payload));

				if (deleteTasks.Count > 50)
				{
					Task.WaitAll(deleteTasks.ToArray());
					deleteTasks.Clear();
				}
			}

			Task.WaitAll(deleteTasks.ToArray());
		}

		/// <summary>
		/// Removes a specific document from the collection.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="documentName">The name of the document.</param>
		/// <returns>A task that represents the asynchronous operation.</returns>
		/// <exception cref="ArgumentNullException">Thrown when the document is null.</exception>
		public virtual async Task RemoveAsync(
			string collectionName,
			string documentName)
		{
			if (documentName == null)
				throw new ArgumentNullException(nameof(documentName));

			collectionName = FixCollectionName(collectionName);

			// retrieve all the document ids
			var result = await PostAsync(
					"docs/search",
					new
					{
						top = 999999,
						select = "id",
						filter = CreateDocumentNameFilter(collectionName, documentName)
					});

			var matches = (dynamic[])result.value;
			var idsToDelete = matches.Select(m => (string)m.id);

			var payload = new DynamicObject();
			payload["@search.action"] = "delete";

			var deleteTasks = new List<Task>();
			foreach (var id in idsToDelete)
			{
				payload["id"] = id;
				deleteTasks.Add(PostAsync("docs/index", payload));

				if (deleteTasks.Count > 50)
				{
					Task.WaitAll(deleteTasks.ToArray());
					deleteTasks.Clear();
				}
			}

			Task.WaitAll(deleteTasks.ToArray());
		}

		/// <summary>
		/// Retrieves a specific document from the collection.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="documentName">The name of the document.</param>
		/// <param name="includeEmbedding">Indicates whether to include the embedding in the retrieval.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains an <see cref="EmbeddedDocument"/> object.</returns>
		public virtual async Task<EmbeddedDocument> RetrieveAsync(
			string collectionName,
			string documentName,
			bool includeEmbedding)
		{
			return
				await ReadDocumentAsync(
					collectionName, documentName, includeEmbedding);
		}

		/// <summary>
		/// Retrieves documents from the collection based on a filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="includeEmbedding">Indicates whether to include the embedding in the retrieval.</param>
		/// <param name="filter">An optional filter predicate for the documents. Default is null.</param>
		/// <returns>A task that represents the asynchronous operation.
		/// The task result contains an array of <see cref="EmbeddedDocument"/> objects.</returns>
		public virtual async Task<EmbeddedDocument[]> RetrieveAsync(
			string collectionName,
			bool includeEmbedding,
			Predicate<EmbeddedDocument> filter = null)
		{
			collectionName = FixCollectionName(collectionName);

			var result = await PostAsync(
				"docs/search",
				new
				{
					top = 10000,
					select = "documentName,chunk,metadata",
					filter = CreateCollectionNameFilter(collectionName, true)
				});

			var list = new List<EmbeddedDocument>();

			var matches = (dynamic[])result.value;
			if (matches.Length > 0)
			{
				for (var i = 0; i < matches.Length; i++)
				{
					var metadata = new Metadata(JSON.Parse(matches[i].metadata));
					var documentName = (string)matches[i].documentName;
					metadata.Remove(EMBEDDING_MODEL);

					var document = new EmbeddedDocument(documentName, metadata);

					if (filter != null && !filter(document))
						continue;

					list.Add(document);
				}
			}

			// now read the chunks and vectors.
			if (includeEmbedding)
			{
				for (var i = 0; i < list.Count; i++)
				{
					var document = list[i];
					document = await ReadDocumentAsync(collectionName, document.Name, true);
					list[i] = document;
				}
			}

			return list.ToArray();
		}

		/// <summary>
		/// Stores a document in the specified collection.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="document">The document to store.</param>
		/// <returns>A task that represents the asynchronous operation.</returns>
		/// <exception cref="ArgumentNullException">Thrown when the document is null.</exception>
		public virtual async Task StoreAsync(
			string collectionName,
			EmbeddedDocument document)
		{
			if (document == null)
				throw new ArgumentNullException(nameof(document));

			collectionName = FixCollectionName(collectionName);

			await PostAsync(
				"docs/index",
				BuildDocumentData(document, collectionName));
		}

		#region Implementation

		//
		private string FixCollectionName(string collectionName)
		{
			if (String.IsNullOrEmpty(collectionName))
			{
				collectionName = "default";
			}
			else
			{
				collectionName = collectionName
					.Replace(" ", "-")
					.Replace(".", "_")
					.Replace("/", "_")
					.Replace("\\", "_").ToLowerInvariant();
			}

			return collectionName;
		}

		//
		private string BuildDocumentId(string collectionName, string documentName)
		{
			documentName = documentName
				.Replace(" ", "-")
				.Replace(".", "_")
				.Replace("/", "_")
				.Replace("\\", "_").ToLowerInvariant();

			return $"{collectionName}={documentName}";
		}

		//
		private static float ScoreToSimilarity(float score)
		{
			return -((1 - score) / score) + 1;
		}

		//
		private string CreateDocumentNameFilter(string collectionName,string documentName, bool includeEmbeddings)
		{
			var docId = BuildDocumentId(collectionName, documentName);
			return
				includeEmbeddings
					? $"documentName eq '{documentName}' and collectionName eq '{collectionName}'"
					: $"documentName eq '{documentName}' and id eq '{docId}-0'";
		}

		//
		private string CreateCollectionNameFilter(string collectionName, bool? master)
		{
			if (master != null)
				return $"collectionName eq '{collectionName}' and master eq {(master.Value ? "true" : "false")}";

			return $"collectionName eq '{collectionName}'";
		}

		//
		private string CreateDocumentNameFilter(string collectionName, params string[] documentNames)
		{
			if (documentNames == null || documentNames.Length == 0)
				return $"collectionName eq '{collectionName}'";

			if (documentNames.Length == 1)
				return $"collectionName eq '{collectionName}' and documentName eq '{documentNames[0]}'";

			var documentNameIn =
				String.Join(",", documentNames.Select(n => $"'{n}'"));

			return
				$"collectionName eq '{collectionName}' and search.in(documentName, [{documentNameIn}])";
		}

		//
		private async Task<EmbeddedDocument> ReadDocumentAsync(
			string collectionName, string documentName, bool includeEmbedding)
		{
			collectionName = FixCollectionName(collectionName);

			var select = includeEmbedding
				? "id,metadata,documentName,chunk,vector"
				: "id,metadata,documentName";

			var result = await PostAsync(
				"docs/search",
				new
				{
					select = select,
					filter = CreateDocumentNameFilter(collectionName, documentName, includeEmbedding)
				});

			if (result.value.Length > 0)
			{
				if (includeEmbedding)
				{
					var values = (dynamic[])result.value;

					var chunks = values.Select(v => (string)v.chunk).ToArray();
					var vectors = values.Select(v => (float[])v.vector).ToArray();
					var metadata = new Metadata(JSON.Parse(values[0].metadata));
					var embeddingModel = (string)metadata[EMBEDDING_MODEL];
					metadata.Remove(EMBEDDING_MODEL);

					return new EmbeddedDocument(
						documentName,
						metadata,
						new Embedding(
							chunks,
							vectors,
							embeddingModel));
				}
				else
				{
					var entry = result.value[0];
					var metadata = new Metadata(JSON.Parse(entry.metadata));
					return new EmbeddedDocument(documentName, metadata);
				}
			}

			return null;
		}

		//
		private object BuildDocumentData(EmbeddedDocument document, string collectionName)
		{
			var embedding = document.GetEmbedding();
			var documentName = document.Name;
			var docId = BuildDocumentId(collectionName, documentName);

			var value = new List<object>();
			for (var i = 0; i < embedding.Vectors.Length; i++)
			{
				var chunk = new DynamicObject();
				chunk["@search.action"] = "mergeOrUpload";
				chunk["id"] = $"{docId}-{i}";
				chunk["chunk"] = embedding.Chunks[i];
				chunk["vector"] = embedding.Vectors[i];
				chunk["documentName"] = documentName;
				chunk["collectionName"] = collectionName;

				if (i == 0)
				{
					var metadata = document.Metadata.Clone();
					metadata[EMBEDDING_MODEL] = embedding.Model;
					chunk["master"] = true;
					chunk["metadata"] = JSON.Stringify(metadata, false);
				}

				value.Add(chunk);
			}

			return new
			{
				value = value
			};
		}

		private async Task<dynamic> PostAsync(string endpoint, object data)
		{
			var client = this.Client;
			var url = $"{this.URL}/{endpoint}?api-version={this.ApiVersion}";

			var payload = JSON.Stringify(data, JSON.SerializerOptions.None);
			var content = new StringContent(payload, Encoding.UTF8, "application/json");

			var response = await client.PostAsync(url, content);

			if (!response.IsSuccessStatusCode)
			{
				var json = JSON.Parse(await response.Content.ReadAsStreamAsync());
				if (json != null)
					throw new Exception(json.error.message);
				else
					response.EnsureSuccessStatusCode();
			}

			var stream = await response.Content.ReadAsStreamAsync();
			return JSON.Parse(stream);
		}

		//
		private async Task<dynamic> GetAsync(string endpoint)
		{
			var client = this.Client;
			var url = $"{this.URL}/{endpoint}?api-version={this.ApiVersion}";

			var response = await client.GetAsync(url);

			if (!response.IsSuccessStatusCode)
			{
				var json = JSON.Parse(await response.Content.ReadAsStreamAsync());
				if (json != null)
					throw new Exception(json.error.message);
				else
					response.EnsureSuccessStatusCode();
			}

			var stream = await response.Content.ReadAsStreamAsync();
			return JSON.Parse(stream);
		}

		#endregion
	}
}