// Copyright 2014 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

#include <memory>
#include <vector>

#include "base/bind.h"
#include "base/memory/ptr_util.h"
#include "base/message_loop/message_loop.h"
#include "base/run_loop.h"
#include "base/threading/thread_task_runner_handle.h"
#include "chromecast/media/base/decrypt_context_impl.h"
#include "chromecast/media/cdm/cast_cdm_context.h"
#include "chromecast/media/cma/pipeline/av_pipeline_client.h"
#include "chromecast/media/cma/pipeline/media_pipeline_impl.h"
#include "chromecast/media/cma/pipeline/video_pipeline_client.h"
#include "chromecast/media/cma/test/frame_generator_for_test.h"
#include "chromecast/media/cma/test/mock_frame_provider.h"
#include "chromecast/media/cma/test/mock_media_pipeline_backend.h"
#include "chromecast/public/media/cast_decoder_buffer.h"
#include "media/base/audio_decoder_config.h"
#include "media/base/media_util.h"
#include "media/base/video_decoder_config.h"
#include "media/cdm/player_tracker_impl.h"
#include "testing/gtest/include/gtest/gtest.h"

using testing::_;
using testing::Invoke;
using testing::NiceMock;
using testing::Return;
using testing::SaveArg;

namespace {
// Total number of frames generated by CodedFrameProvider.
// The first frame has config, while the last one is EOS.
const int kNumFrames = 100;
const int kFrameSize = 512;
const int kFrameDurationUs = 40 * 1000;
const int kLastFrameTimestamp = (kNumFrames - 2) * kFrameDurationUs;
}  // namespace

namespace chromecast {
namespace media {

ACTION_P2(PushBuffer, delegate, buffer_pts) {
  if (arg0->end_of_stream()) {
    base::ThreadTaskRunnerHandle::Get()->PostTask(
        FROM_HERE,
        base::BindOnce(&MediaPipelineBackend::Decoder::Delegate::OnEndOfStream,
                       base::Unretained(*delegate)));
  } else {
    *buffer_pts = arg0->timestamp();
  }
  return MediaPipelineBackend::kBufferSuccess;
}

class CastCdmContextForTest : public CastCdmContext {
 public:
  CastCdmContextForTest() : license_installed_(false) {}
  void SetLicenseInstalled() {
    license_installed_ = true;
    player_tracker_.NotifyNewKey();
  }

  // CastCdmContext implementation:
  int RegisterPlayer(const base::Closure& new_key_cb,
                     const base::Closure& cdm_unset_cb) override {
    return player_tracker_.RegisterPlayer(new_key_cb, cdm_unset_cb);
  }

  void UnregisterPlayer(int registration_id) override {
    return player_tracker_.UnregisterPlayer(registration_id);
  }

  std::unique_ptr<DecryptContextImpl> GetDecryptContext(
      const std::string& key_id,
      const EncryptionScheme& encryption_scheme) override {
    if (license_installed_) {
      return std::unique_ptr<DecryptContextImpl>(
          new DecryptContextImpl(KEY_SYSTEM_CLEAR_KEY));
    } else {
      return std::unique_ptr<DecryptContextImpl>();
    }
  }

  void SetKeyStatus(const std::string& key_id,
                    CastKeyStatus key_status,
                    uint32_t system_code) override {}

  void SetVideoResolution(int width, int height) override {}

 private:
  bool license_installed_;
  base::Closure new_key_cb_;
  ::media::PlayerTrackerImpl player_tracker_;

  DISALLOW_COPY_AND_ASSIGN(CastCdmContextForTest);
};

// Helper class for managing pipeline setup, teardown, feeding data, stop/start
// etc in a simple API for tests to use.
class PipelineHelper {
 public:
  enum Stream { STREAM_AUDIO, STREAM_VIDEO };

  PipelineHelper(bool audio, bool video, bool encrypted)
      : have_audio_(audio),
        have_video_(video),
        encrypted_(encrypted),
        pipeline_backend_(nullptr),
        audio_decoder_delegate_(nullptr),
        video_decoder_delegate_(nullptr) {}

  void Setup() {
    if (encrypted_) {
      cdm_context_.reset(new CastCdmContextForTest());
    }

    pipeline_backend_ = new MockMediaPipelineBackend();
    ON_CALL(*pipeline_backend_, SetPlaybackRate(_)).WillByDefault(Return(true));
    ON_CALL(audio_decoder_, SetConfig(_)).WillByDefault(Return(true));
    ON_CALL(audio_decoder_, PushBuffer(_))
        .WillByDefault(PushBuffer(&audio_decoder_delegate_,
                                  &last_push_pts_[STREAM_AUDIO]));
    ON_CALL(video_decoder_, SetConfig(_)).WillByDefault(Return(true));
    ON_CALL(video_decoder_, PushBuffer(_))
        .WillByDefault(PushBuffer(&video_decoder_delegate_,
                                  &last_push_pts_[STREAM_VIDEO]));

    media_pipeline_ = std::make_unique<MediaPipelineImpl>();
    media_pipeline_->Initialize(kLoadTypeURL,
                                base::WrapUnique(pipeline_backend_));

    if (have_audio_) {
      ::media::AudioDecoderConfig audio_config(
          ::media::kCodecMP3, ::media::kSampleFormatS16,
          ::media::CHANNEL_LAYOUT_STEREO, 44100, ::media::EmptyExtraData(),
          ::media::Unencrypted());
      AvPipelineClient client;
      client.eos_cb = base::Bind(&PipelineHelper::OnEos, base::Unretained(this),
                                 STREAM_AUDIO);
      EXPECT_CALL(*pipeline_backend_, CreateAudioDecoder())
          .Times(1)
          .WillOnce(Return(&audio_decoder_));
      EXPECT_CALL(audio_decoder_, SetDelegate(_))
          .Times(1)
          .WillOnce(SaveArg<0>(&audio_decoder_delegate_));
      ::media::PipelineStatus status = media_pipeline_->InitializeAudio(
          audio_config, client, CreateFrameProvider());
      ASSERT_EQ(::media::PIPELINE_OK, status);
    }
    if (have_video_) {
      std::vector<::media::VideoDecoderConfig> video_configs;
      video_configs.push_back(::media::VideoDecoderConfig(
          ::media::kCodecH264, ::media::H264PROFILE_MAIN,
          ::media::PIXEL_FORMAT_I420, ::media::COLOR_SPACE_UNSPECIFIED,
          ::media::VIDEO_ROTATION_0, gfx::Size(640, 480),
          gfx::Rect(0, 0, 640, 480), gfx::Size(640, 480),
          ::media::EmptyExtraData(), ::media::EncryptionScheme()));
      VideoPipelineClient client;
      client.av_pipeline_client.eos_cb = base::Bind(
          &PipelineHelper::OnEos, base::Unretained(this), STREAM_VIDEO);
      EXPECT_CALL(*pipeline_backend_, CreateVideoDecoder())
          .Times(1)
          .WillOnce(Return(&video_decoder_));
      EXPECT_CALL(video_decoder_, SetDelegate(_))
          .Times(1)
          .WillOnce(SaveArg<0>(&video_decoder_delegate_));
      ::media::PipelineStatus status = media_pipeline_->InitializeVideo(
          video_configs, client, CreateFrameProvider());
      ASSERT_EQ(::media::PIPELINE_OK, status);
    }
  }

  void Start(const base::Closure& eos_cb) {
    eos_cb_ = eos_cb;
    eos_[STREAM_AUDIO] = !media_pipeline_->HasAudio();
    eos_[STREAM_VIDEO] = !media_pipeline_->HasVideo();
    last_push_pts_[STREAM_AUDIO] = std::numeric_limits<int64_t>::min();
    last_push_pts_[STREAM_VIDEO] = std::numeric_limits<int64_t>::min();
    int64_t start_pts = 0;

    EXPECT_CALL(*pipeline_backend_, Initialize())
        .Times(1)
        .WillOnce(Return(true));
    EXPECT_CALL(*pipeline_backend_, Start(start_pts))
        .Times(1)
        .WillOnce(Return(true));

    media_pipeline_->StartPlayingFrom(
        base::TimeDelta::FromMilliseconds(start_pts));
    media_pipeline_->SetPlaybackRate(1.0f);
  }
  void SetCdm() { media_pipeline_->SetCdm(cdm_context_.get()); }
  void Flush(const base::Closure& flush_cb) {
    EXPECT_CALL(*pipeline_backend_, Stop()).Times(1);
    media_pipeline_->Flush(flush_cb);
  }
  void Stop() {
    media_pipeline_.reset();
    base::RunLoop::QuitCurrentWhenIdleDeprecated();
  }
  void SetCdmLicenseInstalled() { cdm_context_->SetLicenseInstalled(); }

  bool have_audio() const { return have_audio_; }
  bool have_video() const { return have_video_; }
  int64_t last_push_pts(Stream stream) const { return last_push_pts_[stream]; }

 private:
  std::unique_ptr<CodedFrameProvider> CreateFrameProvider() {
    std::vector<FrameGeneratorForTest::FrameSpec> frame_specs;
    frame_specs.resize(kNumFrames);
    for (size_t k = 0; k < frame_specs.size() - 1; k++) {
      frame_specs[k].has_config = (k == 0);
      frame_specs[k].timestamp =
          base::TimeDelta::FromMicroseconds(kFrameDurationUs) * k;
      frame_specs[k].size = kFrameSize;
      frame_specs[k].has_decrypt_config = encrypted_;
    }
    frame_specs.back().is_eos = true;

    std::unique_ptr<FrameGeneratorForTest> frame_generator(
        new FrameGeneratorForTest(frame_specs));
    bool provider_delayed_pattern[] = {false, true};
    std::unique_ptr<MockFrameProvider> frame_provider(new MockFrameProvider());
    frame_provider->Configure(
        std::vector<bool>(
            provider_delayed_pattern,
            provider_delayed_pattern + arraysize(provider_delayed_pattern)),
        std::move(frame_generator));
    frame_provider->SetDelayFlush(true);
    return std::move(frame_provider);
  }

  void OnEos(Stream stream) {
    eos_[stream] = true;
    if (eos_[STREAM_AUDIO] && eos_[STREAM_VIDEO] && !eos_cb_.is_null())
      eos_cb_.Run();
  }

  bool have_audio_;
  bool have_video_;
  bool encrypted_;
  bool eos_[2];
  int64_t last_push_pts_[2];
  base::Closure eos_cb_;
  std::unique_ptr<CastCdmContextForTest> cdm_context_;
  MockMediaPipelineBackend* pipeline_backend_;
  NiceMock<MockAudioDecoder> audio_decoder_;
  NiceMock<MockVideoDecoder> video_decoder_;
  MediaPipelineBackend::Decoder::Delegate* audio_decoder_delegate_;
  MediaPipelineBackend::Decoder::Delegate* video_decoder_delegate_;
  std::unique_ptr<MediaPipelineImpl> media_pipeline_;

  DISALLOW_COPY_AND_ASSIGN(PipelineHelper);
};

using AudioVideoTuple = ::testing::tuple<bool, bool>;

class AudioVideoPipelineImplTest
    : public ::testing::TestWithParam<AudioVideoTuple> {
 public:
  AudioVideoPipelineImplTest() {}

 protected:
  void SetUp() override {
    pipeline_helper_.reset(new PipelineHelper(
        ::testing::get<0>(GetParam()), ::testing::get<1>(GetParam()), false));
    pipeline_helper_->Setup();
  }

  base::MessageLoop message_loop_;
  std::unique_ptr<PipelineHelper> pipeline_helper_;

  DISALLOW_COPY_AND_ASSIGN(AudioVideoPipelineImplTest);
};

static void VerifyPlay(PipelineHelper* pipeline_helper) {
  // The decoders must have received the last frame.
  if (pipeline_helper->have_audio())
    EXPECT_EQ(kLastFrameTimestamp,
              pipeline_helper->last_push_pts(PipelineHelper::STREAM_AUDIO));
  if (pipeline_helper->have_video())
    EXPECT_EQ(kLastFrameTimestamp,
              pipeline_helper->last_push_pts(PipelineHelper::STREAM_VIDEO));

  pipeline_helper->Stop();
}

TEST_P(AudioVideoPipelineImplTest, Play) {
  base::Closure verify_task =
      base::Bind(&VerifyPlay, base::Unretained(pipeline_helper_.get()));
  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Start,
                     base::Unretained(pipeline_helper_.get()), verify_task));
  base::RunLoop().Run();
}

static void VerifyFlush(PipelineHelper* pipeline_helper) {
  // The decoders must not have received any frame.
  if (pipeline_helper->have_audio())
    EXPECT_LT(pipeline_helper->last_push_pts(PipelineHelper::STREAM_AUDIO), 0);
  if (pipeline_helper->have_video())
    EXPECT_LT(pipeline_helper->last_push_pts(PipelineHelper::STREAM_VIDEO), 0);

  pipeline_helper->Stop();
}

static void VerifyNotReached() {
  EXPECT_TRUE(false);
}

TEST_P(AudioVideoPipelineImplTest, Flush) {
  base::Closure verify_task =
      base::Bind(&VerifyFlush, base::Unretained(pipeline_helper_.get()));
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::Start,
                                base::Unretained(pipeline_helper_.get()),
                                base::Bind(&VerifyNotReached)));
  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Flush,
                     base::Unretained(pipeline_helper_.get()), verify_task));

  base::RunLoop().Run();
}

TEST_P(AudioVideoPipelineImplTest, FullCycle) {
  base::Closure stop_task = base::Bind(
      &PipelineHelper::Stop, base::Unretained(pipeline_helper_.get()));
  base::Closure eos_cb =
      base::Bind(&PipelineHelper::Flush,
                 base::Unretained(pipeline_helper_.get()), stop_task);

  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Start,
                     base::Unretained(pipeline_helper_.get()), eos_cb));
  base::RunLoop().Run();
};

// Test all three types of pipeline: audio-only, video-only, audio-video.
INSTANTIATE_TEST_CASE_P(
    MediaPipelineImplTests,
    AudioVideoPipelineImplTest,
    ::testing::Values(AudioVideoTuple(true, false),   // Audio only.
                      AudioVideoTuple(false, true),   // Video only.
                      AudioVideoTuple(true, true)));  // Audio and Video.

// These tests verify that the pipeline handles encrypted media playback
// events (in particular, CDM and license installation) correctly.
class EncryptedAVPipelineImplTest : public ::testing::Test {
 public:
  EncryptedAVPipelineImplTest() {}

 protected:
  void SetUp() override {
    pipeline_helper_.reset(new PipelineHelper(true, true, true));
    pipeline_helper_->Setup();
  }

  base::MessageLoop message_loop_;
  std::unique_ptr<PipelineHelper> pipeline_helper_;

  DISALLOW_COPY_AND_ASSIGN(EncryptedAVPipelineImplTest);
};

// Sets a CDM with license already installed before starting the pipeline.
TEST_F(EncryptedAVPipelineImplTest, SetCdmWithLicenseBeforeStart) {
  base::Closure verify_task =
      base::Bind(&VerifyPlay, base::Unretained(pipeline_helper_.get()));
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Start,
                     base::Unretained(pipeline_helper_.get()), verify_task));
  base::RunLoop().Run();
}

// Start the pipeline, then set a CDM with existing license.
TEST_F(EncryptedAVPipelineImplTest, SetCdmWithLicenseAfterStart) {
  base::Closure verify_task =
      base::Bind(&VerifyPlay, base::Unretained(pipeline_helper_.get()));
  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Start,
                     base::Unretained(pipeline_helper_.get()), verify_task));

  base::RunLoop().RunUntilIdle();
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));
  base::RunLoop().Run();
}

// Start the pipeline, set a CDM, and then install the license.
TEST_F(EncryptedAVPipelineImplTest, SetCdmAndInstallLicenseAfterStart) {
  base::Closure verify_task =
      base::Bind(&VerifyPlay, base::Unretained(pipeline_helper_.get()));
  message_loop_.task_runner()->PostTask(
      FROM_HERE,
      base::BindOnce(&PipelineHelper::Start,
                     base::Unretained(pipeline_helper_.get()), verify_task));
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdm,
                                base::Unretained(pipeline_helper_.get())));

  base::RunLoop().RunUntilIdle();
  message_loop_.task_runner()->PostTask(
      FROM_HERE, base::BindOnce(&PipelineHelper::SetCdmLicenseInstalled,
                                base::Unretained(pipeline_helper_.get())));
  base::RunLoop().Run();
}

}  // namespace media
}  // namespace chromecast