Merge HLS segments first by discontinuity then via FFmpeg

HLS playlists where each segment is in an mp4 container seems to corrupt when the EXT-X-MAP is changed out, unless you first merge segments by discontinuity and then merge the merges via FFmpeg (which demuxes all the merged segment continuities and then concatanates them together, probably giving it new init data too).

devine/core/manifests/hls.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 import html
 import logging
 import re
+import shutil
+import subprocess
 import sys
 import time
 from concurrent import futures
@@ -29,7 +31,7 @@ from devine.core.downloaders import downloader
 from devine.core.downloaders import requests as requests_downloader
 from devine.core.drm import DRM_T, ClearKey, Widevine
 from devine.core.tracks import Audio, Subtitle, Tracks, Video
-from devine.core.utilities import is_close_match, try_ensure_utf8
+from devine.core.utilities import get_binary_path, is_close_match, try_ensure_utf8
 
 
 class HLS:
@@ -250,12 +252,28 @@ class HLS:
         range_offset.put(0)
         drm_lock = Lock()
 
+        discontinuities: list[list[segment]] = []
+        discontinuity_index = -1
+        for i, segment in enumerate(master.segments):
+            if i == 0 or segment.discontinuity:
+                discontinuity_index += 1
+                discontinuities.append([])
+            discontinuities[discontinuity_index].append(segment)
+
+        for d_i, discontinuity in enumerate(discontinuities):
+            # each discontinuity is a separate 'file'/encode and must be processed separately
+            discontinuity_save_dir = save_dir / str(d_i).zfill(len(str(len(discontinuities))))
+            discontinuity_save_path = discontinuity_save_dir.with_suffix(Path(discontinuity[0].uri).suffix)
+
             with ThreadPoolExecutor(max_workers=16) as pool:
                 for i, download in enumerate(futures.as_completed((
                     pool.submit(
                         HLS.download_segment,
                         segment=segment,
-                    out_path=(save_dir / str(n).zfill(len(str(len(master.segments))))).with_suffix(".mp4"),
+                        out_path=(
+                            discontinuity_save_dir /
+                            str(s_i).zfill(len(str(len(discontinuity))))
+                        ).with_suffix(Path(segment.uri).suffix),
                         track=track,
                         init_data=init_data,
                         segment_key=segment_key,
@@ -266,7 +284,7 @@ class HLS:
                         session=session,
                         proxy=proxy
                     )
-                for n, segment in enumerate(master.segments)
+                    for s_i, segment in enumerate(discontinuity)
                 ))):
                     try:
                         download_size = download.result()
@@ -307,11 +325,8 @@ class HLS:
                             last_speed_refresh = now
                             download_sizes.clear()
 
-        if DOWNLOAD_LICENCE_ONLY.is_set():
+            with open(discontinuity_save_path, "wb") as f:
-            return
+                for segment_file in sorted(discontinuity_save_dir.iterdir()):
-
-        with open(save_path, "wb") as f:
-            for segment_file in sorted(save_dir.iterdir()):
                     segment_data = segment_file.read_bytes()
                     if isinstance(track, Subtitle):
                         segment_data = try_ensure_utf8(segment_data)
@@ -323,11 +338,29 @@ class HLS:
                                 encode("utf8")
                     f.write(segment_data)
                     segment_file.unlink()
+                shutil.rmtree(discontinuity_save_dir)
+
+        if DOWNLOAD_LICENCE_ONLY.is_set():
+            return
+
+        if isinstance(track, (Video, Audio)):
+            progress(downloaded="Merging")
+            HLS.merge_segments(
+                segments=sorted(list(save_dir.iterdir())),
+                save_path=save_path
+            )
+            shutil.rmtree(save_dir)
+        else:
+            with open(save_path, "wb") as f:
+                for discontinuity_file in sorted(save_dir.iterdir()):
+                    discontinuity_data = discontinuity_file.read_bytes()
+                    f.write(discontinuity_data)
+                    discontinuity_file.unlink()
+            save_dir.rmdir()
 
         progress(downloaded="Downloaded")
 
         track.path = save_path
-        save_dir.rmdir()
 
     @staticmethod
     def download_segment(
@@ -482,6 +515,37 @@ class HLS:
 
         return download_size
 
+    @staticmethod
+    def merge_segments(segments: list[Path], save_path: Path) -> int:
+        """
+        Concatenate Segments by first demuxing with FFmpeg.
+
+        Returns the file size of the merged file.
+        """
+        ffmpeg = get_binary_path("ffmpeg")
+        if not ffmpeg:
+            raise EnvironmentError("FFmpeg executable was not found but is required to merge HLS segments.")
+
+        demuxer_file = segments[0].parent / "ffmpeg_concat_demuxer.txt"
+        demuxer_file.write_text("\n".join([
+            f"file '{segment}'"
+            for segment in segments
+        ]))
+
+        subprocess.check_call([
+            ffmpeg, "-hide_banner",
+            "-loglevel", "panic",
+            "-f", "concat",
+            "-safe", "0",
+            "-i", demuxer_file,
+            "-map", "0",
+            "-c", "copy",
+            save_path
+        ])
+        demuxer_file.unlink()
+
+        return save_path.stat().st_size
+
     @staticmethod
     def get_drm(
         keys: list[Union[m3u8.model.SessionKey, m3u8.model.Key]],