/*
 * Copyright (C) 2014 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 com.google.android.exoplayer.chunk;

import com.google.android.exoplayer.C;
import com.google.android.exoplayer.LoadControl;
import com.google.android.exoplayer.MediaFormat;
import com.google.android.exoplayer.MediaFormatHolder;
import com.google.android.exoplayer.SampleHolder;
import com.google.android.exoplayer.SampleSource;
import com.google.android.exoplayer.TrackInfo;
import com.google.android.exoplayer.TrackRenderer;
import com.google.android.exoplayer.extractor.DefaultTrackOutput;
import com.google.android.exoplayer.upstream.Loader;
import com.google.android.exoplayer.upstream.Loader.Loadable;
import com.google.android.exoplayer.util.Assertions;

import android.os.Handler;
import android.os.SystemClock;

import java.io.IOException;
import java.util.Collections;
import java.util.LinkedList;
import java.util.List;

/**
 * A {@link SampleSource} that loads media in {@link Chunk}s, which are themselves obtained from a
 * {@link ChunkSource}.
 */
public class ChunkSampleSource implements SampleSource, Loader.Callback {

  /**
   * Interface definition for a callback to be notified of {@link ChunkSampleSource} events.
   */
  public interface EventListener extends BaseChunkSampleSourceEventListener {}

  /**
   * The default minimum number of times to retry loading data prior to failing.
   */
  public static final int DEFAULT_MIN_LOADABLE_RETRY_COUNT = 3;

  private static final int STATE_UNPREPARED = 0;
  private static final int STATE_PREPARED = 1;
  private static final int STATE_ENABLED = 2;

  private static final int NO_RESET_PENDING = -1;

  private final int eventSourceId;
  private final LoadControl loadControl;
  private final ChunkSource chunkSource;
  private final ChunkOperationHolder currentLoadableHolder;
  private final LinkedList<BaseMediaChunk> mediaChunks;
  private final List<BaseMediaChunk> readOnlyMediaChunks;
  private final DefaultTrackOutput sampleQueue;
  private final int bufferSizeContribution;
  private final boolean frameAccurateSeeking;
  private final Handler eventHandler;
  private final EventListener eventListener;
  private final int minLoadableRetryCount;

  private int state;
  private long downstreamPositionUs;
  private long lastSeekPositionUs;
  private long pendingResetPositionUs;
  private long lastPerformedBufferOperation;
  private boolean pendingDiscontinuity;

  private Loader loader;
  private boolean loadingFinished;
  private IOException currentLoadableException;
  private boolean currentLoadableExceptionFatal;
  private int currentLoadableExceptionCount;
  private long currentLoadableExceptionTimestamp;
  private long currentLoadStartTimeMs;

  private MediaFormat downstreamMediaFormat;
  private Format downstreamFormat;

  public ChunkSampleSource(ChunkSource chunkSource, LoadControl loadControl,
      int bufferSizeContribution, boolean frameAccurateSeeking) {
    this(chunkSource, loadControl, bufferSizeContribution, frameAccurateSeeking, null, null, 0);
  }

  public ChunkSampleSource(ChunkSource chunkSource, LoadControl loadControl,
      int bufferSizeContribution, boolean frameAccurateSeeking, Handler eventHandler,
      EventListener eventListener, int eventSourceId) {
    this(chunkSource, loadControl, bufferSizeContribution, frameAccurateSeeking, eventHandler,
        eventListener, eventSourceId, DEFAULT_MIN_LOADABLE_RETRY_COUNT);
  }

  public ChunkSampleSource(ChunkSource chunkSource, LoadControl loadControl,
      int bufferSizeContribution, boolean frameAccurateSeeking, Handler eventHandler,
      EventListener eventListener, int eventSourceId, int minLoadableRetryCount) {
    this.chunkSource = chunkSource;
    this.loadControl = loadControl;
    this.bufferSizeContribution = bufferSizeContribution;
    this.frameAccurateSeeking = frameAccurateSeeking;
    this.eventHandler = eventHandler;
    this.eventListener = eventListener;
    this.eventSourceId = eventSourceId;
    this.minLoadableRetryCount = minLoadableRetryCount;
    currentLoadableHolder = new ChunkOperationHolder();
    mediaChunks = new LinkedList<BaseMediaChunk>();
    readOnlyMediaChunks = Collections.unmodifiableList(mediaChunks);
    sampleQueue = new DefaultTrackOutput(loadControl.getAllocator());
    state = STATE_UNPREPARED;
    pendingResetPositionUs = NO_RESET_PENDING;
  }

  @Override
  public boolean prepare(long positionUs) {
    Assertions.checkState(state == STATE_UNPREPARED);
    loader = new Loader("Loader:" + chunkSource.getTrackInfo().mimeType);
    state = STATE_PREPARED;
    return true;
  }

  @Override
  public int getTrackCount() {
    Assertions.checkState(state != STATE_UNPREPARED);
    return 1;
  }

  @Override
  public TrackInfo getTrackInfo(int track) {
    Assertions.checkState(state != STATE_UNPREPARED);
    Assertions.checkState(track == 0);
    return chunkSource.getTrackInfo();
  }

  @Override
  public void enable(int track, long positionUs) {
    Assertions.checkState(state == STATE_PREPARED);
    Assertions.checkState(track == 0);
    state = STATE_ENABLED;
    chunkSource.enable();
    loadControl.register(this, bufferSizeContribution);
    downstreamFormat = null;
    downstreamMediaFormat = null;
    downstreamPositionUs = positionUs;
    lastSeekPositionUs = positionUs;
    restartFrom(positionUs);
  }

  @Override
  public void disable(int track) {
    Assertions.checkState(state == STATE_ENABLED);
    Assertions.checkState(track == 0);
    pendingDiscontinuity = false;
    state = STATE_PREPARED;
    try {
      chunkSource.disable(mediaChunks);
    } finally {
      loadControl.unregister(this);
      if (loader.isLoading()) {
        loader.cancelLoading();
      } else {
        sampleQueue.clear();
        mediaChunks.clear();
        clearCurrentLoadable();
        loadControl.trimAllocator();
      }
    }
  }

  @Override
  public boolean continueBuffering(long positionUs) throws IOException {
    Assertions.checkState(state == STATE_ENABLED);
    downstreamPositionUs = positionUs;
    chunkSource.continueBuffering(positionUs);
    updateLoadControl();

    boolean haveSamples = !sampleQueue.isEmpty();
    if (!haveSamples) {
      maybeThrowLoadableException();
    }
    return loadingFinished || haveSamples;
  }

  @Override
  public int readData(int track, long positionUs, MediaFormatHolder formatHolder,
      SampleHolder sampleHolder, boolean onlyReadDiscontinuity) throws IOException {
    Assertions.checkState(state == STATE_ENABLED);
    Assertions.checkState(track == 0);
    downstreamPositionUs = positionUs;

    if (pendingDiscontinuity) {
      pendingDiscontinuity = false;
      return DISCONTINUITY_READ;
    }

    if (onlyReadDiscontinuity) {
      return NOTHING_READ;
    }

    if (isPendingReset()) {
      maybeThrowLoadableException();
      return NOTHING_READ;
    }

    boolean haveSamples = !sampleQueue.isEmpty();
    BaseMediaChunk currentChunk = mediaChunks.getFirst();
    while (haveSamples && mediaChunks.size() > 1
        && mediaChunks.get(1).getFirstSampleIndex() == sampleQueue.getReadIndex()) {
      mediaChunks.removeFirst();
      currentChunk = mediaChunks.getFirst();
    }

    if (downstreamFormat == null || !downstreamFormat.equals(currentChunk.format)) {
      notifyDownstreamFormatChanged(currentChunk.format, currentChunk.trigger,
          currentChunk.startTimeUs);
      downstreamFormat = currentChunk.format;
    }

    if (haveSamples || currentChunk.isFormatFinal) {
      MediaFormat mediaFormat = currentChunk.getMediaFormat();
      if (mediaFormat != null && !mediaFormat.equals(downstreamMediaFormat, true)) {
        chunkSource.getMaxVideoDimensions(mediaFormat);
        formatHolder.format = mediaFormat;
        formatHolder.drmInitData = currentChunk.getDrmInitData();
        downstreamMediaFormat = mediaFormat;
        return FORMAT_READ;
      }
    }

    if (!haveSamples) {
      if (loadingFinished) {
        return END_OF_STREAM;
      }
      maybeThrowLoadableException();
      return NOTHING_READ;
    }

    if (sampleQueue.getSample(sampleHolder)) {
      boolean decodeOnly = frameAccurateSeeking && sampleHolder.timeUs < lastSeekPositionUs;
      sampleHolder.flags |= decodeOnly ? C.SAMPLE_FLAG_DECODE_ONLY : 0;
      onSampleRead(currentChunk, sampleHolder);
      return SAMPLE_READ;
    }

    maybeThrowLoadableException();
    return NOTHING_READ;
  }

  @Override
  public void seekToUs(long positionUs) {
    Assertions.checkState(state == STATE_ENABLED);
    lastSeekPositionUs = positionUs;
    if ((isPendingReset() ? pendingResetPositionUs : downstreamPositionUs) == positionUs) {
      return;
    }

    downstreamPositionUs = positionUs;
    // If we're not pending a reset, see if we can seek within the sample queue.
    boolean seekInsideBuffer = !isPendingReset() && sampleQueue.skipToKeyframeBefore(positionUs);
    if (seekInsideBuffer) {
      // We succeeded. All we need to do is discard any chunks that we've moved past.
      boolean haveSamples = !sampleQueue.isEmpty();
      while (haveSamples && mediaChunks.size() > 1
          && mediaChunks.get(1).getFirstSampleIndex() <= sampleQueue.getReadIndex()) {
        mediaChunks.removeFirst();
      }
    } else {
      // We failed, and need to restart.
      restartFrom(positionUs);
    }
    // Either way, we need to send a discontinuity to the downstream components.
    pendingDiscontinuity = true;
  }

  private void maybeThrowLoadableException() throws IOException {
    if (currentLoadableException != null && (currentLoadableExceptionFatal
        || currentLoadableExceptionCount > minLoadableRetryCount)) {
      throw currentLoadableException;
    }
    if (sampleQueue.isEmpty() && currentLoadableHolder.chunk == null) {
      IOException chunkSourceException = chunkSource.getError();
      if (chunkSourceException != null) {
        throw chunkSourceException;
      }
    }
  }

  @Override
  public long getBufferedPositionUs() {
    Assertions.checkState(state == STATE_ENABLED);
    if (isPendingReset()) {
      return pendingResetPositionUs;
    } else if (loadingFinished) {
      return TrackRenderer.END_OF_TRACK_US;
    } else {
      long largestParsedTimestampUs = sampleQueue.getLargestParsedTimestampUs();
      return largestParsedTimestampUs == Long.MIN_VALUE ? downstreamPositionUs
          : largestParsedTimestampUs;
    }
  }

  @Override
  public void release() {
    Assertions.checkState(state != STATE_ENABLED);
    if (loader != null) {
      loader.release();
      loader = null;
    }
    state = STATE_UNPREPARED;
  }

  @Override
  public void onLoadCompleted(Loadable loadable) {
    long now = SystemClock.elapsedRealtime();
    long loadDurationMs = now - currentLoadStartTimeMs;
    Chunk currentLoadable = currentLoadableHolder.chunk;
    chunkSource.onChunkLoadCompleted(currentLoadable);
    if (isMediaChunk(currentLoadable)) {
      MediaChunk mediaChunk = (MediaChunk) currentLoadable;
      notifyLoadCompleted(currentLoadable.bytesLoaded(), mediaChunk.type, mediaChunk.trigger,
          mediaChunk.format, mediaChunk.startTimeUs, mediaChunk.endTimeUs, now, loadDurationMs);
      loadingFinished = ((BaseMediaChunk) currentLoadable).isLastChunk;
    } else {
      notifyLoadCompleted(currentLoadable.bytesLoaded(), currentLoadable.type,
          currentLoadable.trigger, currentLoadable.format, -1, -1, now, loadDurationMs);
    }
    clearCurrentLoadable();
    updateLoadControl();
  }

  @Override
  public void onLoadCanceled(Loadable loadable) {
    Chunk currentLoadable = currentLoadableHolder.chunk;
    notifyLoadCanceled(currentLoadable.bytesLoaded());
    clearCurrentLoadable();
    if (state == STATE_ENABLED) {
      restartFrom(pendingResetPositionUs);
    } else {
      sampleQueue.clear();
      mediaChunks.clear();
      clearCurrentLoadable();
      loadControl.trimAllocator();
    }
  }

  @Override
  public void onLoadError(Loadable loadable, IOException e) {
    currentLoadableException = e;
    currentLoadableExceptionCount++;
    currentLoadableExceptionTimestamp = SystemClock.elapsedRealtime();
    notifyLoadError(e);
    chunkSource.onChunkLoadError(currentLoadableHolder.chunk, e);
    updateLoadControl();
  }

  /**
   * Called when a sample has been read. Can be used to perform any modifications necessary before
   * the sample is returned.
   *
   * @param mediaChunk The chunk from which the sample was obtained.
   * @param sampleHolder Holds the read sample.
   */
  protected void onSampleRead(MediaChunk mediaChunk, SampleHolder sampleHolder) {
    // Do nothing.
  }

  private void restartFrom(long positionUs) {
    pendingResetPositionUs = positionUs;
    loadingFinished = false;
    if (loader.isLoading()) {
      loader.cancelLoading();
    } else {
      sampleQueue.clear();
      mediaChunks.clear();
      clearCurrentLoadable();
      updateLoadControl();
    }
  }

  private void clearCurrentLoadable() {
    currentLoadableHolder.chunk = null;
    clearCurrentLoadableException();
  }

  private void clearCurrentLoadableException() {
    currentLoadableException = null;
    currentLoadableExceptionCount = 0;
    currentLoadableExceptionFatal = false;
  }

  private void updateLoadControl() {
    if (currentLoadableExceptionFatal) {
      // We've failed, but we still need to update the control with our current state.
      loadControl.update(this, downstreamPositionUs, -1, false, true);
      return;
    }

    long now = SystemClock.elapsedRealtime();
    long nextLoadPositionUs = getNextLoadPositionUs();
    boolean isBackedOff = currentLoadableException != null;
    boolean loadingOrBackedOff = loader.isLoading() || isBackedOff;

    // If we're not loading or backed off, evaluate the operation if (a) we don't have the next
    // chunk yet and we're not finished, or (b) if the last evaluation was over 2000ms ago.
    if (!loadingOrBackedOff && ((currentLoadableHolder.chunk == null && nextLoadPositionUs != -1)
        || (now - lastPerformedBufferOperation > 2000))) {
      // Perform the evaluation.
      lastPerformedBufferOperation = now;
      currentLoadableHolder.queueSize = readOnlyMediaChunks.size();
      chunkSource.getChunkOperation(readOnlyMediaChunks, pendingResetPositionUs,
          downstreamPositionUs, currentLoadableHolder);
      boolean chunksDiscarded = discardUpstreamMediaChunks(currentLoadableHolder.queueSize);
      // Update the next load position as appropriate.
      if (currentLoadableHolder.chunk == null) {
        // Set loadPosition to -1 to indicate that we don't have anything to load.
        nextLoadPositionUs = -1;
      } else if (chunksDiscarded) {
        // Chunks were discarded, so we need to re-evaluate the load position.
        nextLoadPositionUs = getNextLoadPositionUs();
      }
    }

    // Update the control with our current state, and determine whether we're the next loader.
    boolean nextLoader = loadControl.update(this, downstreamPositionUs, nextLoadPositionUs,
        loadingOrBackedOff, false);

    if (isBackedOff) {
      long elapsedMillis = now - currentLoadableExceptionTimestamp;
      if (elapsedMillis >= getRetryDelayMillis(currentLoadableExceptionCount)) {
        resumeFromBackOff();
      }
      return;
    }

    if (!loader.isLoading() && nextLoader) {
      maybeStartLoading();
    }
  }

  /**
   * Gets the next load time, assuming that the next load starts where the previous chunk ended (or
   * from the pending reset time, if there is one).
   */
  private long getNextLoadPositionUs() {
    if (isPendingReset()) {
      return pendingResetPositionUs;
    } else {
      BaseMediaChunk lastMediaChunk = mediaChunks.getLast();
      return lastMediaChunk.isLastChunk ? -1 : lastMediaChunk.endTimeUs;
    }
  }

  /**
   * Resumes loading.
   * <p>
   * If the {@link ChunkSource} returns a chunk equivalent to the backed off chunk B, then the
   * loading of B will be resumed. In all other cases B will be discarded and the new chunk will
   * be loaded.
   */
  private void resumeFromBackOff() {
    currentLoadableException = null;

    Chunk backedOffChunk = currentLoadableHolder.chunk;
    if (!isMediaChunk(backedOffChunk)) {
      currentLoadableHolder.queueSize = readOnlyMediaChunks.size();
      chunkSource.getChunkOperation(readOnlyMediaChunks, pendingResetPositionUs,
          downstreamPositionUs, currentLoadableHolder);
      discardUpstreamMediaChunks(currentLoadableHolder.queueSize);
      if (currentLoadableHolder.chunk == backedOffChunk) {
        // Chunk was unchanged. Resume loading.
        loader.startLoading(backedOffChunk, this);
      } else {
        // Chunk was changed. Notify that the existing load was canceled.
        notifyLoadCanceled(backedOffChunk.bytesLoaded());
        // Start loading the replacement.
        maybeStartLoading();
      }
      return;
    }

    if (backedOffChunk == mediaChunks.getFirst()) {
      // We're not able to clear the first media chunk, so we have no choice but to continue
      // loading it.
      loader.startLoading(backedOffChunk, this);
      return;
    }

    // The current loadable is the last media chunk. Remove it before we invoke the chunk source,
    // and add it back again afterwards.
    BaseMediaChunk removedChunk = mediaChunks.removeLast();
    Assertions.checkState(backedOffChunk == removedChunk);
    currentLoadableHolder.queueSize = readOnlyMediaChunks.size();
    chunkSource.getChunkOperation(readOnlyMediaChunks, pendingResetPositionUs, downstreamPositionUs,
        currentLoadableHolder);
    mediaChunks.add(removedChunk);

    if (currentLoadableHolder.chunk == backedOffChunk) {
      // Chunk was unchanged. Resume loading.
      loader.startLoading(backedOffChunk, this);
    } else {
      // Chunk was changed. Notify that the existing load was canceled.
      notifyLoadCanceled(backedOffChunk.bytesLoaded());
      // This call will remove and release at least one chunk from the end of mediaChunks. Since
      // the current loadable is the last media chunk, it is guaranteed to be removed.
      discardUpstreamMediaChunks(currentLoadableHolder.queueSize);
      clearCurrentLoadableException();
      maybeStartLoading();
    }
  }

  private void maybeStartLoading() {
    Chunk currentLoadable = currentLoadableHolder.chunk;
    if (currentLoadable == null) {
      // Nothing to load.
      return;
    }
    currentLoadStartTimeMs = SystemClock.elapsedRealtime();
    if (isMediaChunk(currentLoadable)) {
      BaseMediaChunk mediaChunk = (BaseMediaChunk) currentLoadable;
      mediaChunk.init(sampleQueue);
      mediaChunks.add(mediaChunk);
      if (isPendingReset()) {
        pendingResetPositionUs = NO_RESET_PENDING;
      }
      notifyLoadStarted(mediaChunk.dataSpec.length, mediaChunk.type, mediaChunk.trigger,
          mediaChunk.format, mediaChunk.startTimeUs, mediaChunk.endTimeUs);
    } else {
      notifyLoadStarted(currentLoadable.dataSpec.length, currentLoadable.type,
          currentLoadable.trigger, currentLoadable.format, -1, -1);
    }
    loader.startLoading(currentLoadable, this);
  }

  /**
   * Discard upstream media chunks until the queue length is equal to the length specified.
   *
   * @param queueLength The desired length of the queue.
   * @return True if chunks were discarded. False otherwise.
   */
  private boolean discardUpstreamMediaChunks(int queueLength) {
    if (mediaChunks.size() <= queueLength) {
      return false;
    }
    long startTimeUs = 0;
    long endTimeUs = mediaChunks.getLast().endTimeUs;

    BaseMediaChunk removed = null;
    while (mediaChunks.size() > queueLength) {
      removed = mediaChunks.removeLast();
      startTimeUs = removed.startTimeUs;
    }
    sampleQueue.discardUpstreamSamples(removed.getFirstSampleIndex());

    notifyUpstreamDiscarded(startTimeUs, endTimeUs);
    return true;
  }

  private boolean isMediaChunk(Chunk chunk) {
    return chunk instanceof BaseMediaChunk;
  }

  private boolean isPendingReset() {
    return pendingResetPositionUs != NO_RESET_PENDING;
  }

  private long getRetryDelayMillis(long errorCount) {
    return Math.min((errorCount - 1) * 1000, 5000);
  }

  protected final int usToMs(long timeUs) {
    return (int) (timeUs / 1000);
  }

  private void notifyLoadStarted(final long length, final int type, final int trigger,
      final Format format, final long mediaStartTimeUs, final long mediaEndTimeUs) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onLoadStarted(eventSourceId, length, type, trigger, format,
              usToMs(mediaStartTimeUs), usToMs(mediaEndTimeUs));
        }
      });
    }
  }

  private void notifyLoadCompleted(final long bytesLoaded, final int type, final int trigger,
      final Format format, final long mediaStartTimeUs, final long mediaEndTimeUs,
      final long elapsedRealtimeMs, final long loadDurationMs) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onLoadCompleted(eventSourceId, bytesLoaded, type, trigger, format,
              usToMs(mediaStartTimeUs), usToMs(mediaEndTimeUs), elapsedRealtimeMs, loadDurationMs);
        }
      });
    }
  }

  private void notifyLoadCanceled(final long bytesLoaded) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onLoadCanceled(eventSourceId, bytesLoaded);
        }
      });
    }
  }

  private void notifyLoadError(final IOException e) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onLoadError(eventSourceId, e);
        }
      });
    }
  }

  private void notifyUpstreamDiscarded(final long mediaStartTimeUs, final long mediaEndTimeUs) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onUpstreamDiscarded(eventSourceId, usToMs(mediaStartTimeUs),
              usToMs(mediaEndTimeUs));
        }
      });
    }
  }

  private void notifyDownstreamFormatChanged(final Format format, final int trigger,
      final long positionUs) {
    if (eventHandler != null && eventListener != null) {
      eventHandler.post(new Runnable()  {
        @Override
        public void run() {
          eventListener.onDownstreamFormatChanged(eventSourceId, format, trigger,
              usToMs(positionUs));
        }
      });
    }
  }

}
