#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim: ai ts=4 sts=4 et sw=4 nu

"""
    create project on Google Developer console
    Add Youtube Data API v3 to it
    Create credentials (Other non-UI, Public Data)
"""

import os
import json
import locale
import shutil
import datetime
from pathlib import Path
import concurrent.futures
from functools import partial
from gettext import gettext as _

import jinja2
import youtube_dl
from dateutil import parser as dt_parser
from babel.dates import format_date

from .zim import ZimInfo, make_zim_file
from .utils import (
    clean_text,
    resize_image,
    load_json,
    save_json,
    get_slug,
    save_file,
    get_colors,
    is_hex_color,
    get_language_details,
    setlocale,
)
from .youtube import (
    get_channel_json,
    credentials_ok,
    Playlist,
    get_channel_playlists_json,
    get_videos_json,
    get_videos_authors_info,
    save_channel_branding,
    skip_deleted_videos,
)
from .converter import hook_youtube_dl_ffmpeg
from .constants import logger, ROOT_DIR, CHANNEL, PLAYLIST, USER


class Youtube2Zim(object):
    def __init__(
        self,
        collection_type,
        youtube_id,
        api_key,
        video_format,
        low_quality,
        nb_videos_per_page,
        all_subtitles,
        autoplay,
        output_dir,
        no_zim,
        fname,
        debug,
        keep_build_dir,
        skip_download,
        max_concurrency,
        youtube_store,
        language,
        locale_name,
        tags,
        title=None,
        description=None,
        creator=None,
        publisher=None,
        name=None,
        profile_image=None,
        banner_image=None,
        main_color=None,
        secondary_color=None,
        only_test_branding=None,
    ):
        # data-retrieval info
        self.collection_type = collection_type
        self.youtube_id = youtube_id
        self.api_key = api_key

        # video-encoding info
        self.video_format = video_format
        self.low_quality = low_quality

        # options & zim params
        self.nb_videos_per_page = nb_videos_per_page
        self.all_subtitles = all_subtitles
        self.autoplay = autoplay
        self.fname = fname
        self.language = language
        self.tags = [t.strip() for t in tags.split(",")]
        self.title = title
        self.description = description
        self.creator = creator
        self.publisher = publisher
        self.name = name
        self.profile_image = profile_image
        self.banner_image = banner_image
        self.main_color = main_color
        self.secondary_color = secondary_color

        # process-related
        self.output_dir = Path(output_dir).expanduser().resolve()
        self.playlists = []
        self.uploads_playlist_id = None
        self.videos_ids = []
        self.main_channel_id = None  # use for branding
        self.only_test_branding = only_test_branding

        # debug/devel options
        self.no_zim = no_zim
        self.debug = debug
        self.keep_build_dir = keep_build_dir
        self.skip_download = skip_download
        self.max_concurrency = max_concurrency

        self.build_dir = self.output_dir.joinpath("build")

        # store ZIM-related info
        self.zim_info = ZimInfo(
            language=language,
            tags=tags,
            title=title,
            description=description,
            creator=creator,
            publisher=publisher,
            name=name,
        )

        # update youtube credentials store
        youtube_store.update(
            build_dir=self.build_dir, api_key=self.api_key, cache_dir=self.cache_dir
        )

        # set and record locale for translations
        locale_name = locale_name or get_language_details(self.language)["iso-639-1"]
        try:
            self.locale = setlocale(locale_name)
        except locale.Error:
            logger.error(
                f"No locale for {locale_name}. Use --locale to specify it. defaulting to en_US"
            )
            self.locale = setlocale("en")

    @property
    def root_dir(self):
        return ROOT_DIR

    @property
    def templates_dir(self):
        return self.root_dir.joinpath("templates")

    @property
    def assets_src_dir(self):
        return self.templates_dir.joinpath("assets")

    @property
    def assets_dir(self):
        return self.build_dir.joinpath("assets")

    @property
    def channels_dir(self):
        return self.build_dir.joinpath("channels")

    @property
    def cache_dir(self):
        return self.build_dir.joinpath("cache")

    @property
    def videos_dir(self):
        return self.build_dir.joinpath("videos")

    @property
    def profile_path(self):
        return self.build_dir.joinpath("profile.jpg")

    @property
    def banner_path(self):
        return self.build_dir.joinpath("banner.jpg")

    @property
    def is_user(self):
        return self.collection_type == USER

    @property
    def is_channel(self):
        return self.collection_type == CHANNEL

    @property
    def is_playlist(self):
        return self.collection_type == PLAYLIST

    @property
    def is_single_channel(self):
        if self.is_channel or self.is_user:
            return True
        return len(list(set([pl.creator_id for pl in self.playlists]))) == 1

    @property
    def sorted_playlists(self):
        """ sorted list of playlists (by title) but with Uploads one at first if any """
        if len(self.playlists) < 2:
            return self.playlists

        sorted_playlists = sorted(self.playlists, key=lambda x: x.title)
        index = 0
        # make sure our Uploads, special playlist is first
        if self.uploads_playlist_id:
            try:
                index = [
                    index
                    for index, p in enumerate(sorted_playlists)
                    if p.playlist_id == self.uploads_playlist_id
                ][-1]
            except Exception:
                index = 0
        return (
            [sorted_playlists[index]]
            + sorted_playlists[0:index]
            + sorted_playlists[index + 1 :]
        )

    def run(self):
        """ execute the scrapper step by step """
        logger.info(
            f"starting youtube scraper for {self.collection_type}#{self.youtube_id}"
        )

        logger.info("preparing build folder at {}".format(self.build_dir.resolve()))
        if not self.keep_build_dir and self.build_dir.exists():
            shutil.rmtree(self.cache_dir, ignore_errors=True)
            shutil.rmtree(self.build_dir)
        self.make_build_folder()

        logger.info("testing Youtbe credentials")
        if not credentials_ok():
            raise ValueError("Unable to connect to Youtube API v3. check `API_KEY`.")

        # fail early if supplied branding files are missing
        self.check_branding_values()

        logger.info("compute playlists list to retrieve")
        self.extract_playlists()

        logger.info(
            ".. {} playlists:\n   {}".format(
                len(self.playlists),
                "\n   ".join([p.playlist_id for p in self.playlists]),
            )
        )

        if self.only_test_branding:
            logger.info("update general metadata")
            self.update_metadata()

            logger.info("building fake homepage with branding")
            self.make_test_html_file()

            logger.info("done.")
            return

        logger.info("compute list of videos")
        self.extract_videos_list()
        logger.info(".. {} videos.".format(len(self.videos_ids)))

        # download videos (and recompress)
        logger.info(
            f"downloading all videos, subtitles and thumbnails (concurrency={self.max_concurrency})"
        )
        logger.info(f"  format: {self.video_format}")
        logger.info(f"  recompress: {self.low_quality}")
        logger.info(f"  generated-subtitles: {self.all_subtitles}")
        if not self.skip_download:
            succeeded, failed = self.download_video_files(
                max_concurrency=self.max_concurrency
            )
            if failed:
                logger.error(f"{len(failed)} video(s) failed to download: {failed}")
                if len(failed) >= len(succeeded):
                    logger.critical("More than half of videos failed. exiting")
                    raise IOError("Too much videos failed to download")

        logger.info("retrieve channel-info for all videos (author details)")
        get_videos_authors_info(succeeded)

        logger.info("download all author's profile pictures")
        self.download_authors_branding()

        logger.info("update general metadata")
        self.update_metadata()

        logger.info("creating HTML files")
        self.make_html_files(succeeded)

        # make zim file
        if not self.no_zim:
            period = datetime.datetime.now().strftime("%Y-%m")
            self.fname = Path(self.fname if self.fname else f"{self.name}_{period}.zim")
            logger.info("building ZIM file")
            print(self.zim_info.to_zimwriterfs_args())
            make_zim_file(self.build_dir, self.output_dir, self.fname, self.zim_info)
            logger.info("removing HTML folder")
            if not self.keep_build_dir:
                shutil.rmtree(self.build_dir, ignore_errors=True)

        logger.info("all done!")

    def make_build_folder(self):
        """ prepare build folder before we start downloading data """

        # create build folder
        os.makedirs(self.build_dir, exist_ok=True)

        # copy assets
        if self.assets_dir.exists():
            shutil.rmtree(self.assets_dir)
        shutil.copytree(self.assets_src_dir, self.assets_dir)

        # cache folder to store youtube-api results
        self.cache_dir.mkdir(exist_ok=True)

        # make videos placeholder
        self.videos_dir.mkdir(exist_ok=True)

        # make channels placeholder (profile files)
        self.channels_dir.mkdir(exist_ok=True)

    def check_branding_values(self):
        """ checks that user-supplied images and colors are valid (so to fail early)

            Images are checked for existence or downloaded then resized
            Colors are check for validity """

        # skip this step if none of related values were supplied
        if not sum(
            [
                bool(x)
                for x in (
                    self.profile_image,
                    self.banner_image,
                    self.main_color,
                    self.secondary_color,
                )
            ]
        ):
            return
        logger.info("checking your branding files and values")
        if self.profile_image:
            if self.profile_image.startswith("http"):
                save_file(self.profile_image, self.profile_path)
            else:
                if not self.profile_image.exists():
                    raise IOError(
                        f"--profile image could not be found: {self.profile_image}"
                    )
                shutil.move(self.profile_image, self.profile_path)
            resize_image(self.profile_path, width=100, height=100, method="thumbnail")
        if self.banner_image:
            if self.banner_image.startswith("http"):
                save_file(self.banner_image, self.banner_path)
            else:
                if not self.banner_image.exists():
                    raise IOError(
                        f"--banner image could not be found: {self.banner_image}"
                    )
                shutil.move(self.banner_image, self.banner_path)
            resize_image(self.banner_path, width=1060, height=175, method="thumbnail")

        if self.main_color and not is_hex_color(self.main_color):
            raise ValueError(
                f"--main-color is not a valid hex color: {self.main_color}"
            )

        if self.secondary_color and not is_hex_color(self.secondary_color):
            raise ValueError(
                f"--secondary_color-color is not a valid hex color: {self.secondary_color}"
            )

    def extract_playlists(self):
        """ prepare a list of Playlist from user request

            USER: we fetch the hidden channel associate to it
            CHANNEL (and USER): we grab all playlists + `uploads` playlist
            PLAYLIST: we retrieve from the playlist Id(s) """

        if self.is_user or self.is_channel:
            if self.is_user:
                # youtube_id is a Username, fetch actual channelId through channel
                channel_json = get_channel_json(self.youtube_id, for_username=True)
            else:
                # youtube_id is a channelId
                channel_json = get_channel_json(self.youtube_id)

            self.main_channel_id = channel_json["id"]

            # retrieve list of playlists for that channel
            playlist_ids = [
                p["id"] for p in get_channel_playlists_json(self.main_channel_id)
            ]
            # we always include uploads playlist (contains everything)
            playlist_ids += [
                channel_json["contentDetails"]["relatedPlaylists"]["uploads"]
            ]
            self.uploads_playlist_id = playlist_ids[-1]
        elif self.is_playlist:
            playlist_ids = self.youtube_id.split(",")
            self.main_channel_id = Playlist.from_id(playlist_ids[0]).creator_id
        else:
            raise NotImplementedError("unsupported collection_type")

        self.playlists = [
            Playlist.from_id(playlist_id) for playlist_id in list(set(playlist_ids))
        ]

    def extract_videos_list(self):
        all_videos = load_json(self.cache_dir, "videos")
        if all_videos is None:
            all_videos = {}

            # we only return video_ids that we'll use later on. per-playlist JSON stored
            for playlist in self.playlists:
                all_videos.update(
                    {
                        v["contentDetails"]["videoId"]: v
                        for v in filter(
                            skip_deleted_videos, get_videos_json(playlist.playlist_id)
                        )
                    }
                )

            # self.videos_ids = list(set(videos_ids))
            save_json(self.cache_dir, "videos", all_videos)
        self.videos_ids = [*all_videos.keys()]  # unpacking so it's subscriptable

    def download_video_files(self, max_concurrency):

        # prepare options which are shared with every downloader
        options = {
            "cachedir": self.videos_dir,
            "writethumbnail": True,
            "write_all_thumbnails": True,
            "writesubtitles": True,
            "allsubtitles": True,
            "subtitlesformat": "vtt",
            "keepvideo": False,
            "ignoreerrors": False,
            "retries": 20,
            "fragment-retries": 50,
            "skip-unavailable-fragments": True,
            # "external_downloader": "aria2c",
            # "external_downloader_args": ["--max-tries=20", "--retry-wait=30"],
            "outtmpl": str(self.videos_dir.joinpath("%(id)s", "video.%(ext)s")),
            "preferredcodec": self.video_format,
            "format": "{fmt}/best".format(fmt=self.video_format),
            "progress_hooks": [
                partial(hook_youtube_dl_ffmpeg, self.video_format, self.low_quality)
            ],
        }
        if self.all_subtitles:
            options.update({"writeautomaticsub": True})

        # find number of actuall parallel workers
        nb_videos = len(self.videos_ids)
        concurrency = nb_videos if nb_videos < max_concurrency else max_concurrency

        # short-circuit concurency if we have only one thread (can help debug)
        if concurrency <= 1:
            return self.download_video_files_batch(options, self.videos_ids)

        # prepare out videos_ids batches
        def get_slot():
            n = 0
            while True:
                yield n
                n += 1
                if n >= concurrency:
                    n = 0

        batches = [[] for _ in range(0, concurrency)]
        slot = get_slot()
        for video_id in self.videos_ids:
            batches[next(slot)].append(video_id)

        overall_succeeded = []
        overall_failed = []
        # execute the batches concurrently
        with concurrent.futures.ThreadPoolExecutor(max_workers=concurrency) as executor:
            fs = [
                executor.submit(self.download_video_files_batch, options, videos_ids)
                for videos_ids in batches
            ]
            done, not_done = concurrent.futures.wait(
                fs, return_when=concurrent.futures.ALL_COMPLETED
            )

            # we have some `not_done` batches, indicating errors within
            if not_done:
                logger.critical(
                    "Not all video-processing batches completed. Cancelling…"
                )
                for future in not_done:
                    exc = future.exception()
                    if exc:
                        logger.exception(exc)
                        raise exc

            # retrieve our list of successful/failed video_ids
            for future in done:
                succeeded, failed = future.result()
                overall_succeeded += succeeded
                overall_failed += failed

        # remove left-over files for failed downloads
        logger.debug(f"removing left-over files of {len(overall_failed)} failed videos")
        for video_id in overall_failed:
            shutil.rmtree(self.videos_dir.joinpath(video_id), ignore_errors=True)

        return overall_succeeded, overall_failed

    @staticmethod
    def download_video_files_batch(options, videos_ids):

        succeeded = []
        failed = []
        with youtube_dl.YoutubeDL(options) as ydl:
            for video_id in videos_ids:
                try:
                    ydl.download([video_id])
                    succeeded.append(video_id)
                except youtube_dl.utils.DownloadError:
                    failed.append(video_id)
        return succeeded, failed

    def download_authors_branding(self):
        videos_channels_json = load_json(self.cache_dir, "videos_channels")
        uniq_channel_ids = list(
            set([chan["channelId"] for chan in videos_channels_json.values()])
        )
        for channel_id in uniq_channel_ids:
            save_channel_branding(self.channels_dir, channel_id, save_banner=False)

    def update_metadata(self):
        # we use title, description, profile and banner of channel/user
        # or channel of first playlist
        main_channel_json = get_channel_json(self.main_channel_id)
        save_channel_branding(self.channels_dir, self.main_channel_id, save_banner=True)

        # if a single playlist was requested, use if for names;
        # otherwise, use main_channel's details.
        auto_title = (
            self.playlists[0].title
            if self.is_playlist and len(self.playlists) == 1
            else main_channel_json["snippet"]["title"].strip()
        )
        auto_description = (
            clean_text(self.playlists[0].description)
            if self.is_playlist and len(self.playlists) == 1
            else clean_text(main_channel_json["snippet"]["description"])
        )
        self.title = self.title or auto_title or "-"
        self.description = self.description or auto_description or "-"

        if self.creator is None:
            if self.is_single_channel:
                self.creator = _("Youtube Channel “{title}”").format(
                    title=main_channel_json["snippet"]["title"]
                )
            else:
                self.creator = _("Youtube Channels")
        self.publisher = self.publisher or "Kiwix"

        self.tags = self.tags or ["youtube"]
        if "_videos:yes" not in self.tags:
            self.tags.append("_videos:yes")

        self.zim_info.update(
            title=self.title,
            description=self.description,
            creator=self.creator,
            publisher=self.publisher,
            name=self.name,
            tags=self.tags,
        )

        # copy our main_channel branding into /(profile|banner).jpg if not supplied
        if not self.profile_path.exists():
            shutil.copy(
                self.channels_dir.joinpath(self.main_channel_id, "profile.jpg"),
                self.profile_path,
            )
        if not self.banner_path.exists():
            shutil.copy(
                self.channels_dir.joinpath(self.main_channel_id, "banner.jpg"),
                self.banner_path,
            )

        # set colors from images if not supplied
        if self.main_color is None or self.secondary_color is None:
            profile_main, profile_secondary = get_colors(self.profile_path)
        self.main_color = self.main_color or profile_main
        self.secondary_color = self.secondary_color or profile_secondary

        resize_image(
            self.profile_path,
            width=48,
            height=48,
            method="thumbnail",
            to=self.build_dir.joinpath("favicon.jpg"),
        )

    def make_test_html_file(self):

        video_id = "aaaabbbccddd"
        video_dir = self.videos_dir.joinpath(video_id)
        video_dir.mkdir(exist_ok=True)
        shutil.copy(
            self.assets_dir.joinpath("sample.jpg"), video_dir.joinpath("video.jpg")
        )

        env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(str(self.templates_dir)), autoescape=True
        )

        html = env.get_template("home.html").render(
            playlists=self.sorted_playlists,
            video_format=self.video_format,
            title=self.title,
            description=self.description,
            color=self.main_color,
            background_color=self.secondary_color,
            page_label=_("Page {current}/{total}"),
            back_label=_("Back to top"),
        )
        with open(self.build_dir.joinpath("home.html"), "w", encoding="utf-8") as fp:
            fp.write(html)

        with open(self.assets_dir.joinpath("db.js"), "w", encoding="utf-8") as fp:
            fp.write(
                env.get_template("assets/db.js").render(
                    NB_VIDEOS_PER_PAGE=self.nb_videos_per_page
                )
            )

        with open(self.assets_dir.joinpath("data.js"), "w", encoding="utf-8") as fp:
            videos = [
                {
                    "id": video_id,
                    "title": "A Sample Title",
                    "slug": "a-sample-title",
                    "description": "A sample description of video",
                    "subtitles": [],
                    "thumbnail": str(Path("videos").joinpath(video_id, "video.jpg")),
                }
                for _ in range(0, 6)
            ]
            fp.write(
                "var json_{slug} = {json_str};\n".format(
                    slug=self.sorted_playlists[0].slug,
                    json_str=json.dumps(videos, indent=4),
                )
            )

    def make_html_files(self, actual_videos_ids):
        """ make up HTML structure to read the content

        /home.html                                  Homepage

        for each video:
            - <slug-title>.html                     HTML article
            - videos/<videoId>/video.<ext>          video file
            - videos/<videoId>/video.<lang>.vtt     subtititle(s)
            - videos/<videoId>/video.jpg            template
        """

        def is_present(video):
            """ whether this video has actually been succeffuly downloaded """
            return video["contentDetails"]["videoId"] in actual_videos_ids

        def get_subtitles(video_id):
            video_dir = self.videos_dir.joinpath(video_id)
            languages = [
                x.stem.split(".")[1]
                for x in video_dir.iterdir()
                if x.is_file() and x.name.endswith(".vtt")
            ]

            return [get_language_details(language) for language in languages]

        env = jinja2.Environment(
            loader=jinja2.FileSystemLoader(str(self.templates_dir)), autoescape=True
        )

        videos = load_json(self.cache_dir, "videos").values()
        # filter videos so we only include the ones we could retrieve
        videos = list(filter(is_present, videos))
        videos_channels = load_json(self.cache_dir, "videos_channels")
        for video in videos:
            video_id = video["contentDetails"]["videoId"]
            title = video["snippet"]["title"]
            slug = get_slug(title)
            description = video["snippet"]["description"]
            publication_date = dt_parser.parse(
                video["contentDetails"]["videoPublishedAt"]
            )
            author = videos_channels[video_id]
            subtitles = get_subtitles(video_id)
            video_url = f"https://www.youtube.com/watch?v={video_id}"

            html = env.get_template("article.html").render(
                video_id=video_id,
                video_format=self.video_format,
                author=author,
                title=title,
                description=description,
                date=format_date(publication_date, format="medium", locale=self.locale),
                subtitles=subtitles,
                url=video_url,
                channel_id=video["snippet"]["channelId"],
                color=self.main_color,
                background_color=self.secondary_color,
                autoplay=self.autoplay,
            )
            with open(
                self.build_dir.joinpath(f"{slug}.html"), "w", encoding="utf-8"
            ) as fp:
                fp.write(html)

        # build homepage
        html = env.get_template("home.html").render(
            playlists=self.playlists,
            video_format=self.video_format,
            title=self.title,
            description=self.description,
            color=self.main_color,
            background_color=self.secondary_color,
            page_label=_("Page {current}/{total}"),
            back_label=_("Back to top"),
        )
        with open(self.build_dir.joinpath("home.html"), "w", encoding="utf-8") as fp:
            fp.write(html)

        # rewrite app.js including `format`
        with open(self.assets_dir.joinpath("app.js"), "w", encoding="utf-8") as fp:
            fp.write(
                env.get_template("assets/app.js").render(video_format=self.video_format)
            )

        # rewrite app.js including `pagination`
        with open(self.assets_dir.joinpath("db.js"), "w", encoding="utf-8") as fp:
            fp.write(
                env.get_template("assets/db.js").render(
                    NB_VIDEOS_PER_PAGE=self.nb_videos_per_page
                )
            )

        # write list of videos in data.js
        def to_data_js(video):
            return {
                "id": video["contentDetails"]["videoId"],
                "title": video["snippet"]["title"],
                "slug": get_slug(video["snippet"]["title"]),
                "description": video["snippet"]["description"],
                "subtitles": get_subtitles(video["contentDetails"]["videoId"]),
                "thumbnail": str(
                    Path("videos").joinpath(
                        video["contentDetails"]["videoId"], "video.jpg"
                    )
                ),
            }

        with open(self.assets_dir.joinpath("data.js"), "w", encoding="utf-8") as fp:
            # write all playlists as they are
            for playlist in self.playlists:
                # retrieve list of videos for PL
                playlist_videos = load_json(
                    self.cache_dir, f"playlist_{playlist.playlist_id}_videos"
                )
                # filtering-out missing ones (deleted or not downloaded)
                playlist_videos = list(filter(skip_deleted_videos, playlist_videos))
                playlist_videos = list(filter(is_present, playlist_videos))
                # sorting them based on playlist
                playlist_videos.sort(key=lambda v: v["snippet"]["position"])

                fp.write(
                    "var json_{slug} = {json_str};\n".format(
                        slug=playlist.slug,
                        json_str=json.dumps(
                            list(map(to_data_js, playlist_videos)), indent=4
                        ),
                    )
                )

        # write a metadata.json file with some content-related data
        with open(
            self.build_dir.joinpath("metadata.json"), "w", encoding="utf-8"
        ) as fp:
            json.dump({"video_format": self.video_format}, fp, indent=4)
