FastAnime/fastanime/Utility/downloader/downloader.py

import logging
import os
import shutil
import subprocess
import tempfile
from queue import Queue
from threading import Thread

import yt_dlp
from rich import print
from rich.prompt import Confirm
from yt_dlp.utils import sanitize_filename

logger = logging.getLogger(__name__)


class YtDLPDownloader:
    downloads_queue = Queue()

    def _worker(self):
        while True:
            task, args = self.downloads_queue.get()
            try:
                task(*args)
            except Exception as e:
                logger.error(f"Something went wrong {e}")
            self.downloads_queue.task_done()

    def __init__(self):
        self._thread = Thread(target=self._worker)
        self._thread.daemon = True
        self._thread.start()

    def _download_file(
        self,
        url: str,
        anime_title: str,
        episode_title: str,
        download_dir: str,
        silent: bool,
        vid_format: str = "best",
        force_unknown_ext=False,
        verbose=False,
        headers={},
        sub="",
        merge=False,
        clean=False,
        prompt=True,
    ):
        """Helper function that downloads anime given url and path details

        Args:
            url: [TODO:description]
            anime_title: [TODO:description]
            episode_title: [TODO:description]
            download_dir: [TODO:description]
            silent: [TODO:description]
            vid_format: [TODO:description]
        """
        anime_title = sanitize_filename(anime_title)
        episode_title = sanitize_filename(episode_title)
        ydl_opts = {
            # Specify the output path and template
            "http_headers": headers,
            "outtmpl": f"{download_dir}/{anime_title}/{episode_title}.%(ext)s",
            "silent": silent,
            "verbose": verbose,
            "format": vid_format,
            "compat_opts": ("allow-unsafe-ext",) if force_unknown_ext else tuple(),
        }
        urls = [url]
        if sub:
            urls.append(sub)
        vid_path = ""
        sub_path = ""
        for i, url in enumerate(urls):
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
            if not info:
                continue
            if i == 0:
                vid_path = info["requested_downloads"][0]["filepath"]
            else:
                sub_path = info["requested_downloads"][0]["filepath"]
        if sub_path and vid_path and merge:
            self.merge_subtitles(vid_path, sub_path, clean, prompt)

    def merge_subtitles(self, video_path, sub_path, clean, prompt):
        # Extract the directory and filename
        video_dir = os.path.dirname(video_path)
        video_name = os.path.basename(video_path)
        video_name, _ = os.path.splitext(video_name)
        video_name += ".mkv"

        FFMPEG_EXECUTABLE = shutil.which("ffmpeg")
        if not FFMPEG_EXECUTABLE:
            print("[yellow bold]WARNING: [/]FFmpeg not found")
            return
        # Create a temporary directory
        with tempfile.TemporaryDirectory() as temp_dir:
            # Temporary output path in the temporary directory
            temp_output_path = os.path.join(temp_dir, video_name)
            # FFmpeg command to merge subtitles
            command = [
                FFMPEG_EXECUTABLE,
                "-hide_banner",
                "-i",
                video_path,
                "-i",
                sub_path,
                "-c",
                "copy",
                "-map",
                "0",
                "-map",
                "1",
                temp_output_path,
            ]

            # Run the command
            try:
                subprocess.run(command, check=True)

                # Move the file back to the original directory with the original name
                final_output_path = os.path.join(video_dir, video_name)

                if os.path.exists(final_output_path):
                    if not prompt or Confirm.ask(
                        f"File exists({final_output_path}) would you like to overwrite it",
                        default=True,
                    ):
                        # move file to dest
                        os.remove(final_output_path)
                        shutil.move(temp_output_path, final_output_path)
                else:
                    shutil.move(temp_output_path, final_output_path)
                # clean up
                if clean:
                    print("[cyan]Cleaning original files...[/]")
                    os.remove(video_path)
                    os.remove(sub_path)

                print(
                    f"[green bold]Subtitles merged successfully.[/] Output file: {final_output_path}"
                )
            except subprocess.CalledProcessError as e:
                print(f"[red bold]Error[/] during merging subtitles: {e}")
            except Exception as e:
                print(f"[red bold]An error[/] occurred: {e}")

    # WARN: May remove this legacy functionality
    def download_file(self, url: str, title, silent=True):
        """A helper that just does things in the background

        Args:
            title ([TODO:parameter]): [TODO:description]
            silent ([TODO:parameter]): [TODO:description]
            url: [TODO:description]
        """
        self.downloads_queue.put((self._download_file, (url, title, silent)))


downloader = YtDLPDownloader()