From 77a91c1086be84443e72df2535bb159cc35910d3 Mon Sep 17 00:00:00 2001
From: Sergey Petrov <psa@insit.ru>
Date: Mon, 7 Apr 2025 10:10:12 +0500
Subject: [PATCH] Update muxer.

---
 reader/internal/processor/h264-aac_muxer.go | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/reader/internal/processor/h264-aac_muxer.go b/reader/internal/processor/h264-aac_muxer.go
index 175f494..59f9aa2 100644
--- a/reader/internal/processor/h264-aac_muxer.go
+++ b/reader/internal/processor/h264-aac_muxer.go
@@ -31,6 +31,9 @@ type MpegTSMuxer struct {
 	H264Format       *format.H264
 	Mpeg4AudioFormat *format.MPEG4Audio
 
+	idrFrame   [][]byte
+	idrChecker bool
+
 	f               *os.File
 	b               *bufio.Writer
 	w               *mpegts.Writer
@@ -61,6 +64,8 @@ func (e *MpegTSMuxer) Initialize() error {
 
 	e.w = mpegts.NewWriter(e.b, []*mpegts.Track{e.h264Track, e.mpeg4AudioTrack})
 
+	e.idrChecker = false
+
 	return nil
 }
 
@@ -99,6 +104,7 @@ func (e *MpegTSMuxer) WriteH264(pts int64, au [][]byte) error {
 		case h264.NALUTypeIDR:
 			idrPresent = true
 			j++
+			e.idrFrame = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, nalu)
 
 		case h264.NALUTypeNonIDR:
 			nonIDRPresent = true
@@ -111,14 +117,20 @@ func (e *MpegTSMuxer) WriteH264(pts int64, au [][]byte) error {
 
 	au = filteredAU
 
-	if au == nil || (!nonIDRPresent && !idrPresent) {
+	if !e.idrChecker {
+		au = append(e.idrFrame, au...)
+		e.idrChecker = true
+		//mux.frameChecker = true
+	}
+
+	if au == nil || ((!nonIDRPresent && !idrPresent) && e.idrFrame == nil) {
 		return nil
 	}
 
 	// Add SPS and PPS before access unit that contains an IDR.
-	if idrPresent {
-		au = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, au...)
-	}
+	//if idrPresent {
+	//	au = append([][]byte{e.H264Format.SPS, e.H264Format.PPS}, au...)
+	//}
 
 	if e.dtsExtractor == nil {
 		// Skip samples silently until we find one with an IDR.