sample_muxer: added support for WebVTT chapters

Change-Id: Ic5ab8097c0981ef300eadc4a3c151f63b2aad81d
diff --git a/sample_muxer.cpp b/sample_muxer.cpp
index 62659ce..1e84410 100644
--- a/sample_muxer.cpp
+++ b/sample_muxer.cpp
@@ -63,6 +63,8 @@
          "add WebVTT descriptions as metadata track\n");
   printf("  -webvtt-metadata <vttfile>     "
          "add WebVTT subtitles as metadata track\n");
+  printf("  -webvtt-chapters <vttfile>     "
+         "add WebVTT chapters as MKV chapters element\n");
 }
 
 struct MetadataFile {
@@ -98,13 +100,14 @@
     metadata_files_t* metadata_files) {
   int& i = *argv_index;
 
-  enum { kCount = 4 };
+  enum { kCount = 5 };
   struct Arg { const char* name; SampleMuxerMetadata::Kind kind; };
   const Arg args[kCount] = {
     { "-webvtt-subtitles", SampleMuxerMetadata::kSubtitles },
     { "-webvtt-captions", SampleMuxerMetadata::kCaptions },
     { "-webvtt-descriptions", SampleMuxerMetadata::kDescriptions },
-    { "-webvtt-metadata", SampleMuxerMetadata::kMetadata }
+    { "-webvtt-metadata", SampleMuxerMetadata::kMetadata },
+    { "-webvtt-chapters", SampleMuxerMetadata::kChapters }
   };
 
   for (int idx = 0; idx < kCount; ++idx) {
@@ -147,8 +150,8 @@
   uint64 max_cluster_duration = 0;
   uint64 max_cluster_size = 0;
   bool switch_tracks = false;
-  int audio_track_number = 0; // 0 tells muxer to decide.
-  int video_track_number = 0; // 0 tells muxer to decide.
+  int audio_track_number = 0;  // 0 tells muxer to decide.
+  int video_track_number = 0;  // 0 tells muxer to decide.
   bool chunking = false;
   const char* chunk_name = NULL;
 
@@ -291,8 +294,8 @@
   // Set Tracks element attributes
   const mkvparser::Tracks* const parser_tracks = parser_segment->GetTracks();
   unsigned long i = 0;
-  uint64 vid_track = 0; // no track added
-  uint64 aud_track = 0; // no track added
+  uint64 vid_track = 0;  // no track added
+  uint64 aud_track = 0;  // no track added
 
   using mkvparser::Track;
 
@@ -408,6 +411,9 @@
   if (!LoadMetadataFiles(metadata_files, &metadata))
     return EXIT_FAILURE;
 
+  if (!metadata.AddChapters())
+    return EXIT_FAILURE;
+
   // Set Cues element attributes
   mkvmuxer::Cues* const cues = muxer_segment.GetCues();
   cues->set_output_block_number(output_cues_block_number);
@@ -427,8 +433,7 @@
 
     long status = cluster->GetFirst(block_entry);
 
-    if (status)
-    {
+    if (status) {
         printf("\n Could not get first block of cluster.\n");
         return EXIT_FAILURE;
     }
@@ -483,8 +488,7 @@
 
       status = cluster->GetNext(block_entry, block_entry);
 
-      if (status)
-      {
+      if (status) {
           printf("\n Could not get next block of cluster.\n");
           return EXIT_FAILURE;
       }
@@ -498,7 +502,10 @@
   if (!metadata.Write(-1))
     return EXIT_FAILURE;
 
-  muxer_segment.Finalize();
+  if (!muxer_segment.Finalize()) {
+    printf("Finalization of segment failed.\n");
+    return EXIT_FAILURE;
+  }
 
   delete [] data;
   delete parser_segment;
diff --git a/sample_muxer_metadata.cc b/sample_muxer_metadata.cc
index 4066670..c3a2e3f 100644
--- a/sample_muxer_metadata.cc
+++ b/sample_muxer_metadata.cc
@@ -16,6 +16,9 @@
 }
 
 bool SampleMuxerMetadata::Load(const char* file, Kind kind) {
+  if (kind == kChapters)
+    return LoadChapters(file);
+
   mkvmuxer::uint64 track_num;
 
   if (!AddTrack(kind, &track_num)) {
@@ -26,6 +29,21 @@
   return Parse(file, kind, track_num);
 }
 
+bool SampleMuxerMetadata::AddChapters() {
+  typedef cue_list_t::const_iterator iter_t;
+  iter_t i = chapter_cues_.begin();
+  const iter_t j = chapter_cues_.end();
+
+  while (i != j) {
+    const cue_t& chapter = *i++;
+
+    if (!AddChapter(chapter))
+      return false;
+  }
+
+  return true;
+}
+
 bool SampleMuxerMetadata::Write(mkvmuxer::int64 time_ns) {
   typedef cues_set_t::iterator iter_t;
 
@@ -49,6 +67,129 @@
   return true;
 }
 
+bool SampleMuxerMetadata::LoadChapters(const char* file) {
+  if (!chapter_cues_.empty()) {
+    printf("Support for more than one chapters file is not yet implemented\n");
+    return false;
+  }
+
+  cue_list_t cues;
+
+  if (!ParseChapters(file, &cues))
+    return false;
+
+  // TODO(matthewjheaney): support more than one chapters file
+  chapter_cues_.swap(cues);
+
+  return true;
+}
+
+bool SampleMuxerMetadata::ParseChapters(
+    const char* file,
+    cue_list_t* cues_ptr) {
+  cue_list_t& cues = *cues_ptr;
+  cues.clear();
+
+  libwebvtt::VttReader r;
+  int e = r.Open(file);
+
+  if (e) {
+    printf("Unable to open WebVTT file: \"%s\"\n", file);
+    return false;
+  }
+
+  libwebvtt::Parser p(&r);
+  e = p.Init();
+
+  if (e < 0) {  // error
+    printf("Error parsing WebVTT file: \"%s\"\n", file);
+    return false;
+  }
+
+  libwebvtt::Time t;
+  t.hours = -1;
+
+  for (;;) {
+    cue_t c;
+    e = p.Parse(&c);
+
+    if (e < 0) {  // error
+      printf("Error parsing WebVTT file: \"%s\"\n", file);
+      return false;
+    }
+
+    if (e > 0)  // EOF
+      return true;
+
+    if (c.start_time < t) {
+      printf("bad WebVTT cue timestamp (out-of-order)\n");
+      return false;
+    }
+
+    if (c.stop_time < c.start_time) {
+      printf("bad WebVTT cue timestamp (stop < start)\n");
+      return false;
+    }
+
+    t = c.start_time;
+    cues.push_back(c);
+  }
+}
+
+bool SampleMuxerMetadata::AddChapter(const cue_t& cue) {
+  // TODO(matthewjheaney): support language and country
+
+  mkvmuxer::Chapter* const chapter = segment_->AddChapter();
+
+  if (chapter == NULL) {
+    printf("Unable to add chapter\n");
+    return false;
+  }
+
+  if (cue.identifier.empty()) {
+    chapter->set_id(NULL);
+  } else {
+    const char* const id = cue.identifier.c_str();
+    if (!chapter->set_id(id)) {
+      printf("Unable to set chapter id\n");
+      return false;
+    }
+  }
+
+  typedef libwebvtt::presentation_t time_ms_t;
+  const time_ms_t start_time_ms = cue.start_time.presentation();
+  const time_ms_t stop_time_ms = cue.stop_time.presentation();
+
+  enum { kNsPerMs = 1000000 };
+  const mkvmuxer::uint64 start_time_ns = start_time_ms * kNsPerMs;
+  const mkvmuxer::uint64 stop_time_ns = stop_time_ms * kNsPerMs;
+
+  chapter->set_time(*segment_, start_time_ns, stop_time_ns);
+
+  typedef libwebvtt::Cue::payload_t::const_iterator iter_t;
+  iter_t i = cue.payload.begin();
+  const iter_t j = cue.payload.end();
+
+  string title;
+
+  for (;;) {
+    title += *i++;
+
+    if (i == j)
+      break;
+
+    enum { kLF = '\x0A' };
+    title += kLF;
+  }
+
+  if (!chapter->add_string(title.c_str(), NULL, NULL)) {
+    printf("Unable to set chapter title\n");
+    return false;
+  }
+
+  return true;
+}
+
 bool SampleMuxerMetadata::AddTrack(
     Kind kind,
     mkvmuxer::uint64* track_num) {
diff --git a/sample_muxer_metadata.h b/sample_muxer_metadata.h
index a714721..353cf19 100644
--- a/sample_muxer_metadata.h
+++ b/sample_muxer_metadata.h
@@ -21,7 +21,8 @@
     kSubtitles,
     kCaptions,
     kDescriptions,
-    kMetadata
+    kMetadata,
+    kChapters
   };
 
   SampleMuxerMetadata();
@@ -31,10 +32,12 @@
   bool Init(mkvmuxer::Segment* segment);
 
   // Parse the WebVTT file |filename| having the indicated |kind|, and
-  // create a corresponding track in the segment.  Returns false on
-  // error.
+  // create a corresponding track (or chapters element) in the
+  // segment.  Returns false on error.
   bool Load(const char* filename, Kind kind);
 
+  bool AddChapters();
+
   // Write any WebVTT cues whose time is less or equal to |time_ns| as
   // a metadata block in its corresponding track.  If |time_ns| is
   // negative, write all remaining cues. Returns false on error.
@@ -74,6 +77,21 @@
   };
 
   typedef std::multiset<SortableCue> cues_set_t;
+  typedef std::list<cue_t> cue_list_t;
+
+  // Parse the WebVTT cues in the named |file|, returning false on
+  // error.  We handle chapters as a special case, because they are
+  // stored in their own, dedicated level-1 element.
+  bool LoadChapters(const char* file);
+
+  // Parse the WebVTT chapters in |file| to populate |cues|.  Returns
+  // false on error.
+  static bool ParseChapters(const char* file,
+                            cue_list_t* cues);
+
+  // Adds WebVTT cue |chapter| to the chapters element of the output
+  // file's segment element.  Returns false on error.
+  bool AddChapter(const cue_t& chapter);
 
   // Add a metadata track to the segment having the indicated |kind|,
   // returning the |track_num| that has been chosen for this track.
@@ -105,6 +123,10 @@
   // Set of cues ordered by time and then by track number.
   cues_set_t cues_set_;
 
+  // The cues that will be used to populate the Chapters level-1
+  // element of the output file.
+  cue_list_t chapter_cues_;
+
   // Disable copy ctor and copy assign.
   SampleMuxerMetadata(const SampleMuxerMetadata&);
   SampleMuxerMetadata& operator=(const SampleMuxerMetadata&);