/*
 * Copyright (C) 2017 The Android Open Source Project
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package androidx.media3.test.utils;

import static com.google.common.truth.Truth.assertThat;

import android.net.Uri;
import android.os.SystemClock;
import androidx.annotation.Nullable;
import androidx.media3.common.C;
import androidx.media3.common.MimeTypes;
import androidx.media3.common.TrackGroup;
import androidx.media3.common.util.Assertions;
import androidx.media3.common.util.NullableType;
import androidx.media3.common.util.UnstableApi;
import androidx.media3.common.util.Util;
import androidx.media3.datasource.DataSpec;
import androidx.media3.datasource.TransferListener;
import androidx.media3.exoplayer.LoadingInfo;
import androidx.media3.exoplayer.SeekParameters;
import androidx.media3.exoplayer.drm.DrmSessionEventListener;
import androidx.media3.exoplayer.drm.DrmSessionManager;
import androidx.media3.exoplayer.source.CompositeSequenceableLoader;
import androidx.media3.exoplayer.source.LoadEventInfo;
import androidx.media3.exoplayer.source.MediaPeriod;
import androidx.media3.exoplayer.source.MediaSourceEventListener;
import androidx.media3.exoplayer.source.SampleStream;
import androidx.media3.exoplayer.source.SequenceableLoader;
import androidx.media3.exoplayer.source.TrackGroupArray;
import androidx.media3.exoplayer.source.chunk.ChunkSampleStream;
import androidx.media3.exoplayer.trackselection.ExoTrackSelection;
import androidx.media3.exoplayer.upstream.Allocator;
import androidx.media3.exoplayer.upstream.DefaultLoadErrorHandlingPolicy;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * Fake {@link MediaPeriod} that provides tracks from the given {@link TrackGroupArray}. Selecting a
 * track will give the player a {@link ChunkSampleStream}.
 */
@UnstableApi
public class FakeAdaptiveMediaPeriod
    implements MediaPeriod, SequenceableLoader.Callback<ChunkSampleStream<FakeChunkSource>> {

  private static final DataSpec FAKE_DATA_SPEC = new DataSpec(Uri.parse("http://fake.test"));

  private final TrackGroupArray trackGroupArray;
  private final MediaSourceEventListener.EventDispatcher mediaSourceEventDispatcher;
  private final long fakePreparationLoadTaskId;
  private final FakeChunkSource.Factory chunkSourceFactory;
  private final Allocator allocator;
  private final long durationUs;
  @Nullable private final TransferListener transferListener;
  private final List<ChunkSampleStream<FakeChunkSource>> sampleStreams;

  @Nullable private Callback callback;
  private boolean prepared;
  private SequenceableLoader sequenceableLoader;

  public FakeAdaptiveMediaPeriod(
      TrackGroupArray trackGroupArray,
      MediaSourceEventListener.EventDispatcher mediaSourceEventDispatcher,
      Allocator allocator,
      FakeChunkSource.Factory chunkSourceFactory,
      long durationUs,
      @Nullable TransferListener transferListener) {
    this.trackGroupArray = trackGroupArray;
    this.mediaSourceEventDispatcher = mediaSourceEventDispatcher;
    this.chunkSourceFactory = chunkSourceFactory;
    this.allocator = allocator;
    this.durationUs = durationUs;
    this.transferListener = transferListener;
    sampleStreams = new ArrayList<>();
    sequenceableLoader =
        new CompositeSequenceableLoader(
            /* loaders= */ ImmutableList.of(), /* loaderTrackTypes= */ ImmutableList.of());
    fakePreparationLoadTaskId = LoadEventInfo.getNewId();
  }

  /** Releases the media period. */
  public void release() {
    prepared = false;
    for (ChunkSampleStream<FakeChunkSource> sampleStream : sampleStreams) {
      sampleStream.release();
    }
    sampleStreams.clear();
    sequenceableLoader =
        new CompositeSequenceableLoader(
            /* loaders= */ ImmutableList.of(), /* loaderTrackTypes= */ ImmutableList.of());
  }

  @Override
  public void prepare(Callback callback, long positionUs) {
    mediaSourceEventDispatcher.loadStarted(
        new LoadEventInfo(fakePreparationLoadTaskId, FAKE_DATA_SPEC, SystemClock.elapsedRealtime()),
        C.DATA_TYPE_MEDIA,
        C.TRACK_TYPE_UNKNOWN,
        /* trackFormat= */ null,
        C.SELECTION_REASON_UNKNOWN,
        /* trackSelectionData= */ null,
        /* mediaStartTimeUs= */ 0,
        /* mediaEndTimeUs= */ C.TIME_UNSET);
    this.callback = callback;
    prepared = true;
    Util.castNonNull(this.callback).onPrepared(this);
    mediaSourceEventDispatcher.loadCompleted(
        new LoadEventInfo(
            fakePreparationLoadTaskId,
            FAKE_DATA_SPEC,
            FAKE_DATA_SPEC.uri,
            /* responseHeaders= */ ImmutableMap.of(),
            SystemClock.elapsedRealtime(),
            /* loadDurationMs= */ 0,
            /* bytesLoaded= */ 100),
        C.DATA_TYPE_MEDIA,
        C.TRACK_TYPE_UNKNOWN,
        /* trackFormat= */ null,
        C.SELECTION_REASON_UNKNOWN,
        /* trackSelectionData= */ null,
        /* mediaStartTimeUs= */ 0,
        /* mediaEndTimeUs= */ C.TIME_UNSET);
  }

  @Override
  public void maybeThrowPrepareError() throws IOException {
    // Do nothing.
  }

  @Override
  public TrackGroupArray getTrackGroups() {
    assertThat(prepared).isTrue();
    return trackGroupArray;
  }

  @SuppressWarnings({"unchecked"}) // Casting sample streams created by this class.
  @Override
  public long selectTracks(
      @NullableType ExoTrackSelection[] selections,
      boolean[] mayRetainStreamFlags,
      @NullableType SampleStream[] streams,
      boolean[] streamResetFlags,
      long positionUs) {
    assertThat(prepared).isTrue();
    int rendererCount = selections.length;
    for (int i = 0; i < rendererCount; i++) {
      if (streams[i] != null && (selections[i] == null || !mayRetainStreamFlags[i])) {
        ((ChunkSampleStream<FakeChunkSource>) streams[i]).release();
        sampleStreams.remove(streams[i]);
        streams[i] = null;
      }
      if (streams[i] == null && selections[i] != null) {
        ExoTrackSelection selection = selections[i];
        assertThat(selection.length()).isAtLeast(1);
        TrackGroup trackGroup = selection.getTrackGroup();
        assertThat(trackGroupArray.indexOf(trackGroup)).isNotEqualTo(C.INDEX_UNSET);
        int indexInTrackGroup = selection.getIndexInTrackGroup(selection.getSelectedIndex());
        assertThat(indexInTrackGroup).isAtLeast(0);
        assertThat(indexInTrackGroup).isLessThan(trackGroup.length);
        FakeChunkSource chunkSource =
            chunkSourceFactory.createChunkSource(selection, durationUs, transferListener);
        ChunkSampleStream<FakeChunkSource> sampleStream =
            new ChunkSampleStream<>(
                MimeTypes.getTrackType(selection.getSelectedFormat().sampleMimeType),
                /* embeddedTrackTypes= */ null,
                /* embeddedTrackFormats= */ null,
                chunkSource,
                /* callback= */ this,
                allocator,
                positionUs,
                DrmSessionManager.DRM_UNSUPPORTED,
                new DrmSessionEventListener.EventDispatcher(),
                new DefaultLoadErrorHandlingPolicy(/* minimumLoadableRetryCount= */ 3),
                mediaSourceEventDispatcher,
                /* canReportInitialDiscontinuity= */ false);
        streams[i] = sampleStream;
        sampleStreams.add(sampleStream);
        streamResetFlags[i] = true;
      }
    }
    sequenceableLoader =
        new CompositeSequenceableLoader(
            sampleStreams,
            Lists.transform(sampleStreams, s -> ImmutableList.of(s.primaryTrackType)));
    return seekToUs(positionUs);
  }

  @Override
  public void discardBuffer(long positionUs, boolean toKeyframe) {
    for (ChunkSampleStream<FakeChunkSource> sampleStream : sampleStreams) {
      sampleStream.discardBuffer(positionUs, toKeyframe);
    }
  }

  @Override
  public void reevaluateBuffer(long positionUs) {
    sequenceableLoader.reevaluateBuffer(positionUs);
  }

  @Override
  public long readDiscontinuity() {
    assertThat(prepared).isTrue();
    return C.TIME_UNSET;
  }

  @Override
  public long getBufferedPositionUs() {
    assertThat(prepared).isTrue();
    return sequenceableLoader.getBufferedPositionUs();
  }

  @Override
  public long seekToUs(long positionUs) {
    assertThat(prepared).isTrue();
    for (ChunkSampleStream<FakeChunkSource> sampleStream : sampleStreams) {
      sampleStream.seekToUs(positionUs);
    }
    return positionUs;
  }

  @Override
  public long getAdjustedSeekPositionUs(long positionUs, SeekParameters seekParameters) {
    return positionUs;
  }

  @Override
  public long getNextLoadPositionUs() {
    assertThat(prepared).isTrue();
    return sequenceableLoader.getNextLoadPositionUs();
  }

  @Override
  public boolean continueLoading(LoadingInfo loadingInfo) {
    sequenceableLoader.continueLoading(loadingInfo);
    return true;
  }

  @Override
  public boolean isLoading() {
    return sequenceableLoader.isLoading();
  }

  @Override
  public void onContinueLoadingRequested(ChunkSampleStream<FakeChunkSource> source) {
    Assertions.checkStateNotNull(callback).onContinueLoadingRequested(this);
  }
}