using System.Text;

namespace Firefly
{
    public partial class Firefly
    {
        #region String Operations
        /// <summary>
        /// Handles the SET command which sets a key-value pair in the string store.
        /// </summary>
        /// <param name="args">Command arguments in format: "key value"</param>
        /// <returns>OK on success, error if arguments are invalid</returns>
        static byte[] HandleSetCommand(string args)
        {
            string[] parts = args.Split(' ', 2);
            if (parts.Length < 2)
            {
                return Encoding.UTF8.GetBytes("-ERR wrong number of arguments for 'set' command\r\n");
            }

            string key = parts[0];
            string value = parts[1];

            // Store the key-value pair using the sharded helper method
            if (!StringStoreSet(key, value))
            {
                string? existingType = GetKeyType(key);
                if (existingType != null)
                {
                    return Encoding.UTF8.GetBytes($"-ERR key '{key}' already exists as type '{existingType}'\r\n");
                }
                return Encoding.UTF8.GetBytes("-ERR operation failed\r\n");
            }

            return Encoding.UTF8.GetBytes("+OK\r\n");
        }

        /// <summary>
        /// Handles the GET command which retrieves a value from the string store.
        /// </summary>
        /// <param name="args">Command arguments containing the key</param>
        /// <returns>The value associated with the key, or nil if the key doesn't exist</returns>
        static byte[] HandleGetCommand(string args)
        {
            if (string.IsNullOrWhiteSpace(args))
            {
                return Encoding.UTF8.GetBytes("-ERR wrong number of arguments for 'get' command\r\n");
            }

            string key = args.Trim();

            // Retrieve the value using the sharded helper method
            if (StringStoreGet(key, out string? value))
            {
                return Encoding.UTF8.GetBytes($"+{value}\r\n");
            }
            else
            {
                return Encoding.UTF8.GetBytes("$-1\r\n"); // Redis returns nil for non-existent keys
            }
        }

        #region String Store Helpers
        /// <summary>
        /// Checks if a string exists for a given key.
        /// </summary>
        /// <param name="key">The key to check the string for</param>
        /// <returns>True if the string exists, false otherwise</returns>
        private static bool StringStoreExists(string key)
        {
            int shardIndex = GetShardIndex(key);
            return stringStoreShards[shardIndex].ContainsKey(key);
        }

        /// <summary>
        /// Sets a string for a given key.
        /// </summary>
        /// <param name="key">The key to set the string for</param>
        /// <param name="value">The value to set</param>
        /// <returns>True if the string was set, false otherwise</returns>
        private static bool StringStoreSet(string key, string value)
        {
            // Check if key exists in any other store
            if (!EnsureKeyDoesNotExist(key, "string"))
            {
                return false;
            }

            int shardIndex = GetShardIndex(key);
            stringStoreShards[shardIndex][key] = value;

            MarkDataAsModified();
            
            return true;
        }

        /// <summary>
        /// Gets a string for a given key.
        /// </summary>
        /// <param name="key">The key to get the string for</param>
        /// <param name="value">The value</param>
        /// <returns>True if the string was found, false otherwise</returns>
        private static bool StringStoreGet(string key, out string? value)
        {
            int shardIndex = GetShardIndex(key);
            return stringStoreShards[shardIndex].TryGetValue(key, out value);
        }

        /// <summary>
        /// Removes a string for a given key.
        /// </summary>
        /// <param name="key">The key to remove the string for</param>
        /// <returns>True if the string was removed, false otherwise</returns>
        private static bool StringStoreRemove(string key)
        {
            int shardIndex = GetShardIndex(key);
            bool result = stringStoreShards[shardIndex].TryRemove(key, out _);
            
            if (result)
                MarkDataAsModified();
            
            return result;
        }
        #endregion
        #endregion
    }
}