using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using FluentAssertions;
using SqrtSpace.SpaceTime.Linq;
using Xunit;

namespace SqrtSpace.SpaceTime.Tests.Linq;

public class SpaceTimeEnumerableTests
{
    private static IEnumerable<int> GenerateNumbers(int count)
    {
        for (int i = 0; i < count; i++)
        {
            yield return i;
        }
    }

    private static IEnumerable<TestItem> GenerateTestItems(int count)
    {
        var random = new Random(42); // Fixed seed for reproducibility
        for (int i = 0; i < count; i++)
        {
            yield return new TestItem
            {
                Id = i,
                Value = random.Next(1000),
                Category = $"Category{random.Next(10)}",
                Date = DateTime.Today.AddDays(-random.Next(365))
            };
        }
    }

    public class TestItem
    {
        public int Id { get; set; }
        public int Value { get; set; }
        public string Category { get; set; } = "";
        public DateTime Date { get; set; }
    }

    public class OrderByExternalTests
    {
        [Fact]
        public void OrderByExternal_SmallCollection_ReturnsSortedResults()
        {
            // Arrange
            var items = new[] { 5, 2, 8, 1, 9, 3, 7, 4, 6 };

            // Act
            var result = items.OrderByExternal(x => x).ToList();

            // Assert
            result.Should().BeEquivalentTo(new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9 });
        }

        [Fact]
        public void OrderByExternal_LargeCollection_ReturnsSortedResults()
        {
            // Arrange
            var items = GenerateNumbers(10_000).OrderBy(_ => Guid.NewGuid()).ToList();

            // Act
            var result = items.OrderByExternal(x => x).ToList();

            // Assert
            result.Should().BeInAscendingOrder();
            result.Should().HaveCount(10_000);
        }

        [Fact]
        public void OrderByExternal_WithCustomComparer_UsesComparer()
        {
            // Arrange
            var items = new[] { "apple", "Banana", "cherry", "Date" };
            var comparer = StringComparer.OrdinalIgnoreCase;

            // Act
            var result = items.OrderByExternal(x => x, comparer).ToList();

            // Assert
            result.Should().BeEquivalentTo(new[] { "apple", "Banana", "cherry", "Date" });
        }

        [Fact]
        public void OrderByExternal_WithCustomBufferSize_RespectsBufferSize()
        {
            // Arrange
            var items = GenerateNumbers(1000).ToList();

            // Act
            var result = items.OrderByExternal(x => x, bufferSize: 10).ToList();

            // Assert
            result.Should().BeInAscendingOrder();
            result.Should().HaveCount(1000);
        }

        [Fact]
        public void OrderByDescendingExternal_ReturnsDescendingOrder()
        {
            // Arrange
            var items = new[] { 5, 2, 8, 1, 9, 3, 7, 4, 6 };

            // Act
            var result = items.OrderByDescendingExternal(x => x).ToList();

            // Assert
            result.Should().BeEquivalentTo(new[] { 9, 8, 7, 6, 5, 4, 3, 2, 1 });
        }

        [Fact]
        public void OrderByExternal_WithComplexKey_SortsCorrectly()
        {
            // Arrange
            var items = GenerateTestItems(100).ToList();

            // Act
            var result = items.OrderByExternal(x => x.Date)
                             .ThenByExternal(x => x.Value)
                             .ToList();

            // Assert
            result.Should().BeInAscendingOrder(x => x.Date)
                  .And.ThenBeInAscendingOrder(x => x.Value);
        }

        [Fact]
        public void OrderByExternal_EmptyCollection_ReturnsEmpty()
        {
            // Arrange
            var items = Enumerable.Empty<int>();

            // Act
            var result = items.OrderByExternal(x => x).ToList();

            // Assert
            result.Should().BeEmpty();
        }

        [Fact]
        public void OrderByExternal_NullKeySelector_ThrowsException()
        {
            // Arrange
            var items = new[] { 1, 2, 3 };

            // Act & Assert
            var action = () => items.OrderByExternal<int, int>(null!).ToList();
            action.Should().Throw<ArgumentNullException>();
        }
    }

    public class GroupByExternalTests
    {
        [Fact]
        public void GroupByExternal_SimpleGrouping_ReturnsCorrectGroups()
        {
            // Arrange
            var items = new[] { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10 };

            // Act
            var result = items.GroupByExternal(x => x % 3).ToList();

            // Assert
            result.Should().HaveCount(3);
            result.SelectMany(g => g).Should().BeEquivalentTo(items);
            result.Single(g => g.Key == 0).Should().BeEquivalentTo(new[] { 3, 6, 9 });
            result.Single(g => g.Key == 1).Should().BeEquivalentTo(new[] { 1, 4, 7, 10 });
            result.Single(g => g.Key == 2).Should().BeEquivalentTo(new[] { 2, 5, 8 });
        }

        [Fact]
        public void GroupByExternal_WithElementSelector_TransformsElements()
        {
            // Arrange
            var items = GenerateTestItems(100).ToList();

            // Act
            var result = items.GroupByExternal(
                x => x.Category,
                x => x.Value
            ).ToList();

            // Assert
            result.Should().HaveCount(10); // 10 categories
            result.Sum(g => g.Count()).Should().Be(100);
            result.All(g => g.All(v => v.GetType() == typeof(int))).Should().BeTrue();
        }

        [Fact]
        public void GroupByExternal_WithResultSelector_AppliesTransformation()
        {
            // Arrange
            var items = GenerateTestItems(50).ToList();

            // Act
            var result = items.GroupByExternal(
                x => x.Category,
                x => x.Value,
                (key, values) => new
                {
                    Category = key,
                    Sum = values.Sum(),
                    Count = values.Count()
                }
            ).ToList();

            // Assert
            result.Should().HaveCount(10);
            result.Sum(x => x.Count).Should().Be(50);
            result.All(x => x.Sum > 0).Should().BeTrue();
        }

        [Fact]
        public void GroupByExternal_LargeDataset_HandlesCorrectly()
        {
            // Arrange
            var items = GenerateNumbers(10_000).Select(x => new { Id = x, Group = x % 100 });

            // Act
            var result = items.GroupByExternal(x => x.Group).ToList();

            // Assert
            result.Should().HaveCount(100);
            result.All(g => g.Count() == 100).Should().BeTrue();
        }

        [Fact]
        public void GroupByExternal_WithCustomComparer_UsesComparer()
        {
            // Arrange
            var items = new[] { "apple", "Apple", "banana", "Banana", "cherry" };
            var comparer = StringComparer.OrdinalIgnoreCase;

            // Act
            var result = items.GroupByExternal(x => x, comparer).ToList();

            // Assert
            result.Should().HaveCount(3);
            result.Single(g => comparer.Equals(g.Key, "apple")).Count().Should().Be(2);
            result.Single(g => comparer.Equals(g.Key, "banana")).Count().Should().Be(2);
        }

        [Fact]
        public void GroupByExternal_EmptyCollection_ReturnsEmpty()
        {
            // Arrange
            var items = Enumerable.Empty<int>();

            // Act
            var result = items.GroupByExternal(x => x).ToList();

            // Assert
            result.Should().BeEmpty();
        }
    }

    public class DistinctExternalTests
    {
        [Fact]
        public void DistinctExternal_RemovesDuplicates()
        {
            // Arrange
            var items = new[] { 1, 2, 3, 2, 4, 3, 5, 1, 6, 4, 7 };

            // Act
            var result = items.DistinctExternal().ToList();

            // Assert
            result.Should().BeEquivalentTo(new[] { 1, 2, 3, 4, 5, 6, 7 });
        }

        [Fact]
        public void DistinctExternal_WithComparer_UsesComparer()
        {
            // Arrange
            var items = new[] { "apple", "Apple", "banana", "Banana", "cherry" };
            var comparer = StringComparer.OrdinalIgnoreCase;

            // Act
            var result = items.DistinctExternal(comparer).ToList();

            // Assert
            result.Should().HaveCount(3);
        }

        [Fact]
        public void DistinctExternal_LargeDataset_HandlesCorrectly()
        {
            // Arrange
            var items = GenerateNumbers(10_000).Concat(GenerateNumbers(10_000));

            // Act
            var result = items.DistinctExternal().ToList();

            // Assert
            result.Should().HaveCount(10_000);
            result.Should().BeEquivalentTo(Enumerable.Range(0, 10_000));
        }

        [Fact]
        public void DistinctExternal_PreservesFirstOccurrence()
        {
            // Arrange
            var items = new[]
            {
                new TestItem { Id = 1, Value = 100 },
                new TestItem { Id = 2, Value = 200 },
                new TestItem { Id = 1, Value = 300 },
                new TestItem { Id = 3, Value = 400 }
            };

            // Act
            var result = items.DistinctExternal(new TestItemIdComparer()).ToList();

            // Assert
            result.Should().HaveCount(3);
            result.Single(x => x.Id == 1).Value.Should().Be(100); // First occurrence
        }

        private class TestItemIdComparer : IEqualityComparer<TestItem>
        {
            public bool Equals(TestItem? x, TestItem? y)
            {
                if (ReferenceEquals(x, y)) return true;
                if (x is null || y is null) return false;
                return x.Id == y.Id;
            }

            public int GetHashCode(TestItem obj) => obj.Id.GetHashCode();
        }
    }

    public class BatchBySqrtNTests
    {
        [Fact]
        public void BatchBySqrtN_SmallCollection_ReturnsSingleBatch()
        {
            // Arrange
            var items = GenerateNumbers(100).ToList();

            // Act
            var batches = items.BatchBySqrtN().ToList();

            // Assert
            batches.Should().HaveCount(10); // sqrt(100) = 10, so 10 batches of 10
            batches.All(b => b.Count() == 10).Should().BeTrue();
            batches.SelectMany(b => b).Should().BeEquivalentTo(items);
        }

        [Fact]
        public void BatchBySqrtN_LargeCollection_ReturnsOptimalBatches()
        {
            // Arrange
            var items = GenerateNumbers(10_000).ToList();

            // Act
            var batches = items.BatchBySqrtN().ToList();

            // Assert
            var expectedBatchSize = (int)Math.Sqrt(10_000); // 100
            batches.Should().HaveCount(100);
            batches.Take(99).All(b => b.Count() == expectedBatchSize).Should().BeTrue();
            batches.SelectMany(b => b).Should().BeEquivalentTo(items);
        }

        [Fact]
        public void BatchBySqrtN_NonSquareNumber_HandlesRemainder()
        {
            // Arrange
            var items = GenerateNumbers(150).ToList();

            // Act
            var batches = items.BatchBySqrtN().ToList();

            // Assert
            var batchSize = (int)Math.Sqrt(150); // 12
            batches.Should().HaveCount(13); // 12 full batches + 1 partial
            batches.Take(12).All(b => b.Count() == batchSize).Should().BeTrue();
            batches.Last().Count().Should().Be(150 - (12 * batchSize));
        }

        [Fact]
        public void BatchBySqrtN_EmptyCollection_ReturnsNoBatches()
        {
            // Arrange
            var items = Enumerable.Empty<int>();

            // Act
            var batches = items.BatchBySqrtN().ToList();

            // Assert
            batches.Should().BeEmpty();
        }

        [Fact]
        public async Task BatchBySqrtNAsync_ProcessesAsynchronously()
        {
            // Arrange
            var items = GenerateNumbers(1000);

            // Act
            var batchCount = 0;
            var totalItems = 0;
            await foreach (var batch in items.BatchBySqrtNAsync())
            {
                batchCount++;
                totalItems += batch.Count();
            }

            // Assert
            batchCount.Should().BeGreaterThan(1);
            totalItems.Should().Be(1000);
        }
    }

    public class ToCheckpointedListAsyncTests
    {
        [Fact]
        public async Task ToCheckpointedListAsync_SmallCollection_ReturnsAllItems()
        {
            // Arrange
            var items = GenerateNumbers(100);

            // Act
            var result = await items.ToCheckpointedListAsync();

            // Assert
            result.Should().HaveCount(100);
            result.Should().BeEquivalentTo(Enumerable.Range(0, 100));
        }

        [Fact]
        public async Task ToCheckpointedListAsync_WithCheckpointAction_CallsCheckpoint()
        {
            // Arrange
            var items = GenerateNumbers(10_000);
            var checkpointCount = 0;
            var lastCheckpointedCount = 0;

            // Act
            var result = await items.ToCheckpointedListAsync(
                checkpointAction: async (list) =>
                {
                    checkpointCount++;
                    lastCheckpointedCount = list.Count;
                    await Task.Delay(1); // Simulate async work
                });

            // Assert
            result.Should().HaveCount(10_000);
            checkpointCount.Should().BeGreaterThan(0);
            checkpointCount.Should().BeLessThanOrEqualTo(100); // sqrt(10000) = 100
        }

        [Fact]
        public async Task ToCheckpointedListAsync_WithCancellation_ThrowsWhenCancelled()
        {
            // Arrange
            var items = GenerateNumbers(100_000);
            var cts = new CancellationTokenSource();
            
            // Act
            var task = items.ToCheckpointedListAsync(
                checkpointAction: async (list) =>
                {
                    if (list.Count > 5000)
                    {
                        cts.Cancel();
                    }
                    await Task.Delay(1);
                },
                cancellationToken: cts.Token);

            // Assert
            await task.Invoking(t => t).Should().ThrowAsync<OperationCanceledException>();
        }
    }

    public class StreamAsJsonAsyncTests
    {
        [Fact]
        public async Task StreamAsJsonAsync_SerializesCorrectly()
        {
            // Arrange
            var items = GenerateTestItems(10);
            var stream = new System.IO.MemoryStream();

            // Act
            await items.StreamAsJsonAsync(stream);
            stream.Position = 0;
            var json = new System.IO.StreamReader(stream).ReadToEnd();

            // Assert
            json.Should().StartWith("[");
            json.Should().EndWith("]");
            json.Should().Contain("\"Id\"");
            json.Should().Contain("\"Value\"");
            json.Should().Contain("\"Category\"");
        }

        [Fact]
        public async Task StreamAsJsonAsync_EmptyCollection_WritesEmptyArray()
        {
            // Arrange
            var items = Enumerable.Empty<TestItem>();
            var stream = new System.IO.MemoryStream();

            // Act
            await items.StreamAsJsonAsync(stream);
            stream.Position = 0;
            var json = new System.IO.StreamReader(stream).ReadToEnd();

            // Assert
            json.Trim().Should().Be("[]");
        }
    }
}