///////////////////////////////////////////////////////////////////////////////
//
// (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 retrieving embeddings using Chroma.
	/// </summary>
	/// <remarks>
	/// This class provides methods to store, retrieve, query, and remove embedded documents
	/// in a Chroma vector database. It uses an HTTP client service to communicate with the
	/// Chroma API.
	/// </remarks>
	[ApiCategory("Services/IEmbeddingStorageService")]
	public class ChromaEmbeddingStorageService : IEmbeddingStorageService
	{
		// metadata keys used internally
		private const string MASTER = "master";
		private const string DOCUMENT_NAME = "documentName";
		private const string EMBEDDING_MODEL = "embeddingModel";

		/// <summary>
		/// Initializes a new instance of the <see cref="ChromaEmbeddingStorageService"/> class.
		/// </summary>
		/// <param name="url">The URL of the Chroma service. Default is <c>null</c>.</param>
		public ChromaEmbeddingStorageService(string url = null)
		{
			this.URL = url;
		}

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

		/// <summary>
		/// Gets or sets the Authentication method for the Chroma service.
		/// </summary>
		[DefaultValue(null)]
		public virtual string Authentication
		{
			get;
			set;
		}

		/// <summary>
		/// Gets or sets the API key for the Chroma service.
		/// </summary>
		[DefaultValue(null)]
		public virtual string ApiKey
		{
			get;
			set;
		}

		/// <summary>
		/// Gets the HTTP client service used to communicate with the Chroma API.
		/// </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 Chroma service.
		/// </summary>
		/// <returns>The API key for the Chroma service.</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, null, documentName, false)) != null;
		}

		/// <summary>
		/// Queries the specified 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 to compare against.</param>
		/// <param name="topN">The number of top similar documents to return.</param>
		/// <param name="minSimilarity">The minimum similarity threshold for documents to be considered.</param>
		/// <param name="filter">An optional filter predicate to apply to 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>();

			var collectionId = await EnsureCollectionAsync(collectionName, false);

			// 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();

			// perform the NN query, optionally limited to the filtered doc names
			var result = await PostAsync(
				$"collections/{collectionId}/query",
				new
				{
					n_results = topN,
					query_embeddings = new[] {
						query
					},
					where = CreateDocumentNameWhere(filteredNames),
					include = new[] {
						"metadatas",
						"documents",
						"distances"
					}
				});

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

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

				for (int i = 0, l = metadatas[0].Length; i < l; i++)
				{
					var distance = distances[0][i];
					var similarity = (1.0f - distance);
					if (similarity >= minSimilarity)
					{
						var metadata = new Metadata(metadatas[0][i]);
						var documentName = (string)metadata[DOCUMENT_NAME];
						chunks.Add((documentName, documents[0][i], similarity));
					}
				}

				foreach (var g in chunks.GroupBy(c => c.name))
				{
					var document = await ReadDocumentAsync(collectionName, collectionId, 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 the specified collection for a document similar to the given 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 similar documents to retrieve.</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));

			var collectionId = await EnsureCollectionAsync(collectionName, false);

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

			if (query == null)
			{
				var result = await PostAsync(
					$"collections/{collectionId}/get",
					new
					{
						n_results = topN,
						where = CreateDocumentNameWhere(documentName),
						include = new[] {
							"documents"
						}
					});

				var chunks = (string[])result.documents;
				document.SetMatches(
					new Matches(
						chunks,
						Enumerable.Repeat(1f, chunks.Length)));
			}
			else
			{
				var result = await PostAsync(
					$"collections/{collectionId}/query",
					new
					{
						n_results = topN,
						query_embeddings = new[] {
							query
						},
						where = CreateDocumentNameWhere(documentName),
						include = new[] {
							"documents",
							"distances"
						}
					});

				var chunks = (string[][])result.documents;
				var distances = (float[][])result.distances;

				if (distances.Length > 0)
				{
					var take = distances.Length;
					for (var i = 0; i < distances.Length; i++)
					{
						var similarity = (1.0f - (float)distances[i][0]);
						if (similarity < minSimilarity)
						{
							take = i;
							break;
						}
					}

					if (take > 0)
					{
						document.SetMatches(new Matches(
							chunks.Take(take).Select(c => c[0]),
							distances.Take(take).Select(d => (1.0f - d[0]))));
					}
				}
			}

			return document;
		}

		/// <summary>
		/// Removes documents from the specified collection that match the given filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="filter">An optional filter predicate to apply to 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)
		{
			var collectionId = await EnsureCollectionAsync(collectionName, false);

			if (filter == null)
			{
				await PostAsync(
					$"collections/{collectionId}/delete", null);

				return;
			}

			var documents = await RetrieveAsync(collectionName, false, filter);

			await PostAsync(
				$"collections/{collectionId}/delete",
				new
				{
					where = CreateDocumentNameWhere(documents.Select(d => d.Name).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>
		public async Task RemoveAsync(
			string collectionName,
			string documentName)
		{
			var collectionId = await EnsureCollectionAsync(collectionName, false);

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

			await PostAsync(
				$"collections/{collectionId}/delete",
				new
				{
					where = CreateDocumentNameWhere(documentName)
				});
		}

		/// <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 retrieved document.</param>
		/// <returns>A task that represents the asynchronous operation. The task result contains the <see cref="EmbeddedDocument"/>.</returns>
		public virtual async Task<EmbeddedDocument> RetrieveAsync(
			string collectionName,
			string documentName,
			bool includeEmbedding)
		{
			return
				await ReadDocumentAsync(
					collectionName, null, documentName, includeEmbedding);
		}

		/// <summary>
		/// Retrieves all documents from the specified collection that match the given filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="includeEmbedding">Indicates whether to include the embedding in the retrieved documents.</param>
		/// <param name="filter">An optional filter predicate to apply to 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)
		{
			var collectionId = await EnsureCollectionAsync(collectionName, false);

			var result = await PostAsync(
				$"collections/{collectionId}/get",
				new
				{
					include = new[]
					{
						"metadatas"
					},
					where =  CreateMasterWhere(true)
				});

			var list = new List<EmbeddedDocument>();
			for (int i = 0, l = result.metadatas.Length; i < l; i++)
			{
				var metadata = new Metadata(result.metadatas[i]);
				var documentName = (string)metadata[DOCUMENT_NAME];
				metadata.Remove(MASTER, DOCUMENT_NAME, 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, collectionId, 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));

			var collectionId = await EnsureCollectionAsync(collectionName, true);

			await PostAsync(
				$"collections/{collectionId}/add",
				BuildDocumentData(document));
		}

		#region Implementation

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

			return collectionName;
		}

		//
		private object CreateDocumentNameWhere(params string[] names)
		{
			if (names == null || names.Length == 0)
				return null;

			var where = new DynamicObject();
			if (names.Length == 1)
			{
				where[DOCUMENT_NAME] = names[0];
			}
			else
			{

				var filter = new DynamicObject();
				filter["$in"] = names;
				where[DOCUMENT_NAME] = filter;
			}

			return where;
		}

		//
		private object CreateMasterWhere(bool value)
		{
			var where = new DynamicObject();
			where[MASTER] = value;
			return where;
		}

		//
		private async Task<EmbeddedDocument> ReadDocumentAsync(
			string collectionName, string collectionId, string documentName, bool includeEmbedding)
		{
			collectionId = collectionId ?? await EnsureCollectionAsync(collectionName, false);

			var result = await PostAsync(
				$"collections/{collectionId}/get",
				new
				{
					include = includeEmbedding
						? new[] { "metadatas", "embeddings", "documents" }
						: new[] { "metadatas" },

					ids = includeEmbedding
						? null
						: new[] { $"{documentName}:0" },

					where = CreateDocumentNameWhere(documentName)
				});

			if (result.ids.Length > 0)
			{
				if (includeEmbedding)
				{
					var chunks = (string[])result.documents;
					var vectors = (float[][])result.embeddings;
					var metadata = new Metadata(result.metadatas[0]);

					var embeddingModel = (string)metadata[EMBEDDING_MODEL];
					metadata.Remove(MASTER, DOCUMENT_NAME, EMBEDDING_MODEL);

					return new EmbeddedDocument(
						documentName,
						metadata,
						new Embedding(
							chunks,
							vectors,
							embeddingModel));
				}
				else
				{
					var metadata = new Metadata(result.metadatas[0]);
					metadata.Remove(MASTER, DOCUMENT_NAME, EMBEDDING_MODEL);
					return new EmbeddedDocument(documentName, metadata);
				}
			}

			return null;
		}

		//
		private object BuildDocumentData(EmbeddedDocument document)
		{
			var embedding = document.GetEmbedding();
			var documentName = document.Name;

			// repeat all the chunk metadatas, they contain only the document name
			var count = embedding.Chunks.Length;
			var chunk_metadata = new Metadata();
			chunk_metadata[DOCUMENT_NAME] = documentName;
			var metadatas = Enumerable.Repeat((object)chunk_metadata, count).ToArray();

			// only the master metadata is the complete metadata object
			var master_metadata = document.Metadata.Clone();
			master_metadata[MASTER] = true;
			master_metadata[DOCUMENT_NAME] = documentName;
			master_metadata[EMBEDDING_MODEL] = embedding.Model;
			metadatas[0] = master_metadata;

			var ids = Enumerable.Range(0, count).Select(i => $"{documentName}:{i}");

			return new
			{
				ids = ids,
				metadatas = metadatas,
				documents = embedding.Chunks,
				embeddings = embedding.Vectors
			};
		}

		//
		private async Task<string> EnsureCollectionAsync(string collectionName, bool create)
		{
			collectionName = FixCollectionName(collectionName);

			var data = await GetAsync($"collections/{collectionName}");
			var collectionId = (string)data?.id;

			if (collectionId == null && create)
			{
				data = await PostAsync(
					"collections",
					new
					{
						name = collectionName,
						metadata = JSON.Parse("{\"hnsw:space\": \"cosine\"}")
					}
				);
				collectionId = (string)data.id;
			}

			if (String.IsNullOrEmpty(collectionId))
			{
				if (create)
					throw new Exception($"Failed creating collection {collectionName}");

				return null;
			}

			return collectionId;
		}

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

			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.detail[0].msg);
				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}";

			var response = await client.GetAsync(url);

			if (response.StatusCode == System.Net.HttpStatusCode.NotFound)
				return null;

			if (!response.IsSuccessStatusCode)
			{
				var json = JSON.Parse(await response.Content.ReadAsStreamAsync());
				if (json != null)
					throw new Exception(json.detail[0].msg);
				else
					response.EnsureSuccessStatusCode();
			}

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

		#endregion
	}
}