# downloader.py
# This file is part of qarte-5
#    
# Author: Vincent Vande Vyvre <vincent.vandevyvre@oqapy.eu>
# Copyright: 2016-2025 Vincent Vande Vyvre
# Licence: GPL3
# Home page: https://launchpad.net/qarte

import os
import shutil
import contextlib
import urllib.request
import time
import shlex
import subprocess
import re
import string
from urllib.error import URLError
from threading import Thread
import logging
lgg = logging.getLogger(__name__)

import m3u8

from PyQt5.QtCore import QObject, pyqtSignal

from tools import SubtitleConverter
from gui.copywarning import CopyWarning

class Downloader(QObject):
    mergingFinished = pyqtSignal(str)
    def __init__(self, artetv):
        super().__init__()
        self.artetv = artetv
        self.tasks = {}
        self.idxs = []
        self.loaders = []
        self.current_hls = {}
        self.current_downloading = None
        self.is_running = False
        self.is_merging = 0
        self.no_merging = False
        self.mergingFinished.connect(self.on_merging_finished)

    def add_task(self, task):
        idx = task.idx
        self.tasks[idx] = task
        self.idxs.append(idx)
        lgg.info("Task added: %s, tot: %s" % (idx, len(self.tasks)))

    def config_task(self, idx, new):
        pass

    def remove_task(self, idx):
        lgg.info("Remove task %s" % idx)
        if idx >= len(self.idxs):
            return

        try:
            idx = self.idxs.pop(idx)
            self.tasks.pop(idx, None)
        except Exception as why:
            lgg.warning("Error when removing a task:\n\t%s" % why)

    def shift_task(self, src, dst):
        lgg.info("Shift %s to %s. Tasks: %s" %(src, dst, len(self.tasks)))
        try:
            idx = self.idxs.pop(src)
            self.idxs.insert(dst, idx)
        except Exception as why:
            lgg.warning("Error when shifting a task:\n\t%s" % why)

    def abort_task(self):
        self.no_merging = True
        for l in self.loaders:
            l.kill_request = True

        self.loaders = []
        self.is_running = False
        self.is_merging = 0

    def start_downloading(self):
        if not self.idxs:
            self.is_running = False
            self.artetv.on_empty_list()
            return

        self.is_running = True
        idx = self.idxs.pop(0)
        video = self.tasks[idx]
        self.artetv.ui.tv_basket.remove_first()
        if video.quality["is_hls"]:
            self.start_hls_downloading(video)
            return

        self.files_needed = []
        index = video.idx
        m3u8_url = video.quality["uri"]
        base_url = video.quality["base_url"]
        if not m3u8_url.startswith("http"):
            if base_url is None:
                lgg.warning("No base url for item %s" % video.title)
                return

            m3u8_url = os.path.join(base_url, video.quality["uri"])
        target = self.get_video_target(index)
        mp4_url = self.get_mp4_url(m3u8_url)
        self.files_needed.append([mp4_url, target])
        target = self.get_audio_target(index)
        m3u8_url = video.audio["url"]
        if not m3u8_url.startswith("http"):
            m3u8_url = os.path.join(base_url, video.audio["url"])

        mp4_url = self.get_mp4_url(m3u8_url)
        self.files_needed.append([mp4_url, target])
        if video.subtitle:
            target = self.get_subtitle_target(index)
            m3u8_url = video.subtitle["url"]
            if not m3u8_url.startswith("http"):
                m3u8_url = os.path.join(base_url, video.subtitle["url"])

            vtt_url  = self.get_vtt_url(m3u8_url)
            self.files_needed.append([vtt_url, target])

        self.run_loader(video)

    def start_hls_downloading(self, video):
        target = self.get_video_target(video.idx)
        segments = self.get_segments(video.quality["uri"])
        self.current_hls[video.idx] = self.get_video_path(video.title)
        loader = Loader()
        loader.loadingProgress.connect(self.artetv.display_progress)
        loader.loadingFinished.connect(self.on_downloading_hls_finished)
        loader.config_loader_by_segments(segments, target)
        self.loaders = [loader]
        Thread(target=loader.download_by_segments).start()

    def run_loader(self, item):
        video = self.files_needed[0]
        loader = Loader()
        loader.loadingProgress.connect(self.artetv.display_progress)
        loader.loadingFinished.connect(self.on_downloading_finished)
        if not self.initiate_loader(loader, video[0], video[1]):
            return

        # TODO check free space (artetv line 975)

        lgg.info('URL: %s' % video[0])
        self.merged_filename = self.get_final_file_name(item.idx)
        self.video_title = item.fname
        audio = self.files_needed[1]
        loader1 = Loader()
        if not self.initiate_loader(loader1, audio[0], audio[1]):
            return

        lgg.info('Start download audio %s size: %s' 
                    %(item.fname, loader1.size))
        Thread(target=loader1.download).start()
        lgg.info('Start download video %s size: %s' 
                    %(item.fname, loader.size))
        Thread(target=loader.download).start()
        self.loaders = [loader, loader1]

        if len(self.files_needed) > 2:
            subt = self.files_needed[2]
            loader2 = Loader()
            loader2.loadingFinished.connect(self.on_subtitle_loaded)
            if not self.initiate_loader(loader2, subt[0], subt[1]):
                return

            Thread(target=loader2.download).start()

    def initiate_loader(self, engine, url, fname):
        if not url:
            return False

        lgg.info("initiate loader: %s\n\t%s" %(url, fname))
        self.file_size = engine.config_loader(url, fname)
        if isinstance(self.file_size, int):
            return True

        lgg.info('Loading error: %s' % self.file_size)
        return False

    def get_segments(self, url):
        try:
            data = m3u8.load(url)
        except Exception as why:
            lgg.warning("Can't read segments: %s" % why)
            return False

        return data.segments

    def on_downloading_finished(self, fname):
        """Rename the video downloaded.

        Args:
        fname -- temporary file name
        """
        self.is_running = False
        self.artetv.on_downloading_finished(fname, len(self.tasks))
        if not fname:
            self.artetv.handle_downloading_error(self.loaders[0])

        else:
            if self.no_merging:
                # Downloading was canceled
                return

            if self.artetv.merge_files:
                self.merge_video_audio(fname)

    def on_downloading_hls_finished(self, fname):
        """Rename the video downloaded.

        Args:
        fname -- temporary file name
        """
        lgg.info("Downloading %s finished" % fname)
        self.is_running = False
        if not fname:
            self.artetv.handle_downloading_error(self.loaders[0])

        else:
            idx = os.path.basename(fname).split("v")[0]
            dest = self.current_hls.pop(idx)
            try:
                shutil.move(fname, dest)
                lgg.info('File renamed: %s' % os.path.basename(dest))
            except Exception as why:
                lgg.warning("Can't rename %s" % dest)
                lgg.warning("Reason: %s" % why)

            self.artetv.on_downloading_finished(os.path.basename(dest), len(self.task))
            self.start_downloading()

    def get_mp4_url(self, url):
        lgg.info("Get m3u8 file 0: \n   %s" % url)
        try:
            data = m3u8.load(url)
            base = data.base_uri
            if not base.endswith("/"):
                base += "/"

            if type(data.segment_map) == dict:
                return base + data.segment_map["uri"]

            return base + data.segment_map[0].uri
        except Exception as why:
            lgg.info("Error when reading m3u8:\n\t%s" % why)
            return False

    def get_vtt_url(self, url):
        lgg.info("Get m3u8 file 1: \n   %s" % url)
        try:
            data = m3u8.load(url)
            base = data.base_uri
            if not base.endswith("/"):
                base += "/"

            return base + data.files[0]
        except Exception as why:
            lgg.info("Error when reading m3u8:\n\t%s" % why)
            return False

    def get_video_target(self, fname):
        return os.path.join(self.artetv.temp_dir, fname+"video.mp4")

    def get_audio_target(self, fname):
        return os.path.join(self.artetv.temp_dir, fname+"audio.mp4")

    def get_subtitle_target(self, fname):
        return os.path.join(self.artetv.temp_dir, fname+"subtitle.vtt")

    def get_final_file_name(self, fname):
        return os.path.join(self.artetv.temp_dir, fname+"final.mp4")

    def on_subtitle_loaded(self, fname):
        if not os.path.isfile(fname):
            lgg.warning("File %s not found !" % fname)
            self.files_needed.pop(2)
            return

        inf = self.files_needed[2][1]
        converter = SubtitleConverter(inf)
        outf = inf.replace(".vtt", ".srt")
        converter.vtt_to_srt()
        self.files_needed[2][1] = outf

    def merge_video_audio(self, fname):
        lgg.info("Merge %s files" % len(self.files_needed))
        for i in self.files_needed[:2]:
            if not os.path.isfile(i[1]):
                lgg.warning("File %s not found !" % i[1])
                return

        com = False
        if len(self.files_needed) == 1:
            self.mergingFinished.emit(self.merged_filename)
            self.start_downloading()
            return

        if len(self.files_needed) == 2:
            com = 'ffmpeg -y -i %s -i %s -c:v copy -c:a aac "%s"' \
                    % (self.files_needed[0][1], self.files_needed[1][1],
                        self.merged_filename)

        elif len(self.files_needed) == 3 and os.path.isfile(
                                            self.files_needed[2][1]):
            com = self.build_command()

        self.is_merging += 1
        def merge(*args):
            cmd = shlex.split(args[0])
            r = subprocess.Popen(cmd, universal_newlines=True, 
                                 stdout=subprocess.PIPE, 
                                 stderr=subprocess.STDOUT).communicate()
            self.is_loading = False
            self.mergingFinished.emit(args[1])

        if com:
            Thread(target=merge, args=(com, self.merged_filename)).start()

        else:
            self.mergingFinished.emit(self.merged_filename)

        self.start_downloading()

    def build_command(self):
        video = self.files_needed[0][1]
        audio = self.files_needed[1][1]
        subt = self.files_needed[2][1]
        self.merged_filename = self.merged_filename.replace(".mp4", ".mkv")
        final = self.merged_filename
        cmd = 'ffmpeg -y -i "%s" -i "%s" -i "%s"' %(video, audio, subt)
        cmd += " -map 0 -map 1 -metadata:s:1 language=mul -metadata:s:1"
        cmd += " handler='mul' -metadata:s:1 title='mul' -map 2 -metadata:s:2"
        cmd += " language=fra -metadata:s:2 handler='fra (VO)'" 
        cmd += " -metadata:s:2 title='fra (VO)'"
        cmd += " -disposition:s -default -disposition 0 -c:v copy -c:a copy -c:s copy"
        cmd += ' "%s"' % final
        return cmd

    def on_merging_finished(self, fname):
        lgg.info("Merging finished, file: %s" % fname)
        ext = os.path.splitext(fname)[1]
        if not os.path.isfile(fname):
            lgg.warning("Merging failed ! file not present")
            self.is_merging -= 1
            return

        else:
            idx = os.path.split(fname)[1].rstrip("final.mp4kv")
            dest = self.get_video_path(self.tasks[idx].fname, ext)
            if fname.endswith("mkv"):
                dest = dest.replace(".mp4", ".mkv")

            try:
                shutil.move(fname, dest)
                lgg.info('File renamed: %s' % os.path.basename(dest))
            except OSError as why:
                if why.errno == 22:
                    lgg.warning("ERROR: %s" % why.strerror)
                    self.call_filename_error(orig, dest)
            except Exception as why:
                lgg.warning("Can't rename %s" % dest)
                lgg.warning("Reason: %s" % why)

        self.is_merging -= 1
        self.artetv.on_merging_finished(self.tasks[idx], dest)
            
    def get_video_path(self, title, ext=".mp4"):
        """Return the path/name of the video file.

        Args:
        title -- the video title
        """
        lgg.info("set path for: %s, %s" % (title, ext))
        fld = self.artetv.get_video_dir()
        path = os.path.join(fld, title + ext)
        count = 1
        while os.path.isfile(path):
            path = os.path.join(fld, title + "(%s)%s" %(count, ext))
            count += 1

        return path

    def call_filename_error(self, orig, target):
        path, base = os.path.split(target)
        fname, ext = os.path.splitext(base)
        self.new_name = ""
        warn = CopyWarning(fname, self)
        rep = warn.exec_()
        if not self.new_name:
            self.new_name = fname
        dest = os.path.join(path, self.new_name + ext)
        self.finalise(orig, dest)


class Loader(QObject):
    loadingProgress = pyqtSignal(int, int, int)
    loadingFinished = pyqtSignal(str)
    def __init__(self):
        super().__init__()
        self.is_alive = False

    def config_loader(self, url, filename):
        self.url = url
        self.filename = filename
        self.kill_request = False
        self.last_error = ''
        try:
            with contextlib.closing(urllib.request.urlopen(url, None, 6)) as fp:
                headers = fp.info()
                if "content-length" in headers:
                    self.size = int(headers["Content-Length"])
                else:
                    self.size = 0
        except URLError as why:
            self.size = why

        except Exception as why:
            self.size = why

        return self.size

    def download(self):
        """A hacked version of urllib.request.urlretrieve

        """
        self.is_alive = True
        size = -1
        try:
            with contextlib.closing(urllib.request.urlopen(self.url, 
                                    None, 6)) as fp:
                bs = 1024 * 16
                percent = 1
                blocknum = 0
                amount = 0
                blocks_percent = 0
                begin_at = time.time()
                if self.size:
                    blocks_percent = int(self.size / bs / 100)

                with open(self.filename, 'wb') as target:
                    while True:
                        block = fp.read(bs)
                        if not block or self.kill_request:
                            break

                        target.write(block)
                        blocknum += 1
                        amount += bs
                        if blocknum == blocks_percent:
                            elapsed = time.time() - begin_at
                            speed = int(amount / elapsed)
                            remain = int(self.size - amount)
                            self.loadingProgress.emit(percent, speed, remain)
                            percent += 1
                            blocknum = 0
        except URLError as why:
            self.last_error = why
            self.filename = ''

        except Exception as why:
            self.last_error = why
            self.filename = ''

        self.is_alive = False
        self.loadingFinished.emit(self.filename)

    def config_loader_by_segments(self, segments, filename):
        lgg.info("Configure downloading %s segments" % len(segments))
        self.segments = segments
        self.filename = filename
        self.kill_request = False
        self.last_error = ''

    def download_by_segments(self):
        """Downloader for protocol hls.

        """
        with open(self.filename, "wb") as target:
            self.is_alive = True
            lenght = len(self.segments)
            loaded = 0
            begin_at = time.time()
            blocks_percent = int(lenght/100)
            percent = 0
            amount = 0
            outf = self.filename
            for segment in self.segments:
                if self.kill_request:
                    break

                try:
                    with contextlib.closing(urllib.request.urlopen(
                                    segment.absolute_uri, None, 6)) as fp:
                        block = fp.read()
                        target.write(block)
                        loaded += 1
                        amount += len(block)
                        if loaded >= blocks_percent:
                            speed = amount / (time.time() - begin_at)
                            percent += 1
                            loaded = 0
                            self.loadingProgress.emit(percent, speed, 0)
                except Exception as why:
                    lgg.info("Downloading error: %s" % why)
                    self.last_error = why
                    outf = ""
                    break

        self.is_alive = False
        self.loadingFinished.emit(outf)

