package com.futo.platformplayer.sabr;

import android.util.Pair;
import android.util.SparseIntArray;

import androidx.annotation.IntDef;
import androidx.annotation.Nullable;

import androidx.media3.common.C;
import androidx.media3.common.Format;
import androidx.media3.common.util.UnstableApi;
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.CompositeSequenceableLoaderFactory;
import androidx.media3.exoplayer.source.EmptySampleStream;
import androidx.media3.exoplayer.source.MediaPeriod;
import androidx.media3.exoplayer.source.MediaSourceEventListener.EventDispatcher;
import androidx.media3.exoplayer.source.SampleStream;
import androidx.media3.exoplayer.source.SequenceableLoader;
import androidx.media3.common.TrackGroup;
import androidx.media3.exoplayer.source.TrackGroupArray;
import androidx.media3.exoplayer.source.chunk.ChunkSampleStream;
import androidx.media3.exoplayer.source.chunk.ChunkSampleStream.EmbeddedSampleStream;
import com.futo.platformplayer.sabr.PlayerEmsgHandler.PlayerEmsgCallback;
import com.futo.platformplayer.sabr.PlayerEmsgHandler.PlayerTrackEmsgHandler;
import com.futo.platformplayer.sabr.SabrChunkSource.Factory;
import com.futo.platformplayer.sabr.manifest.AdaptationSet;
import com.futo.platformplayer.sabr.manifest.EventStream;
import com.futo.platformplayer.sabr.manifest.Period;
import com.futo.platformplayer.sabr.manifest.Representation;
import com.futo.platformplayer.sabr.manifest.SabrManifest;
import androidx.media3.exoplayer.trackselection.ExoTrackSelection;
import androidx.media3.exoplayer.upstream.Allocator;
import androidx.media3.exoplayer.upstream.LoadErrorHandlingPolicy;
import androidx.media3.exoplayer.upstream.LoaderErrorThrower;
import androidx.media3.datasource.TransferListener;
import androidx.media3.common.MimeTypes;
import androidx.media3.common.util.Util;

import java.io.IOException;
import java.lang.annotation.Documented;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.IdentityHashMap;
import java.util.List;
import java.util.regex.Matcher;

@UnstableApi
final class SabrMediaPeriod implements MediaPeriod, SequenceableLoader.Callback<ChunkSampleStream<SabrChunkSource>>, ChunkSampleStream.ReleaseCallback<SabrChunkSource> {
    /* package */ final int id;
    private final Factory chunkSourceFactory;
    @Nullable
    private final TransferListener transferListener;
    private final LoadErrorHandlingPolicy loadErrorHandlingPolicy;
    private final EventDispatcher eventDispatcher;
    private final long elapsedRealtimeOffsetMs;
    private final LoaderErrorThrower manifestLoaderErrorThrower;
    private final TrackGroupArray trackGroups;
    private final TrackGroupInfo[] trackGroupInfos;
    private final Allocator allocator;
    private final CompositeSequenceableLoaderFactory compositeSequenceableLoaderFactory;
    private final PlayerEmsgHandler playerEmsgHandler;
    private final IdentityHashMap<ChunkSampleStream<SabrChunkSource>, PlayerTrackEmsgHandler>
            trackEmsgHandlerBySampleStream;

    private @Nullable Callback callback;
    private ChunkSampleStream<SabrChunkSource>[] sampleStreams;
    private SequenceableLoader compositeSequenceableLoader;
    private EventSampleStream[] eventSampleStreams;
    private SabrManifest manifest;
    private int periodIndex;
    private List<EventStream> eventStreams;
    private boolean notifiedReadingStarted;
    private final DrmSessionManager drmSessionManager;
    private final DrmSessionEventListener.EventDispatcher drmEventDispatcher;

    public SabrMediaPeriod(
            int id,
            SabrManifest manifest,
            int periodIndex,
            SabrChunkSource.Factory chunkSourceFactory,
            @Nullable TransferListener transferListener,
            LoadErrorHandlingPolicy loadErrorHandlingPolicy,
            EventDispatcher eventDispatcher,
            long elapsedRealtimeOffsetMs,
            LoaderErrorThrower manifestLoaderErrorThrower,
            Allocator allocator,
            CompositeSequenceableLoaderFactory compositeSequenceableLoaderFactory,
            PlayerEmsgCallback playerEmsgCallback) {
        this.id = id;
        this.manifest = manifest;
        this.periodIndex = periodIndex;
        this.chunkSourceFactory = chunkSourceFactory;
        this.transferListener = transferListener;
        this.loadErrorHandlingPolicy = loadErrorHandlingPolicy;
        this.eventDispatcher = eventDispatcher;
        this.elapsedRealtimeOffsetMs = elapsedRealtimeOffsetMs;
        this.manifestLoaderErrorThrower = manifestLoaderErrorThrower;
        this.allocator = allocator;
        this.compositeSequenceableLoaderFactory = compositeSequenceableLoaderFactory;
        playerEmsgHandler = new PlayerEmsgHandler(manifest, playerEmsgCallback, allocator);
        sampleStreams = newSampleStreamArray(0);
        eventSampleStreams = new EventSampleStream[0];
        trackEmsgHandlerBySampleStream = new IdentityHashMap<>();
        compositeSequenceableLoader =
                compositeSequenceableLoaderFactory.createCompositeSequenceableLoader(sampleStreams);
        Period period = manifest.getPeriod(periodIndex);
        Pair<TrackGroupArray, TrackGroupInfo[]> result = buildTrackGroups(period.adaptationSets);
        trackGroups = result.first;
        trackGroupInfos = result.second;
        this.drmSessionManager = DrmSessionManager.DRM_UNSUPPORTED;
        this.drmEventDispatcher = new DrmSessionEventListener.EventDispatcher();
    }

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

    @Override
    public void prepare(Callback callback, long positionUs) {
        this.callback = callback;
        callback.onPrepared(this);
    }

    @Override
    public void maybeThrowPrepareError() throws IOException {
        manifestLoaderErrorThrower.maybeThrowError();
    }

    @Override
    public TrackGroupArray getTrackGroups() {
        return trackGroups;
    }

    @Override
    public long selectTracks(
            @Nullable ExoTrackSelection[] selections,
            boolean[] mayRetainStreamFlags,
            @Nullable SampleStream[] streams,
            boolean[] streamResetFlags,
            long positionUs) {
        int[] streamIndexToTrackGroupIndex = getStreamIndexToTrackGroupIndex(selections);
        releaseDisabledStreams(selections, mayRetainStreamFlags, streams);
        releaseOrphanEmbeddedStreams(selections, streams, streamIndexToTrackGroupIndex);
        selectNewStreams(
                selections, streams, streamResetFlags, positionUs, streamIndexToTrackGroupIndex);

        ArrayList<ChunkSampleStream<SabrChunkSource>> sampleStreamList = new ArrayList<>();
        ArrayList<EventSampleStream> eventSampleStreamList = new ArrayList<>();
        for (SampleStream sampleStream : streams) {
            if (sampleStream instanceof ChunkSampleStream) {
                @SuppressWarnings("unchecked")
                ChunkSampleStream<SabrChunkSource> stream =
                        (ChunkSampleStream<SabrChunkSource>) sampleStream;
                sampleStreamList.add(stream);
            } else if (sampleStream instanceof EventSampleStream) {
                eventSampleStreamList.add((EventSampleStream) sampleStream);
            }
        }
        sampleStreams = newSampleStreamArray(sampleStreamList.size());
        sampleStreamList.toArray(sampleStreams);
        eventSampleStreams = new EventSampleStream[eventSampleStreamList.size()];
        eventSampleStreamList.toArray(eventSampleStreams);

        compositeSequenceableLoader =
                compositeSequenceableLoaderFactory.createCompositeSequenceableLoader(sampleStreams);
        return positionUs;
    }

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

    @Override
    public long readDiscontinuity() {
        if (!notifiedReadingStarted) {
            notifiedReadingStarted = true;
        }
        return C.TIME_UNSET;
    }

    @Override
    public long seekToUs(long positionUs) {
        for (ChunkSampleStream<SabrChunkSource> sampleStream : sampleStreams) {
            sampleStream.seekToUs(positionUs);
        }
        for (EventSampleStream sampleStream : eventSampleStreams) {
            sampleStream.seekToUs(positionUs);
        }
        return positionUs;
    }

    @Override
    public long getAdjustedSeekPositionUs(long positionUs, SeekParameters seekParameters) {
        for (ChunkSampleStream<SabrChunkSource> sampleStream : sampleStreams) {
            if (sampleStream.primaryTrackType == C.TRACK_TYPE_VIDEO) {
                return sampleStream.getAdjustedSeekPositionUs(positionUs, seekParameters);
            }
        }
        return positionUs;
    }

    @Override
    public long getBufferedPositionUs() {
        return compositeSequenceableLoader.getBufferedPositionUs();
    }

    @Override
    public long getNextLoadPositionUs() {
        return compositeSequenceableLoader.getNextLoadPositionUs();
    }

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

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

    // SequenceableLoader.Callback implementation.
    @Override
    public void onContinueLoadingRequested(ChunkSampleStream<SabrChunkSource> source) {
        callback.onContinueLoadingRequested(this);
    }

    @Override
    public synchronized void onSampleStreamReleased(ChunkSampleStream<SabrChunkSource> stream) {
        PlayerTrackEmsgHandler trackEmsgHandler = trackEmsgHandlerBySampleStream.remove(stream);
        if (trackEmsgHandler != null) {
            trackEmsgHandler.release();
        }
    }

    /**
     * Updates the {@link SabrManifest} and the index of this period in the manifest.
     *
     * @param manifest The updated manifest.
     * @param periodIndex the new index of this period in the updated manifest.
     */
    public void updateManifest(SabrManifest manifest, int periodIndex) {
        this.manifest = manifest;
        this.periodIndex = periodIndex;
        playerEmsgHandler.updateManifest(manifest);
        if (sampleStreams != null) {
            for (ChunkSampleStream<SabrChunkSource> sampleStream : sampleStreams) {
                sampleStream.getChunkSource().updateManifest(manifest, periodIndex);
            }
            callback.onContinueLoadingRequested(this);
        }
    }

    public void release() {
        playerEmsgHandler.release();
        for (ChunkSampleStream<SabrChunkSource> sampleStream : sampleStreams) {
            sampleStream.release(this);
        }
        callback = null;
    }


    @SuppressWarnings("unchecked")
    private static ChunkSampleStream<SabrChunkSource>[] newSampleStreamArray(int length) {
        return new ChunkSampleStream[length];
    }

    private static Pair<TrackGroupArray, TrackGroupInfo[]> buildTrackGroups(
            List<AdaptationSet> adaptationSets) {
        int[][] groupedAdaptationSetIndices = getGroupedAdaptationSetIndices(adaptationSets);

        int primaryGroupCount = groupedAdaptationSetIndices.length;
        boolean[] primaryGroupHasEventMessageTrackFlags = new boolean[primaryGroupCount];
        Format[][] primaryGroupCea608TrackFormats = new Format[primaryGroupCount][];
        int totalEmbeddedTrackGroupCount =
                identifyEmbeddedTracks(
                        primaryGroupCount,
                        adaptationSets,
                        groupedAdaptationSetIndices,
                        primaryGroupHasEventMessageTrackFlags,
                        primaryGroupCea608TrackFormats);

        int totalGroupCount = primaryGroupCount + totalEmbeddedTrackGroupCount;
        TrackGroup[] trackGroups = new TrackGroup[totalGroupCount];
        TrackGroupInfo[] trackGroupInfos = new TrackGroupInfo[totalGroupCount];

        int trackGroupCount =
                buildPrimaryAndEmbeddedTrackGroupInfos(
                        adaptationSets,
                        groupedAdaptationSetIndices,
                        primaryGroupCount,
                        primaryGroupHasEventMessageTrackFlags,
                        primaryGroupCea608TrackFormats,
                        trackGroups,
                        trackGroupInfos);

        return Pair.create(new TrackGroupArray(trackGroups), trackGroupInfos);
    }

    private static int[][] getGroupedAdaptationSetIndices(List<AdaptationSet> adaptationSets) {
        int adaptationSetCount = adaptationSets.size();
        SparseIntArray idToIndexMap = new SparseIntArray(adaptationSetCount);
        for (int i = 0; i < adaptationSetCount; i++) {
            idToIndexMap.put(adaptationSets.get(i).id, i);
        }

        int[][] groupedAdaptationSetIndices = new int[adaptationSetCount][];
        boolean[] adaptationSetUsedFlags = new boolean[adaptationSetCount];

        int groupCount = 0;
        for (int i = 0; i < adaptationSetCount; i++) {
            if (adaptationSetUsedFlags[i]) {
                // This adaptation set has already been included in a group.
                continue;
            }
            adaptationSetUsedFlags[i] = true;
            groupedAdaptationSetIndices[groupCount++] = new int[] {i};
        }

        return groupCount < adaptationSetCount
                ? Arrays.copyOf(groupedAdaptationSetIndices, groupCount) : groupedAdaptationSetIndices;
    }

    /**
     * Iterates through list of primary track groups and identifies embedded tracks.
     *
     * @param primaryGroupCount The number of primary track groups.
     * @param adaptationSets The list of {@link AdaptationSet} of the current DASH period.
     * @param groupedAdaptationSetIndices The indices of {@link AdaptationSet} that belongs to the
     *     same primary group, grouped in primary track groups order.
     * @param primaryGroupHasEventMessageTrackFlags An output array to be filled with flags indicating
     *     whether each of the primary track groups contains an embedded event message track.
     * @param primaryGroupCea608TrackFormats An output array to be filled with track formats for
     *     CEA-608 tracks embedded in each of the primary track groups.
     * @return Total number of embedded track groups.
     */
    private static int identifyEmbeddedTracks(
            int primaryGroupCount,
            List<AdaptationSet> adaptationSets,
            int[][] groupedAdaptationSetIndices,
            boolean[] primaryGroupHasEventMessageTrackFlags,
            Format[][] primaryGroupCea608TrackFormats) {
        int numEmbeddedTrackGroups = 0;
        for (int i = 0; i < primaryGroupCount; i++) {
            primaryGroupCea608TrackFormats[i] =
                    getCea608TrackFormats(adaptationSets, groupedAdaptationSetIndices[i]);
            if (primaryGroupCea608TrackFormats[i].length != 0) {
                numEmbeddedTrackGroups++;
            }
        }
        return numEmbeddedTrackGroups;
    }

    private static int buildPrimaryAndEmbeddedTrackGroupInfos(
            List<AdaptationSet> adaptationSets,
            int[][] groupedAdaptationSetIndices,
            int primaryGroupCount,
            boolean[] primaryGroupHasEventMessageTrackFlags,
            Format[][] primaryGroupCea608TrackFormats,
            TrackGroup[] trackGroups,
            TrackGroupInfo[] trackGroupInfos) {
        int trackGroupCount = 0;
        for (int i = 0; i < primaryGroupCount; i++) {
            int[] adaptationSetIndices = groupedAdaptationSetIndices[i];
            List<Representation> representations = new ArrayList<>();
            for (int adaptationSetIndex : adaptationSetIndices) {
                representations.addAll(adaptationSets.get(adaptationSetIndex).representations);
            }
            Format[] formats = new Format[representations.size()];
            for (int j = 0; j < formats.length; j++) {
                formats[j] = representations.get(j).format;
            }

            AdaptationSet firstAdaptationSet = adaptationSets.get(adaptationSetIndices[0]);
            int primaryTrackGroupIndex = trackGroupCount++;
            int eventMessageTrackGroupIndex =
                    primaryGroupHasEventMessageTrackFlags[i] ? trackGroupCount++ : C.INDEX_UNSET;
            int cea608TrackGroupIndex =
                    primaryGroupCea608TrackFormats[i].length != 0 ? trackGroupCount++ : C.INDEX_UNSET;

            trackGroups[primaryTrackGroupIndex] = new TrackGroup(formats);
            trackGroupInfos[primaryTrackGroupIndex] =
                    TrackGroupInfo.primaryTrack(
                            firstAdaptationSet.type,
                            adaptationSetIndices,
                            primaryTrackGroupIndex,
                            eventMessageTrackGroupIndex,
                            cea608TrackGroupIndex);
            if (eventMessageTrackGroupIndex != C.INDEX_UNSET) {
                Format format =
                        new Format.Builder()
                                .setId(firstAdaptationSet.id + ":emsg")
                                .setSampleMimeType(MimeTypes.APPLICATION_EMSG)
                                .build();
                trackGroups[eventMessageTrackGroupIndex] = new TrackGroup(format);
                trackGroupInfos[eventMessageTrackGroupIndex] =
                        TrackGroupInfo.embeddedEmsgTrack(adaptationSetIndices, primaryTrackGroupIndex);
            }
            if (cea608TrackGroupIndex != C.INDEX_UNSET) {
                trackGroups[cea608TrackGroupIndex] = new TrackGroup(primaryGroupCea608TrackFormats[i]);
                trackGroupInfos[cea608TrackGroupIndex] =
                        TrackGroupInfo.embeddedCea608Track(adaptationSetIndices, primaryTrackGroupIndex);
            }
        }
        return trackGroupCount;
    }

    private static Format[] getCea608TrackFormats(
            List<AdaptationSet> adaptationSets, int[] adaptationSetIndices) {
        return new Format[0];
    }

    private int[] getStreamIndexToTrackGroupIndex(ExoTrackSelection[] selections) {
        int[] streamIndexToTrackGroupIndex = new int[selections.length];
        for (int i = 0; i < selections.length; i++) {
            if (selections[i] != null) {
                streamIndexToTrackGroupIndex[i] = trackGroups.indexOf(selections[i].getTrackGroup());
            } else {
                streamIndexToTrackGroupIndex[i] = C.INDEX_UNSET;
            }
        }
        return streamIndexToTrackGroupIndex;
    }

    private void releaseDisabledStreams(ExoTrackSelection[] selections, boolean[] mayRetainStreamFlags, SampleStream[] streams) {
        for (int i = 0; i < selections.length; i++) {
            if (selections[i] == null || !mayRetainStreamFlags[i]) {
                if (streams[i] instanceof ChunkSampleStream) {
                    @SuppressWarnings("unchecked")
                    ChunkSampleStream<SabrChunkSource> stream =
                            (ChunkSampleStream<SabrChunkSource>) streams[i];
                    stream.release(this);
                } else if (streams[i] instanceof ChunkSampleStream.EmbeddedSampleStream) {
                    ((EmbeddedSampleStream) streams[i]).release();
                }
                streams[i] = null;
            }
        }
    }

    private void releaseOrphanEmbeddedStreams(ExoTrackSelection[] selections, SampleStream[] streams, int[] streamIndexToTrackGroupIndex) {
        for (int i = 0; i < selections.length; i++) {
            if (streams[i] instanceof EmptySampleStream || streams[i] instanceof EmbeddedSampleStream) {
                // We need to release an embedded stream if the corresponding primary stream is released.
                int primaryStreamIndex = getPrimaryStreamIndex(i, streamIndexToTrackGroupIndex);
                boolean mayRetainStream;
                if (primaryStreamIndex == C.INDEX_UNSET) {
                    // If the corresponding primary stream is not selected, we may retain an existing
                    // EmptySampleStream.
                    mayRetainStream = streams[i] instanceof EmptySampleStream;
                } else {
                    // If the corresponding primary stream is selected, we may retain the embedded stream if
                    // the stream's parent still matches.
                    mayRetainStream =
                            (streams[i] instanceof EmbeddedSampleStream)
                                    && ((EmbeddedSampleStream) streams[i]).parent == streams[primaryStreamIndex];
                }
                if (!mayRetainStream) {
                    if (streams[i] instanceof EmbeddedSampleStream) {
                        ((EmbeddedSampleStream) streams[i]).release();
                    }
                    streams[i] = null;
                }
            }
        }
    }

    private int getPrimaryStreamIndex(int embeddedStreamIndex, int[] streamIndexToTrackGroupIndex) {
        int embeddedTrackGroupIndex = streamIndexToTrackGroupIndex[embeddedStreamIndex];
        if (embeddedTrackGroupIndex == C.INDEX_UNSET) {
            return C.INDEX_UNSET;
        }
        int primaryTrackGroupIndex = trackGroupInfos[embeddedTrackGroupIndex].primaryTrackGroupIndex;
        for (int i = 0; i < streamIndexToTrackGroupIndex.length; i++) {
            int trackGroupIndex = streamIndexToTrackGroupIndex[i];
            if (trackGroupIndex == primaryTrackGroupIndex
                    && trackGroupInfos[trackGroupIndex].trackGroupCategory
                    == TrackGroupInfo.CATEGORY_PRIMARY) {
                return i;
            }
        }
        return C.INDEX_UNSET;
    }

    private void selectNewStreams(ExoTrackSelection[] selections, SampleStream[] streams, boolean[] streamResetFlags, long positionUs, int[] streamIndexToTrackGroupIndex) {
        // Create newly selected primary and event streams.
        for (int i = 0; i < selections.length; i++) {
            ExoTrackSelection selection = selections[i];
            if (selection == null) {
                continue;
            }
            if (streams[i] == null) {
                // Create new stream for selection.
                streamResetFlags[i] = true;
                int trackGroupIndex = streamIndexToTrackGroupIndex[i];
                TrackGroupInfo trackGroupInfo = trackGroupInfos[trackGroupIndex];
                if (trackGroupInfo.trackGroupCategory == TrackGroupInfo.CATEGORY_PRIMARY) {
                    streams[i] = buildSampleStream(trackGroupInfo, selection, positionUs);
                } else if (trackGroupInfo.trackGroupCategory == TrackGroupInfo.CATEGORY_MANIFEST_EVENTS) {
                    EventStream eventStream = eventStreams.get(trackGroupInfo.eventStreamGroupIndex);
                    Format format = selection.getTrackGroup().getFormat(0);
                    streams[i] = new EventSampleStream(eventStream, format, manifest.dynamic);
                }
            } else if (streams[i] instanceof ChunkSampleStream) {
                // Update selection in existing stream.
                @SuppressWarnings("unchecked")
                ChunkSampleStream<SabrChunkSource> stream = (ChunkSampleStream<SabrChunkSource>) streams[i];
                stream.getChunkSource().updateTrackSelection(selection);
            }
        }
        // Create newly selected embedded streams from the corresponding primary stream. Note that this
        // second pass is needed because the primary stream may not have been created yet in a first
        // pass if the index of the primary stream is greater than the index of the embedded stream.
        for (int i = 0; i < selections.length; i++) {
            if (streams[i] == null && selections[i] != null) {
                int trackGroupIndex = streamIndexToTrackGroupIndex[i];
                TrackGroupInfo trackGroupInfo = trackGroupInfos[trackGroupIndex];
                if (trackGroupInfo.trackGroupCategory == TrackGroupInfo.CATEGORY_EMBEDDED) {
                    int primaryStreamIndex = getPrimaryStreamIndex(i, streamIndexToTrackGroupIndex);
                    if (primaryStreamIndex == C.INDEX_UNSET) {
                        // If an embedded track is selected without the corresponding primary track, create an
                        // empty sample stream instead.
                        streams[i] = new EmptySampleStream();
                    } else {
                        streams[i] =
                                ((ChunkSampleStream) streams[primaryStreamIndex])
                                        .selectEmbeddedTrack(positionUs, trackGroupInfo.trackType);
                    }
                }
            }
        }
    }

    private ChunkSampleStream<SabrChunkSource> buildSampleStream(
            TrackGroupInfo trackGroupInfo,
            ExoTrackSelection selection,
            long positionUs) {

        int embeddedTrackCount = 0;
        boolean enableEventMessageTrack =
                trackGroupInfo.embeddedEventMessageTrackGroupIndex != C.INDEX_UNSET;
        TrackGroup embeddedEventMessageTrackGroup = null;
        if (enableEventMessageTrack) {
            embeddedEventMessageTrackGroup =
                    trackGroups.get(trackGroupInfo.embeddedEventMessageTrackGroupIndex);
            embeddedTrackCount++;
        }
        boolean enableCea608Tracks =
                trackGroupInfo.embeddedCea608TrackGroupIndex != C.INDEX_UNSET;
        TrackGroup embeddedCea608TrackGroup = null;
        if (enableCea608Tracks) {
            embeddedCea608TrackGroup =
                    trackGroups.get(trackGroupInfo.embeddedCea608TrackGroupIndex);
            embeddedTrackCount += embeddedCea608TrackGroup.length;
        }

        Format[] embeddedTrackFormats = new Format[embeddedTrackCount];
        int[] embeddedTrackTypes = new int[embeddedTrackCount];
        embeddedTrackCount = 0;

        if (enableEventMessageTrack) {
            embeddedTrackFormats[embeddedTrackCount] =
                    embeddedEventMessageTrackGroup.getFormat(0);
            embeddedTrackTypes[embeddedTrackCount] = C.TRACK_TYPE_METADATA;
            embeddedTrackCount++;
        }

        List<Format> embeddedCea608TrackFormats = new ArrayList<>();
        if (enableCea608Tracks) {
            for (int i = 0; i < embeddedCea608TrackGroup.length; i++) {
                embeddedTrackFormats[embeddedTrackCount] =
                        embeddedCea608TrackGroup.getFormat(i);
                embeddedTrackTypes[embeddedTrackCount] = C.TRACK_TYPE_TEXT;
                embeddedCea608TrackFormats.add(embeddedTrackFormats[embeddedTrackCount]);
                embeddedTrackCount++;
            }
        }

        PlayerTrackEmsgHandler trackPlayerEmsgHandler =
                manifest.dynamic && enableEventMessageTrack
                        ? playerEmsgHandler.newPlayerTrackEmsgHandler()
                        : null;

        SabrChunkSource chunkSource =
                chunkSourceFactory.createSabrChunkSource(
                        manifestLoaderErrorThrower,
                        manifest,
                        periodIndex,
                        trackGroupInfo.adaptationSetIndices,
                        selection,
                        trackGroupInfo.trackType,
                        elapsedRealtimeOffsetMs,
                        enableEventMessageTrack,
                        embeddedCea608TrackFormats,
                        trackPlayerEmsgHandler,
                        transferListener);

        ChunkSampleStream<SabrChunkSource> stream =
                new ChunkSampleStream<>(
                        trackGroupInfo.trackType,
                        embeddedTrackTypes,
                        embeddedTrackFormats,
                        chunkSource,
                        /* callback= */ this,
                        allocator,
                        positionUs,
                        drmSessionManager,
                        drmEventDispatcher,
                        loadErrorHandlingPolicy,
                        eventDispatcher,
                        /* canReportInitialDiscontinuity= */ true,
                        /* downloadExecutor= */ null);

        synchronized (this) {
            // The map is also accessed on the loading thread so synchronize access.
            trackEmsgHandlerBySampleStream.put(stream, trackPlayerEmsgHandler);
        }
        return stream;
    }

    private static final class TrackGroupInfo {

        @Documented
        @Retention(RetentionPolicy.SOURCE)
        @IntDef({CATEGORY_PRIMARY, CATEGORY_EMBEDDED, CATEGORY_MANIFEST_EVENTS})
        public @interface TrackGroupCategory {}

        /**
         * A normal track group that has its samples drawn from the stream.
         * For example: a video Track Group or an audio Track Group.
         */
        private static final int CATEGORY_PRIMARY = 0;

        /**
         * A track group whose samples are embedded within one of the primary streams. For example: an
         * EMSG track has its sample embedded in emsg atoms in one of the primary streams.
         */
        private static final int CATEGORY_EMBEDDED = 1;

        /**
         * A track group that has its samples listed explicitly in the DASH manifest file.
         * For example: an EventStream track has its sample (Events) included directly in the DASH
         * manifest file.
         */
        private static final int CATEGORY_MANIFEST_EVENTS = 2;

        public final int[] adaptationSetIndices;
        public final int trackType;
        @TrackGroupCategory public final int trackGroupCategory;

        public final int eventStreamGroupIndex;
        public final int primaryTrackGroupIndex;
        public final int embeddedEventMessageTrackGroupIndex;
        public final int embeddedCea608TrackGroupIndex;

        public static TrackGroupInfo primaryTrack(
                int trackType,
                int[] adaptationSetIndices,
                int primaryTrackGroupIndex,
                int embeddedEventMessageTrackGroupIndex,
                int embeddedCea608TrackGroupIndex) {
            return new TrackGroupInfo(
                    trackType,
                    CATEGORY_PRIMARY,
                    adaptationSetIndices,
                    primaryTrackGroupIndex,
                    embeddedEventMessageTrackGroupIndex,
                    embeddedCea608TrackGroupIndex,
                    /* eventStreamGroupIndex= */ -1);
        }

        public static TrackGroupInfo embeddedEmsgTrack(int[] adaptationSetIndices,
                                                       int primaryTrackGroupIndex) {
            return new TrackGroupInfo(
                    C.TRACK_TYPE_METADATA,
                    CATEGORY_EMBEDDED,
                    adaptationSetIndices,
                    primaryTrackGroupIndex,
                    C.INDEX_UNSET,
                    C.INDEX_UNSET,
                    /* eventStreamGroupIndex= */ -1);
        }

        public static TrackGroupInfo embeddedCea608Track(int[] adaptationSetIndices,
                                                         int primaryTrackGroupIndex) {
            return new TrackGroupInfo(
                    C.TRACK_TYPE_TEXT,
                    CATEGORY_EMBEDDED,
                    adaptationSetIndices,
                    primaryTrackGroupIndex,
                    C.INDEX_UNSET,
                    C.INDEX_UNSET,
                    /* eventStreamGroupIndex= */ -1);
        }

        public static TrackGroupInfo mpdEventTrack(int eventStreamIndex) {
            return new TrackGroupInfo(
                    C.TRACK_TYPE_METADATA,
                    CATEGORY_MANIFEST_EVENTS,
                    new int[0],
                    /* primaryTrackGroupIndex= */ -1,
                    C.INDEX_UNSET,
                    C.INDEX_UNSET,
                    eventStreamIndex);
        }

        private TrackGroupInfo(
                int trackType,
                @TrackGroupCategory int trackGroupCategory,
                int[] adaptationSetIndices,
                int primaryTrackGroupIndex,
                int embeddedEventMessageTrackGroupIndex,
                int embeddedCea608TrackGroupIndex,
                int eventStreamGroupIndex) {
            this.trackType = trackType;
            this.adaptationSetIndices = adaptationSetIndices;
            this.trackGroupCategory = trackGroupCategory;
            this.primaryTrackGroupIndex = primaryTrackGroupIndex;
            this.embeddedEventMessageTrackGroupIndex = embeddedEventMessageTrackGroupIndex;
            this.embeddedCea608TrackGroupIndex = embeddedCea608TrackGroupIndex;
            this.eventStreamGroupIndex = eventStreamGroupIndex;
        }
    }
}
