From 506ba0f61511e735b42ec3fb833b32831e6f6ea2 Mon Sep 17 00:00:00 2001
From: rlaphoenix <rlaphoenix@pm.me>
Date: Fri, 16 Feb 2024 16:49:42 +0000
Subject: [PATCH] HLS: Only merge relevant segments on discontinuity

---
 devine/core/manifests/hls.py | 43 ++++++++++++++++++++++++------------
 1 file changed, 29 insertions(+), 14 deletions(-)

diff --git a/devine/core/manifests/hls.py b/devine/core/manifests/hls.py
index de4ac0c..b0b3577 100644
--- a/devine/core/manifests/hls.py
+++ b/devine/core/manifests/hls.py
@@ -364,20 +364,35 @@ class HLS:
 
                 return decrypted_path
 
-            def merge_discontinuity():
-                """Merge all files in the segment save directory so far."""
-                files = list(sorted(segment_save_dir.iterdir()))
+            def merge_discontinuity(include_this_segment: bool):
+                """
+                Merge all segments of the discontinuity.
 
-                to_dir = segment_save_dir.parent
-                to_path = to_dir / f"{str(discon_i).zfill(name_len)}{files[-1].suffix}"
+                All segment files for this discontinuity must already be downloaded and
+                already decrypted (if it needs to be decrypted).
 
-                merge(
-                    to=to_path,
-                    via=files,
-                    delete=True,
-                    include_map_data=True
-                )
-                segment_save_dir.rmdir()
+                Parameters:
+                    include_this_segment: Whether to include the current segment in the
+                        list of segments to merge and decrypt. This should be False if
+                        decrypting on EXT-X-KEY changes, or True when decrypting on the
+                        last segment.
+                """
+                last_segment_i = max(0, i - int(not include_this_segment))
+
+                files = [
+                    file
+                    for file in sorted(segment_save_dir.iterdir())
+                    if int(file.stem.replace("_decrypted", "").split("-")[-1]) <= last_segment_i
+                ]
+                if files:
+                    to_dir = segment_save_dir.parent
+                    to_path = to_dir / f"{str(discon_i).zfill(name_len)}{files[-1].suffix}"
+                    merge(
+                        to=to_path,
+                        via=files,
+                        delete=True,
+                        include_map_data=True
+                    )
 
             if isinstance(track, Subtitle):
                 segment_data = try_ensure_utf8(segment_file_path.read_bytes())
@@ -392,7 +407,7 @@ class HLS:
             if segment.discontinuity and i != 0:
                 if encryption_data:
                     decrypt(include_this_segment=False)
-                merge_discontinuity()
+                merge_discontinuity(include_this_segment=False)
 
                 discon_i += 1
                 range_offset = 0  # TODO: Should this be reset or not?
@@ -446,7 +461,7 @@ class HLS:
                 # required as it won't end with EXT-X-DISCONTINUITY nor a new key
                 if encryption_data:
                     decrypt(include_this_segment=True)
-                merge_discontinuity()
+                merge_discontinuity(include_this_segment=True)
 
             progress(advance=1)