segmenter.h

// Copyright 2015 Google Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd

#ifndef MEDIA_FORMATS_WEBM_SEGMENTER_H_
#define MEDIA_FORMATS_WEBM_SEGMENTER_H_

#include <memory>
#include "packager/media/base/status.h"
#include "packager/media/formats/webm/encryptor.h"
#include "packager/media/formats/webm/mkv_writer.h"
#include "packager/media/formats/webm/seek_head.h"
#include "packager/third_party/libwebm/src/mkvmuxer.hpp"

namespace shaka {
namespace media {

struct MuxerOptions;

class AudioStreamInfo;
class KeySource;
class MediaSample;
class StreamInfo;
class MuxerListener;
class ProgressListener;
class StreamInfo;
class VideoStreamInfo;

namespace webm {

class Segmenter {
 public:
  explicit Segmenter(const MuxerOptions& options);
  virtual ~Segmenter();

  Status Initialize(std::unique_ptr<MkvWriter> writer,
                    StreamInfo* info,
                    ProgressListener* progress_listener,
                    MuxerListener* muxer_listener,
                    KeySource* encryption_key_source,
                    uint32_t max_sd_pixels,
                    uint32_t max_hd_pixels,
                    uint32_t max_uhd1_pixels,
                    double clear_lead_in_seconds);

  Status Finalize();

  Status AddSample(std::shared_ptr<MediaSample> sample);

  virtual Status FinalizeSegment(uint64_t start_timescale,
                                 uint64_t duration_timescale,
                                 bool is_subsegment) = 0;

  virtual bool GetInitRangeStartAndEnd(uint64_t* start, uint64_t* end) = 0;

  virtual bool GetIndexRangeStartAndEnd(uint64_t* start, uint64_t* end) = 0;

  float GetDuration() const;

 protected:
  uint64_t FromBMFFTimescale(uint64_t time_timescale);
  uint64_t FromWebMTimecode(uint64_t time_webm_timecode);
  Status WriteSegmentHeader(uint64_t file_size, MkvWriter* writer);
  Status SetCluster(uint64_t start_webm_timecode,
                    uint64_t position,
                    MkvWriter* writer);

  void UpdateProgress(uint64_t progress);
  void set_progress_target(uint64_t target) { progress_target_ = target; }

  const MuxerOptions& options() const { return options_; }
  mkvmuxer::Cluster* cluster() { return cluster_.get(); }
  mkvmuxer::Cues* cues() { return &cues_; }
  MuxerListener* muxer_listener() { return muxer_listener_; }
  StreamInfo* info() { return info_; }
  SeekHead* seek_head() { return &seek_head_; }

  int track_id() const { return track_id_; }
  uint64_t segment_payload_pos() const { return segment_payload_pos_; }

  virtual Status DoInitialize(std::unique_ptr<MkvWriter> writer) = 0;
  virtual Status DoFinalize() = 0;

 private:
  Status CreateVideoTrack(VideoStreamInfo* info);
  Status CreateAudioTrack(AudioStreamInfo* info);
  Status InitializeEncryptor(KeySource* key_source, uint32_t max_sd_pixels,
                             uint32_t max_hd_pixels, uint32_t max_uhd1_pixels);

  // Writes the previous frame to the file.
  Status WriteFrame(bool write_duration);

  // This is called when there needs to be a new (sub)segment.
  // In single-segment mode, a Cluster is a segment and there is no subsegment.
  // In multi-segment mode, a new file is a segment and the clusters in the file
  // are subsegments.
  virtual Status NewSegment(uint64_t start_timescale, bool is_subsegment) = 0;

  // Store the previous sample so we know which one is the last frame.
  std::shared_ptr<MediaSample> prev_sample_;
  // The reference frame timestamp; used to populate the ReferenceBlock element
  // when writing non-keyframe BlockGroups.
  uint64_t reference_frame_timestamp_ = 0;

  const MuxerOptions& options_;
  std::unique_ptr<Encryptor> encryptor_;
  double clear_lead_ = 0;
  // Encryption is enabled only after clear_lead_.
  bool enable_encryption_ = false;

  std::unique_ptr<mkvmuxer::Cluster> cluster_;
  mkvmuxer::Cues cues_;
  SeekHead seek_head_;
  mkvmuxer::SegmentInfo segment_info_;
  mkvmuxer::Tracks tracks_;

  StreamInfo* info_ = nullptr;
  MuxerListener* muxer_listener_ = nullptr;
  ProgressListener* progress_listener_ = nullptr;
  uint64_t progress_target_ = 0;
  uint64_t accumulated_progress_ = 0;
  uint64_t first_timestamp_ = 0;
  int64_t sample_duration_ = 0;
  // The position (in bytes) of the start of the Segment payload in the init
  // file.  This is also the size of the header before the SeekHead.
  uint64_t segment_payload_pos_ = 0;

  // Indicate whether a new segment needed to be created, which is always true
  // in the beginning.
  bool new_segment_ = true;
  // Indicate whether a new subsegment needed to be created.
  bool new_subsegment_ = false;
  int track_id_ = 0;

  DISALLOW_COPY_AND_ASSIGN(Segmenter);
};

}  // namespace webm
}  // namespace media
}  // namespace shaka

#endif  // MEDIA_FORMATS_WEBM_SEGMENTER_H_