import os, re, click

from typing import List, Self

from ffx.track_type import TrackType
from ffx.iso_language import IsoLanguage

from ffx.track_disposition import TrackDisposition
from ffx.track_codec import TrackCodec

from ffx.track_descriptor import TrackDescriptor
from ffx.logging_utils import get_ffx_logger


class MediaDescriptor:
    """This class represents the structural content of a media file including streams and metadata"""

    CONTEXT_KEY = "context"

    TAGS_KEY = "tags"
    TRACKS_KEY = "tracks"

    TRACK_DESCRIPTOR_LIST_KEY = "track_descriptors"
    ATTACHMENT_DESCRIPTOR_LIST_KEY = "attachment_descriptors"
    CLEAR_TAGS_FLAG_KEY = "clear_tags"

    FFPROBE_DISPOSITION_KEY = "disposition"
    FFPROBE_TAGS_KEY = "tags"
    FFPROBE_CODEC_TYPE_KEY = "codec_type"

    #407 remove as well
    EXCLUDED_MEDIA_TAGS = ["creation_time"]

    SEASON_EPISODE_STREAM_LANGUAGE_DISPOSITIONS_MATCH = '[sS]([0-9]+)[eE]([0-9]+)_([0-9]+)_([a-z]{3})(?:_([A-Z]{3}))*'
    STREAM_LANGUAGE_DISPOSITIONS_MATCH = '([0-9]+)_([a-z]{3})(?:_([A-Z]{3}))*'

    SUBTITLE_FILE_EXTENSION = 'vtt'

    def __init__(self, **kwargs):

        if MediaDescriptor.CONTEXT_KEY in kwargs.keys():
            if type(kwargs[MediaDescriptor.CONTEXT_KEY]) is not dict:
                raise TypeError(
                    f"MediaDescriptor.__init__(): Argument {MediaDescriptor.CONTEXT_KEY} is required to be of type dict"
                )
            self.__context = kwargs[MediaDescriptor.CONTEXT_KEY]
            self.__logger = self.__context['logger']
        else:
            self.__context = {}
            self.__logger = get_ffx_logger()

        if MediaDescriptor.TAGS_KEY in kwargs.keys():
            if type(kwargs[MediaDescriptor.TAGS_KEY]) is not dict:
                raise TypeError(
                    f"MediaDescriptor.__init__(): Argument {MediaDescriptor.TAGS_KEY} is required to be of type dict"
                )
            self.__mediaTags = kwargs[MediaDescriptor.TAGS_KEY]
        else:
            self.__mediaTags = {}

        if MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY in kwargs.keys():
            if (
                type(kwargs[MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY]) is not list
            ):  # Use List typehint for TrackDescriptor as well if it works
                raise TypeError(
                    f"MediaDescriptor.__init__(): Argument {MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY} is required to be of type list"
                )
            for d in kwargs[MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY]:
                if type(d) is not TrackDescriptor:
                    raise TypeError(
                        f"TrackDesciptor.__init__(): All elements of argument list {MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY} are required to be of type TrackDescriptor"
                    )
            self.__trackDescriptors: List[TrackDescriptor] = kwargs[MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY]
        else:
            self.__trackDescriptors: List[TrackDescriptor] = []

    def setTrackLanguage(self, language: str, index: int, trackType: TrackType = None):

        trackLanguage = IsoLanguage.findThreeLetter(language)
        if trackLanguage == IsoLanguage.UNDEFINED:
            self.__logger.warning('MediaDescriptor.setTrackLanguage(): Parameter language does not contain a registered '
                                    + f"ISO 639 3-letter language code, skipping to set language for"
                                    + str('' if trackType is None else trackType.label()) + f"track {index}")

        trackList = self.getTrackDescriptors(trackType=trackType)

        if index < 0 or index > len(trackList) - 1:
            self.__logger.warning(f"MediaDescriptor.setTrackLanguage(): Parameter index ({index}) is "
                                    + f"out of range of {'' if trackType is None else trackType.label()}track list")

        td: TrackDescriptor = trackList[index]
        td.setLanguage(trackLanguage)

        return


    def setTrackTitle(self, title: str, index: int, trackType: TrackType = None):

        trackList = self.getTrackDescriptors(trackType=trackType)

        if index < 0 or index > len(trackList) - 1:
            self.__logger.error(f"MediaDescriptor.setTrackTitle(): Parameter index ({index}) is "
                                    + f"out of range of {'' if trackType is None else trackType.label()}track list")
            raise click.Abort()

        td: TrackDescriptor = trackList[index]
        td.setTitle(title)


    def setDefaultSubTrack(self, trackType: TrackType, subIndex: int):
        # for t in self.getAllTrackDescriptors():
        for t in self.getTrackDescriptors():
            if t.getType() == trackType:
                t.setDispositionFlag(
                    TrackDisposition.DEFAULT, t.getSubIndex() == int(subIndex)
                )

    def setForcedSubTrack(self, trackType: TrackType, subIndex: int):
        # for t in self.getAllTrackDescriptors():
        for t in self.getTrackDescriptors():
            if t.getType() == trackType:
                t.setDispositionFlag(
                    TrackDisposition.FORCED, t.getSubIndex() == int(subIndex)
                )

    def checkConfiguration(self):

        videoTracks = self.getVideoTracks()
        audioTracks = self.getAudioTracks()
        subtitleTracks = self.getSubtitleTracks()
        
        if len([v for v in videoTracks if v.getDispositionFlag(TrackDisposition.DEFAULT)]) > 1:
            raise ValueError('More than one default video track')
        if len([a for a in audioTracks if a.getDispositionFlag(TrackDisposition.DEFAULT)]) > 1:
            raise ValueError('More than one default audio track')
        if len([s for s in subtitleTracks if s.getDispositionFlag(TrackDisposition.DEFAULT)]) > 1:
            raise ValueError('More than one default subtitle track')

        if len([v for v in videoTracks if v.getDispositionFlag(TrackDisposition.FORCED)]) > 1:
            raise ValueError('More than one forced video track')
        if len([a for a in audioTracks if a.getDispositionFlag(TrackDisposition.FORCED)]) > 1:
            raise ValueError('More than one forced audio track')
        if len([s for s in subtitleTracks if s.getDispositionFlag(TrackDisposition.FORCED)]) > 1:
            raise ValueError('More than one forced subtitle track')

        trackDescriptors = videoTracks + audioTracks + subtitleTracks
        sourceIndices = [
            t.getSourceIndex() for t in trackDescriptors
        ]
        if len(set(sourceIndices)) < len(trackDescriptors):
            raise ValueError('Multiple streams originating from the same source stream')


    def applyOverrides(self, overrides: dict):

        if 'languages' in overrides.keys():
            for trackIndex in overrides['languages'].keys():
                self.setTrackLanguage(overrides['languages'][trackIndex], trackIndex)

        if 'titles' in overrides.keys():
            for trackIndex in overrides['titles'].keys():
                self.setTrackTitle(overrides['titles'][trackIndex], trackIndex)

        if 'forced_video' in overrides.keys():
            sti = int(overrides['forced_video'])
            self.setForcedSubTrack(TrackType.VIDEO, sti)
            self.setDefaultSubTrack(TrackType.VIDEO, sti)

        elif 'default_video' in overrides.keys():
            sti = int(overrides['default_video'])
            self.setDefaultSubTrack(TrackType.VIDEO, sti)

        if 'forced_audio' in overrides.keys():
            sti = int(overrides['forced_audio'])
            self.setForcedSubTrack(TrackType.AUDIO, sti)
            self.setDefaultSubTrack(TrackType.AUDIO, sti)

        elif 'default_audio' in overrides.keys():
            sti = int(overrides['default_audio'])
            self.setDefaultSubTrack(TrackType.AUDIO, sti)

        if 'forced_subtitle' in overrides.keys():
            sti = int(overrides['forced_subtitle'])
            self.setForcedSubTrack(TrackType.SUBTITLE, sti)
            self.setDefaultSubTrack(TrackType.SUBTITLE, sti)

        elif 'default_subtitle' in overrides.keys():
            sti = int(overrides['default_subtitle'])
            self.setDefaultSubTrack(TrackType.SUBTITLE, sti)

        if 'stream_order' in overrides.keys():
            self.rearrangeTrackDescriptors(overrides['stream_order'])


    def applySourceIndices(self, sourceMediaDescriptor: Self):
        # sourceTrackDescriptors = sourceMediaDescriptor.getAllTrackDescriptors()
        sourceTrackDescriptors = sourceMediaDescriptor.getTrackDescriptors()

        numTrackDescriptors = len(self.__trackDescriptors)
        if len(sourceTrackDescriptors) != numTrackDescriptors:
            raise ValueError('MediaDescriptor.applySourceIndices (): Number of track descriptors does not match')

        for trackIndex in range(numTrackDescriptors):
            self.__trackDescriptors[trackIndex].setSourceIndex(sourceTrackDescriptors[trackIndex].getSourceIndex())


    def rearrangeTrackDescriptors(self, newOrder: List[int]):
        if len(newOrder) != len(self.__trackDescriptors):
            raise ValueError('Length of list with reordered indices does not match number of track descriptors')
        reorderedTrackDescriptors = []
        for oldIndex in newOrder:
            reorderedTrackDescriptors.append(self.__trackDescriptors[oldIndex])
        self.__trackDescriptors = reorderedTrackDescriptors
        self.reindexSubIndices()
        self.reindexIndices()


    @classmethod
    def fromFfprobe(cls, context, formatData, streamData):

        kwargs = {}

        kwargs[MediaDescriptor.CONTEXT_KEY] = context

        if MediaDescriptor.FFPROBE_TAGS_KEY in formatData.keys():
            kwargs[MediaDescriptor.TAGS_KEY] = formatData[
                MediaDescriptor.FFPROBE_TAGS_KEY
            ]

        kwargs[MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY] = []

        # TODO: Evtl obsolet
        subIndexCounters = {}

        for streamObj in streamData:

            ffprobeCodecType = streamObj[MediaDescriptor.FFPROBE_CODEC_TYPE_KEY]
            trackType = TrackType.fromLabel(ffprobeCodecType)

            if trackType != TrackType.UNKNOWN:

                if trackType not in subIndexCounters.keys():
                    subIndexCounters[trackType] = 0

                kwargs[MediaDescriptor.TRACK_DESCRIPTOR_LIST_KEY].append(
                    TrackDescriptor.fromFfprobe(
                        streamObj, subIndex=subIndexCounters[trackType]
                    )
                )
                subIndexCounters[trackType] += 1

        return cls(**kwargs)

    def getTags(self):
        return self.__mediaTags


    def sortSubIndices(
        self, descriptors: List[TrackDescriptor]
    ) -> List[TrackDescriptor]:
        subIndex = 0
        for d in descriptors:
            d.setSubIndex(subIndex)
            subIndex += 1
        return descriptors

    def reindexSubIndices(self, trackDescriptors: list = []):
        tdList = trackDescriptors if trackDescriptors else self.__trackDescriptors
        subIndexCounter = {}
        for td in tdList:
            trackType = td.getType()
            if trackType not in subIndexCounter.keys():
                subIndexCounter[trackType] = 0
            td.setSubIndex(subIndexCounter[trackType])
            subIndexCounter[trackType] += 1

    def sortIndices(
        self, descriptors: List[TrackDescriptor]
    ) -> List[TrackDescriptor]:
        index = 0
        for d in descriptors:
            d.setIndex(index)
            index += 1
        return descriptors

    def reindexIndices(self, trackDescriptors: list = []):
        tdList = trackDescriptors if trackDescriptors else self.__trackDescriptors
        for trackIndex in range(len(tdList)):
            tdList[trackIndex].setIndex(trackIndex)


    # def getAllTrackDescriptors(self):
    #     """Returns all track descriptors sorted by type: video, audio then subtitles"""                        
    #     return self.getVideoTracks() + self.getAudioTracks() + self.getSubtitleTracks()


    def getTrackDescriptors(self,
                            trackType: TrackType = None) -> List[TrackDescriptor]:

        if trackType is None:
            return self.__trackDescriptors

        descriptorList = []
        for td in self.__trackDescriptors:
            if td.getType() == trackType:
                descriptorList.append(td)
        
        return descriptorList


    def getVideoTracks(self) -> List[TrackDescriptor]:
        return [v for v in self.__trackDescriptors if v.getType() == TrackType.VIDEO]

    def getAudioTracks(self) -> List[TrackDescriptor]:
        return [a for a in self.__trackDescriptors if a.getType() == TrackType.AUDIO]

    def getSubtitleTracks(self) -> List[TrackDescriptor]:
        return [
            s
            for s in self.__trackDescriptors
            if s.getType() == TrackType.SUBTITLE
        ]

    def getAttachmentTracks(self) -> List[TrackDescriptor]:
        return [
            s
            for s in self.__trackDescriptors
            if s.getType() == TrackType.ATTACHMENT
        ]


    def getImportFileTokens(self, use_sub_index: bool = True):
        """Generate ffmpeg import options for external stream files"""

        importFileTokens = []

        td: TrackDescriptor
        for td in self.__trackDescriptors:

            importedFilePath = td.getExternalSourceFilePath()

            if importedFilePath:

                self.__logger.info(f"Substituting subtitle stream #{td.getIndex()} "
                                    + f"({td.getType().label()}:{td.getSubIndex()}) "
                                    + f"with import from file {td.getExternalSourceFilePath()}")

                importFileTokens += [
                    "-i",
                    importedFilePath,
                ]

        return importFileTokens


    def getInputMappingTokens(self,
                                use_sub_index: bool = True,
                                only_video: bool = False,
                                sourceMediaDescriptor: Self = None):
        """Tracks must be reordered for source index order"""

        inputMappingTokens = []
        
        sortedTrackDescriptors = sorted(self.__trackDescriptors, key=lambda d: d.getIndex())
        sourceTrackDescriptorsByIndex = {
            td.getIndex(): td
            for td in (
                sourceMediaDescriptor.getTrackDescriptors()
                if sourceMediaDescriptor is not None
                else sortedTrackDescriptors
            )
        }

        # raise click.ClickException(' '.join([f"\nindex={td.getIndex()} subIndex={td.getSubIndex()} srcIndex={td.getSourceIndex()} type={td.getType().label()}" for td in self.__trackDescriptors]))

        filePointer = 1
        for trackIndex in range(len(sortedTrackDescriptors)):

            td: TrackDescriptor = sortedTrackDescriptors[trackIndex]

            #HINT: Attached thumbnails are not supported by .webm container format
            if td.getCodec() != TrackCodec.PNG:

                sourceTrackDescriptor = sourceTrackDescriptorsByIndex.get(td.getSourceIndex())
                if sourceTrackDescriptor is None:
                    raise ValueError(f"No source track descriptor found for source index {td.getSourceIndex()}")

                stdi = sourceTrackDescriptor.getIndex()
                stdsi = sourceTrackDescriptor.getSubIndex()

                trackType = td.getType()
                trackCodec = td.getCodec()

                if (trackType != TrackType.ATTACHMENT
                        and (trackType == TrackType.VIDEO or not only_video)):


                    importedFilePath = td.getExternalSourceFilePath()

                    if use_sub_index:

                        if importedFilePath:

                            inputMappingTokens += [
                                "-map",
                                f"{filePointer}:{trackType.indicator()}:0",
                            ]
                            filePointer += 1

                        else:

                            if not trackCodec in [TrackCodec.PGS, TrackCodec.VOBSUB]:
                                inputMappingTokens += [
                                    "-map",
                                    f"0:{trackType.indicator()}:{stdsi}",
                                ]

                    else:
                        if not trackCodec in [TrackCodec.PGS, TrackCodec.VOBSUB]:
                            inputMappingTokens += ["-map", f"0:{stdi}"]

        if sourceMediaDescriptor:
            fontDescriptors = [ftd for ftd in sourceMediaDescriptor.getAttachmentTracks()
                                if ftd.getCodec() == TrackCodec.TTF]
        else:
            fontDescriptors = [ftd for ftd in self.__trackDescriptors
                                if ftd.getType() == TrackType.ATTACHMENT
                                and ftd.getCodec() == TrackCodec.TTF]

        for ad in sorted(fontDescriptors, key=lambda d: d.getIndex()):
            inputMappingTokens += ["-map", f"0:{ad.getIndex()}"]

        return inputMappingTokens


    def searchSubtitleFiles(self, searchDirectory, prefix):

        sesld_match = re.compile(f"{prefix}_{MediaDescriptor.SEASON_EPISODE_STREAM_LANGUAGE_DISPOSITIONS_MATCH}")
        sld_match = re.compile(f"{prefix}_{MediaDescriptor.STREAM_LANGUAGE_DISPOSITIONS_MATCH}")

        subtitleFileDescriptors = []

        for subtitleFilename in os.listdir(searchDirectory):
            if subtitleFilename.startswith(prefix) and subtitleFilename.endswith(
                "." + MediaDescriptor.SUBTITLE_FILE_EXTENSION
            ):

                sesld_result = sesld_match.search(subtitleFilename)
                sld_result = None if not sesld_result is None else sld_match.search(subtitleFilename)
                
                if not sesld_result is None:

                    subtitleFilePath = os.path.join(searchDirectory, subtitleFilename)
                    if os.path.isfile(subtitleFilePath):

                        subtitleFileDescriptor = {}
                        subtitleFileDescriptor["path"] = subtitleFilePath
                        subtitleFileDescriptor["season"] = int(sesld_result.group(1))
                        subtitleFileDescriptor["episode"] = int(sesld_result.group(2))
                        subtitleFileDescriptor["index"] = int(sesld_result.group(3))
                        subtitleFileDescriptor["language"] = sesld_result.group(4)

                        dispSet = set()
                        dispCaptGroups = sesld_result.groups()
                        numCaptGroups = len(dispCaptGroups)
                        if numCaptGroups > 4:
                            for groupIndex in range(numCaptGroups - 4):
                                disp = TrackDisposition.fromIndicator(dispCaptGroups[groupIndex + 4])
                                if disp is not None:
                                    dispSet.add(disp)
                        subtitleFileDescriptor["disposition_set"] = dispSet

                        subtitleFileDescriptors.append(subtitleFileDescriptor)

                if not sld_result is None:

                    subtitleFilePath = os.path.join(searchDirectory, subtitleFilename)
                    if os.path.isfile(subtitleFilePath):

                        subtitleFileDescriptor = {}
                        subtitleFileDescriptor["path"] = subtitleFilePath
                        subtitleFileDescriptor["index"] = int(sld_result.group(1))
                        subtitleFileDescriptor["language"] = sld_result.group(2)

                        dispSet = set()
                        dispCaptGroups = sld_result.groups()
                        numCaptGroups = len(dispCaptGroups)
                        if numCaptGroups > 2:
                            for groupIndex in range(numCaptGroups - 2):
                                disp = TrackDisposition.fromIndicator(dispCaptGroups[groupIndex + 2])
                                if disp is not None:
                                    dispSet.add(disp)
                        subtitleFileDescriptor["disposition_set"] = dispSet

                        subtitleFileDescriptors.append(subtitleFileDescriptor)


        self.__logger.debug(f"searchSubtitleFiles(): Available subtitle files {subtitleFileDescriptors}")

        return subtitleFileDescriptors


    def importSubtitles(
        self,
        searchDirectory,
        prefix,
        season: int = -1,
        episode: int = -1,
        preserve_dispositions: bool = False,
    ):
        
        # click.echo(f"Season: {season} Episode: {episode}")
        self.__logger.debug(f"importSubtitles(): Season: {season} Episode: {episode}")

        availableFileSubtitleDescriptors = self.searchSubtitleFiles(searchDirectory, prefix)
        
        self.__logger.debug(f"importSubtitles(): availableFileSubtitleDescriptors: {availableFileSubtitleDescriptors}")

        subtitleTracks = self.getSubtitleTracks()
        
        self.__logger.debug(f"importSubtitles(): subtitleTracks: {[s.getIndex() for s in subtitleTracks]}")

        matchingSubtitleFileDescriptors = (
            sorted(
                [
                    d
                    for d in availableFileSubtitleDescriptors
                    if ((season == -1 and episode == -1)
                        or (
                            d.get("season") == int(season)
                            and d.get("episode") == int(episode)
                        ))
                ],
                key=lambda d: d["index"],
            )
            if availableFileSubtitleDescriptors
            else []
        )

        self.__logger.debug(f"importSubtitles(): matchingSubtitleFileDescriptors: {matchingSubtitleFileDescriptors}")

        for msfd in matchingSubtitleFileDescriptors:
            matchingSubtitleTrackDescriptor = [s for s in subtitleTracks if s.getIndex() == msfd["index"]]
            if matchingSubtitleTrackDescriptor:
                # click.echo(f"Found matching subtitle file {msfd["path"]}\n")
                self.__logger.debug(f"importSubtitles(): Found matching subtitle file {msfd['path']}")
                matchingTrack = matchingSubtitleTrackDescriptor[0]
                matchingTrack.setExternalSourceFilePath(msfd["path"])

                # Prefer metadata coming from the external single-track source when
                # it is provided explicitly by the filename contract.
                matchingTrack.getTags()["language"] = msfd["language"]
                if msfd["disposition_set"] and not preserve_dispositions:
                    matchingTrack.setDispositionSet(msfd["disposition_set"])


    def getConfiguration(self, label: str = ''):
        yield f"--- {label if label else 'MediaDescriptor '+str(id(self))} {' '.join([str(k)+'='+str(v) for k,v in self.__mediaTags.items()])}"
        # for td in self.getAllTrackDescriptors():
        for td in self.getTrackDescriptors():
            yield (f"{td.getIndex()}:{td.getType().indicator()}:{td.getSubIndex()} "
                   + '|'.join([d.indicator() for d in td.getDispositionSet()])
                   + ' ' + ' '.join([str(k)+'='+str(v) for k,v in td.getTags().items()]))
