From 2c5a2ccd3fa5b8381faa3f4eccb687fdcb07a712 Mon Sep 17 00:00:00 2001
From: Nhan Trong Nguyen <49386581+NhanNguyen700@users.noreply.github.com>
Date: Fri, 3 Nov 2023 16:34:15 +0700
Subject: [PATCH] Simple support for WebVTT style (#95)

* Simple support for WebVTT style

* fixup! Simple support for WebVTT style

* [webvtt/test] Add cases of STYLES

* [webvtt] Use long condition instead of function

---------

Co-authored-by: Nhan Nguyen <nhan.trong.nguyen@edgeware.tv>
---
 subtitles.go             |  1 +
 testdata/example-in.vtt  | 14 ++++++++++++++
 testdata/example-out.vtt | 15 +++++++++++++++
 webvtt.go                | 35 ++++++++++++++++++++++++++++++++---
 4 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/subtitles.go b/subtitles.go
index 3225979..7219e8f 100644
--- a/subtitles.go
+++ b/subtitles.go
@@ -243,6 +243,7 @@ type StyleAttributes struct {
 	WebVTTRegionAnchor   string
 	WebVTTScroll         string
 	WebVTTSize           string
+	WebVTTStyles         []string
 	WebVTTVertical       string
 	WebVTTViewportAnchor string
 	WebVTTWidth          string
diff --git a/testdata/example-in.vtt b/testdata/example-in.vtt
index e4e0d6a..9a95149 100644
--- a/testdata/example-in.vtt
+++ b/testdata/example-in.vtt
@@ -7,6 +7,20 @@ STYLE
 ::cue(b) {
   color: peachpuff;
 }
+::cue(c) {
+  color: white;
+}
+
+STYLE
+::cue(a) {
+  color: red;
+}
+::cue(d) {
+  color: red;
+
+  background-image: linear-gradient(to bottom, dimgray, lightgray);
+}
+
 
 Region: id=fred width=40% lines=3 regionanchor=0%,100% viewportanchor=10%,90% scroll=up
 Region: id=bill width=40% lines=3 regionanchor=100%,100% viewportanchor=90%,90% scroll=up
diff --git a/testdata/example-out.vtt b/testdata/example-out.vtt
index 98cd730..0eb8974 100644
--- a/testdata/example-out.vtt
+++ b/testdata/example-out.vtt
@@ -1,5 +1,20 @@
 WEBVTT
 
+STYLE
+::cue(b) {
+color: peachpuff;
+}
+::cue(c) {
+color: white;
+}
+::cue(a) {
+color: red;
+}
+::cue(d) {
+color: red;
+background-image: linear-gradient(to bottom, dimgray, lightgray);
+}
+
 Region: id=bill lines=3 regionanchor=100%,100% scroll=up viewportanchor=90%,90% width=40%
 Region: id=fred lines=3 regionanchor=0%,100% scroll=up viewportanchor=10%,90% width=40%
 
diff --git a/webvtt.go b/webvtt.go
index b153469..11133c3 100644
--- a/webvtt.go
+++ b/webvtt.go
@@ -23,6 +23,7 @@ const (
 	webvttBlockNameRegion         = "region"
 	webvttBlockNameStyle          = "style"
 	webvttBlockNameText           = "text"
+	webvttDefaultStyleID          = "astisub-webvtt-default-style-id"
 	webvttTimeBoundariesSeparator = " --> "
 	webvttTimestampMap            = "X-TIMESTAMP-MAP"
 )
@@ -109,6 +110,7 @@ func ReadFromWebVTT(i io.Reader) (o *Subtitles, err error) {
 	var comments []string
 	var index int
 	var timeOffset time.Duration
+	var webVTTStyles *StyleAttributes
 
 	for scanner.Scan() {
 		// Fetch line
@@ -122,8 +124,14 @@ func ReadFromWebVTT(i io.Reader) (o *Subtitles, err error) {
 			comments = append(comments, strings.TrimPrefix(line, "NOTE "))
 		// Empty line
 		case len(line) == 0:
-			// Reset block name
-			blockName = ""
+			// Reset block name, if we are not in the middle of CSS.
+			// If we are in STYLE block and the CSS is empty or we meet the right brace at the end of last line,
+			// then we are not in CSS and can switch to parse next WebVTT block.
+			if blockName != webvttBlockNameStyle || webVTTStyles == nil ||
+				len(webVTTStyles.WebVTTStyles) == 0 ||
+				strings.HasSuffix(webVTTStyles.WebVTTStyles[len(webVTTStyles.WebVTTStyles)-1], "}") {
+				blockName = ""
+			}
 		// Region
 		case strings.HasPrefix(line, "Region: "):
 			// Add region styles
@@ -162,6 +170,15 @@ func ReadFromWebVTT(i io.Reader) (o *Subtitles, err error) {
 		// Style
 		case strings.HasPrefix(line, "STYLE"):
 			blockName = webvttBlockNameStyle
+
+			if _, ok := o.Styles[webvttDefaultStyleID]; !ok {
+				webVTTStyles = &StyleAttributes{}
+				o.Styles[webvttDefaultStyleID] = &Style{
+					InlineStyle: webVTTStyles,
+					ID:          webvttDefaultStyleID,
+				}
+			}
+
 		// Time boundaries
 		case strings.Contains(line, webvttTimeBoundariesSeparator):
 			// Set block name
@@ -257,7 +274,7 @@ func ReadFromWebVTT(i io.Reader) (o *Subtitles, err error) {
 			case webvttBlockNameComment:
 				comments = append(comments, line)
 			case webvttBlockNameStyle:
-				// TODO Do something with the style
+				webVTTStyles.WebVTTStyles = append(webVTTStyles.WebVTTStyles, line)
 			case webvttBlockNameText:
 				// Parse line
 				if l := parseTextWebVTT(line); len(l.Items) > 0 {
@@ -360,11 +377,23 @@ func (s Subtitles) WriteToWebVTT(o io.Writer) (err error) {
 	var c []byte
 	c = append(c, []byte("WEBVTT\n\n")...)
 
+	var style []string
+	for _, s := range s.Styles {
+		if s.InlineStyle != nil {
+			style = append(style, s.InlineStyle.WebVTTStyles...)
+		}
+	}
+
+	if len(style) > 0 {
+		c = append(c, []byte(fmt.Sprintf("STYLE\n%s\n\n", strings.Join(style, "\n")))...)
+	}
+
 	// Add regions
 	var k []string
 	for _, region := range s.Regions {
 		k = append(k, region.ID)
 	}
+
 	sort.Strings(k)
 	for _, id := range k {
 		c = append(c, []byte("Region: id="+s.Regions[id].ID)...)