From 513e3b34cc957c07f0f44552704ec1578946bc2d Mon Sep 17 00:00:00 2001
From: Adam Zombor <azombor@jwplayer.com>
Date: Thu, 8 Aug 2024 12:58:31 +0100
Subject: [PATCH] Add new SetMediaSegments for Playlist Segments overwrite

---
 writer.go      | 11 +++++++++++
 writer_test.go | 28 ++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/writer.go b/writer.go
index 05f6f5d..cca0df3 100644
--- a/writer.go
+++ b/writer.go
@@ -394,6 +394,17 @@ func (p *MediaPlaylist) InsertSegments(segments []*MediaSegment, seqID uint64) e
 	return nil
 }
 
+// SetMediaSegments takes in []*MediaSegment and sets the playlist segments to it
+// This operation does reset playlist cache
+func (p *MediaPlaylist) SetMediaSegments(segments []*MediaSegment) {
+	p.Segments = segments
+
+	p.count = uint(len(segments))
+	p.tail = p.count
+	p.capacity = p.count
+	p.buf.Reset()
+}
+
 // last returns the previously written segment's index
 func (p *MediaPlaylist) last() uint {
 	if p.tail == 0 {
diff --git a/writer_test.go b/writer_test.go
index 3c8ed16..cdf3b5e 100644
--- a/writer_test.go
+++ b/writer_test.go
@@ -1159,6 +1159,34 @@ func TestInsertSegments(t *testing.T) {
 	})
 }
 
+func TestSetMediaSegments(t *testing.T) {
+	p, e := NewMediaPlaylist(3, 4)
+	if e != nil {
+		t.Fatalf("Create media playlist failed: %s", e)
+	}
+	p.Close()
+	if e = p.Append("test01.ts", 5.0, ""); e != nil {
+		t.Errorf("Add 1st segment to a media playlist failed: %s", e)
+	}
+	if e = p.Append("test02.ts", 5.0, ""); e != nil {
+		t.Errorf("Add 2nd segment to a media playlist failed: %s", e)
+	}
+
+	newSegments := []*MediaSegment{
+		{SeqId: 1, URI: "test03.ts"},
+		{SeqId: 2, URI: "test04.ts"},
+	}
+
+	p.SetMediaSegments(newSegments)
+
+	require.Equal(t, 2, len(p.Segments))
+	require.Equal(t, uint(2), p.capacity)
+	require.Equal(t, uint(2), p.tail)
+	require.Equal(t, uint(2), p.count)
+	require.Equal(t, "test03.ts", p.Segments[0].URI)
+	require.Equal(t, "test04.ts", p.Segments[1].URI)
+}
+
 // Create new media playlist
 // Add three segments to media playlist
 // Set gap tag for the 2nd segment.