using System.Diagnostics;
using System.Net.Sockets;
using System.Text;

#pragma warning disable IDE0130 // Namespace does not match folder structure
namespace Firefly
#pragma warning restore IDE0130 // Namespace does not match folder structure
{
    /// <summary>
    /// Test client for verifying pipelining and batching functionality
    /// </summary>
    public class PipelineTest
    {
        private static string Host = "127.0.0.1";
        private static int Port = 6379;
        private static string Password = "";
        private static int CommandCount = 10000;
        private static int BatchSize = 100;
        private static int BufferSize = 1024 * 1024; // 1MB buffer
        private static int ClientCount = 1;

        /// <summary>
        /// Runs the pipeline test
        /// </summary>
        public static async Task RunTest(string[] args)
        {
            // Parse command line arguments
            ParseArguments(args);

            Console.WriteLine("\nStarting Pipeline Test...");
            Console.WriteLine($"Host: {Host}");
            Console.WriteLine($"Port: {Port}");
            Console.WriteLine($"Clients: {ClientCount}");
            if (!string.IsNullOrEmpty(Password))
            {
                Console.WriteLine("Authentication: Enabled");
            }
            Console.WriteLine($"Command Count: {CommandCount}");
            Console.WriteLine($"Batch Size: {BatchSize}");
            Console.WriteLine($"Buffer Size: {BufferSize / 1024}KB\n");

            // Test 1: Sequential commands (no pipelining)
            await RunSequentialTest();

            // Test 2: Pipelined commands
            await RunPipelinedTest();

            // Test 3: Batched commands
            await RunBatchedTest();

            Console.WriteLine("\nPipeline Test completed.");
        }

        private static void ParseArguments(string[] args)
        {
            for (int i = 0; i < args.Length; i++)
            {
                string arg = args[i].ToLower();

                switch (arg)
                {
                    case "--host":
                    case "-h":
                        if (i + 1 < args.Length)
                        {
                            Host = args[i + 1];
                            i++;
                        }
                        break;

                    case "--port":
                    case "-p":
                        if (i + 1 < args.Length && int.TryParse(args[i + 1], out int port))
                        {
                            Port = port;
                            i++;
                        }
                        break;

                    case "--password":
                    case "-pass":
                        if (i + 1 < args.Length)
                        {
                            Password = args[i + 1];
                            i++;
                        }
                        break;

                    case "--commands":
                    case "-c":
                        if (i + 1 < args.Length && int.TryParse(args[i + 1], out int commands))
                        {
                            CommandCount = commands;
                            i++;
                        }
                        break;

                    case "--batch":
                    case "-b":
                        if (i + 1 < args.Length && int.TryParse(args[i + 1], out int batch))
                        {
                            BatchSize = batch;
                            i++;
                        }
                        break;

                    case "--buffer":
                        if (i + 1 < args.Length && int.TryParse(args[i + 1], out int buffer))
                        {
                            BufferSize = buffer * 1024; // Convert KB to bytes
                            i++;
                        }
                        break;

                    case "--clients":
                        if (i + 1 < args.Length && int.TryParse(args[i + 1], out int clients))
                        {
                            ClientCount = clients;
                            i++;
                        }
                        break;
                }
            }
        }

        private static async Task Authenticate(NetworkStream stream)
        {
            if (!string.IsNullOrEmpty(Password))
            {
                string authCommand = $"AUTH {Password}\r\n";
                byte[] authData = Encoding.UTF8.GetBytes(authCommand);
                await stream.WriteAsync(authData);
                
                var buffer = new byte[BufferSize];
                int bytesRead = await stream.ReadAsync(buffer);
                string response = Encoding.UTF8.GetString(buffer, 0, bytesRead);
                
                if (!response.StartsWith("+OK"))
                {
                    throw new Exception($"Authentication failed: {response}");
                }
            }
        }

        /// <summary>
        /// Read responses from the stream until we have the expected count
        /// </summary>
        private static async Task<int> ReadResponses(NetworkStream stream, byte[] buffer, int expectedCount, string testType)
        {
            var responseBuilder = new StringBuilder();
            int totalResponses = 0;
            int retryCount = 0;
            const int maxRetries = 10;

            while (totalResponses < expectedCount && retryCount < maxRetries)
            {
                if (stream.DataAvailable)
                {
                    int bytesRead = await stream.ReadAsync(buffer);
                    if (bytesRead > 0)
                    {
                        string response = Encoding.UTF8.GetString(buffer, 0, bytesRead);
                        responseBuilder.Append(response);
                        
                        // Count complete responses
                        int newResponses = response.Split(["\r\n"], StringSplitOptions.RemoveEmptyEntries).Length;
                        totalResponses += newResponses;

                        if (totalResponses % 50 == 0)
                        {
                            Console.WriteLine($"{testType}: Received {totalResponses}/{expectedCount} responses");
                        }

                        retryCount = 0; // Reset retry count on successful read
                    }
                }
                else
                {
                    await Task.Delay(100); // Wait a bit for more data
                    retryCount++;
                }
            }

            if (totalResponses < expectedCount)
            {
                Console.WriteLine($"Warning: Only received {totalResponses}/{expectedCount} responses");
            }

            return totalResponses;
        }

        /// <summary>
        /// Test with sequential commands (no pipelining)
        /// </summary>
        private static async Task RunSequentialTest()
        {
            Console.WriteLine("=== Sequential Test ===");
            var stopwatch = new Stopwatch();
            stopwatch.Start();

            var tasks = new List<Task>();
            var clientResults = new List<(int clientId, int responses, double opsPerSecond)>();

            for (int clientId = 0; clientId < ClientCount; clientId++)
            {
                tasks.Add(Task.Run(async () =>
                {
                    using var client = new TcpClient();
                    await client.ConnectAsync(Host, Port);
                    using var stream = client.GetStream();
                    var buffer = new byte[BufferSize];
                    int totalResponses = 0;

                    // Authenticate if password is set
                    await Authenticate(stream);

                    for (int i = 0; i < CommandCount; i++)
                    {
                        string key = $"seq:client{clientId}:key:{i}";
                        string value = $"value:{i}";

                        // SET command
                        string setCommand = $"SET {key} {value}\r\n";
                        byte[] setData = Encoding.UTF8.GetBytes(setCommand);
                        await stream.WriteAsync(setData);
                        totalResponses += await ReadResponses(stream, buffer, 1, $"Sequential-Client{clientId}");

                        // GET command
                        string getCommand = $"GET {key}\r\n";
                        byte[] getData = Encoding.UTF8.GetBytes(getCommand);
                        await stream.WriteAsync(getData);
                        totalResponses += await ReadResponses(stream, buffer, 1, $"Sequential-Client{clientId}");

                        if (i % 10 == 0)
                        {
                            Console.WriteLine($"Sequential-Client{clientId}: Processed {i}/{CommandCount} commands");
                        }
                    }

                    clientResults.Add((clientId, totalResponses, (CommandCount * 2.0) / (stopwatch.ElapsedMilliseconds / 1000.0)));
                }));
            }

            await Task.WhenAll(tasks);
            stopwatch.Stop();

            // Print individual client results
            foreach (var (clientId, responses, opsPerSecond) in clientResults)
            {
                Console.WriteLine($"Sequential-Client{clientId}:");
                Console.WriteLine($"  Total responses: {responses}");
                Console.WriteLine($"  Operations per second: {opsPerSecond:F2}");
            }

            // Print aggregate results
            double totalOpsPerSecond = clientResults.Sum(r => r.opsPerSecond);
            Console.WriteLine($"\nSequential Test Summary:");
            Console.WriteLine($"Total operations per second: {totalOpsPerSecond:F2}");
            Console.WriteLine($"Average operations per second per client: {totalOpsPerSecond / ClientCount:F2}\n");
        }

        /// <summary>
        /// Test with pipelined commands
        /// </summary>
        private static async Task RunPipelinedTest()
        {
            Console.WriteLine("=== Pipelined Test ===");
            var stopwatch = new Stopwatch();
            stopwatch.Start();

            var tasks = new List<Task>();
            var clientResults = new List<(int clientId, int responses, double opsPerSecond)>();

            for (int clientId = 0; clientId < ClientCount; clientId++)
            {
                tasks.Add(Task.Run(async () =>
                {
                    using var client = new TcpClient();
                    await client.ConnectAsync(Host, Port);
                    using var stream = client.GetStream();
                    var buffer = new byte[BufferSize];

                    // Authenticate if password is set
                    await Authenticate(stream);

                    // Enter pipeline mode
                    string pipelineCommand = "PIPELINE\r\n";
                    byte[] pipelineData = Encoding.UTF8.GetBytes(pipelineCommand);
                    await stream.WriteAsync(pipelineData);
                    await ReadResponses(stream, buffer, 1, $"Pipeline-Client{clientId}");

                    Console.WriteLine($"Pipeline-Client{clientId}: Sending commands...");

                    // Send all commands
                    var allCommands = new StringBuilder();
                    for (int i = 0; i < CommandCount; i++)
                    {
                        string key = $"pipe:client{clientId}:key:{i}";
                        string value = $"value:{i}";
                        allCommands.Append($"SET {key} {value}\r\n");
                        allCommands.Append($"GET {key}\r\n");
                    }

                    byte[] commandsData = Encoding.UTF8.GetBytes(allCommands.ToString());
                    await stream.WriteAsync(commandsData);

                    // Execute pipeline
                    string execCommand = "EXEC\r\n";
                    byte[] execData = Encoding.UTF8.GetBytes(execCommand);
                    await stream.WriteAsync(execData);

                    // Read all responses
                    int totalResponses = await ReadResponses(stream, buffer, CommandCount * 2 + 1, $"Pipeline-Client{clientId}");
                    clientResults.Add((clientId, totalResponses, (CommandCount * 2.0) / (stopwatch.ElapsedMilliseconds / 1000.0)));
                }));
            }

            await Task.WhenAll(tasks);
            stopwatch.Stop();

            // Print individual client results
            foreach (var (clientId, responses, opsPerSecond) in clientResults)
            {
                Console.WriteLine($"Pipeline-Client{clientId}:");
                Console.WriteLine($"  Total responses: {responses}");
                Console.WriteLine($"  Operations per second: {opsPerSecond:F2}");
            }

            // Print aggregate results
            double totalOpsPerSecond = clientResults.Sum(r => r.opsPerSecond);
            Console.WriteLine($"\nPipelined Test Summary:");
            Console.WriteLine($"Total operations per second: {totalOpsPerSecond:F2}");
            Console.WriteLine($"Average operations per second per client: {totalOpsPerSecond / ClientCount:F2}\n");
        }

        /// <summary>
        /// Test with batched commands
        /// </summary>
        private static async Task RunBatchedTest()
        {
            Console.WriteLine("=== Batched Test ===");
            var stopwatch = new Stopwatch();
            stopwatch.Start();

            var tasks = new List<Task>();
            var clientResults = new List<(int clientId, int responses, double opsPerSecond)>();

            for (int clientId = 0; clientId < ClientCount; clientId++)
            {
                tasks.Add(Task.Run(async () =>
                {
                    using var client = new TcpClient();
                    await client.ConnectAsync(Host, Port);
                    using var stream = client.GetStream();
                    var buffer = new byte[BufferSize];
                    int totalResponses = 0;

                    // Authenticate if password is set
                    await Authenticate(stream);

                    for (int batch = 0; batch < CommandCount; batch += BatchSize)
                    {
                        int currentBatchSize = Math.Min(BatchSize, CommandCount - batch);
                        var batchCommands = new StringBuilder();

                        for (int i = 0; i < currentBatchSize; i++)
                        {
                            int index = batch + i;
                            string key = $"batch:client{clientId}:key:{index}";
                            string value = $"value:{index}";
                            batchCommands.Append($"SET {key} {value}\r\n");
                            batchCommands.Append($"GET {key}\r\n");
                        }

                        byte[] batchData = Encoding.UTF8.GetBytes(batchCommands.ToString());
                        await stream.WriteAsync(batchData);

                        int responses = await ReadResponses(stream, buffer, currentBatchSize * 2, $"Batch-Client{clientId}");
                        totalResponses += responses;

                        Console.WriteLine($"Batch-Client{clientId}: Processed batch {batch / BatchSize + 1}/{Math.Ceiling((double)CommandCount / BatchSize)} " +
                                        $"({responses} responses)");
                    }

                    clientResults.Add((clientId, totalResponses, (CommandCount * 2.0) / (stopwatch.ElapsedMilliseconds / 1000.0)));
                }));
            }

            await Task.WhenAll(tasks);
            stopwatch.Stop();

            // Print individual client results
            foreach (var (clientId, responses, opsPerSecond) in clientResults)
            {
                Console.WriteLine($"Batch-Client{clientId}:");
                Console.WriteLine($"  Total responses: {responses}");
                Console.WriteLine($"  Operations per second: {opsPerSecond:F2}");
            }

            // Print aggregate results
            double totalOpsPerSecond = clientResults.Sum(r => r.opsPerSecond);
            Console.WriteLine($"\nBatched Test Summary:");
            Console.WriteLine($"Total operations per second: {totalOpsPerSecond:F2}");
            Console.WriteLine($"Average operations per second per client: {totalOpsPerSecond / ClientCount:F2}\n");
        }
    }
} 