monitor

#!/usr/bin/env python3

import collections
import json
import os
import pathlib
import re
import signal
import subprocess
import sys
import time
import urllib.parse

import arrow
import matplotlib.pyplot as plt
plt.rcParams.update({'font.family': 'Consolas', 'font.size': 14})
import requests
import youtube_dl

from lib import auth
from lib.config import INDEX, METADATA_DIR, PLOTS_DIR, VIEWERS_DIR, logger


global youtube
HERE = pathlib.Path(__file__).resolve().parent


class LivestreamEndedError(Exception):
    def __init__(self, video_id, end_time):
        super().__init__()
        self.video_id = video_id
        self.end_time = end_time

    def __str__(self):
        return f'{self.video_id} has ended on {self.end_time}'


def log_exception(exc, msg):
    logger.error(f'{msg}: {exc.__class__.__name__}: {exc}')


def get_current_live_video_id(channel_id):
    try:
        ytdl = youtube_dl.YoutubeDL(params=dict(
            quiet=True,
            youtube_include_dash_manifest=False,
        ))

        # Completely silence youtube-dl
        def no_op(*args, **kwargs):
            pass

        def silent_raise(*args, **kwargs):
            raise RuntimeError

        ytdl.report_warning = no_op
        ytdl.report_error = silent_raise

        info = ytdl.extract_info(f'https://youtube.com/channel/{channel_id}/live', download=False)
        video_id = info.get('id')
        if not isinstance(video_id, str):
            return None
        else:
            if not re.match(r'^[a-zA-Z0-9_-]{11}$', video_id):
                logger.warning(f'youtube-dl returned invalid video id: {video_id}')
                return None
            return video_id
    except Exception:
        return None


def get_live_viewership(video_id):
    try:
        response = youtube.videos().list(part='liveStreamingDetails', id=video_id).execute()
        live_details = response['items'][0]['liveStreamingDetails']
        if 'concurrentViewers' not in live_details:
            if 'actualEndTime' in live_details:
                raise LivestreamEndedError(video_id, live_details['actualEndTime'])
            elif 'actualStartTime' in live_details:
                logger.error(f'{video_id} is live, but concurrentViewers is not available')
                return None
            else:
                logger.error(f'{video_id} is not live')
                return None
        return live_details['concurrentViewers']
    except Exception as exc:
        if isinstance(exc, LivestreamEndedError):
            raise exc
        log_exception(exc, 'failed to extract concurrentViewers of livestream')
        return None


def get_start_time(video_id):
    try:
        response = youtube.videos().list(part='liveStreamingDetails', id=video_id).execute()
        return (arrow.get(response['items'][0]['liveStreamingDetails']['actualStartTime'])
                .to('Asia/Shanghai'))
    except Exception as exc:
        log_exception(exc, 'failed to extract concurrentViewers of livestream')
        return None


def plot(video_id, logfile):
    start_time = get_start_time(video_id)
    date = start_time.strftime('%Y%m%d')
    start_time_str = start_time.strftime('%Y-%m-%dT%H:%M%z')
    data = []
    peak_offset = 0
    peak_count = 0
    peak_offset_seconds = 0
    with open(logfile) as fp:
        for line in fp:
            timestamp, count = line.split()
            timestamp = int(timestamp)
            count = int(count)
            offset = int((timestamp - start_time.timestamp) / 60)
            data.append((offset, count))
            if count > peak_count:
                peak_count = count
                peak_offset = offset
                peak_offset_seconds = int(timestamp - start_time.timestamp)
    plt.figure(figsize=(15, 10))
    plt.plot(*zip(*data), linewidth=1, color='green')
    plt.xlim(xmin=0, xmax=max(offset for offset, _ in data))
    plt.ylim(ymin=0)
    plt.axvline(x=peak_offset, ymax=peak_count / plt.ylim()[1], linewidth=1, color='red')
    plt.text(peak_offset + 1, peak_count + 1, f'Peak: {peak_count}', color='red')
    plt.xlabel('Time since start (minutes)')
    plt.ylabel('Concurrent viewer count')
    plt.title(f'Live viewership curve for youtu.be/{video_id}\nStart time: {start_time_str}')
    for ext in ('svg', 'png'):
        path = PLOTS_DIR / f'{date}-{video_id}.{ext}'
        plt.savefig(os.fspath(path), dpi=120, bbox_inches='tight')
        logger.info(f'Generated image: {path}')
        optimizer = 'svgo' if ext == 'svg' else 'optipng'
        subprocess.run([optimizer, path], stdin=subprocess.DEVNULL,
                       stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    logger.info(f'Peak: {peak_count} at https://youtu.be/{video_id}?t={peak_offset_seconds}')
    return peak_count, peak_offset_seconds


def postprocess(video_id, logfile):
    try:
        response = youtube.videos().list(part='liveStreamingDetails', id=video_id).execute()
        live_details = response['items'][0]['liveStreamingDetails']
        start_time = arrow.get(live_details['actualStartTime']).to('Asia/Shanghai')
        end_time = arrow.get(live_details['actualEndTime']).to('Asia/Shanghai')
        date = start_time.strftime('%Y%m%d')

        # Livestream shorter than 15 minutes
        duration = end_time.timestamp - start_time.timestamp
        if duration < 900:
            logger.warning(f'{video_id} lasted for only {duration} seconds, '
                           f'shorter than the postprocessing threshold of 15 minutes')
            return

        peak_count, peak_offset_seconds = plot(video_id, logfile)

        with open(INDEX, 'a') as fp:
            print(date, video_id, file=fp)

        respobj = requests.get('https://snh48live.org/schedule/json/').json()
        events = list(reversed(respobj['past'])) + respobj['scheduled']
        related_events = []
        for event in events:
            event_start_time = arrow.get(event['datetime'])
            # Allow the livestream to be 30 minutes late
            if start_time.timestamp - 1800 <= event_start_time.timestamp <= end_time.timestamp:
                related_events.append(collections.OrderedDict([
                    ('datetime', event['datetime']),
                    ('title', event['title']),
                    ('subtitle', event['subtitle']),
                    ('thumbnail_url', urllib.parse.urljoin('https://snh48live.org/',
                                                           event['local_thumbnail_url'])),
                ]))
        event_names = [f'{e["title"]} {e["subtitle"]}' for e in related_events]
        logger.info('Events identified: %s' % '; '.join(event_names))
        data = collections.OrderedDict([
            ('video_id', video_id),
            ('start_time', str(start_time)),
            ('end_time', str(end_time)),
            ('events', related_events),
            ('plot', f'{date}-{video_id}.svg'),
            ('peak_viewers', peak_count),
            ('peak_offset', peak_offset_seconds),
        ])
        metadata_file = METADATA_DIR / f'{date}-{video_id}.json'
        with open(metadata_file, 'w') as fp:
            fp.write(json.dumps(data, indent=2, ensure_ascii=False))
        logger.info(f'Wrote metadata file: {metadata_file}')

        # Generate website
        subprocess.run([HERE / 'freeze'])
    except Exception as exc:
        log_exception(exc, f'failed to postprocess {viedo_id}')


def main():
    parser = auth.ArgumentParser()
    subparsers = parser.add_subparsers()
    parser.add_argument('--channel', default='UClwRU9iNX7UbzyuVzvZTSkA', help='channel id')
    plot_parser = subparsers.add_parser('plot', help='plot with existing data')
    plot_parser.add_argument('video_ids', nargs='+', metavar='video_id')
    plot_parser.set_defaults(plot=True)
    args = parser.parse_args()

    global youtube
    youtube = auth.get_youtube_client('youtube.readonly', args=args)

    if hasattr(args, 'plot'):
        retval = 0
        for video_id in args.video_ids:
            try:
                logfile = list(VIEWERS_DIR.glob(f'*-{video_id}.log'))[0]
            except IndexError:
                logger.error(f'video id {video_id} not found')
                retval = 1
            else:
                plot(video_id, logfile)
        sys.exit(retval)

    video_id = None
    try:
        while True:
            if not video_id:
                video_id = get_current_live_video_id(args.channel)
                if not video_id:
                    logger.info('No livestream at the moment')
                    time.sleep(60)
                    continue
                logger.info(f'Found livestream: https://youtu.be/{video_id}')
                start_time = get_start_time(video_id) or arrow.get().to('Asia/Shanghai')
                date = start_time.strftime('%Y%m%d')
                logfile = VIEWERS_DIR / f'{date}-{video_id}.log'
                with open(logfile, 'a') as fp:
                    pass

                chat_recorder = subprocess.Popen([HERE / 'record-chat', '--', video_id],
                                                 stdin=subprocess.DEVNULL)

            time.sleep(60 - time.time() % 60)
            now = int(time.time())
            try:
                viewership = get_live_viewership(video_id)
            except LivestreamEndedError as exc:
                logger.info(exc)
                chat_recorder.send_signal(signal.SIGINT)
                chat_recorder.wait(5)
                chat_recorder.kill()
                postprocess(video_id, logfile)
                video_id = None
                logfile = None
            else:
                if viewership is not None:
                    print(video_id, arrow.get(now).to('Asia/Shanghai'), viewership)
                    with open(logfile, 'a') as fp:
                        print(now, viewership, file=fp)
    except KeyboardInterrupt:
        pass


if __name__ == '__main__':
    main()