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

		private int _dimension;

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

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

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

		/// <summary>
		/// Gets or sets the maximum batch size used to store document vectors into the database.
		/// </summary>
		[DefaultValue(100)]
		protected virtual int MaxBatchSize
		{
			get;
			set;
		} = 100;

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

		/// <summary>
		/// Queries the specified collection for documents similar to the given 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 similar documents to retrieve.</param>
		/// <param name="minSimilarity">The minimum similarity threshold.</param>
		/// <param name="filter">An optional filter to apply to the documents. Default is <c>null</c>.</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 = FixNamespaceName(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();

			// perform the NN query, optionally limited to the filtered doc names
			var result = await PostAsync(
				"query",
				new
				{
					topK = topN,
					vector = query,
					filter = CreateDocumentNameFilter(filteredNames),
					@namespace = collectionName,
					includeValues = false,
					includeMetadata = true
				});

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

			var matches = (dynamic[])result.matches;
			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 = (float)matches[i].score;
					if (similarity >= minSimilarity)
					{
						var metadata = matches[i].metadata;
						var chunk = (string)metadata[DOCUMENT_CHUNK];
						var documentName = (string)metadata[DOCUMENT_NAME];

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

			collectionName = FixNamespaceName(collectionName);

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

			if (query == null)
			{
				var result = await PostAsync(
					"query",
					new
					{
						topK = topN,
						vector = CreateEmptyVector(collectionName),
						filter = CreateDocumentNameFilter(documentName),
						includeMetadata = true
					});

				var matches = (dynamic[])result.matches;
				if (matches.Length > 0)
				{
					var ids = matches.Select(m => (string)m.id);
					var chunks = matches.Select(m => (string)m.metadata[DOCUMENT_CHUNK]).ToArray();

					// must sort by ids, there is no way to get the vectors in order otherwise.
					Array.Sort(ids.ToArray(), chunks);

					return document.SetMatches(
						new Matches(
							chunks,
							Enumerable.Repeat(1f, chunks.Length)));
				}
			}
			else
			{
				var result = await PostAsync(
					"query",
					new
					{
						topK = topN,
						vector = query,
						@namespace = collectionName,
						filter = CreateDocumentNameFilter(documentName),
						includeMetadata = true
					});

				var matches = (dynamic[])result.matches;
				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 = (float)match.score;
						if (similarity < minSimilarity)
							break;

						chunks.Add(((string)match.metadata[DOCUMENT_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 specified collection that match the given filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="filter">An optional filter to apply to the documents. Default is <c>null</c>.</param>
		/// <returns>A task that represents the asynchronous operation.</returns>
		public virtual async Task RemoveAsync(
			string collectionName,
			Predicate<EmbeddedDocument> filter = null)
		{
			collectionName = FixNamespaceName(collectionName);

			if (filter == null)
			{
				await PostAsync(
					"vectors/delete",
					new
					{
						deleteAll = true,
						@namespace = collectionName,
					});

				return;
			}

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

			await PostAsync(
				"vectors/delete",
				new
				{
					deleteAll = true,
					@namespace = collectionName,
					filter = CreateDocumentNameFilter(documents.Select(d => d.Name).ToArray())
				});
		}

		/// <summary>
		/// Removes a document from 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.</returns>
		public virtual async Task RemoveAsync(
			string collectionName,
			string documentName)
		{
			collectionName = FixNamespaceName(collectionName);

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

			await PostAsync(
				"vectors/delete",
				new
				{
					deleteAll = true,
					@namespace = collectionName,
					filter = CreateDocumentNameFilter(documentName)
				});
		}

		/// <summary>
		/// Retrieves a document from the specified collection.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="documentName">The name of the document.</param>
		/// <param name="includeEmbedding">A value indicating whether to include the embedding in the result.</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 specified collection that match the given filter.
		/// </summary>
		/// <param name="collectionName">The name of the collection.</param>
		/// <param name="includeEmbedding">A value indicating whether to include the embedding in the result.</param>
		/// <param name="filter">An optional filter to apply to the documents. Default is <c>null</c>.</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 = FixNamespaceName(collectionName);

			var result = await PostAsync(
				"query",
				new
				{
					topK = 10000,
					vector = CreateEmptyVector(collectionName),
					@namespace = collectionName,
					includeMetadata = true,
					includeValues = false,
					filter = CreateMasterFilter(true)
				});

			var list = new List<EmbeddedDocument>();

			var matches = (dynamic[])result.matches;
			if (matches.Length > 0)
			{
				for (var i = 0; i < matches.Length; i++)
				{
					var metadata = new Metadata(matches[i].metadata);
					var documentName = (string)metadata[DOCUMENT_NAME];
					metadata.Remove(MASTER, DOCUMENT_NAME, EMBEDDING_MODEL, DOCUMENT_CHUNK);

					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 <c>null</c>.</exception>
		public virtual async Task StoreAsync(
			string collectionName,
			EmbeddedDocument document)
		{
			if (document == null)
				throw new ArgumentNullException(nameof(document));

			collectionName = FixNamespaceName(collectionName);

			var batches = BuildDocumentData(document, collectionName, this.MaxBatchSize);
			for (int i = 0; i < batches.Count; i++)
			{
				await PostAsync("vectors/upsert", batches[i]);
			}
		}

		#region Implementation

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

			return collectionName;
		}

		//
		private object CreateDocumentNameFilter(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 CreateMasterFilter(bool value)
		{
			var where = new DynamicObject();
			where[MASTER] = value;
			return where;
		}

		//
		private float[] CreateEmptyVector(string collectionName)
		{
			var response = PostAsync("describe_index_stats", null).Result;
			var dimension = (int)response.dimension;
			return Enumerable.Repeat(0.0f, dimension).ToArray();
		}

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

			var result =
				includeEmbedding

				? await PostAsync(
					"query",
					new
					{
						topK = 10000,
						vector = CreateEmptyVector(collectionName),
						@namespace = collectionName,
						includeMetadata = true,
						includeValues = true,
						filter = CreateDocumentNameFilter(documentName)
					})

				: await PostAsync(
					"query",
					new
					{
						topK = 1,
						@namespace = collectionName,
						id = $"{documentName}:0",
						includeMetadata = true,
						includeValues = false,
						filter = CreateDocumentNameFilter(documentName)
					});

			var matches = (dynamic[])result.matches;
			if (matches.Length > 0)
			{
				if (includeEmbedding)
				{
					var ids = matches.Select(m => (string)m.id);
					var vectors = matches.Select(m => (float[])m.values).ToArray();
					var chunks = matches.Select(m => (string)m.metadata[DOCUMENT_CHUNK]).ToArray();

					// must sort by ids, there is no way to get the vectors in order otherwise.
					Array.Sort(ids.ToArray(), vectors);
					Array.Sort(ids.ToArray(), chunks);

					var metadata = new Metadata(matches[0].metadata);
					var embedding = new Embedding(chunks, vectors, (string)metadata[EMBEDDING_MODEL]);
					metadata.Remove(MASTER, DOCUMENT_NAME, EMBEDDING_MODEL, DOCUMENT_CHUNK);

					return new EmbeddedDocument(documentName, metadata, embedding);
				}
				else
				{
					var metadata = new Metadata(matches[0].metadata);
					metadata.Remove(MASTER, DOCUMENT_NAME, EMBEDDING_MODEL, DOCUMENT_CHUNK);
					return new EmbeddedDocument(documentName, metadata);
				}
			}

			return null;
		}

		//
		private IList<object> BuildDocumentData(EmbeddedDocument document, string collectionName, int maxBatchSize)
		{
			var embedding = document.GetEmbedding();

			var vectors = new List<object>();
			var batches = new List<object>();

			var metadata = document.Metadata;
			var documentName = document.Name;
			var count = embedding.Chunks.Length;

			for (var i = 0; i < count; i++)
			{
				if (i == 0)
				{
					metadata = metadata.Clone();
					metadata[MASTER] = true;
					metadata[DOCUMENT_NAME] = documentName;
					metadata[EMBEDDING_MODEL] = embedding.Model;
				}
				else
				{
					metadata = new Metadata();
					metadata[DOCUMENT_NAME] = documentName;
				}

				metadata[DOCUMENT_CHUNK] = embedding.Chunks[i];

				vectors.Add(new
				{
					id = $"{documentName}:{i}",
					values = embedding.Vectors[i],
					metadata = metadata
				});

				if (vectors.Count >= maxBatchSize)
				{
					batches.Add(new
					{
						vectors = vectors,
						@namespace = collectionName
					});

					vectors = new List<object>();
				}
			}

			if (vectors.Count > 0)
			{
				batches.Add(new
				{
					vectors = vectors,
					@namespace = collectionName
				});
			}

			return batches;
		}

		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 reader = new StreamReader(await response.Content.ReadAsStreamAsync());
				throw new Exception(reader.ReadToEnd());
			}

			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.IsSuccessStatusCode)
			{
				var reader = new StreamReader(await response.Content.ReadAsStreamAsync());
				throw new Exception(reader.ReadToEnd());
			}

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

		#endregion
	}
}