#AUTOGENERATED! DO NOT EDIT! File to edit: dev/07_scraping.ipynb (unless otherwise specified).

__all__ = ['fetch_img_from_url', 'load_praw_credentials', 'is_imgur_page', 'fetch_imgur_url',
           'scrape_subreddit_images_1k', 'scrape_subreddit_images', 'remove_corrupt_images']

#Cell
from bs4 import BeautifulSoup
import cv2
import glob
import logging
import os
import pandas as pd
from pathlib import Path
from PIL import Image
import praw
import re
import requests

from accio.s3tool import S3tool
from .utils import IMG_EXT

#Cell
def fetch_img_from_url(url, out_path, verbose=False):
    """Given a URL, fetch an image and download it to the specified path.

    Parameters
    ----------
    url: str
        Location of image online.
    out_path: str
        Path to download the image to.
    verbose: bool
        If True, prints a message alerting the user when the image could not be
        retrieved.

    Returns
    -------
    bool: Specifies whether image was successfully retrieved.
    """
    try:
        with requests.get(url, stream=True, timeout=10) as r:
            if r.status_code != 200:
                if verbose: print(f'STATUS CODE ERROR: {url}')
                return False

            # Write bytes to file chunk by chunk.
            with open(out_path, 'wb') as f:
                for chunk in r.iter_content(256):
                    f.write(chunk)

    # Any time url cannot be accessed, don't care about exact error.
    except Exception as e:
        if verbose: print(e)
        return False

    return True

#Cell
def load_praw_credentials(bucket='gg-datascience',
                          s3_path='hmamin/praw_creds.csv'):
    """Load credentials for dune script owned by hdmamin user on reddit.

    Returns
    -------
    dict[str, str]: Keyword arguments that can be passed directly to
        praw.Reddit to instantiate a new instance.
    """
    tmp_dir = '/tmp'
    tmp_file = os.path.join(tmp_dir, 'praw_creds.csv')
    s3 = S3tool()
    s3.download_file(bucket, s3_path, tmp_dir, maintain_s3_partitions=False)

    # Read credentials into dict and delete temporary file.
    with open(tmp_file, 'r') as f:
        cols, vals = [line.strip().split(',') for line in f]
    os.remove(tmp_file)
    return dict(zip(cols, vals))

#Cell
def is_imgur_page(url):
    """Check if a URL is a link to an imgur html page. Use as validation before
    trying to fetch an image directly.

    Returns
    -------
    bool
    """
    url = url.lower()
    return 'imgur' in url and not any(ext in url.split('.')[-1]
                                      for ext in IMG_EXT)

#Cell
def fetch_imgur_url(url):
    """Fetch and extract an image url from an imgur page. Returns None if an
    image is not found.
    """
    try:
        r = requests.get(url)

    # If url throws error, see if it is one of the links formatted like this:
    # ex: http://[Imgur](http://i.imgur.com/1gmLlya)
    # If so, recursive call on part in parens. Otherwise, return None.
    except:
        parens_links = re.findall(r'\((.+)\)', url)
        if parens_links:
            return fetch_imgur_url(parens_links[0])
        else:
            return

    if r.status_code != 200: return
    img_links = BeautifulSoup(r.content, features='lxml')\
        .find_all('link', rel='image_src')
    try:
        return img_links[0]['href']
    except:
        return

#Cell
def scrape_subreddit_images_1k(sub_name, out_dir, n=None, verbose=True,
                               print_freq=50):
    """Scrape subreddit using PRAW. Limited to 1,000 posts per call, and
    haven't found way to use this iteratively to get all posts.

    Parameters
    ----------
    sub_name: str
        Name of subreddit.
    out_dir: str
        Name of output image directory.
    n: int
        Max number of posts to scrape. If None, scrapes as many as the
        subreddit has (or until Praw limit is reached).
    verbose: bool

    print_freq: n
        Frequency with which to print message with number of posts scraped.
        Makes monitoring easier.

    Returns
    -------
    pd.DataFrame: Two columns consisting of image urls and post urls.
    """
    img_dir = os.path.join(out_dir, 'test', sub_name)
    os.makedirs(img_dir, exist_ok=True)
    creds = load_praw_credentials()
    reddit = praw.Reddit(**creds)
    sub = reddit.subreddit(sub_name)

    # Scrape top level image from each post.
    posts = []
    for i, post in enumerate(sub.top(limit=n), 1):
        if i % print_freq == 0: print(f'Scraping post {i}...')
        url = post.url
        if is_imgur_page(url): url = fetch_imgur_url(url)
        if not url: continue

        # Check image is fetched, then check file is not empty/corrupt.
        fname = re.sub('\?\d*$', '', url.strip('/')).split('/')[-1]
        img_path = os.path.join(img_dir, fname)
        if fetch_img_from_url(url, img_path, verbose):
            # Animated images can't be opened by cv2, which is used by
            # ImageSequence. In future Pillow may be better choice.
            try:
                img = cv2.imread(img_path)
                assert img is not None, 'Image must not be None.'
            except Exception as e:
                if verbose: print(f'IMAGE OPEN ERROR: {fname}')
                os.remove(img_path)
                continue
            posts.append((url, post.permalink))

    # Save csv with image and post urls.
    df = pd.DataFrame(posts, columns=['url', 'post_url'])
    df['post_url'] = 'https://www.reddit.com' + df.post_url
    df.to_csv(os.path.join(out_dir, 'urls.csv'), index=False)
    return df

#Cell
def scrape_subreddit_images(sub, root_dir=None, steps=None, step_size=1_000,
                            after_utc=None, verbose=True, print_freq=50):
    """Scrape reddit images using the pushshift api. This is sometimes
    necessary because PRAW limits us to 1,000 images. Here, we still have to
    collect <=1000 images per api call but we can sort and filter by timestamp
    to accumulate all posts in the subreddit.

    Parameters
    ----------
    sub: str
        Name of subreddit to scrape.
    root_dir: str
        Output directory to save results in. This will contain a subdirectory
        'images' containing images and a text file 'urls.csv' in the root.
        If None is specified, defaults to {sub}_images/.
    steps: int or None
        Number of pushshift api calls to make, where each call contains
        `step_size` posts. Default is None which lets us get all posts in the
        subreddit.
    step_size: int
        Number of posts per api call. Max available is 1,000, but when testing
        we may want to use a smaller number.
    verbose: bool
        If True, print error messages when applicable. This is separate from
        the print_freq functionality, which is concerned with updates rather
        than errors.
    print_freq: int
        Frequency with which to print messages mentioning which image is being
        scraped.
    """
    root_dir = root_dir or sub + '_images'
    img_dir = os.path.join(root_dir, 'test', sub)
    os.makedirs(img_dir, exist_ok=True)

    # Use logger to track last timestamp scraped in case process fails midway.
    logging.basicConfig(handlers=[
            logging.FileHandler(os.path.join(root_dir, 'scrape.log'),
            mode='w')
        ],
        format='%(message)s',
        level=logging.INFO
    )
    logger = logging.getLogger()

    curr_url = push_url = (f'http://api.pushshift.io/reddit/search/submission/'
                           f'?subreddit={sub}&size={step_size}&sort_type='
                           f'created_utc&sort=asc')
    if after_utc: latest = after_utc

    metadata = []
    i = 1
    n_scraped = 0
    while True:
        if i > 1 or after_utc: curr_url = f'{push_url}&after={latest}'
        if verbose: print(f'\nPushshift call #{i}...')

        # Get dict of posts from pushshift.
        r = requests.get(curr_url)
        if r.status_code != 200: break
        data = r.json()['data']
        # When no more posts are found, data is an empty list.
        if not data: break

        # Iterate over each post returned from 1 api call.
        latest = data[-1]['created_utc']
        for post in data:
            img_url = post['url']
            if is_imgur_page(img_url): img_url = fetch_imgur_url(img_url)
            if not img_url: continue

            # Confirm image is fetched, then confirm file is not empty/corrupt.
            fname = re.sub('\?\d*$', '', url.strip('/')).split('/')[-1]
            img_path = os.path.join(img_dir, fname)
            if fetch_img_from_url(img_url, img_path, verbose):
                # Animated images can't be opened by cv2, which is used by
                # ImageSequence. In future Pillow may be better choice.
                try:
                    img = cv2.imread(img_path)
                    assert img is not None, 'Image must not be None.'
                except Exception as e:
                    if verbose: print(f'IMAGE OPEN ERROR: {str(e)} ({fname})')
                    os.remove(img_path)
                    continue

                # Only store for images that were successfully saved.
                metadata.append((post['full_link'], img_url))
                logger.info(post['created_utc'])
                if n_scraped % print_freq == 0:
                    print(f'>>> Scraping image #{n_scraped}: {img_url}')
                n_scraped += 1

        if i == steps: break
        i += 1

    df_kwargs = dict(path_or_buf=os.path.join(root_dir, 'urls.csv'),
                     index=False)
    if after_utc:
        df_kwargs['mode'] = 'a'
        df_kwargs['header'] = False
    df = pd.DataFrame(metadata, columns=['post_url', 'img_url'])
    df.to_csv(**df_kwargs)
    return df

#Cell
def remove_corrupt_images(dir_, mode='move', cv=True, verbose=True):
    """This can be used to remove empty, corrupt, or animated images
    from a directory. Note that some images can be opened by Pillow but
    not cv2, which is what ImageSequence uses to load images.

    Parameters
    ----------
    dir_: str
        Directory of images.
    mode: str
        'move' to move images 1 directory up (safer). 'remove' to remove
        images.
    cv: bool
        True to use cv2, False to use Pillow. Pillow can open animated gifs
        and cv2 can't. Keep in mind spellotape.dl.ImageSequence uses cv2.

    Returns
    -------
    int: Number of images moved or deleted.
    """
    def test_cv2(path):
        img = cv2.imread(path)
        assert img is not None, 'Image must not be None.'

    def test_pillow(path):
        img = Image.open(path)

    open_img = test_cv2 if cv else test_pillow
    changed = 0
    for i, path in enumerate(glob.glob(os.path.join(dir_, '*'))):
        if i % 50 == 0: print(i)
        try:
            open_img(path)
        except Exception as e:
            if verbose: print('error:', e, 'path:', path)
            if mode == 'move':
                path_obj = Path(path)
                os.rename(path,
                          os.path.join(path_obj.parents[1], path_obj.name))
            elif mode == 'remove':
                os.remove(path)
            changed += 1

    return changed