/*
 * Copyright (C) 2019 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.android.tv.tuner.exoplayer2.buffer;

import android.media.MediaFormat;
import android.os.ConditionVariable;
import android.support.annotation.NonNull;
import android.support.annotation.VisibleForTesting;
import android.util.ArrayMap;
import android.util.Log;
import android.util.Pair;

import com.android.tv.common.SoftPreconditions;
import com.android.tv.common.util.CommonUtils;
import com.android.tv.tuner.exoplayer2.SampleExtractor;
import com.google.android.exoplayer2.Format;
import com.google.android.exoplayer2.decoder.DecoderInputBuffer;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.ConcurrentModificationException;
import java.util.LinkedList;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.SortedMap;
import java.util.TreeMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Reads and writes the {@link SampleChunk} objects during playback and DVR. I/O operations are
 * handled by {@link StorageManager}.
 *
 * <p>The buffer manager is enabled for DVR and it can be disabled for playback, while running, if
 * the write throughput to the associated external storage is detected to be lower than a threshold
 * {@code MINIMUM_DISK_WRITE_SPEED_MBPS}". This leads to restarting playback flow.
 */
public class BufferManager {
    private static final String TAG = "BufferManager";
    private static final boolean DEBUG = false;

    // Constants for the disk write speed checking
    private static final long MINIMUM_WRITE_SIZE_FOR_SPEED_CHECK =
            10L * 1024 * 1024; // Checks for every 10M disk write
    private static final int MINIMUM_SAMPLE_SIZE_FOR_SPEED_CHECK = 15 * 1024;
    private static final int MAXIMUM_SPEED_CHECK_COUNT = 5; // Checks only 5 times
    private static final int MINIMUM_DISK_WRITE_SPEED_MBPS = 3; // 3 Megabytes per second

    private final SampleChunk.SampleChunkCreator mSampleChunkCreator;
    // Maps from track name to a map which maps from starting position to {@link SampleChunk}.
    private final Map<String, SortedMap<Long, Pair<SampleChunk, Integer>>> mChunkMap =
            new ArrayMap<>();
    private final Map<String, Long> mStartPositionMap = new ArrayMap<>();
    private final Map<String, ChunkEvictedListener> mEvictListeners = new ArrayMap<>();
    private final StorageManager mStorageManager;
    private long mBufferSize = 0;
    private final EvictChunkQueueMap mPendingDelete = new EvictChunkQueueMap();
    private final SampleChunk.ChunkCallback mChunkCallback =
            new SampleChunk.ChunkCallback() {
                @Override
                public void onChunkWrite(SampleChunk chunk) {
                    mBufferSize += chunk.getSize();
                }

                @Override
                public void onChunkDelete(SampleChunk chunk) {
                    mBufferSize -= chunk.getSize();
                }
            };

    private int mMinSampleSizeForSpeedCheck = MINIMUM_SAMPLE_SIZE_FOR_SPEED_CHECK;
    private long mTotalWriteSize;
    private long mTotalWriteTimeNs;
    private float mWriteBandwidth = 0.0f;
    private final AtomicInteger mSpeedCheckCount = new AtomicInteger();

    public interface ChunkEvictedListener {
        /**
         * Listener for when {@link SampleChunk} is removed from track.
         *
         * @param createdTimeMs creation time of the evicted chunk.
         */
        void onChunkEvicted(long createdTimeMs);
    }
    /** Handles I/O between BufferManager and {@link SampleExtractor}. */
    public interface SampleBuffer {

        /**
         * Initializes SampleBuffer.
         *
         * @param ids track identifiers for storage read/write.
         * @param formats meta-data for each track.
         * @throws IOException if an I/O error occurs.
         */
        void init(
                @NonNull List<String> ids,
                @NonNull List<Format> formats)
                throws IOException;

        /** Selects the track {@code index} for reading sample data. */
        void selectTrack(int index);

        /**
         * Deselects the track at {@code index}, so that no more samples will be read from the
         * track.
         */
        void deselectTrack(int index);

        /**
         * Writes sample to storage.
         *
         * @param index track index
         * @param sample sample to write at storage
         * @param conditionVariable notifies the completion of writing sample.
         * @throws IOException if an I/O error occurs.
         */
        void writeSample(int index, DecoderInputBuffer sample, ConditionVariable conditionVariable)
                throws IOException;

        /** Checks whether storage write speed is slow. */
        boolean isWriteSpeedSlow(int sampleSize, long writeDurationNs);

        /**
         * Handles when write speed is slow.
         */
        void handleWriteSpeedSlow();

        /** Sets the flag when EoS was reached. */
        void setEos();

        /**
         * Reads the next sample in the track at index {@code track} into {@code DecoderInputBuffer}
         * returning {@link com.google.android.exoplayer.SampleSource#SAMPLE_READ} if it is
         * available. If the next sample is not available, returns {@link
         * com.google.android.exoplayer.SampleSource#NOTHING_READ}.
         */
        int readSample(int index, DecoderInputBuffer outSample);

        /** Seeks to the specified time in microseconds. */
        void seekTo(long positionUs);

        /** Returns whether there is buffered data. */
        boolean continueLoading(long positionUs);

        /**
         * Cleans up and releases everything.
         *
         * @throws IOException if an I/O error occurs.
         */
        void release() throws IOException;
    }

    /** A Track format which will be loaded and saved from the permanent storage for recordings. */
    public static class TrackFormat {

        /**
         * The track id for the specified track. The track id will be used as a track identifier for
         * recordings.
         */
        public final String trackId;

        /** The {@link MediaFormat} for the specified track. */
        // TODO: Refactor to Format.
        public final MediaFormat mediaFormat;

        /**
         * Creates TrackFormat.
         *
         * @param trackId Track id
         * @param mediaFormat Media mediaFormat of track
         */
        public TrackFormat(String trackId, MediaFormat mediaFormat) {
            this.trackId = trackId;
            this.mediaFormat = mediaFormat;
        }
    }

    /** A Holder for a sample position which will be loaded from the index file for recordings. */
    public static class PositionHolder {

        /**
         * The current sample position in microseconds. The position is identical to the
         * PTS(presentation time stamp) of the sample.
         */
        public final long positionUs;

        /** Base sample position for the current {@link SampleChunk}. */
        public final long basePositionUs;

        /** The file offset for the current sample in the current {@link SampleChunk}. */
        public final int offset;

        /**
         * Creates a holder for a specific position in the recording.
         *
         * @param positionUs Position in the recording
         * @param basePositionUs Position of base sample
         * @param offset Offset in the recording
         */
        public PositionHolder(long positionUs, long basePositionUs, int offset) {
            this.positionUs = positionUs;
            this.basePositionUs = basePositionUs;
            this.offset = offset;
        }
    }

    /** Storage configuration and policy manager for {@link BufferManager} */
    public interface StorageManager {

        /**
         * Provides eligible storage directory for {@link BufferManager}.
         *
         * @return a directory to save buffer(chunks) and meta files
         */
        File getBufferDir();

        /**
         * Informs whether the storage is used for persistent use. (eg. dvr recording/play)
         *
         * @return {@code true} if stored files are persistent
         */
        boolean isPersistent();

        /**
         * Informs whether the storage usage exceeds pre-determined size.
         *
         * @param bufferSize the current total usage of Storage in bytes.
         * @param pendingDelete the current storage usage which will be deleted in near future by
         *     bytes
         * @return {@code true} if it reached pre-determined max size
         */
        boolean reachedStorageMax(long bufferSize, long pendingDelete);

        /**
         * Informs whether the storage has enough remained space.
         *
         * @param pendingDelete the current storage usage which will be deleted in near future by
         *     bytes
         * @return {@code true} if it has enough space
         */
        boolean hasEnoughBuffer(long pendingDelete);

        /**
         * Reads track name & {@link MediaFormat} from storage.
         *
         * @param isAudio {@code true} if it is for audio track
         * @return {@link List} of TrackFormat
         */
        List<TrackFormat> readTrackInfoFiles(boolean isAudio);

        /**
         * Reads key sample positions for each written sample from storage.
         *
         * @param trackId track name
         * @return indexes of the specified track
         * @throws IOException if an I/O error occurs.
         */
        ArrayList<PositionHolder> readIndexFile(String trackId) throws IOException;

        /**
         * Writes track information to storage.
         *
         * @param formatList {@list List} of TrackFormat
         * @param isAudio {@code true} if it is for audio track
         * @throws IOException if an I/O error occurs.
         */
        void writeTrackInfoFiles(List<TrackFormat> formatList, boolean isAudio) throws IOException;

        /**
         * Writes index file to storage.
         *
         * @param trackName track name
         * @param index {@link SampleChunk} container
         * @throws IOException if an I/O error occurs.
         */
        void writeIndexFile(String trackName, SortedMap<Long, Pair<SampleChunk, Integer>> index)
                throws IOException;

        /**
         * Writes to index file to storage.
         *
         * @param trackName track name
         * @param size size of sample
         * @param position position in micro seconds
         * @param sampleChunk {@link SampleChunk} chunk to be added
         * @param offset offset
         * @throws IOException if an I/O error occurs.
         */
        void updateIndexFile(
                String trackName, int size, long position, SampleChunk sampleChunk, int offset)
                throws IOException;
    }

    private static class EvictChunkQueueMap {
        private final Map<String, LinkedList<SampleChunk>> mEvictMap = new ArrayMap<>();
        private long mSize;

        private void init(String key) {
            mEvictMap.put(key, new LinkedList<>());
        }

        private void add(String key, SampleChunk chunk) {
            LinkedList<SampleChunk> queue = mEvictMap.get(key);
            if (queue != null) {
                mSize += chunk.getSize();
                queue.add(chunk);
            }
        }

        private SampleChunk poll(String key, long startPositionUs) {
            LinkedList<SampleChunk> queue = mEvictMap.get(key);
            if (queue != null) {
                SampleChunk chunk = queue.peek();
                if (chunk != null && chunk.getStartPositionUs() < startPositionUs) {
                    mSize -= chunk.getSize();
                    return queue.poll();
                }
            }
            return null;
        }

        private long getSize() {
            return mSize;
        }

        private void release() {
            for (Map.Entry<String, LinkedList<SampleChunk>> entry : mEvictMap.entrySet()) {
                for (SampleChunk chunk : entry.getValue()) {
                    SampleChunk.IoState.release(chunk, true);
                }
            }
            mEvictMap.clear();
            mSize = 0;
        }
    }

    public BufferManager(StorageManager storageManager) {
        this(storageManager, new SampleChunk.SampleChunkCreator());
    }

    public BufferManager(
            StorageManager storageManager, SampleChunk.SampleChunkCreator sampleChunkCreator) {
        mStorageManager = storageManager;
        mSampleChunkCreator = sampleChunkCreator;
    }

    public void registerChunkEvictedListener(String id, ChunkEvictedListener listener) {
        mEvictListeners.put(id, listener);
    }

    public void unregisterChunkEvictedListener(String id) {
        mEvictListeners.remove(id);
    }

    private static String getFileName(String id, long positionUs) {
        return String.format(Locale.ENGLISH, "%s_%016x.chunk", id, positionUs);
    }

    /**
     * Creates a new {@link SampleChunk} for caching samples if it is needed.
     *
     * @param id the name of the track
     * @param positionUs current position to write a sample in micro seconds.
     * @param inputBufferPool {@link InputBufferPool} for the fast creation of samples.
     * @param currentChunk the current {@link SampleChunk} to write, {@code null} when to create a
     *     new {@link SampleChunk}.
     * @param currentOffset the current offset to write.
     * @return returns the created {@link SampleChunk}.
     * @throws IOException if an I/O error occurs.
     */
    public SampleChunk createNewWriteFileIfNeeded(
            String id,
            long positionUs,
            InputBufferPool inputBufferPool,
            SampleChunk currentChunk,
            int currentOffset,
            boolean updateIndexFile)
            throws IOException {
        if (!maybeEvictChunk()) {
            throw new IOException("Not enough storage space");
        }
        SortedMap<Long, Pair<SampleChunk, Integer>> map = mChunkMap.get(id);
        if (map == null) {
            map = new TreeMap<>();
            mChunkMap.put(id, map);
            mStartPositionMap.put(id, positionUs);
            mPendingDelete.init(id);
        }
        if (currentChunk == null) {
            File file = new File(mStorageManager.getBufferDir(), getFileName(id, positionUs));
            SampleChunk sampleChunk =
                    mSampleChunkCreator.createSampleChunk(
                            inputBufferPool, file, positionUs, mChunkCallback);
            map.put(positionUs, Pair.create(sampleChunk, 0));
            if (updateIndexFile) {
                mStorageManager.updateIndexFile(id, map.size(), positionUs, sampleChunk, 0);
            }
            return sampleChunk;
        } else {
            map.put(positionUs, Pair.create(currentChunk, currentOffset));
            if (updateIndexFile) {
                mStorageManager.updateIndexFile(
                        id, map.size(), positionUs, currentChunk, currentOffset);
            }
            return null;
        }
    }

    /**
     * Loads a track using {@link BufferManager.StorageManager}.
     *
     * @param trackId the name of the track.
     * @param inputBufferPool {@link InputBufferPool} for the fast creation of samples.
     * @throws IOException if an I/O error occurs.
     */
    public void loadTrackFromStorage(String trackId, InputBufferPool inputBufferPool)
            throws IOException {
        ArrayList<PositionHolder> keyPositions = mStorageManager.readIndexFile(trackId);
        long startPositionUs = keyPositions.size() > 0 ? keyPositions.get(0).positionUs : 0;

        SortedMap<Long, Pair<SampleChunk, Integer>> map = mChunkMap.get(trackId);
        if (map == null) {
            map = new TreeMap<>();
            mChunkMap.put(trackId, map);
            mStartPositionMap.put(trackId, startPositionUs);
            mPendingDelete.init(trackId);
        }
        SampleChunk chunk = null;
        long basePositionUs = -1;
        for (PositionHolder position : keyPositions) {
            if (position.basePositionUs != basePositionUs) {
                chunk =
                        mSampleChunkCreator.loadSampleChunkFromFile(
                                inputBufferPool,
                                mStorageManager.getBufferDir(),
                                getFileName(trackId, position.positionUs),
                                position.positionUs,
                                mChunkCallback,
                                chunk);
                basePositionUs = position.basePositionUs;
            }
            map.put(position.positionUs, Pair.create(chunk, position.offset));
        }
    }

    /**
     * Finds a {@link SampleChunk} for the specified track name and the position.
     *
     * @param id the name of the track.
     * @param positionUs the position.
     * @return returns the found {@link SampleChunk}.
     */
    public Pair<SampleChunk, Integer> getReadFile(String id, long positionUs) {
        SortedMap<Long, Pair<SampleChunk, Integer>> map = mChunkMap.get(id);
        if (map == null) {
            return null;
        }
        Pair<SampleChunk, Integer> ret;
        SortedMap<Long, Pair<SampleChunk, Integer>> headMap = map.headMap(positionUs + 1);
        if (!headMap.isEmpty()) {
            ret = headMap.get(headMap.lastKey());
        } else {
            ret = map.get(map.firstKey());
        }
        return ret;
    }

    /**
     * Evicts chunks which are ready to be evicted for the specified track
     *
     * @param id the specified track
     * @param earlierThanPositionUs the start position of the {@link SampleChunk} should be earlier
     *     than
     */
    public void evictChunks(String id, long earlierThanPositionUs) {
        SampleChunk chunk;
        while ((chunk = mPendingDelete.poll(id, earlierThanPositionUs)) != null) {
            SampleChunk.IoState.release(chunk, !mStorageManager.isPersistent());
        }
    }

    /**
     * Returns the start position of the specified track in micro seconds.
     *
     * @param id the specified track
     */
    public long getStartPositionUs(String id) {
        Long ret = mStartPositionMap.get(id);
        return ret == null ? 0 : ret;
    }

    private boolean maybeEvictChunk() {
        long pendingDelete = mPendingDelete.getSize();
        while (mStorageManager.reachedStorageMax(mBufferSize, pendingDelete)
                || !mStorageManager.hasEnoughBuffer(pendingDelete)) {
            if (mStorageManager.isPersistent()) {
                // Since chunks are persistent, we cannot evict chunks.
                return false;
            }
            SortedMap<Long, Pair<SampleChunk, Integer>> earliestChunkMap = null;
            SampleChunk earliestChunk = null;
            String earliestChunkId = null;
            for (Map.Entry<String, SortedMap<Long, Pair<SampleChunk, Integer>>> entry :
                    mChunkMap.entrySet()) {
                SortedMap<Long, Pair<SampleChunk, Integer>> map = entry.getValue();
                if (map.isEmpty()) {
                    continue;
                }
                SampleChunk chunk = map.get(map.firstKey()).first;
                if (earliestChunk == null
                        || chunk.getCreatedTimeMs() < earliestChunk.getCreatedTimeMs()) {
                    earliestChunkMap = map;
                    earliestChunk = chunk;
                    earliestChunkId = entry.getKey();
                }
            }
            if (earliestChunk == null) {
                break;
            }
            mPendingDelete.add(earliestChunkId, earliestChunk);
            earliestChunkMap.remove(earliestChunk.getStartPositionUs());
            if (DEBUG) {
                Log.d(
                        TAG,
                        String.format(
                                "bufferSize = %d; pendingDelete = %b; "
                                        + "earliestChunk size = %d; %s@%d (%s)",
                                mBufferSize,
                                pendingDelete,
                                earliestChunk.getSize(),
                                earliestChunkId,
                                earliestChunk.getStartPositionUs(),
                                CommonUtils.toIsoDateTimeString(earliestChunk.getCreatedTimeMs())));
            }
            ChunkEvictedListener listener = mEvictListeners.get(earliestChunkId);
            if (listener != null) {
                listener.onChunkEvicted(earliestChunk.getCreatedTimeMs());
            }
            pendingDelete = mPendingDelete.getSize();
        }
        for (Map.Entry<String, SortedMap<Long, Pair<SampleChunk, Integer>>> entry :
                mChunkMap.entrySet()) {
            SortedMap<Long, Pair<SampleChunk, Integer>> map = entry.getValue();
            if (map.isEmpty()) {
                continue;
            }
            mStartPositionMap.put(entry.getKey(), map.firstKey());
        }
        return true;
    }

    /**
     * Reads track information which includes {@link MediaFormat}.
     *
     * @return returns all track information which is found by {@link BufferManager.StorageManager}.
     * @throws IOException if an I/O error occurs.
     */
    public List<TrackFormat> readTrackInfoFiles() throws IOException {
        List<TrackFormat> trackFormatList = new ArrayList<>();
        trackFormatList.addAll(mStorageManager.readTrackInfoFiles(false));
        trackFormatList.addAll(mStorageManager.readTrackInfoFiles(true));
        if (trackFormatList.isEmpty()) {
            throw new IOException("No track information to load");
        }
        return trackFormatList;
    }

    /**
     * Writes track information and index information for all tracks.
     *
     * @param audios list of audio track information
     * @param videos list of audio track information
     * @throws IOException if an I/O error occurs.
     */
    public void writeMetaFiles(List<TrackFormat> audios, List<TrackFormat> videos)
            throws IOException {
        if (audios.isEmpty() && videos.isEmpty()) {
            throw new IOException("No track information to save");
        }
        if (!audios.isEmpty()) {
            mStorageManager.writeTrackInfoFiles(audios, true);
            for (TrackFormat trackFormat : audios) {
                SortedMap<Long, Pair<SampleChunk, Integer>> map =
                        mChunkMap.get(trackFormat.trackId);
                if (map == null) {
                    throw new IOException("Audio track index missing");
                }
                mStorageManager.writeIndexFile(trackFormat.trackId, map);
            }
        }
        if (!videos.isEmpty()) {
            mStorageManager.writeTrackInfoFiles(videos, false);
            for (TrackFormat trackFormat : videos) {
                SortedMap<Long, Pair<SampleChunk, Integer>> map =
                        mChunkMap.get(trackFormat.trackId);
                if (map == null) {
                    throw new IOException("Video track index missing");
                }
                mStorageManager.writeIndexFile(trackFormat.trackId, map);
            }
        }
    }

    /**
     * Writes track information for all tracks.
     *
     * @param audios list of audio track information
     * @param videos list of audio track information
     * @throws IOException if an I/O error occurs.
     */
    public void writeMetaFilesOnly(List<TrackFormat> audios, List<TrackFormat> videos)
            throws IOException {
        if (audios.isEmpty() && videos.isEmpty()) {
            throw new IOException("No track information to save");
        }
        if (!audios.isEmpty()) {
            mStorageManager.writeTrackInfoFiles(audios, true);
        }
        if (!videos.isEmpty()) {
            mStorageManager.writeTrackInfoFiles(videos, false);
        }
    }

    /** Releases all the resources. */
    public void release() {
        try {
            mPendingDelete.release();
            for (Map.Entry<String, SortedMap<Long, Pair<SampleChunk, Integer>>> entry :
                    mChunkMap.entrySet()) {
                SampleChunk toRelease = null;
                for (Pair<SampleChunk, Integer> positions : entry.getValue().values()) {
                    if (toRelease != positions.first) {
                        toRelease = positions.first;
                        SampleChunk.IoState.release(toRelease, !mStorageManager.isPersistent());
                    }
                }
            }
            mChunkMap.clear();
        } catch (ConcurrentModificationException | NullPointerException e) {
            // TODO: remove this after it it confirmed that race condition issues are resolved.
            // b/32492258, b/32373376
            SoftPreconditions.checkState(
                    false, "Exception on BufferManager#release: ", e.toString());
        }
    }

    private void resetWriteStat(float writeBandwidth) {
        mWriteBandwidth = writeBandwidth;
        mTotalWriteSize = 0;
        mTotalWriteTimeNs = 0;
    }

    /** Adds a disk write sample size to calculate the average disk write bandwidth. */
    public void addWriteStat(long size, long timeNs) {
        if (size >= mMinSampleSizeForSpeedCheck) {
            mTotalWriteSize += size;
            mTotalWriteTimeNs += timeNs;
        }
    }

    /**
     * Returns if the average disk write bandwidth is slower than threshold {@code
     * MINIMUM_DISK_WRITE_SPEED_MBPS}.
     */
    public boolean isWriteSlow() {
        if (mTotalWriteSize < MINIMUM_WRITE_SIZE_FOR_SPEED_CHECK) {
            return false;
        }

        // Checks write speed for only MAXIMUM_SPEED_CHECK_COUNT times to ignore outliers
        // by temporary system overloading during the playback.
        if (mSpeedCheckCount.get() > MAXIMUM_SPEED_CHECK_COUNT) {
            return false;
        }
        mSpeedCheckCount.incrementAndGet();
        float megabytePerSecond = calculateWriteBandwidth();
        resetWriteStat(megabytePerSecond);
        if (DEBUG) {
            Log.d(TAG, "Measured disk write performance: " + megabytePerSecond + "MBps");
        }
        return megabytePerSecond < MINIMUM_DISK_WRITE_SPEED_MBPS;
    }

    /**
     * Returns recent write bandwidth in MBps. If recent bandwidth is not available, returns {float
     * -1.0f}.
     */
    public float getWriteBandwidth() {
        return mWriteBandwidth == 0.0f ? -1.0f : mWriteBandwidth;
    }

    private float calculateWriteBandwidth() {
        if (mTotalWriteTimeNs == 0) {
            return -1;
        }
        return ((float) mTotalWriteSize * 1000 / mTotalWriteTimeNs);
    }

    /**
     * Returns if {@link BufferManager} has checked the write speed, which is suitable for
     * Trickplay.
     */
    @VisibleForTesting
    public boolean hasSpeedCheckDone() {
        return mSpeedCheckCount.get() > 0;
    }

    /**
     * Sets minimum sample size for write speed check.
     *
     * @param sampleSize minimum sample size for write speed check.
     */
    @VisibleForTesting
    public void setMinimumSampleSizeForSpeedCheck(int sampleSize) {
        mMinSampleSizeForSpeedCheck = sampleSize;
    }
}