From 9ad8126d6095b1402da9e9f0e5755065ffc6f738 Mon Sep 17 00:00:00 2001
From: Tyrrrz <1935960+Tyrrrz@users.noreply.github.com>
Date: Mon, 13 Feb 2023 22:07:18 +0200
Subject: [PATCH] Allow passing category IDs to the `export` command in CLI

Closes #798
---
 .../Commands/Base/ExportCommandBase.cs        | 27 ++++++++++++++++++-
 .../Commands/ExportChannelsCommand.cs         |  2 +-
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/DiscordChatExporter.Cli/Commands/Base/ExportCommandBase.cs b/DiscordChatExporter.Cli/Commands/Base/ExportCommandBase.cs
index a237f93..3cf7164 100644
--- a/DiscordChatExporter.Cli/Commands/Base/ExportCommandBase.cs
+++ b/DiscordChatExporter.Cli/Commands/Base/ExportCommandBase.cs
@@ -15,6 +15,7 @@ using DiscordChatExporter.Core.Exporting;
 using DiscordChatExporter.Core.Exporting.Filtering;
 using DiscordChatExporter.Core.Exporting.Partitioning;
 using DiscordChatExporter.Core.Utils;
+using DiscordChatExporter.Core.Utils.Extensions;
 using Gress;
 
 namespace DiscordChatExporter.Cli.Commands.Base;
@@ -233,12 +234,36 @@ public abstract class ExportCommandBase : TokenCommandBase
     protected async ValueTask ExecuteAsync(IConsole console, IReadOnlyList<Snowflake> channelIds)
     {
         var cancellationToken = console.RegisterCancellationHandler();
+
+        await console.Output.WriteLineAsync("Resolving channel(s)...");
+
         var channels = new List<Channel>();
+        var guildChannelMap = new Dictionary<Snowflake, IReadOnlyList<Channel>>();
 
         foreach (var channelId in channelIds)
         {
             var channel = await Discord.GetChannelAsync(channelId, cancellationToken);
-            channels.Add(channel);
+
+            // Unwrap categories
+            if (channel.Kind == ChannelKind.GuildCategory)
+            {
+                var guildChannels =
+                    guildChannelMap.GetValueOrDefault(channel.GuildId) ??
+                    await Discord.GetGuildChannelsAsync(channel.GuildId, cancellationToken);
+
+                foreach (var guildChannel in guildChannels)
+                {
+                    if (guildChannel.Category.Id == channel.Id)
+                        channels.Add(guildChannel);
+                }
+
+                // Cache the guild channels to avoid redundant work
+                guildChannelMap[channel.GuildId] = guildChannels;
+            }
+            else
+            {
+                channels.Add(channel);
+            }
         }
 
         await ExecuteAsync(console, channels);
diff --git a/DiscordChatExporter.Cli/Commands/ExportChannelsCommand.cs b/DiscordChatExporter.Cli/Commands/ExportChannelsCommand.cs
index d43d80a..b216432 100644
--- a/DiscordChatExporter.Cli/Commands/ExportChannelsCommand.cs
+++ b/DiscordChatExporter.Cli/Commands/ExportChannelsCommand.cs
@@ -15,7 +15,7 @@ public class ExportChannelsCommand : ExportCommandBase
     [CommandOption(
         "channel",
         'c',
-        Description = "Channel ID(s)."
+        Description = "Channel ID(s). If provided with a category ID, all channels in that category will be exported."
     )]
     public required IReadOnlyList<Snowflake> ChannelIds { get; init; }