#!/usr/bin/env python

import re
import json
from urllib import parse
from html import unescape
from xml.dom.minidom import parseString

from lulu.common import (
    match1,
    url_size,
    urls_size,
    get_content,
    mime_to_container,
    parse_query_param,
    download_url_ffmpeg,
)
from lulu.util import log
from lulu.config import YOUTUBE_CODECS
from lulu.extractor import VideoExtractor


class YouTube(VideoExtractor):
    name = 'YouTube'

    # YouTube media encoding options, in descending quality order.
    # http://en.wikipedia.org/wiki/YouTube#Quality_and_codecs.
    # Retrieved July 17, 2014.
    stream_types = YOUTUBE_CODECS

    def decipher(js, s):
        def tr_js(code):
            code = re.sub(r'function', r'def', code)
            code = re.sub(r'(\W)(as|if|in|is|or)\(', r'\1_\2(', code)
            code = re.sub(r'\$', '_dollar', code)
            code = re.sub(r'\{', r':\n\t', code)
            code = re.sub(r'\}', r'\n', code)
            code = re.sub(r'var\s+', r'', code)
            code = re.sub(r'(\w+).join\(""\)', r'"".join(\1)', code)
            code = re.sub(r'(\w+).length', r'len(\1)', code)
            code = re.sub(r'(\w+).slice\((\w+)\)', r'\1[\2:]', code)
            code = re.sub(
                r'(\w+).splice\((\w+),(\w+)\)', r'del \1[\2:\2+\3]', code
            )
            code = re.sub(r'(\w+).split\(""\)', r'list(\1)', code)
            return code

        js = js.replace('\n', ' ')
        f1 = match1(js, r'"signature",([$\w]+)\(\w+\.\w+\)')
        f1def = match1(
            js, r'function %s(\(\w+\)\{[^\{]+\})' % re.escape(f1)
        ) or match1(js, r'\W%s=function(\(\w+\)\{[^\{]+\})' % re.escape(f1))
        f1def = re.sub(r'([$\w]+\.)([$\w]+\(\w+,\d+\))', r'\2', f1def)
        f1def = 'function %s%s' % (f1, f1def)
        code = tr_js(f1def)
        f2s = set(re.findall(r'([$\w]+)\(\w+,\d+\)', f1def))
        for f2 in f2s:
            f2e = re.escape(f2)
            f2def = re.search(
                r'[^$\w]%s:function\((\w+,\w+)\)(\{[^\{\}]+\})' % f2e, js
            )
            if f2def:
                f2def = 'function {}({}){}'.format(
                    f2e, f2def.group(1), f2def.group(2)
                )
            else:
                f2def = re.search(
                    r'[^$\w]%s:function\((\w+)\)(\{[^\{\}]+\})' % f2e, js
                )
                f2def = 'function {}({},b){}'.format(
                    f2e, f2def.group(1), f2def.group(2)
                )
            f2 = re.sub(r'(\W)(as|if|in|is|or)\(', r'\1_\2(', f2)
            f2 = re.sub(r'\$', '_dollar', f2)
            code = code + 'global %s\n' % f2 + tr_js(f2def)

        f1 = re.sub(r'(as|if|in|is|or)', r'_\1', f1)
        f1 = re.sub(r'\$', '_dollar', f1)
        code = code + 'sig=%s(s)' % f1
        exec(code, globals(), locals())
        return locals()['sig']

    def get_url_from_vid(vid):
        return 'https://youtu.be/{}'.format(vid)

    def get_vid_from_url(url):
        """Extracts video ID from URL.
        """
        return match1(url, r'youtu\.be/([^?/]+)') or \
            match1(url, r'youtube\.com/embed/([^/?]+)') or \
            match1(url, r'youtube\.com/v/([^/?]+)') or \
            match1(url, r'youtube\.com/watch/([^/?]+)') or \
            parse_query_param(url, 'v') or \
            parse_query_param(parse_query_param(url, 'u'), 'v')

    def get_playlist_id_from_url(url):
        """Extracts playlist ID from URL.
        """
        return parse_query_param(url, 'list') or \
            parse_query_param(url, 'p')

    def download_playlist_by_url(self, url, **kwargs):
        self.url = url

        playlist_id = self.__class__.get_playlist_id_from_url(self.url)
        if playlist_id is None:
            log.wtf('[Failed] Unsupported URL pattern.')

        video_page = get_content(
            'https://www.youtube.com/playlist?list={}'.format(playlist_id)
        )
        from html.parser import HTMLParser
        videos = sorted(
            [HTMLParser().unescape(video) for video in re.findall(
                r'<a href="(/watch\?[^"]+)"', video_page
            ) if parse_query_param(video, 'index')],
            key=lambda video: parse_query_param(video, 'index')
        )

        # Parse browse_ajax page for more videos to load
        load_more_href = match1(
            video_page, r'data-uix-load-more-href="([^"]+)"'
        )
        while load_more_href:
            browse_ajax = get_content(
                'https://www.youtube.com/{}'.format(load_more_href)
            )
            browse_data = json.loads(browse_ajax)
            load_more_widget_html = browse_data['load_more_widget_html']
            content_html = browse_data['content_html']
            vs = set(re.findall(r'href="(/watch\?[^"]+)"', content_html))
            videos += sorted(
                [HTMLParser().unescape(video) for video in list(vs)
                    if parse_query_param(video, 'index')]
            )
            load_more_href = match1(
                load_more_widget_html, r'data-uix-load-more-href="([^"]+)"'
            )

        self.title = re.search(
            r'<meta name="title" content="([^"]+)"', video_page
        ).group(1)
        self.p_playlist()
        for video in videos:
            vid = parse_query_param(video, 'v')
            index = parse_query_param(video, 'index')
            self.__class__().download_by_url(
                self.__class__.get_url_from_vid(vid), index=index, **kwargs
            )

    def prepare(self, **kwargs):
        assert self.url or self.vid

        if not self.vid and self.url:
            self.vid = self.__class__.get_vid_from_url(self.url)

            if self.vid is None:
                self.download_playlist_by_url(self.url, **kwargs)
                exit(0)

        video_info = parse.parse_qs(
            get_content(
                'https://www.youtube.com/get_video_info?video_id={}'.format(
                    self.vid
                )
            )
        )

        ytplayer_config = None
        if 'status' not in video_info:
            log.wtf('[Failed] Unknown status.')
        elif video_info['status'] == ['ok']:
            if 'use_cipher_signature' not in video_info \
                    or video_info['use_cipher_signature'] == ['False']:
                self.title = parse.unquote_plus(video_info['title'][0])

                # Parse video page (for DASH)
                video_page = get_content(
                    'https://www.youtube.com/watch?v={}'.format(self.vid)
                )
                try:
                    ytplayer_config = json.loads(re.search(
                        'ytplayer.config\s*=\s*([^\n]+?});', video_page
                    ).group(1))
                    self.html5player = 'https://www.youtube.com{}'.format(
                        ytplayer_config['assets']['js']
                    )
                    # Workaround: get_video_info returns bad s. Why?
                    stream_list = ytplayer_config['args'][
                        'url_encoded_fmt_stream_map'
                    ].split(',')
                except Exception:
                    stream_list = video_info[
                        'url_encoded_fmt_stream_map'
                    ][0].split(',')
                    self.html5player = None

            else:
                # Parse video page instead
                video_page = get_content(
                    'https://www.youtube.com/watch?v={}'.format(self.vid)
                )
                ytplayer_config = json.loads(re.search(
                    'ytplayer.config\s*=\s*([^\n]+?});', video_page
                ).group(1))

                self.title = ytplayer_config['args']['title']
                self.html5player = 'https://www.youtube.com{}'.format(
                    ytplayer_config['assets']['js']
                )
                stream_list = ytplayer_config['args'][
                    'url_encoded_fmt_stream_map'
                ].split(',')

        elif video_info['status'] == ['fail']:
            if video_info['errorcode'] == ['150']:
                video_page = get_content(
                    'https://www.youtube.com/watch?v={}'.format(self.vid)
                )
                try:
                    ytplayer_config = json.loads(re.search(
                        'ytplayer.config\s*=\s*([^\n]+});ytplayer', video_page
                    ).group(1))
                except Exception:
                    msg = re.search(
                        'class="message">([^<]+)<', video_page
                    ).group(1)
                    log.wtf('[Failed] "%s"' % msg.strip())

                if 'title' in ytplayer_config['args']:
                    # 150 Restricted from playback on certain sites
                    # Parse video page instead
                    self.title = ytplayer_config['args']['title']
                    self.html5player = 'https://www.youtube.com{}'.format(
                        ytplayer_config['assets']['js']
                    )
                    stream_list = ytplayer_config['args'][
                        'url_encoded_fmt_stream_map'
                    ].split(',')
                else:
                    log.wtf(
                        '[Error] The uploader has not made this video '
                        'available in your country.'
                    )

            elif video_info['errorcode'] == ['100']:
                log.wtf(
                    '[Failed] This video does not exist.',
                    exit_code=int(video_info['errorcode'][0])
                )

            else:
                log.wtf(
                    '[Failed] %s' % video_info['reason'][0],
                    exit_code=int(video_info['errorcode'][0])
                )

        else:
            log.wtf('[Failed] Invalid status.')

        # YouTube Live
        if ytplayer_config and (
            ytplayer_config['args'].get('livestream') == '1' or
            ytplayer_config['args'].get('live_playback') == '1'
        ):
            hlsvp = ytplayer_config['args']['hlsvp']

            if 'info_only' in kwargs and kwargs['info_only']:
                return
            else:
                download_url_ffmpeg(hlsvp, self.title, 'mp4')
                exit(0)

        for stream in stream_list:
            metadata = parse.parse_qs(stream)
            stream_itag = metadata['itag'][0]
            self.streams[stream_itag] = {
                'itag': metadata['itag'][0],
                'url': metadata['url'][0],
                'sig': metadata['sig'][0] if 'sig' in metadata else None,
                's': metadata['s'][0] if 's' in metadata else None,
                'quality': metadata['quality'][0],
                'type': metadata['type'][0],
                'mime': metadata['type'][0].split(';')[0],
                'container': mime_to_container(
                    metadata['type'][0].split(';')[0]
                ),
            }

        # Prepare caption tracks
        try:
            caption_tracks = json.loads(
                ytplayer_config['args']['player_response']
            )['captions']['playerCaptionsTracklistRenderer']['captionTracks']
            for ct in caption_tracks:
                ttsurl, lang = ct['baseUrl'], ct['languageCode']

                tts_xml = parseString(get_content(ttsurl))
                transcript = tts_xml.getElementsByTagName('transcript')[0]
                texts = transcript.getElementsByTagName('text')
                srt = ''
                seq = 0
                for text in texts:
                    if text.firstChild is None:
                        continue  # empty element
                    seq += 1
                    start = float(text.getAttribute('start'))
                    if text.getAttribute('dur'):
                        dur = float(text.getAttribute('dur'))
                    else:
                        dur = 1.0  # could be ill-formed XML
                    finish = start + dur
                    m, s = divmod(start, 60)
                    h, m = divmod(m, 60)
                    start = '{:0>2}:{:0>2}:{:06.3f}'.format(
                        int(h), int(m), s
                    ).replace('.', ',')
                    m, s = divmod(finish, 60)
                    h, m = divmod(m, 60)
                    finish = '{:0>2}:{:0>2}:{:06.3f}'.format(
                        int(h), int(m), s
                    ).replace('.', ',')
                    content = unescape(text.firstChild.nodeValue)

                    srt += '{}\n'.format(str(seq))
                    srt += '{} --> {}\n'.format(start, finish)
                    srt += '{}\n\n'.format(content)

                self.caption_tracks[lang] = srt
        except Exception:
            pass

        # Prepare DASH streams
        try:
            dashmpd = ytplayer_config['args']['dashmpd']
            dash_xml = parseString(get_content(dashmpd))
            for aset in dash_xml.getElementsByTagName('AdaptationSet'):
                mimeType = aset.getAttribute('mimeType')
                if mimeType == 'audio/mp4':
                    rep = aset.getElementsByTagName('Representation')[-1]
                    burls = rep.getElementsByTagName('BaseURL')
                    dash_mp4_a_url = burls[0].firstChild.nodeValue
                    dash_mp4_a_size = burls[0].getAttribute('yt:contentLength')
                    if not dash_mp4_a_size:
                        try:
                            dash_mp4_a_size = url_size(dash_mp4_a_url)
                        except Exception:
                            continue
                elif mimeType == 'audio/webm':
                    rep = aset.getElementsByTagName('Representation')[-1]
                    burls = rep.getElementsByTagName('BaseURL')
                    dash_webm_a_url = burls[0].firstChild.nodeValue
                    dash_webm_a_size = burls[0].getAttribute(
                        'yt:contentLength'
                    )
                    if not dash_webm_a_size:
                        try:
                            dash_webm_a_size = url_size(dash_webm_a_url)
                        except Exception:
                            continue
                elif mimeType == 'video/mp4':
                    for rep in aset.getElementsByTagName('Representation'):
                        w = int(rep.getAttribute('width'))
                        h = int(rep.getAttribute('height'))
                        itag = rep.getAttribute('id')
                        burls = rep.getElementsByTagName('BaseURL')
                        dash_url = burls[0].firstChild.nodeValue
                        dash_size = burls[0].getAttribute('yt:contentLength')
                        if not dash_size:
                            try:
                                dash_size = url_size(dash_url)
                            except Exception:
                                continue
                        self.dash_streams[itag] = {
                            'quality': '{}x{}'.format(w, h),
                            'itag': itag,
                            'type': mimeType,
                            'mime': mimeType,
                            'container': 'mp4',
                            'src': [dash_url, dash_mp4_a_url],
                            'size': int(dash_size) + int(dash_mp4_a_size)
                        }
                elif mimeType == 'video/webm':
                    for rep in aset.getElementsByTagName('Representation'):
                        w = int(rep.getAttribute('width'))
                        h = int(rep.getAttribute('height'))
                        itag = rep.getAttribute('id')
                        burls = rep.getElementsByTagName('BaseURL')
                        dash_url = burls[0].firstChild.nodeValue
                        dash_size = burls[0].getAttribute('yt:contentLength')
                        if not dash_size:
                            try:
                                dash_size = url_size(dash_url)
                            except Exception:
                                continue
                        self.dash_streams[itag] = {
                            'quality': '%sx%s' % (w, h),
                            'itag': itag,
                            'type': mimeType,
                            'mime': mimeType,
                            'container': 'webm',
                            'src': [dash_url, dash_webm_a_url],
                            'size': int(dash_size) + int(dash_webm_a_size)
                        }
        except Exception:
            # VEVO
            if not self.html5player:
                return
            self.js = get_content(self.html5player)
            if 'adaptive_fmts' in ytplayer_config['args']:
                streams = [
                    dict(
                        [
                            (
                                i.split('=')[0],
                                parse.unquote(i.split('=')[1])
                            ) for i in afmt.split('&')
                        ]
                    ) for afmt in ytplayer_config['args'][
                        'adaptive_fmts'
                    ].split(',')
                ]
                for stream in streams:  # get over speed limiting
                    stream['url'] += '&ratebypass=yes'
                for stream in streams:  # audio
                    if stream['type'].startswith('audio/mp4'):
                        dash_mp4_a_url = stream['url']
                        if 's' in stream:
                            sig = self.__class__.decipher(self.js, stream['s'])
                            dash_mp4_a_url += '&signature={}'.format(sig)
                        dash_mp4_a_size = stream['clen']
                    elif stream['type'].startswith('audio/webm'):
                        dash_webm_a_url = stream['url']
                        if 's' in stream:
                            sig = self.__class__.decipher(self.js, stream['s'])
                            dash_webm_a_url += '&signature={}'.format(sig)
                        dash_webm_a_size = stream['clen']
                for stream in streams:  # video
                    if 'size' in stream:
                        if stream['type'].startswith('video/mp4'):
                            mimeType = 'video/mp4'
                            dash_url = stream['url']
                            if 's' in stream:
                                sig = self.__class__.decipher(
                                    self.js, stream['s']
                                )
                                dash_url += '&signature={}'.format(sig)
                            dash_size = stream['clen']
                            itag = stream['itag']
                            self.dash_streams[itag] = {
                                'quality': stream['size'],
                                'itag': itag,
                                'type': mimeType,
                                'mime': mimeType,
                                'container': 'mp4',
                                'src': [dash_url, dash_mp4_a_url],
                                'size': int(dash_size) + int(dash_mp4_a_size)
                            }
                        elif stream['type'].startswith('video/webm'):
                            mimeType = 'video/webm'
                            dash_url = stream['url']
                            if 's' in stream:
                                sig = self.__class__.decipher(
                                    self.js, stream['s']
                                )
                                dash_url += '&signature={}'.format(sig)
                            dash_size = stream['clen']
                            itag = stream['itag']
                            self.dash_streams[itag] = {
                                'quality': stream['size'],
                                'itag': itag,
                                'type': mimeType,
                                'mime': mimeType,
                                'container': 'webm',
                                'src': [dash_url, dash_webm_a_url],
                                'size': int(dash_size) + int(dash_webm_a_size)
                            }

    def extract(self, **kwargs):
        if not self.streams_sorted:
            # No stream is available
            return

        if 'stream_id' in kwargs and kwargs['stream_id']:
            # Extract the stream
            stream_id = kwargs['stream_id']
            if stream_id not in self.streams \
                    and stream_id not in self.dash_streams:
                log.e('[Error] Invalid video format.')
                log.e(
                    'Run \'-i\' command with no specific video format to '
                    'view all available formats.'
                )
                exit(2)
        else:
            # Extract stream with the best quality
            stream_id = self.streams_sorted[0]['itag']

        if stream_id in self.streams:
            src = self.streams[stream_id]['url']
            if self.streams[stream_id]['sig'] is not None:
                sig = self.streams[stream_id]['sig']
                src += '&signature={}'.format(sig)
            elif self.streams[stream_id]['s'] is not None:
                if not hasattr(self, 'js'):
                    self.js = get_content(self.html5player)
                s = self.streams[stream_id]['s']
                sig = self.__class__.decipher(self.js, s)
                src += '&signature={}'.format(sig)

            self.streams[stream_id]['src'] = [src]
            self.streams[stream_id]['size'] = urls_size(
                self.streams[stream_id]['src']
            )


site = YouTube()
download = site.download_by_url
download_playlist = site.download_playlist_by_url
