"""Script to run image capture screenshots for state data pages.
"""

from datetime import datetime
import io
import json
import os
from pytz import timezone
import sys
import time

from argparse import ArgumentParser, RawDescriptionHelpFormatter
import boto3
from loguru import logger
import pandas as pd
import requests


parser = ArgumentParser(
    description=__doc__,
    formatter_class=RawDescriptionHelpFormatter)

parser.add_argument(
    '--temp-dir',
    default='/tmp/public-cache',
    help='Local temp dir for snapshots')

parser.add_argument(
    '--s3-bucket',
    default='covid-data-archive',
    help='S3 bucket name')

parser.add_argument('--states',
    default='',
    help='Comma-separated list of state 2-letter names. If present, will only screenshot those.')

parser.add_argument('--push-to-s3', dest='push_to_s3', action='store_true', default=False,
    help='Push screenshots to S3')

parser.add_argument('--replace-most-recent-snapshot', action='store_true', default=False,
    help='If present, will first delete the most recent snapshot for the state before saving ' 
         'new screenshot to S3')

parser.add_argument('--phantomjscloud-key', default='',
    help='API key for PhantomJScloud, used for browser image capture')


# This is a set of (state name, suffix) tuples where we expect a PDF that we should just download.
# Otherwise, the normal behavior is to download the website as a .png image.
PDF_FILES = {('TN', 'secondary')}


class S3Backup():

    def __init__(self, bucket_name):
        self.s3 = boto3.resource('s3')
        self.bucket_name = bucket_name
        self.bucket = self.s3.Bucket(self.bucket_name)

    # uploads image (PNG) file with specified name
    def upload_file(self, local_path, s3_path):
        if local_path.endswith('.png'):
            extra_args = {'ContentType': 'image/png'}
        elif local_path.endswith('.pdf'):
            extra_args = {'ContentType': 'application/pdf', 'ContentDisposition': 'inline'}

        self.s3.meta.client.upload_file(local_path, self.bucket_name, s3_path, ExtraArgs=extra_args)

    # delete most recent snapshot with filename containing <state>-<suffix> or <state> if no suffix
    def delete_most_recent_snapshot(self, state, suffix=''):
        state_with_modifier = state if len(suffix) == 0 else '%s-%s' % (state, suffix)
        state_file_keys = [
            file.key for file in self.bucket.objects.all() if state_with_modifier in file.key]
        most_recent_state_key = sorted(state_file_keys, reverse=True)[0]
        self.s3.Object(self.bucket_name, most_recent_state_key).delete()


class Screenshotter():

    def __init__(self, local_dir, s3_backup, phantomjscloud_key):
        self.phantomjs_url = 'https://phantomjscloud.com/api/browser/v2/%s/' % phantomjscloud_key
        self.local_dir = local_dir
        self.s3_backup = s3_backup

    # makes a PhantomJSCloud call to data_url and saves the output to specified path
    def save_url_image_to_path(self, state, data_url, path):
        """Saves URL image from data_url to the specified path.

        Parameters
        ----------
        state : str
            Two-letter abbreviation of the state or territory. Used for special-casing sizes, etc.

        data_url : str
            URL of data site to save

        path : str
            Local path to which to save .png screenshot of data_url
        """
        logger.info(f"Retrieving {data_url}")

        # if we're just saving a PDF, don't use phantomjscloud
        if path.endswith('.pdf'):
            logger.info(f"Downloading PDF from {data_url}")
            response = requests.get(data_url)
            if response.status_code == 200:
                with open(path, 'wb') as f:
                    f.write(response.content)
                return
            else:
                logger.error(f'Response status code: {response.status_code}')
                raise ValueError(f'Could not retrieve URL: {data_url}')

        data = {
            'url': data_url,
            'renderType': 'png',
        }

        # PhantomJScloud gets the page length wrong for some states, need to set those manually
        if state in ['ID', 'PA', 'CA']:
            logger.info(f"using larger viewport for state {state}")
            data['renderSettings'] = {'viewport': {'width': 1400, 'height': 3000}}
        elif state == 'NE':
            # really huge viewport for some reason
            logger.info(f"using huge viewport for state {state}")
            data['renderSettings'] = {'viewport': {'width': 1400, 'height': 5000}}
        elif state == 'IN':
            # even huger viewport
            logger.info(f"using huger viewport for state {state}")
            data['renderSettings'] = {'viewport': {'width': 1400, 'height': 8500}}
            data['overseerScript'] = 'page.manualWait(); \
                                      await page.waitForSelector("#prefix-dismissButton"); \
                                      page.click("#prefix-dismissButton"); \
                                      await page.waitForFunction(()=>document.querySelector("#main-content").textContent!==""); \
                                      page.done();'
        elif state == 'UT':
            # Utah dashboard doesn't render in phantomjscloud unless I set clipRectangle
            data['renderSettings'] = {'clipRectangle': {'width': 1400, 'height': 3000}}

        elif state == 'TX':
            # TX dashboard takes forever and a half
            data['overseerScript'] = """page.manualWait();
                                      await page.waitForDelay(10000);
                                      page.done();"""

        # for the CDC testing tab, need to do clicking magic
        elif state == 'CDC' and 'testing' in path:
            # try clicking on a tab somewhere there
            logger.info(f"Custom CDC logic")
            data['overseerScript'] = """page.manualWait();
                                      await page.waitForSelector("[data-tabname='tabAllLabs']");
                                      page.click("[data-tabname='tabAllLabs']", {delay: 150});
                                      await page.waitForFunction(()=>document.querySelector("#mainContent_Title").textContent=="United States Laboratory Testing");
                                      await page.waitForDelay(1000);
                                      page.done();"""
        logger.info('Posting request...')
        response = requests.post(self.phantomjs_url, json.dumps(data))
        logger.info('Done.')

        if response.status_code == 200:
            with open(path, 'wb') as f:
                f.write(response.content)
        else:
            logger.error(f'Response status code: {response.status_code}')
            raise ValueError(f'Could not retrieve URL: {data_url}')

    @staticmethod
    def timestamped_filename(state, suffix=''):
        # we default to .png files except for some state/suffix pairs
        if (state, suffix) in PDF_FILES:
            fileext = 'pdf'
        else:
            fileext = 'png'
        state_with_modifier = state if len(suffix) == 0 else '%s-%s' % (state, suffix)
        timestamp = datetime.now(timezone('US/Eastern')).strftime("%Y%m%d-%H%M%S")
        return "%s-%s.%s" % (state_with_modifier, timestamp, fileext)

    @staticmethod
    def get_s3_path(state, suffix=''):
        filename = Screenshotter.timestamped_filename(state, suffix)
        # CDC goes into its own top-level folder to not mess with state_screenshots
        if state == 'CDC':
            return os.path.join(state, filename)
        else:
            return os.path.join('state_screenshots', state, filename)

    def get_local_path(self, state, suffix=''):
        # basename will be e.g. 'CA' if no suffix, or 'CA-secondary' if suffix is 'secondary'
        filename = Screenshotter.timestamped_filename(state, suffix)
        return os.path.join(self.local_dir, filename)

    def screenshot(self, state, data_url, suffix='',
            backup_to_s3=False, replace_most_recent_snapshot=False):
        """Screenshots state data site.

        Parameters
        ----------
        state : str
            Two-letter abbreviation of the state or territory

        data_url : str
            URL containing data site to screenshot

        suffix : str
            If present, will be used in the resulting local and S3 filename (STATE_suffix)

        backup_to_s3 : bool
            If true, will push to S3. If false, backup will be only local

        replace_most_recent_snapshot : bool
            If true, will delete the most recent S3 snapshot corresponding to this state
            (and suffix) before pushing the one taken during this run. Largely used to manually
            replace faulty state screenshots as we learn about them.
        """

        logger.info(f"Screenshotting {state} {suffix} from {data_url}")
        local_path = self.get_local_path(state, suffix)
        self.save_url_image_to_path(state, data_url, local_path)
        if backup_to_s3:
            s3_path = self.get_s3_path(state, suffix)
            if replace_most_recent_snapshot:
                logger.info(f"    3a. first delete the most recent snapshot")
                self.s3_backup.delete_most_recent_snapshot(state, suffix)
            logger.info(f"    4. push to s3 at {s3_path}")
            self.s3_backup.upload_file(local_path, s3_path)


def main(args_list=None):
    if args_list is None:
        args_list = sys.argv[1:]
    args = parser.parse_args(args_list)

    s3 = S3Backup(bucket_name=args.s3_bucket)
    screenshotter = Screenshotter(
        local_dir=args.temp_dir, s3_backup=s3, phantomjscloud_key=args.phantomjscloud_key)

    # get states info from API
    url = 'https://covidtracking.com/api/states/info.csv'
    content = requests.get(url).content
    state_info_df = pd.read_csv(io.StringIO(content.decode('utf-8')))
    
    failed_states = []

    # screenshot state images
    if args.states:
        logger.info(f'Snapshotting states {args.states}')
        states_list = args.states.split(',')
        state_info_df = state_info_df[state_info_df.state.isin(states_list)]

    for idx, r in state_info_df.iterrows():
        state = r["state"]
        data_url = r["covid19Site"]
        secondary_data_url = r["covid19SiteSecondary"]
        try:
            screenshotter.screenshot(
                state, data_url,
                backup_to_s3=args.push_to_s3,
                replace_most_recent_snapshot=args.replace_most_recent_snapshot)
            if not pd.isnull(secondary_data_url):
                screenshotter.screenshot(
                    state, secondary_data_url, suffix='secondary',
                    backup_to_s3=args.push_to_s3,
                    replace_most_recent_snapshot=args.replace_most_recent_snapshot)
        except:
            failed_states.append(state)

    if failed_states:
        logger.error(f"Failed states for this run: {','.join(failed_states)}")
    else:
        logger.info("All required states successfully screenshotted")

    # special-case: screenshot CDC "US Cases" and "US COVID Testing" tabs
    cdc_link = 'https://www.cdc.gov/covid-data-tracker/'
    screenshotter.screenshot('CDC', cdc_link, suffix='testing', backup_to_s3=args.push_to_s3)
    screenshotter.screenshot('CDC', cdc_link, suffix='cases', backup_to_s3=args.push_to_s3)


if __name__ == "__main__":
    main()
