import contextlib
import os
from pathlib import Path
from typing import Dict
from typing import Iterable
from typing import Iterator
from typing import List
from typing import Optional
from typing import Set
from typing import Tuple

from yt_dlp.utils import RejectedVideoReached

from ytdl_sub.config.overrides import Overrides
from ytdl_sub.downloaders.source_plugin import SourcePlugin
from ytdl_sub.downloaders.source_plugin import SourcePluginExtension
from ytdl_sub.downloaders.url.validators import MultiUrlValidator
from ytdl_sub.downloaders.url.validators import UrlThumbnailListValidator
from ytdl_sub.downloaders.url.validators import UrlValidator
from ytdl_sub.downloaders.ytdl_options_builder import YTDLOptionsBuilder
from ytdl_sub.downloaders.ytdlp import YTDLP
from ytdl_sub.entries.entry import Entry
from ytdl_sub.entries.entry_parent import EntryParent
from ytdl_sub.entries.script.variable_definitions import VARIABLES
from ytdl_sub.entries.script.variable_definitions import VariableDefinitions
from ytdl_sub.utils.file_handler import FileHandler
from ytdl_sub.utils.logger import Logger
from ytdl_sub.utils.script import ScriptUtils
from ytdl_sub.utils.thumbnail import ThumbnailTypes
from ytdl_sub.utils.thumbnail import download_and_convert_url_thumbnail
from ytdl_sub.utils.thumbnail import try_convert_download_thumbnail
from ytdl_sub.ytdl_additions.enhanced_download_archive import EnhancedDownloadArchive

v: VariableDefinitions = VARIABLES

download_logger = Logger.get(name="downloader")


class URLDownloadState:
    def __init__(self, entries_total: int):
        self.entries_total = entries_total
        self.entries_downloaded = 0


class UrlDownloaderBasePluginExtension(SourcePluginExtension[MultiUrlValidator]):
    def _match_entry_to_url_validator(self, entry: Entry) -> UrlValidator:
        """
        Handle matching a URL to its original validator. This is for .info.json updates
        when older entries have missing variables
        """
        input_url_idx = entry.get(v.ytdl_sub_input_url_index, int)
        entry_input_url = entry.get(v.ytdl_sub_input_url, str)

        if 0 <= input_url_idx < len(self.plugin_options.urls.list):
            validator = self.plugin_options.urls.list[input_url_idx]
            if self.overrides.apply_formatter(validator.url) == entry_input_url:
                return validator

        # Match the first validator based on the URL, if one exists
        for validator in self.plugin_options.urls.list:
            if self.overrides.apply_formatter(validator.url) == entry_input_url:
                return validator

        # Return the first validator if none exist
        return self.plugin_options.urls.list[0]


class UrlDownloaderThumbnailPlugin(UrlDownloaderBasePluginExtension):
    def __init__(
        self,
        options: MultiUrlValidator,
        overrides: Overrides,
        enhanced_download_archive: EnhancedDownloadArchive,
    ):
        super().__init__(
            options=options,
            overrides=overrides,
            enhanced_download_archive=enhanced_download_archive,
        )
        self._thumbnails_downloaded: Set[str] = set()

    def _download_parent_thumbnails(
        self,
        thumbnail_list_info: UrlThumbnailListValidator,
        entry: Entry,
        parent: EntryParent,
    ) -> None:
        """
        Downloads and moves channel avatar and banner images to the output directory.
        """
        for thumbnail_info in thumbnail_list_info.list:
            thumbnail_name = self.overrides.apply_formatter(thumbnail_info.name, entry=entry)
            thumbnail_id = self.overrides.apply_formatter(thumbnail_info.uid)

            # If the thumbnail name is an empty string, completely ignore trying to download it
            if not thumbnail_name:
                continue

            # If latest entry, always update the thumbnail on each entry
            if thumbnail_id == ThumbnailTypes.LATEST_ENTRY:

                # always save in dry-run even if it doesn't exist...
                if self.is_dry_run or entry.is_thumbnail_downloaded():
                    self.save_file(
                        file_name=entry.get_download_thumbnail_name(),
                        output_file_name=thumbnail_name,
                        copy_file=True,
                    )
                    self._thumbnails_downloaded.add(thumbnail_name)
                continue

            # If not latest entry and the thumbnail has already been downloaded, then skip
            if thumbnail_name in self._thumbnails_downloaded:
                continue

            if (thumbnail_url := parent.get_thumbnail_url(thumbnail_id=thumbnail_id)) is None:
                download_logger.debug("Failed to find thumbnail id '%s'", thumbnail_id)
                continue

            if download_and_convert_url_thumbnail(
                thumbnail_url=thumbnail_url,
                output_thumbnail_path=str(Path(self.working_directory) / thumbnail_name),
            ):
                self.save_file(file_name=thumbnail_name)
                self._thumbnails_downloaded.add(thumbnail_name)
            else:
                download_logger.debug("Failed to download thumbnail id '%s'", thumbnail_id)

    def _download_url_thumbnails(self, collection_url: UrlValidator, entry: Entry):
        """
        After all media entries have been downloaded, post processed, and moved to the output
        directory, run this function. This lets the downloader add any extra files directly to the
        output directory, for things like YT channel image, banner.
        """
        if source_metadata := entry.get(v.source_metadata, dict):
            self._download_parent_thumbnails(
                thumbnail_list_info=collection_url.source_thumbnails,
                entry=entry,
                parent=EntryParent(source_metadata, working_directory=self.working_directory),
            )

        if playlist_metadata := entry.get(v.playlist_metadata, dict):
            self._download_parent_thumbnails(
                thumbnail_list_info=collection_url.playlist_thumbnails,
                entry=entry,
                parent=EntryParent(playlist_metadata, working_directory=self.working_directory),
            )

    def modify_entry(self, entry: Entry) -> Optional[Entry]:
        """
        Use the entry to download thumbnails (or move if LATEST_ENTRY).
        In addition, convert the entry thumbnail to jpg
        """
        # We always convert entry thumbnails to jpgs, and is performed here to be done
        # as early as possible in the plugin pipeline (downstream plugins depend on it being jpg)
        if not self.is_dry_run:
            try_convert_download_thumbnail(entry=entry)

        self._download_url_thumbnails(
            collection_url=self._match_entry_to_url_validator(entry=entry),
            entry=entry,
        )
        return entry


class UrlDownloaderCollectionVariablePlugin(UrlDownloaderBasePluginExtension):
    def __init__(
        self,
        options: MultiUrlValidator,
        overrides: Overrides,
        enhanced_download_archive: EnhancedDownloadArchive,
    ):
        super().__init__(
            options=options,
            overrides=overrides,
            enhanced_download_archive=enhanced_download_archive,
        )
        self._thumbnails_downloaded: Set[str] = set()

    def modify_entry_metadata(self, entry: Entry) -> Optional[Entry]:
        """
        Add collection variables to the entry
        """
        collection_url = self._match_entry_to_url_validator(entry=entry)
        entry.add(collection_url.variables.dict_with_format_strings)

        return entry


class MultiUrlDownloader(SourcePlugin[MultiUrlValidator]):
    """
    Class that interacts with ytdl to perform the download of metadata and content,
    and should translate that to list of Entry objects.
    """

    plugin_options_type = MultiUrlValidator
    plugin_extensions = [UrlDownloaderThumbnailPlugin, UrlDownloaderCollectionVariablePlugin]

    @classmethod
    def ytdl_option_defaults(cls) -> Dict:
        """
        .. code-block:: yaml

           ytdl_options:
             ignoreerrors: True  # ignore errors like hidden videos, age restriction, etc
        """
        return {"ignoreerrors": True}

    def __init__(
        self,
        options: MultiUrlValidator,
        enhanced_download_archive: EnhancedDownloadArchive,
        download_ytdl_options: YTDLOptionsBuilder,
        metadata_ytdl_options: YTDLOptionsBuilder,
        overrides: Overrides,
    ):
        """
        Parameters
        ----------
        options
            Options validator for this downloader
        enhanced_download_archive
            Download archive
        download_ytdl_options
            YTDL options builder for downloading media
        metadata_ytdl_options
            YTDL options builder for downloading metadata
        overrides
            Override variables
        """
        super().__init__(
            options=options,
            enhanced_download_archive=enhanced_download_archive,
            download_ytdl_options=download_ytdl_options,
            metadata_ytdl_options=metadata_ytdl_options,
            overrides=overrides,
        )
        self._downloaded_entries: Set[str] = set()
        self._url_state: Optional[URLDownloadState] = None

    def download_ytdl_options(self, url_idx: Optional[int] = None) -> Dict:
        """
        Returns
        -------
        YTLD options dict for downloading
        """
        return (
            self._download_ytdl_options_builder.clone()
            .add(self.ytdl_option_defaults(), before=True)
            .add(
                (
                    self.plugin_options.urls.list[url_idx].ytdl_options.dict
                    if url_idx is not None
                    else None
                ),
                before=True,
            )
            .to_dict()
        )

    def metadata_ytdl_options(self, ytdl_option_overrides: Dict) -> Dict:
        """
        Returns
        -------
        YTDL options dict for fetching metadata
        """

        return (
            self._metadata_ytdl_options_builder.clone()
            .add(self.ytdl_option_defaults(), before=True)
            .add(ytdl_option_overrides, before=True)
            .to_dict()
        )

    @property
    def is_dry_run(self) -> bool:
        """
        Returns
        -------
        True if dry-run is enabled. False otherwise.
        """
        return self.download_ytdl_options().get("skip_download", False)

    @property
    def is_entry_thumbnails_enabled(self) -> bool:
        """
        Returns
        -------
        True if entry thumbnails should be downloaded. False otherwise.
        """
        return self.download_ytdl_options().get("writethumbnail", False)

    ###############################################################################################
    # DOWNLOAD FUNCTIONS

    def _is_downloaded(self, entry: Entry) -> bool:
        return entry.ytdl_uid() in self._downloaded_entries

    def _mark_downloaded(self, entry: Entry) -> None:
        self._downloaded_entries.add(entry.ytdl_uid())

    @property
    def collection(self) -> MultiUrlValidator:
        """Return the download options collection"""
        return self.plugin_options

    @contextlib.contextmanager
    def _separate_download_archives(self, clear_info_json_files: bool = False):
        """
        Separate download archive writing between collection urls. This is so break_on_existing
        does not break when downloading from subset urls.

        Parameters
        ----------
        clear_info_json_files
            Whether to delete info.json files after yield
        """
        archive_path = self.download_ytdl_options().get("download_archive", "")
        backup_archive_path = f"{archive_path}.backup"

        # If archive path exists, maintain download archive is enable
        if archive_file_exists := archive_path and os.path.isfile(archive_path):
            archive_file_exists = True

            # If a backup exists, it's the one prior to any downloading, use that.
            if os.path.isfile(backup_archive_path):
                FileHandler.copy(src_file_path=backup_archive_path, dst_file_path=archive_path)
            # If not, create the backup
            else:
                FileHandler.copy(src_file_path=archive_path, dst_file_path=backup_archive_path)

        yield

        # If an archive path did not exist at first, but now exists, delete it
        if not archive_file_exists and os.path.isfile(archive_path):
            FileHandler.delete(file_path=archive_path)
        # If the archive file did exist, restore the backup
        elif archive_file_exists:
            FileHandler.move(src_file_path=backup_archive_path, dst_file_path=archive_path)

        # Clear info json files if true
        if clear_info_json_files:
            info_json_files = [
                Path(self.working_directory) / path
                for path in os.listdir(self.working_directory)
                if path.endswith(".info.json")
            ]
            for info_json_file in info_json_files:
                FileHandler.delete(info_json_file)

    def _extract_entry_info_with_retry(self, entry: Entry) -> Entry:
        download_entry_dict = YTDLP.extract_info_with_retry(
            ytdl_options_overrides=self.download_ytdl_options(
                url_idx=entry.get(v.ytdl_sub_input_url_index, int)
            ),
            is_downloaded_fn=None if self.is_dry_run else entry.is_downloaded,
            is_thumbnail_downloaded_fn=(
                None
                if (self.is_dry_run or not self.is_entry_thumbnails_enabled)
                else entry.is_thumbnail_downloaded_via_ytdlp
            ),
            url=entry.webpage_url,
        )
        return Entry(
            download_entry_dict,
            working_directory=self.working_directory,
        )

    def _iterate_child_entries(
        self, entries: List[Entry], download_reversed: bool
    ) -> Iterator[Entry]:
        # Iterate a list of entries, and delete the entries after yielding
        entries_to_iter: List[Optional[Entry]] = entries

        indices = list(range(len(entries_to_iter)))
        if download_reversed:
            indices = reversed(indices)

        for idx in indices:
            self._url_state.entries_downloaded += 1

            if self._is_downloaded(entries_to_iter[idx]):
                download_logger.info(
                    "Already downloaded entry %d/%d: %s",
                    self._url_state.entries_downloaded,
                    self._url_state.entries_total,
                    entries_to_iter[idx].title,
                )
                entries_to_iter[idx] = None
                continue

            yield entries_to_iter[idx]
            self._mark_downloaded(entries_to_iter[idx])

            entries_to_iter[idx] = None

    def _iterate_parent_entry(
        self, parent: EntryParent, download_reversed: bool
    ) -> Iterator[Entry]:
        yield from self._iterate_child_entries(
            entries=parent.entry_children(), download_reversed=download_reversed
        )

        # Recursion the parent's parent entries
        for parent_child in reversed(parent.parent_children()):
            yield from self._iterate_parent_entry(
                parent=parent_child, download_reversed=download_reversed
            )

    def _download_url_metadata(
        self, url: str, include_sibling_metadata: bool, ytdl_options_overrides: Dict
    ) -> Tuple[List[EntryParent], List[Entry]]:
        """
        Downloads only info.json files and forms EntryParent trees
        """
        with self._separate_download_archives():
            entry_dicts = YTDLP.extract_info_via_info_json(
                working_directory=self.working_directory,
                ytdl_options_overrides=ytdl_options_overrides,
                log_prefix_on_info_json_dl="Downloading metadata for",
                url=url,
            )

        parents = EntryParent.from_entry_dicts(
            url=url,
            entry_dicts=entry_dicts,
            working_directory=self.working_directory,
            include_sibling_metadata=include_sibling_metadata,
        )
        orphans = EntryParent.from_entry_dicts_with_no_parents(
            parents=parents,
            entry_dicts=entry_dicts,
            working_directory=self.working_directory,
        )

        return parents, orphans

    def _iterate_entries(
        self,
        parents: List[EntryParent],
        orphans: List[Entry],
        download_reversed: bool,
    ) -> Iterator[Entry]:
        """
        Downloads the leaf entries from EntryParent trees
        """
        # Delete info json files afterwards so other collection URLs do not use them
        with self._separate_download_archives(clear_info_json_files=True):
            for parent in parents:
                yield from self._iterate_parent_entry(
                    parent=parent, download_reversed=download_reversed
                )

            yield from self._iterate_child_entries(
                entries=orphans, download_reversed=download_reversed
            )

    def _download_metadata(self, url: str, validator: UrlValidator) -> Iterable[Entry]:
        metadata_ytdl_options = self.metadata_ytdl_options(
            ytdl_option_overrides=validator.ytdl_options.to_native_dict(self.overrides)
        )
        download_reversed = ScriptUtils.bool_formatter_output(
            self.overrides.apply_formatter(validator.download_reverse)
        )
        include_sibling_metadata = ScriptUtils.bool_formatter_output(
            self.overrides.apply_formatter(validator.include_sibling_metadata)
        )

        parents, orphan_entries = self._download_url_metadata(
            url=url,
            include_sibling_metadata=include_sibling_metadata,
            ytdl_options_overrides=metadata_ytdl_options,
        )

        # TODO: Encapsulate this logic into its own class
        self._url_state = URLDownloadState(
            entries_total=sum(parent.num_children() for parent in parents) + len(orphan_entries)
        )

        download_logger.info("Beginning downloads for %s", url)
        yield from self._iterate_entries(
            parents=parents,
            orphans=orphan_entries,
            download_reversed=download_reversed,
        )

    def download_metadata(self) -> Iterable[Entry]:
        """The function to perform the download of all media entries"""
        # download the bottom-most urls first since they are top-priority
        for idx, url_validator in reversed(list(enumerate(self.collection.urls.list))):
            # URLs can be empty. If they are, then skip
            if not (url := self.overrides.apply_formatter(url_validator.url)):
                continue

            for entry in self._download_metadata(url=url, validator=url_validator):
                entry.initialize_script(self.overrides).add(
                    {
                        v.ytdl_sub_input_url: url,
                        v.ytdl_sub_input_url_index: idx,
                        v.ytdl_sub_input_url_count: len(self.collection.urls.list),
                    }
                )

                yield entry

    def download(self, entry: Entry) -> Optional[Entry]:
        """
        Parameters
        ----------
        entry
            Entry to download

        Returns
        -------
        The entry that was downloaded successfully

        Raises
        ------
        RejectedVideoReached
          If a video was rejected and was not from match_filter
        """
        download_logger.info(
            "Downloading entry %d/%d: %s",
            self._url_state.entries_downloaded,
            self._url_state.entries_total,
            entry.title,
        )

        # Match-filters can be applied at the download stage. If the download is rejected,
        # then return None
        try:
            download_entry = self._extract_entry_info_with_retry(entry=entry)
        except RejectedVideoReached:
            download_logger.info("Entry rejected by download match-filter, skipping ..")
            return None

        upload_date_idx = self._enhanced_download_archive.mapping.get_num_entries_with_upload_date(
            upload_date_standardized=entry.get(v.upload_date_standardized, str)
        )
        download_idx = self._enhanced_download_archive.num_entries

        return entry.add_injected_variables(
            download_entry=download_entry,
            download_idx=download_idx,
            upload_date_idx=upload_date_idx,
        )
