using System;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using Microsoft.Azure.Cosmos;
using Microsoft.Azure.Cosmos.Fluent;
using MoreLinq.Extensions;
using Newtonsoft.Json;

namespace CosmosLongPartitionKey
{
    class Program
    {
        public const string DatabaseName = "testcosmosclient";
        public const int Throughput = 1200;
        public const string LocalConnectionString = "AccountEndpoint=https://localhost:8081/;AccountKey=C2y6yDjf5/R+ob0N8A7Cgv30VRDJIWEHLM+4QDU5DE2nQ9nDuVTqobD4b8mGGyPMbIZnqyMsEcaGQy67XIw/Jw==;";
        public const string DefaultKey = "objectKey";
        public const string TestCollection = "testcollection";
        private static readonly Random Random = new Random();
        
        static void Main(string[] args)
        {
            var cosmosClientBuilder = new CosmosClientBuilder(LocalConnectionString)
                .WithConnectionModeDirect()
                .WithThrottlingRetryOptions(TimeSpan.FromMinutes(2), 20);

            var cosmosClient = cosmosClientBuilder.Build();

            CreateDb(cosmosClient);

            var databaseResponse = cosmosClient.CreateDatabaseIfNotExistsAsync(DatabaseName, Throughput);
            databaseResponse.Wait();
            var database = databaseResponse.Result.Database;

            database.DefineContainer(TestCollection, $"/{DefaultKey}")
                .WithUniqueKey().Path($"/{DefaultKey}").Attach().CreateIfNotExistsAsync().Wait();

            var container = cosmosClient.GetContainer(DatabaseName, TestCollection);

            const int numRecords = 10000;
            var queryKeys = new List<string>();
            var contacts = new List<Contact>();

            for (var c = 0; c < 6; c++)
            {
                contacts.Add(new Contact
                {
                    Identifier = $"u{RandomString(32)}",
                    Type = "phone"
                });
            }

            contacts.Add(new Contact
            {
                Identifier = $"0{Random.Next(400000000, 600000000)}",
                Type = "phone"
            });

            for (var i = 0; i < numRecords; i++)
            {
                var contact1 = contacts[Random.Next(0, contacts.Count)];
                var contact2 = contacts[Random.Next(0, contacts.Count)];
                var when = DateTime.UtcNow.AddMinutes(Random.Next(1, 100000));

                var testData = new TestCollectionObject
                {
                    Id = Guid.NewGuid(),
                    ObjectKey = RandomObjectKey(contact1, contact2, when),
                    Text = RandomString(Random.Next(10, 1000)),
                    Text2 = RandomString(Random.Next(10, 1000)),
                };

                WriteDocument(container, testData);

                const int keysToQuery = numRecords / 500;
                if (i % keysToQuery == 0) queryKeys.Add(testData.ObjectKey);
            }

            try
            {
                var results = container
                    .GetItemLinqQueryable<TestCollectionObject>(true, requestOptions: RunInParallelOptions())
                    .Where(r => queryKeys.Contains(r.ObjectKey))
                    .ToList(); // ERROR OCCURS WHEN QUERY IS EXECUTED

                Console.WriteLine($"[\"{string.Join("\", \n\"", results.Select(r => r.ObjectKey))}\"]");
            }
            catch (Exception e)
            {
                Console.WriteLine(e.Message);
                Console.WriteLine(e.StackTrace);
            }

            // Try again in smaller batches
            foreach (var keyBatch in queryKeys.Batch(5))
            {
                try
                {
                    var keyBatchList = keyBatch.ToList();

                    var results = container
                        .GetItemLinqQueryable<TestCollectionObject>(true, requestOptions: RunInParallelOptions())
                        .Where(r => keyBatchList.Contains(r.ObjectKey))
                        .ToList(); // ERROR STILL OCCURS WHEN QUERY IS EXECUTED
                }
                catch (Exception e)
                {
                    Console.WriteLine(e.Message);
                    Console.WriteLine(e.StackTrace);
                    Console.WriteLine($"[\"{string.Join("\", \n\"", keyBatch)}\"]");
                }
            }

            Console.WriteLine("Hit any key to finish");
            Console.ReadKey();
        }

        private static void CreateDb(CosmosClient cosmosClient)
        {
            try
            {
                cosmosClient.GetDatabase(DatabaseName).DeleteAsync().Wait();
            }
            catch (Exception e)
            {
                if (e.InnerException is CosmosException)
                {
                    // fall through
                }
                else throw;
            }
        }

        private static void WriteDocument(Container container, TestCollectionObject testData)
        {
            try
            {
                container.CreateItemAsync(testData, requestOptions: null).Wait();
            }
            catch (CosmosException e)
            {
                if (e.StatusCode != HttpStatusCode.Conflict) throw;
            }
        }

        private static string RandomObjectKey(Contact contact1, Contact contact2, DateTime when)
        {
            return $"message~{contact1}~{contact2}~Chat~{when:yyyyMMddHHmmssK}";
        }

        private static string RandomString(int length)
        {
            const string chars = "abcdef0123456789";
            return new string(Enumerable.Repeat(chars, length)
                .Select(s => s[Random.Next(s.Length)]).ToArray());
        }

        private static QueryRequestOptions RunInParallelOptions()
        {
            return new QueryRequestOptions
            {
                MaxItemCount = -1,
                MaxBufferedItemCount = -1,
                MaxConcurrency = -1
            };
        }
    }
    public class Contact
    {
        [JsonProperty("identifier")]
        public string Identifier { get; set; }

        [JsonProperty("type")]
        public string Type { get; set; }

        public override string ToString()
        {
            return $"{Type}~{Identifier}";
        }
    }

    public class TestCollectionObject
    {
        [JsonProperty("id")]
        public Guid Id { get; set; }
        [JsonProperty(Program.DefaultKey)]
        public string ObjectKey { get; set; }
        [JsonProperty("text")]
        public string Text { get; set; }
        [JsonProperty("text2")]
        public string Text2 { get; set; }
    }
}

  at Microsoft.Azure.Cosmos.IRoutingMapProviderExtensions.<TryGetOverlappingRangesAsync>d__3.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.CosmosQueryClientCore.<GetTargetPartitionKeyRangesAsync>d__13.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.Query.CosmosQueryExecutionContextFactory.<GetTargetPartitionKeyRangesAsync>d__14.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.Query.CosmosQueryExecutionContextFactory.<CreateItemQueryExecutionContextAsync>d__12.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.Query.CosmosQueryExecutionContextFactory.<ExecuteNextAsync>d__11.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.Query.QueryIterator.<ReadNextAsync>d__4.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.FeedIteratorCore`1.<ReadNextAsync>d__5.MoveNext()
   at System.Runtime.ExceptionServices.ExceptionDispatchInfo.Throw()
   at System.Runtime.CompilerServices.TaskAwaiter.HandleNonSuccessAndDebuggerNotification(Task task)
   at System.Runtime.CompilerServices.TaskAwaiter`1.GetResult()
   at Microsoft.Azure.Cosmos.Linq.CosmosLinqQuery`1.<GetEnumerator>d__20.MoveNext()
   at System.Collections.Generic.List`1.AddEnumerable(IEnumerable`1 enumerable)
   at System.Linq.Enumerable.ToList[TSource](IEnumerable`1 source)
   at CosmosLongPartitionKey.Program.Main(String[] args) in C:\\Users\\<XXX>\\CosmosLongPartitionKey\\Program.cs:line 130