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 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_; }
  uint64_t cluster_length_in_time_scale() const {
    return cluster_length_in_time_scale_;
  }

  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 subsegment.  This does nothing
  // in single-segment mode.  In multi-segment mode this creates a new Cluster
  // element.
  virtual Status NewSubsegment(uint64_t start_timescale) = 0;
  // This is called when there needs to be a new segment.  In single-segment
  // mode, this creates a new Cluster element.  In multi-segment mode this
  // creates a new output file.
  virtual Status NewSegment(uint64_t start_timescale) = 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_;

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

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

  StreamInfo* info_;
  MuxerListener* muxer_listener_;
  ProgressListener* progress_listener_;
  uint64_t progress_target_;
  uint64_t accumulated_progress_;
  uint64_t first_timestamp_;
  int64_t sample_duration_;
  // 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_;

  // Durations in timescale.
  uint64_t cluster_length_in_time_scale_;
  uint64_t segment_length_in_time_scale_;

  int track_id_;

  DISALLOW_COPY_AND_ASSIGN(Segmenter);
};

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

#endif  // MEDIA_FORMATS_WEBM_SEGMENTER_H_