lay295 · ScrubN · Jul 15, 2024 · Jul 12, 2024 · Jul 14, 2024 · Jul 14, 2024
diff --git a/TwitchDownloaderCore/ChatDownloader.cs b/TwitchDownloaderCore/ChatDownloader.cs
@@ -42,12 +42,12 @@ public ChatDownloader(ChatDownloadOptions chatDownloadOptions, ITaskProgress pro
                 "TwitchDownloader");
         }
 
-        private static async Task<List<Comment>> DownloadSection(double videoStart, double videoEnd, string videoId, IProgress<int> progress, ChatFormat format, CancellationToken cancellationToken)
+        private static async Task<List<Comment>> DownloadSection(Range downloadRange, string videoId, IProgress<int> progress, ChatFormat format, CancellationToken cancellationToken)
         {
             var comments = new List<Comment>();
-            //GQL only wants ints
-            videoStart = Math.Floor(videoStart);
-            double videoDuration = videoEnd - videoStart;
+            int videoStart = downloadRange.Start.Value;
+            int videoEnd = downloadRange.End.Value;
+            int videoDuration = videoEnd - videoStart;
             double latestMessage = videoStart - 1;
             bool isFirst = true;
             string cursor = "";
@@ -104,25 +104,26 @@ private static async Task<List<Comment>> DownloadSection(double videoStart, doub
                 }
 
                 var convertedComments = ConvertComments(commentResponse[0].data.video, format);
-                comments.EnsureCapacity(Math.Min(0, comments.Capacity + convertedComments.Count));
                 foreach (var comment in convertedComments)
                 {
-                    if (latestMessage < videoEnd && comment.content_offset_seconds > videoStart)
+                    if (comment.content_offset_seconds >= videoStart && comment.content_offset_seconds < videoEnd)
+                    {
                         comments.Add(comment);
+                    }
 
-                    latestMessage = comment.content_offset_seconds;
+                    if (comment.content_offset_seconds > latestMessage)
+                    {
+                        latestMessage = comment.content_offset_seconds;
+                    }
                 }
 
                 if (!commentResponse[0].data.video.comments.pageInfo.hasNextPage)
                     break;
 
                 cursor = commentResponse[0].data.video.comments.edges.Last().cursor;
 
-                if (progress != null)
-                {
-                    int percent = (int)Math.Floor((latestMessage - videoStart) / videoDuration * 100);
-                    progress.Report(percent);
-                }
+                var percent = (int)Math.Floor((latestMessage - videoStart) / videoDuration * 100);
+                progress.Report(percent);
 
                 if (isFirst)
                 {
@@ -276,43 +277,8 @@ private async Task DownloadAsyncImpl(FileInfo outputFileInfo, FileStream outputF
             DownloadType downloadType = downloadOptions.Id.All(char.IsDigit) ? DownloadType.Video : DownloadType.Clip;
 
             var (chatRoot, connectionCount) = await InitChatRoot(downloadType);
-            var videoStart = chatRoot.video.start;
-            var videoEnd = chatRoot.video.end;
-            var videoId = chatRoot.video.id;
-            var videoDuration = videoEnd - videoStart;
-
-            var downloadTasks = new List<Task<List<Comment>>>(connectionCount);
-            var percentages = new int[connectionCount];
-
-            double chunk = videoDuration / connectionCount;
-            for (int i = 0; i < connectionCount; i++)
-            {
-                int tc = i;
-
-                var taskProgress = new Progress<int>(percent =>
-                {
-                    percentages[tc] = Math.Clamp(percent, 0, 100);
-
-                    var reportPercent = percentages.Sum() / connectionCount;
-                    _progress.ReportProgress(reportPercent);
-                });
-
-                double start = videoStart + chunk * i;
-                downloadTasks.Add(DownloadSection(start, start + chunk, videoId, taskProgress, downloadOptions.DownloadFormat, cancellationToken));
-            }
-
-            _progress.SetTemplateStatus("Downloading {0}%", 0);
-            await Task.WhenAll(downloadTasks);
-
-            var sortedComments = new List<Comment>(downloadTasks.Count);
-            foreach (var commentTask in downloadTasks)
-            {
-                sortedComments.AddRange(commentTask.Result);
-            }
 
-            sortedComments.Sort(new CommentOffsetComparer());
-
-            chatRoot.comments = sortedComments.DistinctBy(x => x._id).ToList();
+            chatRoot.comments = await DownloadComments(cancellationToken, chatRoot.video, connectionCount);
 
             if (downloadOptions.EmbedData && (downloadOptions.DownloadFormat is ChatFormat.Json or ChatFormat.Html))
             {
@@ -326,7 +292,7 @@ private async Task DownloadAsyncImpl(FileInfo outputFileInfo, FileStream outputF
                 await BackfillUserInfo(chatRoot);
             }
 
-            _progress.SetStatus("Writing output file");
+            _progress.SetStatus("Writing Output File");
             switch (downloadOptions.DownloadFormat)
             {
                 case ChatFormat.Json:
@@ -440,6 +406,49 @@ private async Task DownloadAsyncImpl(FileInfo outputFileInfo, FileStream outputF
             return (chatRoot, connectionCount);
         }
 
+        private async Task<List<Comment>> DownloadComments(CancellationToken cancellationToken, Video video, int connectionCount)
+        {
+            _progress.SetTemplateStatus("Downloading {0}%", 0);
+
+            var videoStart = (int)Math.Floor(video.start);
+            var videoEnd = (int)Math.Ceiling(video.end) + 1; // Exclusive end
+            var videoDuration = videoEnd - videoStart;
+
+            var downloadTasks = new List<Task<List<Comment>>>(connectionCount);
+            var percentages = new int[connectionCount];
+
+            var chunkSize = (int)Math.Ceiling(videoDuration / (double)connectionCount);
+            for (var i = 0; i < connectionCount; i++)
+            {
+                var tc = i;
+
+                var taskProgress = new Progress<int>(percent =>
+                {
+                    percentages[tc] = Math.Clamp(percent, 0, 100);
+
+                    var reportPercent = percentages.Sum() / connectionCount;
+                    _progress.ReportProgress(reportPercent);
+                });
+
+                var start = videoStart + chunkSize * i;
+                var end = Math.Min(videoEnd, start + chunkSize);
+                var downloadRange = new Range(start, end);
+                downloadTasks.Add(DownloadSection(downloadRange, video.id, taskProgress, downloadOptions.DownloadFormat, cancellationToken));
+            }
+
+            await Task.WhenAll(downloadTasks);
+
+            _progress.ReportProgress(100);
+
+            var commentList = downloadTasks
+                .SelectMany(task => task.Result)
+                .ToHashSet(new CommentIdEqualityComparer())
+                .ToList();
+
+            commentList.Sort(new CommentOffsetComparer());
+            return commentList;
+        }
+
         private async Task EmbedImages(ChatRoot chatRoot, CancellationToken cancellationToken)
         {
             _progress.SetTemplateStatus("Downloading Embed Images {0}%", 0);

diff --git a/TwitchDownloaderCore/ChatUpdater.cs b/TwitchDownloaderCore/ChatUpdater.cs
@@ -431,26 +431,18 @@ private async Task ChatTrimEndingTask(CancellationToken cancellationToken)
             ChatRoot newChatRoot = await ChatJson.DeserializeAsync(inputFile, getComments: true, onlyFirstAndLastComments: false, getEmbeds: false, cancellationToken);
 
             // Append the new comment section
-            SortedSet<Comment> commentsSet = new SortedSet<Comment>(new CommentOffsetComparer());
-            foreach (var comment in newChatRoot.comments)
-            {
-                if (comment.content_offset_seconds < downloadOptions.TrimEndingTime && comment.content_offset_seconds >= downloadOptions.TrimBeginningTime)
-                {
-                    commentsSet.Add(comment);
-                }
-            }
+            var commentsSet = newChatRoot.comments.ToHashSet(new CommentIdEqualityComparer());
 
             lock (_trimChatRootLock)
             {
+                commentsSet.EnsureCapacity(commentsSet.Count + chatRoot.comments.Count);
                 foreach (var comment in chatRoot.comments)
                 {
                     commentsSet.Add(comment);
                 }
 
-                List<Comment> comments = commentsSet.DistinctBy(x => x._id).ToList();
-                commentsSet.Clear();
-
-                chatRoot.comments = comments;
+                chatRoot.comments = commentsSet.ToList();
+                chatRoot.comments.Sort(new CommentOffsetComparer());
             }
         }
 

diff --git a/TwitchDownloaderCore/Tools/CommentIdEqualityComparer.cs b/TwitchDownloaderCore/Tools/CommentIdEqualityComparer.cs
@@ -0,0 +1,18 @@
+using System.Collections.Generic;
+using TwitchDownloaderCore.TwitchObjects;
+
+namespace TwitchDownloaderCore.Tools
+{
+    public class CommentIdEqualityComparer : IEqualityComparer<Comment>
+    {
+        public bool Equals(Comment x, Comment y)
+        {
+            if (x is null) return y is null;
+            if (y is null) return false;
+
+            return x._id.Equals(y._id);
+        }
+
+        public int GetHashCode(Comment obj) => obj._id.GetHashCode();
+    }
+}
diff --git a/TwitchDownloaderCore/Tools/FilenameService.cs b/TwitchDownloaderCore/Tools/FilenameService.cs
@@ -35,27 +35,27 @@ public static string GetFilename(string template, [AllowNull] string title, [All
             if (template.Contains("{trim_start_custom="))
             {
                 var trimStartRegex = new Regex("{trim_start_custom=\"(.*?)\"}");
-                ReplaceCustomWithFormattable(stringBuilder, trimStartRegex, trimStart);
+                ReplaceCustomWithFormattable(stringBuilder, trimStartRegex, trimStart, TimeSpanHFormat.ReusableInstance);
             }
 
             if (template.Contains("{trim_end_custom="))
             {
                 var trimEndRegex = new Regex("{trim_end_custom=\"(.*?)\"}");
-                ReplaceCustomWithFormattable(stringBuilder, trimEndRegex, trimEnd);
+                ReplaceCustomWithFormattable(stringBuilder, trimEndRegex, trimEnd, TimeSpanHFormat.ReusableInstance);
             }
 
             if (template.Contains("{length_custom="))
             {
                 var lengthRegex = new Regex("{length_custom=\"(.*?)\"}");
-                ReplaceCustomWithFormattable(stringBuilder, lengthRegex, videoLength);
+                ReplaceCustomWithFormattable(stringBuilder, lengthRegex, videoLength, TimeSpanHFormat.ReusableInstance);
             }
 
             var fileName = stringBuilder.ToString();
             var additionalSubfolders = GetTemplateSubfolders(ref fileName);
             return Path.Combine(Path.Combine(additionalSubfolders), ReplaceInvalidFilenameChars(fileName));
         }
 
-        private static void ReplaceCustomWithFormattable(StringBuilder sb, Regex regex, IFormattable formattable, IFormatProvider formatProvider = null)
+        private static void ReplaceCustomWithFormattable<TFormattable>(StringBuilder sb, Regex regex, TFormattable formattable, [AllowNull] IFormatProvider formatProvider = null) where TFormattable : IFormattable
         {
             do
             {
@@ -66,8 +66,12 @@ private static void ReplaceCustomWithFormattable(StringBuilder sb, Regex regex,
                     break;
 
                 var formatString = match.Groups[1].Value;
+                var formattedString = formatProvider?.GetFormat(typeof(ICustomFormatter)) is ICustomFormatter customFormatter
+                    ? customFormatter.Format(formatString, formattable, formatProvider)
+                    : formattable.ToString(formatString, formatProvider);
+
                 sb.Remove(match.Groups[0].Index, match.Groups[0].Length);
-                sb.Insert(match.Groups[0].Index, ReplaceInvalidFilenameChars(formattable.ToString(formatString, formatProvider)));
+                sb.Insert(match.Groups[0].Index, ReplaceInvalidFilenameChars(formattedString));
             } while (true);
         }
 

diff --git a/TwitchDownloaderCore/TsMerger.cs b/TwitchDownloaderCore/TsMerger.cs
@@ -34,7 +34,7 @@ public async Task MergeAsync(CancellationToken cancellationToken)
 
             try
             {
-                await MergeAsyncImpl(outputFs, cancellationToken);
+                await MergeAsyncImpl(outputFileInfo, outputFs, cancellationToken);
             }
             catch
             {
@@ -46,7 +46,7 @@ public async Task MergeAsync(CancellationToken cancellationToken)
             }
         }
 
-        private async Task MergeAsyncImpl(FileStream outputFs, CancellationToken cancellationToken)
+        private async Task MergeAsyncImpl(FileInfo outputFileInfo, FileStream outputFs, CancellationToken cancellationToken)
         {
             var isM3U8 = false;
             var isFirst = true;
@@ -74,9 +74,7 @@ private async Task MergeAsyncImpl(FileStream outputFs, CancellationToken cancell
 
             _progress.SetTemplateStatus("Combining Parts {0}% [2/2]", 0);
 
-            await CombineVideoParts(fileList, outputFs, cancellationToken);
-
-            _progress.ReportProgress(100);
+            await CombineVideoParts(fileList, outputFileInfo, outputFs, cancellationToken);
         }
 
         private async Task VerifyVideoParts(IReadOnlyCollection<string> fileList, CancellationToken cancellationToken)
@@ -100,6 +98,8 @@ private async Task VerifyVideoParts(IReadOnlyCollection<string> fileList, Cancel
                 cancellationToken.ThrowIfCancellationRequested();
             }
 
+            _progress.ReportProgress(100);
+
             if (failedParts.Count != 0)
             {
                 if (failedParts.Count == fileList.Count)
@@ -131,9 +131,9 @@ private static async Task<bool> VerifyVideoPart(string filePath)
             return true;
         }
 
-        private async Task CombineVideoParts(IReadOnlyCollection<string> fileList, FileStream outputStream, CancellationToken cancellationToken)
+        private async Task CombineVideoParts(IReadOnlyCollection<string> fileList, FileInfo outputFileInfo, FileStream outputStream, CancellationToken cancellationToken)
         {
-            DriveInfo outputDrive = DriveHelper.GetOutputDrive(mergeOptions.OutputFile);
+            DriveInfo outputDrive = DriveHelper.GetOutputDrive(outputFileInfo.FullName);
 
             int partCount = fileList.Count;
             int doneCount = 0;
@@ -153,6 +153,8 @@ private async Task CombineVideoParts(IReadOnlyCollection<string> fileList, FileS
 
                 cancellationToken.ThrowIfCancellationRequested();
             }
+
+            _progress.ReportProgress(100);
         }
     }
 }